From 1695d4d5e9a4859db75f9c25a2c29751148a9c47 Mon Sep 17 00:00:00 2001
From: Christopher Kolstad <chriswk@getunleash.io>
Date: Fri, 24 Nov 2023 10:38:45 +0100
Subject: [PATCH] feat: Added the possibility to add a custom
 startupExceptionHandler

---
 .../repository/FeatureRepository.java         |  18 ++-
 .../io/getunleash/util/UnleashConfig.java     |  43 +++++--
 .../io/getunleash/DefaultUnleashTest.java     |  63 +++++-----
 .../repository/FeatureRepositoryTest.java     | 110 ++++++++++--------
 4 files changed, 141 insertions(+), 93 deletions(-)

diff --git a/src/main/java/io/getunleash/repository/FeatureRepository.java b/src/main/java/io/getunleash/repository/FeatureRepository.java
index 36e54c33a..dfeb8a61a 100644
--- a/src/main/java/io/getunleash/repository/FeatureRepository.java
+++ b/src/main/java/io/getunleash/repository/FeatureRepository.java
@@ -95,9 +95,15 @@ private void initCollections(UnleashScheduledExecutor executor) {
         }
 
         if (unleashConfig.isSynchronousFetchOnInitialisation()) {
-            updateFeatures(e -> {
-                throw e;
-            }).run();
+            if (this.unleashConfig.getStartupExceptionHandler() != null) {
+                updateFeatures(this.unleashConfig.getStartupExceptionHandler()).run();
+            } else {
+                updateFeatures(
+                                e -> {
+                                    throw e;
+                                })
+                        .run();
+            }
         }
 
         if (!unleashConfig.isDisablePolling()) {
@@ -128,7 +134,11 @@ private Runnable updateFeatures(final Consumer<UnleashException> handler) {
                         featureBackupHandler.write(featureCollection);
                     } else if (response.getStatus() == ClientFeaturesResponse.Status.UNAVAILABLE) {
                         if (!ready && unleashConfig.isSynchronousFetchOnInitialisation()) {
-                            throw new UnleashException(String.format("Could not initialize Unleash, got response code %d", response.getHttpStatusCode()), null);
+                            throw new UnleashException(
+                                    String.format(
+                                            "Could not initialize Unleash, got response code %d",
+                                            response.getHttpStatusCode()),
+                                    null);
                         }
                         if (ready) {
                             throttler.handleHttpErrorCodes(response.getHttpStatusCode());
diff --git a/src/main/java/io/getunleash/util/UnleashConfig.java b/src/main/java/io/getunleash/util/UnleashConfig.java
index 53428b1cc..601200b82 100644
--- a/src/main/java/io/getunleash/util/UnleashConfig.java
+++ b/src/main/java/io/getunleash/util/UnleashConfig.java
@@ -5,6 +5,7 @@
 import io.getunleash.CustomHttpHeadersProvider;
 import io.getunleash.DefaultCustomHttpHeadersProviderImpl;
 import io.getunleash.UnleashContextProvider;
+import io.getunleash.UnleashException;
 import io.getunleash.event.NoOpSubscriber;
 import io.getunleash.event.UnleashSubscriber;
 import io.getunleash.lang.Nullable;
@@ -14,19 +15,15 @@
 import io.getunleash.strategy.Strategy;
 import java.io.File;
 import java.math.BigInteger;
-import java.net.Authenticator;
-import java.net.HttpURLConnection;
-import java.net.InetAddress;
-import java.net.InetSocketAddress;
-import java.net.PasswordAuthentication;
-import java.net.Proxy;
-import java.net.URI;
-import java.net.UnknownHostException;
+import java.net.*;
 import java.nio.charset.StandardCharsets;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.time.Duration;
-import java.util.*;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.function.Consumer;
 
 public class UnleashConfig {
 
@@ -71,6 +68,7 @@ public class UnleashConfig {
     @Nullable private final Strategy fallbackStrategy;
     @Nullable private final ToggleBootstrapProvider toggleBootstrapProvider;
     @Nullable private final Proxy proxy;
+    @Nullable private final Consumer<UnleashException> startupExceptionHandler;
 
     private UnleashConfig(
             @Nullable URI unleashAPI,
@@ -101,7 +99,8 @@ private UnleashConfig(
             @Nullable Strategy fallbackStrategy,
             @Nullable ToggleBootstrapProvider unleashBootstrapProvider,
             @Nullable Proxy proxy,
-            @Nullable Authenticator proxyAuthenticator) {
+            @Nullable Authenticator proxyAuthenticator,
+            @Nullable Consumer<UnleashException> startupExceptionHandler) {
 
         if (appName == null) {
             throw new IllegalStateException("You are required to specify the unleash appName");
@@ -165,6 +164,7 @@ private UnleashConfig(
         this.metricSenderFactory = metricSenderFactory;
         this.clientSpecificationVersion =
                 UnleashProperties.getProperty("client.specification.version");
+        this.startupExceptionHandler = startupExceptionHandler;
     }
 
     public static Builder builder() {
@@ -334,6 +334,11 @@ public UnleashFeatureFetcherFactory getUnleashFeatureFetcherFactory() {
         return this.unleashFeatureFetcherFactory;
     }
 
+    @Nullable
+    public Consumer<UnleashException> getStartupExceptionHandler() {
+        return startupExceptionHandler;
+    }
+
     static class SystemProxyAuthenticator extends Authenticator {
         @Override
         protected @Nullable PasswordAuthentication getPasswordAuthentication() {
@@ -427,6 +432,8 @@ public static class Builder {
         private @Nullable Proxy proxy;
         private @Nullable Authenticator proxyAuthenticator;
 
+        private @Nullable Consumer<UnleashException> startupExceptionHandler;
+
         private static String getHostname() {
             String hostName = System.getProperty("hostname");
             if (hostName == null || hostName.isEmpty()) {
@@ -657,6 +664,19 @@ public Builder apiKey(String apiKey) {
             return this;
         }
 
+        /**
+         * Used to handle exceptions when starting up synchronously. Allows user the option to
+         * choose how errors should be handled.
+         *
+         * @param startupExceptionHandler - a lambda taking the Exception and doing what it wants to
+         *     the system.
+         */
+        public Builder startupExceptionHandler(
+                @Nullable Consumer<UnleashException> startupExceptionHandler) {
+            this.startupExceptionHandler = startupExceptionHandler;
+            return this;
+        }
+
         public UnleashConfig build() {
             return new UnleashConfig(
                     unleashAPI,
@@ -688,7 +708,8 @@ public UnleashConfig build() {
                     fallbackStrategy,
                     toggleBootstrapProvider,
                     proxy,
-                    proxyAuthenticator);
+                    proxyAuthenticator,
+                    startupExceptionHandler);
         }
 
         public String getDefaultSdkVersion() {
diff --git a/src/test/java/io/getunleash/DefaultUnleashTest.java b/src/test/java/io/getunleash/DefaultUnleashTest.java
index d0bb1328c..beebaa196 100644
--- a/src/test/java/io/getunleash/DefaultUnleashTest.java
+++ b/src/test/java/io/getunleash/DefaultUnleashTest.java
@@ -14,13 +14,11 @@
 import io.getunleash.event.EventDispatcher;
 import io.getunleash.event.UnleashReady;
 import io.getunleash.event.UnleashSubscriber;
-import io.getunleash.integration.TestDefinition;
 import io.getunleash.metric.UnleashMetricService;
 import io.getunleash.repository.*;
 import io.getunleash.strategy.DefaultStrategy;
 import io.getunleash.strategy.Strategy;
 import io.getunleash.util.UnleashConfig;
-
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.util.*;
@@ -40,10 +38,10 @@ class DefaultUnleashTest {
 
     @RegisterExtension
     static WireMockExtension serverMock =
-        WireMockExtension.newInstance()
-            .configureStaticDsl(true)
-            .options(wireMockConfig().dynamicPort().dynamicHttpsPort())
-            .build();
+            WireMockExtension.newInstance()
+                    .configureStaticDsl(true)
+                    .options(wireMockConfig().dynamicPort().dynamicHttpsPort())
+                    .build();
 
     @BeforeEach
     public void setup() {
@@ -258,49 +256,51 @@ public void synchronous_fetch_on_initialisation_fails_on_initialization() {
 
     @ParameterizedTest
     @ValueSource(ints = {401, 403, 404, 500})
-    public void synchronous_fetch_on_initialisation_fails_on_non_200_response(int code) throws URISyntaxException {
+    public void synchronous_fetch_on_initialisation_fails_on_non_200_response(int code)
+            throws URISyntaxException {
         mockUnleashAPI(code);
         IsReadyTestSubscriber readySubscriber = new IsReadyTestSubscriber();
         UnleashConfig config =
-            UnleashConfig.builder()
-                .unleashAPI(new URI("http://localhost:" + serverMock.getPort() + "/api/"))
-                .appName("wrong_upstream")
-                .apiKey("default:development:1234567890123456")
-                .instanceId("non-200")
-                .synchronousFetchOnInitialisation(true)
-                .subscriber(readySubscriber)
-                .build();
+                UnleashConfig.builder()
+                        .unleashAPI(new URI("http://localhost:" + serverMock.getPort() + "/api/"))
+                        .appName("wrong_upstream")
+                        .apiKey("default:development:1234567890123456")
+                        .instanceId("non-200")
+                        .synchronousFetchOnInitialisation(true)
+                        .subscriber(readySubscriber)
+                        .build();
 
         assertThatThrownBy(() -> new DefaultUnleash(config)).isInstanceOf(UnleashException.class);
         assertThat(readySubscriber.ready).isFalse();
     }
 
     @Test
-    public void synchronous_fetch_on_initialisation_switches_to_ready_on_200() throws URISyntaxException {
+    public void synchronous_fetch_on_initialisation_switches_to_ready_on_200()
+            throws URISyntaxException {
         mockUnleashAPI(200);
         IsReadyTestSubscriber readySubscriber = new IsReadyTestSubscriber();
         UnleashConfig config =
-            UnleashConfig.builder()
-                .unleashAPI(new URI("http://localhost:" + serverMock.getPort() + "/api/"))
-                .appName("wrong_upstream")
-                .apiKey("default:development:1234567890123456")
-                .instanceId("with-success-response")
-                .synchronousFetchOnInitialisation(true)
-                .subscriber(readySubscriber)
-                .build();
+                UnleashConfig.builder()
+                        .unleashAPI(new URI("http://localhost:" + serverMock.getPort() + "/api/"))
+                        .appName("wrong_upstream")
+                        .apiKey("default:development:1234567890123456")
+                        .instanceId("with-success-response")
+                        .synchronousFetchOnInitialisation(true)
+                        .subscriber(readySubscriber)
+                        .build();
         new DefaultUnleash(config);
         assertThat(readySubscriber.ready).isTrue();
     }
 
     private void mockUnleashAPI(int featuresStatusCode) {
         stubFor(
-            get(urlEqualTo("/api/client/features"))
-                .withHeader("Accept", equalTo("application/json"))
-                .willReturn(
-                    aResponse()
-                        .withStatus(featuresStatusCode)
-                        .withHeader("Content-Type", "application/json")
-                        .withBody("{\"features\": []}")));
+                get(urlEqualTo("/api/client/features"))
+                        .withHeader("Accept", equalTo("application/json"))
+                        .willReturn(
+                                aResponse()
+                                        .withStatus(featuresStatusCode)
+                                        .withHeader("Content-Type", "application/json")
+                                        .withBody("{\"features\": []}")));
         stubFor(post(urlEqualTo("/api/client/register")).willReturn(aResponse().withStatus(200)));
     }
 
@@ -344,6 +344,7 @@ public void client_identifier_handles_api_key_being_null() {
 
     private static class IsReadyTestSubscriber implements UnleashSubscriber {
         public boolean ready = false;
+
         public void onReady(UnleashReady unleashReady) {
             this.ready = true;
         }
diff --git a/src/test/java/io/getunleash/repository/FeatureRepositoryTest.java b/src/test/java/io/getunleash/repository/FeatureRepositoryTest.java
index c0641ef3e..2986eea51 100644
--- a/src/test/java/io/getunleash/repository/FeatureRepositoryTest.java
+++ b/src/test/java/io/getunleash/repository/FeatureRepositoryTest.java
@@ -1,17 +1,15 @@
 package io.getunleash.repository;
 
+import static io.getunleash.repository.FeatureToggleResponse.Status.*;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.*;
+
 import io.getunleash.*;
 import io.getunleash.event.EventDispatcher;
 import io.getunleash.lang.Nullable;
 import io.getunleash.util.UnleashConfig;
 import io.getunleash.util.UnleashScheduledExecutor;
-import org.jetbrains.annotations.NotNull;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
-import org.junit.jupiter.params.ParameterizedTest;
-import org.junit.jupiter.params.provider.ValueSource;
-import org.mockito.ArgumentCaptor;
-
 import java.io.File;
 import java.io.IOException;
 import java.net.URISyntaxException;
@@ -21,11 +19,13 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
-
-import static io.getunleash.repository.FeatureToggleResponse.Status.*;
-import static org.assertj.core.api.Assertions.assertThat;
-import static org.junit.jupiter.api.Assertions.*;
-import static org.mockito.Mockito.*;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.jetbrains.annotations.NotNull;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.mockito.ArgumentCaptor;
 
 public class FeatureRepositoryTest {
     FeatureBackupHandlerFile backupHandler;
@@ -106,14 +106,14 @@ public void feature_toggles_should_be_updated() {
 
         when(backupHandler.read()).thenReturn(simpleFeatureCollection(false));
 
-
         FeatureRepository featureRepository =
                 new FeatureRepository(config, backupHandler, executor, fetcher, bootstrapHandler);
         // run the toggleName fetcher callback
         verify(executor).setInterval(runnableArgumentCaptor.capture(), anyLong(), anyLong());
         verify(fetcher, times(0)).fetchFeatures();
 
-        ClientFeaturesResponse response = new ClientFeaturesResponse(CHANGED, simpleFeatureCollection(true));
+        ClientFeaturesResponse response =
+                new ClientFeaturesResponse(CHANGED, simpleFeatureCollection(true));
         when(fetcher.fetchFeatures()).thenReturn(response);
         runnableArgumentCaptor.getValue().run();
 
@@ -184,9 +184,7 @@ public void should_not_perform_synchronous_fetch_on_initialisation() {
         when(backupHandler.read()).thenReturn(new FeatureCollection());
 
         FeatureCollection featureCollection = populatedFeatureCollection(null);
-        ClientFeaturesResponse response =
-                new ClientFeaturesResponse(
-                        CHANGED, featureCollection);
+        ClientFeaturesResponse response = new ClientFeaturesResponse(CHANGED, featureCollection);
 
         when(fetcher.fetchFeatures()).thenReturn(response);
 
@@ -250,15 +248,35 @@ public void should_not_read_bootstrap_if_backup_was_found()
 
         when(toggleBootstrapProvider.read()).thenReturn(fileToString(file));
 
-        when(backupHandler.read())
-                .thenReturn(
-                    getFeatureCollection());
+        when(backupHandler.read()).thenReturn(getFeatureCollection());
 
         new FeatureRepository(
                 config, backupHandler, new EventDispatcher(config), fetcher, bootstrapHandler);
         verify(toggleBootstrapProvider, times(0)).read();
     }
 
+    @Test
+    public void shouldCallStartupExceptionHandlerIfStartupFails() {
+        ToggleBootstrapProvider toggleBootstrapProvider = mock(ToggleBootstrapProvider.class);
+        UnleashScheduledExecutor executor = mock(UnleashScheduledExecutor.class);
+        AtomicBoolean failed = new AtomicBoolean(false);
+        UnleashConfig config =
+                UnleashConfig.builder()
+                        .synchronousFetchOnInitialisation(true)
+                        .startupExceptionHandler(
+                                (e) -> {
+                                    failed.set(true);
+                                })
+                        .appName("test-sync-update")
+                        .scheduledExecutor(executor)
+                        .unleashAPI("http://localhost:8080")
+                        .toggleBootstrapProvider(toggleBootstrapProvider)
+                        .build();
+
+        Unleash unleash = new DefaultUnleash(config);
+        assertThat(failed).isTrue();
+    }
+
     @ParameterizedTest
     @ValueSource(ints = {403, 404})
     public void should_increase_to_max_interval_when_code(int code)
@@ -460,35 +478,31 @@ private String fileToString(File f) throws IOException {
     @NotNull
     private FeatureCollection simpleFeatureCollection(boolean enabled) {
         return populatedFeatureCollection(
-            null,
-            new FeatureToggle(
-                "toggleFetcherCalled",
-                enabled,
-                Arrays.asList(new ActivationStrategy("custom", null))));
+                null,
+                new FeatureToggle(
+                        "toggleFetcherCalled",
+                        enabled,
+                        Arrays.asList(new ActivationStrategy("custom", null))));
     }
 
     @NotNull
     private FeatureCollection getFeatureCollection() {
         return populatedFeatureCollection(
-            Arrays.asList(
-                new Segment(
-                    1,
-                    "some-name",
-                    Arrays.asList(
-                        new Constraint(
-                            "some-context",
-                            Operator.IN,
-                            "some-value")))),
-            new FeatureToggle(
-                "toggleFeatureName1",
-                true,
-                Collections.singletonList(
-                    new ActivationStrategy("custom", null))),
-            new FeatureToggle(
-                "toggleFeatureName2",
-                true,
-                Collections.singletonList(
-                    new ActivationStrategy("custom", null))));
+                Arrays.asList(
+                        new Segment(
+                                1,
+                                "some-name",
+                                Arrays.asList(
+                                        new Constraint(
+                                                "some-context", Operator.IN, "some-value")))),
+                new FeatureToggle(
+                        "toggleFeatureName1",
+                        true,
+                        Collections.singletonList(new ActivationStrategy("custom", null))),
+                new FeatureToggle(
+                        "toggleFeatureName2",
+                        true,
+                        Collections.singletonList(new ActivationStrategy("custom", null))));
     }
 
     private class TestRunner {
@@ -506,15 +520,17 @@ public TestRunner() {
 
         private void ensureInitialized() {
             if (!initialized) {
-                verify(executor).setInterval(runnableArgumentCaptor.capture(), anyLong(), anyLong());
+                verify(executor)
+                        .setInterval(runnableArgumentCaptor.capture(), anyLong(), anyLong());
                 initialized = true;
             }
         }
 
-        public void assertThatFetchesAndReceives(FeatureToggleResponse.Status status, int statusCode) {
+        public void assertThatFetchesAndReceives(
+                FeatureToggleResponse.Status status, int statusCode) {
             ensureInitialized();
             when(fetcher.fetchFeatures())
-                .thenReturn(new ClientFeaturesResponse(status, statusCode));
+                    .thenReturn(new ClientFeaturesResponse(status, statusCode));
             runnableArgumentCaptor.getValue().run();
             verify(fetcher, times(++count)).fetchFeatures();
         }