fix: Improve client initialization handling. (#27)

kinyoklion · web-flow · commit cf0ed18a91a3 · 2024-06-06T14:52:00.000-07:00
diff --git a/src/main/java/com/launchdarkly/openfeature/serverprovider/Provider.java b/src/main/java/com/launchdarkly/openfeature/serverprovider/Provider.java
@@ -11,8 +11,9 @@
 import dev.openfeature.sdk.*;
 
 import java.io.IOException;
-import java.time.temporal.ChronoUnit;
 import java.util.Collections;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Future;
 
 /**
  * An OpenFeature {@link FeatureProvider} which enables the use of the LaunchDarkly Server-Side SDK for Java
@@ -48,6 +49,8 @@ public String getName() {
 
     private ProviderState state = ProviderState.NOT_READY;
 
+    private final Object stateLock = new Object();
+
     /**
      * Create a provider with the specified SDK and default configuration.
      * <p>
@@ -128,7 +131,9 @@ public ProviderEvaluation<Value> getObjectEvaluation(String key, Value defaultVa
 
     @Override
     public ProviderState getState() {
-        return state;
+        synchronized (state) {
+            return state;
+        }
     }
 
     @Override
@@ -139,50 +144,72 @@ public void initialize(EvaluationContext evaluationContext) throws Exception {
             state = ProviderState.READY;
         }
 
+        var completer = new CompletableFuture<Boolean>();
+
         client.getFlagTracker().addFlagChangeListener(detail -> {
             emitProviderConfigurationChanged(
                 ProviderEventDetails.builder().flagsChanged(Collections.singletonList(detail.getKey())).build());
         });
         // Listen for future status changes.
         client.getDataSourceStatusProvider().addStatusListener((res) -> {
-            switch (res.getState()) {
-                // We will not re-enter INITIALIZING, but it is here to make the switch exhaustive.
-                case INITIALIZING: {
-                }
-                break;
-                case INTERRUPTED: {
-                    state = ProviderState.STALE;
-                    var message = res.getLastError() != null ? res.getLastError().getMessage() : "encountered an unknown error";
-                    emitProviderStale(ProviderEventDetails.builder().message(message).build());
-                }
-                break;
-                case VALID: {
+            handleDataSourceStatus(res, completer);
+        });
+
+        if(state == ProviderState.READY) {
+            return;
+        }
+
+        handleDataSourceStatus(client.getDataSourceStatusProvider().getStatus(), completer);
+        var successfullyInitialized = completer.get();
+
+        if(!successfullyInitialized) {
+            throw new RuntimeException("Failed to initialize LaunchDarkly client.");
+        }
+    }
+
+    private void handleDataSourceStatus(DataSourceStatusProvider.Status res, CompletableFuture<Boolean> completer) {
+        switch (res.getState()) {
+            // We will not re-enter INITIALIZING, but it is here to make the switch exhaustive.
+            case INITIALIZING: {
+            }
+            break;
+            case INTERRUPTED: {
+                setState(ProviderState.STALE);
+
+                var message = res.getLastError() != null ? res.getLastError().getMessage() : "encountered an unknown error";
+                emitProviderStale(ProviderEventDetails.builder().message(message).build());
+            }
+            break;
+            case VALID: {
+                boolean emit = false;
+                synchronized (stateLock) {
                     // If we are ready, then we don't want to emit it again. Other conditions we may be updating the
                     // reason we are stale or interrupted, so we want to emit an event each time.
                     if (state != ProviderState.READY) {
-                        state = ProviderState.READY;
-                        emitProviderReady(ProviderEventDetails.builder().build());
+                        emit = true;
+                        setState(ProviderState.READY);
                     }
                 }
-                break;
-                case OFF: {
-                    // Currently there is not a shutdown state.
-                    // Our client/provider cannot be restarted, so we just go to error.
-                    state = ProviderState.ERROR;
-                    emitProviderError(ProviderEventDetails.builder().message("Provider shutdown").build());
+
+                if (emit) {
+                    completer.complete(true);
+                    emitProviderReady(ProviderEventDetails.builder().build());
                 }
             }
-        });
-        if (state == ProviderState.READY) {
-            return;
+            break;
+            case OFF: {
+                // Currently there is not a shutdown state.
+                // Our client/provider cannot be restarted, so we just go to error.
+                setState(ProviderState.ERROR);
+                completer.complete(false);
+                emitProviderError(ProviderEventDetails.builder().message("Provider shutdown").build());
+            }
         }
+    }
 
-        boolean initialized = client.getDataSourceStatusProvider().waitFor(DataSourceStatusProvider.State.VALID,
-            ChronoUnit.FOREVER.getDuration());
-
-        if (!initialized) {
-            // Here we throw an exception for the OpenFeature SDK, which will handle emitting an event.
-            throw new RuntimeException("Failed to initialize LaunchDarkly client.");
+    private void setState(ProviderState state) {
+        synchronized (stateLock) {
+            this.state = state;
         }
     }
 
diff --git a/src/test/java/com/launchdarkly/openfeature/serverprovider/LifeCycleTest.java b/src/test/java/com/launchdarkly/openfeature/serverprovider/LifeCycleTest.java
@@ -1,17 +1,101 @@
 package com.launchdarkly.openfeature.serverprovider;
 
+import com.launchdarkly.sdk.server.Components;
 import com.launchdarkly.sdk.server.LDConfig;
 import com.launchdarkly.sdk.server.interfaces.DataSourceStatusProvider;
+import com.launchdarkly.sdk.server.subsystems.ClientContext;
+import com.launchdarkly.sdk.server.subsystems.ComponentConfigurer;
+import com.launchdarkly.sdk.server.subsystems.DataSource;
+import com.launchdarkly.sdk.server.subsystems.DataSourceUpdateSink;
 import dev.openfeature.sdk.ImmutableContext;
 import dev.openfeature.sdk.OpenFeatureAPI;
 import dev.openfeature.sdk.ProviderEvent;
 import dev.openfeature.sdk.ProviderState;
+import dev.openfeature.sdk.exceptions.GeneralError;
 import org.junit.jupiter.api.Test;
 
+import java.io.IOException;
+import java.time.Duration;
+import java.time.LocalDateTime;
+import java.time.ZoneOffset;
+import java.time.temporal.ChronoUnit;
+import java.util.Timer;
+import java.util.TimerTask;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class DelayedDataSource implements DataSource {
+    private Duration startDelay;
+    private boolean willError;
+    private boolean initialized = false;
+    private Object lock = new Object();
+    DataSourceUpdateSink sink;
+
+    DelayedDataSource(Duration delay, boolean error, DataSourceUpdateSink sink) {
+        startDelay = delay;
+        willError = error;
+        this.sink = sink;
+    }
+
+    public Future<Void> start() {
+        var future = new CompletableFuture<Void>();
+        var timer = new Timer();
+        timer.schedule(new TimerTask() {
+            @Override
+            public void run() {
+                if (!willError) {
+                    sink.updateStatus(DataSourceStatusProvider.State.VALID, null);
+                    synchronized (lock) {
+                        initialized = true;
+                    }
+                } else {
+                    sink.updateStatus(DataSourceStatusProvider.State.OFF,
+                        new DataSourceStatusProvider.ErrorInfo(
+                            DataSourceStatusProvider.ErrorKind.NETWORK_ERROR,
+                            404,
+                            "bad",
+                            LocalDateTime.now().toInstant(ZoneOffset.UTC)));
+                }
+                future.complete(null);
+            }
+        }, startDelay.toMillis());
+
+        return future;
+    }
+
+    public boolean isInitialized() {
+        synchronized (lock) {
+            return initialized;
+        }
+    }
+
+    public void close() throws IOException {
+    }
+}
+
+class DelayedDataSourceFactory implements ComponentConfigurer<DataSource> {
+    private Duration startDelay;
+    private boolean willError;
+
+    DelayedDataSourceFactory(Duration delay, boolean error) {
+        startDelay = delay;
+        willError = error;
+    }
+
+    @Override
+    public DataSource build(ClientContext clientContext) {
+        return new DelayedDataSource(startDelay, willError, clientContext.getDataSourceUpdateSink());
+    }
+}
 
 /**
  * Tests in this suite use a real client instance and the public constructor.
@@ -54,16 +138,18 @@ public void canShutdownAnOfflineClient() {
     }
 
     @Test
-    public void itEmitsReadyEvents() {
+    public void itEmitsReadyEvents() throws ExecutionException, InterruptedException, TimeoutException {
         var provider = new Provider("fake-key", new LDConfig.Builder()
             .offline(true).build());
 
         var readyCount = new AtomicInteger();
         var errorCount = new AtomicInteger();
         var staleCount = new AtomicInteger();
+        CompletableFuture<Boolean> gotReadyEvent = new CompletableFuture<>();
 
         OpenFeatureAPI.getInstance().on(ProviderEvent.PROVIDER_READY, (detail) -> {
             readyCount.getAndIncrement();
+            gotReadyEvent.complete(true);
         });
 
         OpenFeatureAPI.getInstance().on(ProviderEvent.PROVIDER_STALE, (detail) -> {
@@ -76,10 +162,67 @@ public void itEmitsReadyEvents() {
 
         OpenFeatureAPI.getInstance().setProviderAndWait(provider);
 
-        OpenFeatureAPI.getInstance().shutdown();
-
+        assertTrue(gotReadyEvent.get(1000, TimeUnit.MILLISECONDS));
         assertEquals(1, readyCount.get());
         assertEquals(0, staleCount.get());
         assertEquals(0, errorCount.get());
+
+        OpenFeatureAPI.getInstance().shutdown();
+    }
+
+    @Test
+    public void itCanHandleClientThatIsNotInitializedImmediately() throws Exception {
+        var config = new LDConfig.Builder()
+            .startWait(Duration.ZERO)
+            .dataSource(new DelayedDataSourceFactory(Duration.ofMillis(100), false))
+            .events(Components.noEvents())
+            .build();
+        var provider = new Provider("fake-key", config);
+        assertEquals(ProviderState.NOT_READY, provider.getState());
+
+        var readyCount = new AtomicInteger();
+
+        OpenFeatureAPI.getInstance().on(ProviderEvent.PROVIDER_READY, (detail) -> {
+            readyCount.getAndIncrement();
+        });
+
+        OpenFeatureAPI.getInstance().setProviderAndWait(provider);
+
+        OpenFeatureAPI.getInstance().shutdown();
+
+        assertEquals(ProviderState.READY, provider.getState());
+        assertEquals(1, readyCount.get());
+    }
+
+    @Test
+    public void itCanHandleClientThatIsNotInitializedImmediatelyAndErrors() throws Exception {
+        var config = new LDConfig.Builder()
+            .startWait(Duration.ZERO)
+            .dataSource(new DelayedDataSourceFactory(Duration.ofMillis(100), true))
+            .events(Components.noEvents())
+            .build();
+        var provider = new Provider("fake-key", config);
+        assertEquals(ProviderState.NOT_READY, provider.getState());
+
+        CompletableFuture<Boolean> gotErrorEvent = new CompletableFuture<>();
+
+        OpenFeatureAPI.getInstance().on(ProviderEvent.PROVIDER_ERROR, (detail) -> {
+            gotErrorEvent.complete(true);
+        });
+
+        GeneralError error = null;
+        try {
+            OpenFeatureAPI.getInstance().setProviderAndWait(provider);
+        } catch (GeneralError e) {
+            error = e;
+        }
+
+        assertNotNull(error);
+
+        assertEquals(ProviderState.ERROR, provider.getState());
+
+        assertTrue(gotErrorEvent.get(1000, TimeUnit.MILLISECONDS));
+
+        OpenFeatureAPI.getInstance().shutdown();
     }
 }
