fix: prevent race conditions during repository shutdown

Signed-off-by: Nicklas Lundin <nicklasl@spotify.com>

Author: nicklasl

src/main/java/dev/openfeature/sdk/ProviderRepository.java

@@ -11,6 +11,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
@@ -24,6 +25,7 @@ class ProviderRepository {
     private final Map<String, FeatureProviderStateManager> stateManagers = new ConcurrentHashMap<>();
     private final AtomicReference<FeatureProviderStateManager> defaultStateManger =
             new AtomicReference<>(new FeatureProviderStateManager(new NoOpProvider()));
+    private final AtomicBoolean isShuttingDown = new AtomicBoolean(false);
     private final ExecutorService taskExecutor =
             Executors.newCachedThreadPool(new ConfigurableThreadFactory("openfeature-provider-thread", true));
     private final Object registerStateManagerLock = new Object();
@@ -159,6 +161,10 @@ private void prepareAndInitializeProvider(
             Consumer<FeatureProvider> afterShutdown,
             BiConsumer<FeatureProvider, OpenFeatureError> afterError,
             boolean waitForInit) {
+        if (isShuttingDown.get()) {
+            throw new IllegalStateException("Provider cannot be set while repository is shutting down");
+        }
+
         final FeatureProviderStateManager newStateManager;
         final FeatureProviderStateManager oldStateManager;
 
@@ -255,16 +261,27 @@ private void shutdownProvider(FeatureProviderStateManager manager) {
     }
 
     private void shutdownProvider(FeatureProvider provider) {
-        taskExecutor.submit(() -> {
+        try {
+            taskExecutor.submit(() -> {
+                try {
+                    provider.shutdown();
+                } catch (Exception e) {
+                    log.error(
+                            "Exception when shutting down feature provider {}",
+                            provider.getClass().getName(),
+                            e);
+                }
+            });
+        } catch (java.util.concurrent.RejectedExecutionException e) {
             try {
                 provider.shutdown();
-            } catch (Exception e) {
+            } catch (Exception ex) {
                 log.error(
                         "Exception when shutting down feature provider {}",
                         provider.getClass().getName(),
-                        e);
+                        ex);
             }
-        });
+        }
     }
 
     /**
@@ -273,6 +290,10 @@ private void shutdownProvider(FeatureProvider provider) {
      * including the default feature provider.
      */
     public void shutdown() {
+        if (isShuttingDown.getAndSet(true)) {
+            return;
+        }
+
         Stream.concat(Stream.of(this.defaultStateManger.get()), this.stateManagers.values().stream())
                 .distinct()
                 .forEach(this::shutdownProvider);

src/test/java/dev/openfeature/sdk/ProviderRepositoryTest.java

@@ -4,6 +4,7 @@
 import static dev.openfeature.sdk.testutils.stubbing.ConditionStubber.doDelayResponse;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.awaitility.Awaitility.await;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
@@ -363,6 +364,104 @@ void shouldNotHangIndefinitelyOnShutdown() {
                             return true;
                         });
             }
+
+            @Test
+            @DisplayName("should handle shutdown during provider initialization")
+            void shouldHandleShutdownDuringProviderInitialization() throws Exception {
+                FeatureProvider slowInitProvider = createMockedProvider();
+                AtomicBoolean initStarted = new AtomicBoolean(false);
+                AtomicBoolean shutdownCalled = new AtomicBoolean(false);
+
+                doAnswer(invocation -> {
+                            initStarted.set(true);
+                            Thread.sleep(500);
+                            return null;
+                        })
+                        .when(slowInitProvider)
+                        .initialize(any());
+
+                doAnswer(invocation -> {
+                            shutdownCalled.set(true);
+                            return null;
+                        })
+                        .when(slowInitProvider)
+                        .shutdown();
+
+                providerRepository.setProvider(slowInitProvider, mockAfterSet(), mockAfterInit(), mockAfterShutdown(), mockAfterError(), false);
+
+                await().atMost(Duration.ofSeconds(1)).untilTrue(initStarted);
+
+                // Call shutdown while initialization is in progress
+                assertThatCode(() -> providerRepository.shutdown()).doesNotThrowAnyException();
+
+                await().atMost(Duration.ofSeconds(1)).untilTrue(shutdownCalled);
+                verify(slowInitProvider, times(1)).shutdown();
+            }
+
+            @Test
+            @DisplayName("should handle provider replacement during shutdown")
+            void shouldHandleProviderReplacementDuringShutdown() throws Exception {
+                FeatureProvider oldProvider = createMockedProvider();
+                FeatureProvider newProvider = createMockedProvider();
+                AtomicBoolean oldProviderShutdownCalled = new AtomicBoolean(false);
+
+                doAnswer(invocation -> {
+                            oldProviderShutdownCalled.set(true);
+                            return null;
+                        })
+                        .when(oldProvider)
+                        .shutdown();
+
+                providerRepository.setProvider(oldProvider, mockAfterSet(), mockAfterInit(), mockAfterShutdown(), mockAfterError(), true);
+
+                // Replace provider (this will trigger old provider shutdown in background)
+                providerRepository.setProvider(newProvider, mockAfterSet(), mockAfterInit(), mockAfterShutdown(), mockAfterError(), false);
+
+                assertThatCode(() -> providerRepository.shutdown()).doesNotThrowAnyException();
+
+                await().atMost(Duration.ofSeconds(1)).untilTrue(oldProviderShutdownCalled);
+                verify(oldProvider, times(1)).shutdown();
+                verify(newProvider, times(1)).shutdown();
+            }
+
+            @Test
+            @DisplayName("should prevent adding providers after shutdown has started")
+            void shouldPreventAddingProvidersAfterShutdownHasStarted() {
+                FeatureProvider provider = createMockedProvider();
+                setFeatureProvider(provider);
+
+                providerRepository.shutdown();
+
+                FeatureProvider newProvider = createMockedProvider();
+                assertThatThrownBy(() -> providerRepository.setProvider(
+                                newProvider, mockAfterSet(), mockAfterInit(), mockAfterShutdown(), mockAfterError(), false))
+                        .isInstanceOf(IllegalStateException.class)
+                        .hasMessageContaining("shutting down");
+            }
+
+            @Test
+            @DisplayName("should handle concurrent shutdown calls gracefully")
+            void shouldHandleConcurrentShutdownCallsGracefully() {
+                FeatureProvider provider = createMockedProvider();
+                setFeatureProvider(provider);
+
+                assertThatCode(() -> {
+                            Thread t1 = new Thread(() -> providerRepository.shutdown());
+                            Thread t2 = new Thread(() -> providerRepository.shutdown());
+                            Thread t3 = new Thread(() -> providerRepository.shutdown());
+
+                            t1.start();
+                            t2.start();
+                            t3.start();
+
+                            t1.join(TIMEOUT);
+                            t2.join(TIMEOUT);
+                            t3.join(TIMEOUT);
+                        })
+                        .doesNotThrowAnyException();
+
+                verify(provider, times(1)).shutdown();
+            }
         }
     }