Make IdempotencyConfig thread-safe to avoid Lamdba context leaks in config singleton.

Author: phipag

powertools-idempotency/powertools-idempotency-core/src/main/java/software/amazon/lambda/powertools/idempotency/IdempotencyConfig.java

@@ -34,7 +34,7 @@ public class IdempotencyConfig {
     private final boolean throwOnNoIdempotencyKey;
     private final String hashFunction;
     private final BiFunction<Object, DataRecord, Object> responseHook;
-    private Context lambdaContext;
+    private final InheritableThreadLocal<Context> lambdaContext = new InheritableThreadLocal<>();
 
     private IdempotencyConfig(String eventKeyJMESPath, String payloadValidationJMESPath,
             boolean throwOnNoIdempotencyKey, boolean useLocalCache, int localCacheMaxItems,
@@ -87,11 +87,11 @@ public String getHashFunction() {
     }
 
     public Context getLambdaContext() {
-        return lambdaContext;
+        return lambdaContext.get();
     }
 
     public void setLambdaContext(Context lambdaContext) {
-        this.lambdaContext = lambdaContext;
+        this.lambdaContext.set(lambdaContext);
     }
 
     public BiFunction<Object, DataRecord, Object> getResponseHook() {

powertools-idempotency/powertools-idempotency-core/src/test/java/software/amazon/lambda/powertools/idempotency/PowertoolsIdempotencyTest.java

@@ -165,4 +165,72 @@ void firstCall_withExplicitIdempotencyKey_shouldPutInStore() {
         verify(store).saveSuccess(any(), resultCaptor.capture(), any());
         assertThat(resultCaptor.getValue()).isEqualTo(basket);
     }
+
+    @Test
+    void concurrentInvocations_shouldNotLeakContext() throws Exception {
+        Idempotency.config()
+                .withPersistenceStore(store)
+                .configure();
+
+        PowertoolsIdempotencyMultiArgFunction function = new PowertoolsIdempotencyMultiArgFunction();
+
+        // GIVEN
+        int threadCount = 10;
+        Thread[] threads = new Thread[threadCount];
+        Context[] capturedContexts = new Context[threadCount];
+        int[] capturedRemainingTimes = new int[threadCount];
+        boolean[] success = new boolean[threadCount];
+
+        // WHEN - Multiple threads call handleRequest with different contexts
+        for (int i = 0; i < threadCount; i++) {
+            final int threadIndex = i;
+            final int expectedTime = (i + 1) * 2000; // 2000, 4000, 6000, ..., 20000
+
+            final Context threadContext = new TestLambdaContext() {
+                @Override
+                public int getRemainingTimeInMillis() {
+                    return expectedTime;
+                }
+            };
+
+            threads[i] = new Thread(() -> {
+                try {
+                    Product p = new Product(threadIndex, "product" + threadIndex, 10);
+                    function.handleRequest(p, threadContext);
+
+                    // Capture the context that was actually stored in ThreadLocal by this thread
+                    Context captured = Idempotency.getInstance().getConfig().getLambdaContext();
+                    capturedContexts[threadIndex] = captured;
+                    capturedRemainingTimes[threadIndex] = captured != null ? captured.getRemainingTimeInMillis() : -1;
+                    success[threadIndex] = true;
+                } catch (Exception e) {
+                    success[threadIndex] = false;
+                }
+            });
+        }
+
+        // Start all threads
+        for (Thread thread : threads) {
+            thread.start();
+        }
+
+        // Wait for all threads to complete
+        for (Thread thread : threads) {
+            thread.join();
+        }
+
+        // THEN - All threads should complete successfully
+        for (boolean result : success) {
+            assertThat(result).isTrue();
+        }
+
+        // THEN - Each thread should have captured its own context (no leakage)
+        for (int i = 0; i < threadCount; i++) {
+            int expectedTime = (i + 1) * 2000;
+            assertThat(capturedRemainingTimes[i])
+                    .as("Thread %d should have remaining time %d", i, expectedTime)
+                    .isEqualTo(expectedTime);
+            assertThat(capturedContexts[i]).as("Thread %d should have non-null context", i).isNotNull();
+        }
+    }
 }

powertools-idempotency/powertools-idempotency-core/src/test/java/software/amazon/lambda/powertools/idempotency/internal/IdempotencyAspectTest.java

@@ -194,7 +194,7 @@ void secondCall_notExpired_shouldNotGetFromStoreIfPresentOnIdempotencyException(
                 "Test message",
                 new RuntimeException("Test Cause"),
                 dr))
-                .when(store).saveInProgress(any(), any(), any());
+                        .when(store).saveInProgress(any(), any(), any());
 
         // WHEN
         IdempotencyEnabledFunction function = new IdempotencyEnabledFunction();
@@ -538,4 +538,73 @@ void idempotencyOnSubMethodVoid_shouldThrowException() {
                 IdempotencyConfigurationException.class);
     }
 
+    @Test
+    void concurrentInvocations_shouldNotLeakContext() throws Exception {
+        Idempotency.config()
+                .withPersistenceStore(store)
+                .configure();
+
+        // Use IdempotencyInternalFunction which calls registerLambdaContext
+        IdempotencyInternalFunction function = new IdempotencyInternalFunction(true);
+
+        // GIVEN
+        int threadCount = 10;
+        Thread[] threads = new Thread[threadCount];
+        Context[] capturedContexts = new Context[threadCount];
+        int[] capturedRemainingTimes = new int[threadCount];
+        boolean[] success = new boolean[threadCount];
+
+        // WHEN - Multiple threads call handleRequest with different contexts
+        for (int i = 0; i < threadCount; i++) {
+            final int threadIndex = i;
+            final int expectedTime = (i + 1) * 1000; // 1000, 2000, 3000, ..., 10000
+
+            final Context threadContext = new TestLambdaContext() {
+                @Override
+                public int getRemainingTimeInMillis() {
+                    return expectedTime;
+                }
+            };
+
+            threads[i] = new Thread(() -> {
+                try {
+                    Product p = new Product(threadIndex, "product" + threadIndex, 10);
+                    function.handleRequest(p, threadContext);
+
+                    // Capture the context that was actually stored in ThreadLocal by this thread
+                    Context captured = Idempotency.getInstance().getConfig().getLambdaContext();
+                    capturedContexts[threadIndex] = captured;
+                    capturedRemainingTimes[threadIndex] = captured != null ? captured.getRemainingTimeInMillis() : -1;
+                    success[threadIndex] = true;
+                } catch (Exception e) {
+                    success[threadIndex] = false;
+                }
+            });
+        }
+
+        // Start all threads
+        for (Thread thread : threads) {
+            thread.start();
+        }
+
+        // Wait for all threads to complete
+        for (Thread thread : threads) {
+            thread.join();
+        }
+
+        // THEN - All threads should complete successfully
+        for (boolean result : success) {
+            assertThat(result).isTrue();
+        }
+
+        // THEN - Each thread should have captured its own context (no leakage)
+        for (int i = 0; i < threadCount; i++) {
+            int expectedTime = (i + 1) * 1000;
+            assertThat(capturedRemainingTimes[i])
+                    .as("Thread %d should have remaining time %d", i, expectedTime)
+                    .isEqualTo(expectedTime);
+            assertThat(capturedContexts[i]).as("Thread %d should have non-null context", i).isNotNull();
+        }
+    }
+
 }