diff --git a/openfeature-provider-local/src/main/java/com/spotify/confidence/SwapWasmResolverApi.java b/openfeature-provider-local/src/main/java/com/spotify/confidence/SwapWasmResolverApi.java index 51ae0876..2615e80b 100644 --- a/openfeature-provider-local/src/main/java/com/spotify/confidence/SwapWasmResolverApi.java +++ b/openfeature-provider-local/src/main/java/com/spotify/confidence/SwapWasmResolverApi.java @@ -1,5 +1,6 @@ package com.spotify.confidence; +import com.dylibso.chicory.wasm.ChicoryException; import com.spotify.confidence.flags.resolver.v1.MaterializationMap; import com.spotify.confidence.flags.resolver.v1.ResolveWithStickyRequest; import com.spotify.confidence.flags.resolver.v1.ResolveWithStickyResponse; @@ -17,6 +18,7 @@ class SwapWasmResolverApi implements ResolverApi { private final AtomicReference wasmResolverApiRef = new AtomicReference<>(); private final StickyResolveStrategy stickyResolveStrategy; private final WasmFlagLogger flagLogger; + private final String accountId; public SwapWasmResolverApi( WasmFlagLogger flagLogger, @@ -25,6 +27,7 @@ public SwapWasmResolverApi( StickyResolveStrategy stickyResolveStrategy) { this.stickyResolveStrategy = stickyResolveStrategy; this.flagLogger = flagLogger; + this.accountId = accountId; // Create initial instance final WasmResolveApi initialInstance = new WasmResolveApi(flagLogger); @@ -36,7 +39,9 @@ public SwapWasmResolverApi( public void updateStateAndFlushLogs(byte[] state, String accountId) { // Create new instance with updated state final WasmResolveApi newInstance = new WasmResolveApi(flagLogger); - newInstance.setResolverState(state, accountId); + if (state != null) { + newInstance.setResolverState(state, accountId); + } // Get current instance before switching final WasmResolveApi oldInstance = wasmResolverApiRef.getAndSet(newInstance); @@ -191,6 +196,9 @@ public ResolveFlagsResponse resolve(ResolveFlagsRequest request) { return instance.resolve(request); } catch (IsClosedException e) { return resolve(request); + } catch (ChicoryException ce) { + updateStateAndFlushLogs(null, accountId); + throw ce; } } } diff --git a/openfeature-provider-local/src/test/java/com/spotify/confidence/SwapWasmResolverApiTest.java b/openfeature-provider-local/src/test/java/com/spotify/confidence/SwapWasmResolverApiTest.java new file mode 100644 index 00000000..3328b883 --- /dev/null +++ b/openfeature-provider-local/src/test/java/com/spotify/confidence/SwapWasmResolverApiTest.java @@ -0,0 +1,110 @@ +package com.spotify.confidence; + +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.dylibso.chicory.wasm.ChicoryException; +import com.spotify.confidence.shaded.flags.resolver.v1.ResolveFlagsRequest; +import com.spotify.confidence.shaded.flags.resolver.v1.ResolveFlagsResponse; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; + +public class SwapWasmResolverApiTest { + + @Test + public void testChicoryExceptionTriggersStateReset() throws Exception { + // Create a mock WasmFlagLogger + final WasmFlagLogger mockLogger = mock(WasmFlagLogger.class); + + // Use valid test state bytes and accountId + final byte[] initialState = ResolveTest.exampleStateBytes; + final String accountId = "test-account-id"; + + // Create a spy of SwapWasmResolverApi to verify method calls + final SwapWasmResolverApi swapApi = + spy( + new SwapWasmResolverApi( + mockLogger, initialState, accountId, mock(ResolverFallback.class))); + + // Create a mock WasmResolveApi that throws ChicoryException + final WasmResolveApi mockWasmApi = mock(WasmResolveApi.class); + when(mockWasmApi.resolve(any(ResolveFlagsRequest.class))) + .thenThrow(new ChicoryException("WASM runtime error")); + + // Replace the internal WasmResolveApi with our mock + final var field = SwapWasmResolverApi.class.getDeclaredField("wasmResolverApiRef"); + field.setAccessible(true); + @SuppressWarnings("unchecked") + final AtomicReference ref = + (AtomicReference) field.get(swapApi); + ref.set(mockWasmApi); + + // Create a test resolve request + final ResolveFlagsRequest request = + ResolveFlagsRequest.newBuilder() + .addFlags("test-flag") + .setClientSecret("test-secret") + .build(); + + // Execute the resolve and expect ChicoryException to be thrown + assertThrows(ChicoryException.class, () -> swapApi.resolve(request)); + + // Verify that updateStateAndFlushLogs was called with null state and the accountId + verify(swapApi, atLeastOnce()).updateStateAndFlushLogs(null, accountId); + } + + @Test + public void testIsClosedExceptionTriggersRetry() throws Exception { + // Create a mock WasmFlagLogger + final WasmFlagLogger mockLogger = mock(WasmFlagLogger.class); + + // Use valid test state bytes and accountId + final byte[] initialState = ResolveTest.exampleStateBytes; + final String accountId = "test-account-id"; + + // Create a spy of SwapWasmResolverApi to verify method calls + final SwapWasmResolverApi swapApi = + spy( + new SwapWasmResolverApi( + mockLogger, initialState, accountId, mock(ResolverFallback.class))); + + // Create a mock WasmResolveApi that throws IsClosedException on first call, then succeeds + final WasmResolveApi mockWasmApi = mock(WasmResolveApi.class); + final ResolveFlagsResponse mockResponse = ResolveFlagsResponse.newBuilder().build(); + when(mockWasmApi.resolve(any(ResolveFlagsRequest.class))) + .thenThrow(new IsClosedException()) + .thenReturn(mockResponse); + + // Replace the internal WasmResolveApi with our mock + final var field = SwapWasmResolverApi.class.getDeclaredField("wasmResolverApiRef"); + field.setAccessible(true); + @SuppressWarnings("unchecked") + final AtomicReference ref = + (AtomicReference) field.get(swapApi); + ref.set(mockWasmApi); + + // Create a test resolve request + final ResolveFlagsRequest request = + ResolveFlagsRequest.newBuilder() + .addFlags("flags/flag-1") + .setClientSecret(TestBase.secret.getSecret()) + .build(); + + // Call resolve - it should retry when IsClosedException is thrown and succeed on second attempt + final ResolveFlagsResponse response = swapApi.resolve(request); + + // Verify response is not null + assertNotNull(response); + + // Verify that the mock WasmResolveApi.resolve was called twice (first threw exception, second + // succeeded) + verify(mockWasmApi, times(2)).resolve(request); + } +}