fix(idempotency): ensure unique idempotency keys for multiple methods with same key (#1124)

hjgraca · web-flow · commit fe53e3a0bc07 · 2026-01-29T10:02:10.000Z
diff --git a/libraries/src/AWS.Lambda.Powertools.Idempotency/Persistence/BasePersistenceStore.cs b/libraries/src/AWS.Lambda.Powertools.Idempotency/Persistence/BasePersistenceStore.cs
@@ -2,6 +2,7 @@
 using System.Security.Cryptography;
 using System.Text;
 using System.Text.Json;
+using System.Threading;
 using System.Threading.Tasks;
 using AWS.Lambda.Powertools.Common;
 using AWS.Lambda.Powertools.Idempotency.Exceptions;
@@ -34,9 +35,15 @@ public abstract class BasePersistenceStore : IPersistenceStore
     private IdempotencyOptions _idempotencyOptions = null!;
 
     /// <summary>
-    /// Function name
+    /// Base function name (Lambda function name from environment or "testFunction")
     /// </summary>
-    private string _functionName;
+    private string _baseFunctionName = null!;
+    
+    /// <summary>
+    /// Full function name including method name (used for key generation)
+    /// This is stored per-call via AsyncLocal to support multiple idempotent methods
+    /// </summary>
+    private readonly AsyncLocal<string> _fullFunctionName = new();
 
     /// <summary>
     /// Boolean to indicate whether or not payload validation is enabled
@@ -58,27 +65,26 @@ public abstract class BasePersistenceStore : IPersistenceStore
     public void Configure(IdempotencyOptions idempotencyOptions, string functionName, string keyPrefix)
     {
         // Fast path - already configured
-        if (_isConfigured) return;
+        if (_isConfigured)
+        {
+            // Even if already configured, we need to set the full function name for this call
+            // This supports multiple idempotent methods in the same Lambda
+            SetFullFunctionName(functionName, keyPrefix);
+            return;
+        }
         
         lock (_configureLock)
         {
             // Double-check pattern
-            if (_isConfigured) return;
-            
-            if (!string.IsNullOrEmpty(keyPrefix))
-            {
-                _functionName = keyPrefix;
-            }
-            else
+            if (_isConfigured)
             {
-                var funcEnv = Environment.GetEnvironmentVariable(Constants.LambdaFunctionNameEnv);
-
-                _functionName = funcEnv ?? "testFunction";
-                if (!string.IsNullOrWhiteSpace(functionName))
-                {
-                    _functionName += "." + functionName;
-                }
+                SetFullFunctionName(functionName, keyPrefix);
+                return;
             }
+            
+            // Set the base function name (Lambda function name from environment)
+            var funcEnv = Environment.GetEnvironmentVariable(Constants.LambdaFunctionNameEnv);
+            _baseFunctionName = funcEnv ?? "testFunction";
 
             _idempotencyOptions = idempotencyOptions;
 
@@ -94,6 +100,33 @@ public void Configure(IdempotencyOptions idempotencyOptions, string functionName
             }
             
             _isConfigured = true;
+            
+            // Set the full function name for this call
+            SetFullFunctionName(functionName, keyPrefix);
+        }
+    }
+    
+    /// <summary>
+    /// Sets the full function name for the current call context.
+    /// This method is called for each idempotent method invocation to ensure
+    /// the correct method name is used in the idempotency key.
+    /// </summary>
+    /// <param name="functionName">The name of the decorated method</param>
+    /// <param name="keyPrefix">Optional custom key prefix</param>
+    private void SetFullFunctionName(string functionName, string keyPrefix)
+    {
+        if (!string.IsNullOrEmpty(keyPrefix))
+        {
+            _fullFunctionName.Value = keyPrefix;
+        }
+        else
+        {
+            var fullName = _baseFunctionName;
+            if (!string.IsNullOrWhiteSpace(functionName))
+            {
+                fullName += "." + functionName;
+            }
+            _fullFunctionName.Value = fullName;
         }
     }
 
@@ -105,27 +138,24 @@ internal void Configure(IdempotencyOptions options, string functionName, string
         LRUCache<string, DataRecord> cache)
     {
         // Fast path - already configured
-        if (_isConfigured) return;
+        if (_isConfigured)
+        {
+            SetFullFunctionName(functionName, keyPrefix);
+            return;
+        }
         
         lock (_configureLock)
         {
             // Double-check pattern
-            if (_isConfigured) return;
-            
-            if (!string.IsNullOrEmpty(keyPrefix))
-            {
-                _functionName = keyPrefix;
-            }
-            else
+            if (_isConfigured)
             {
-                var funcEnv = Environment.GetEnvironmentVariable(Constants.LambdaFunctionNameEnv);
-
-                _functionName = funcEnv ?? "testFunction";
-                if (!string.IsNullOrWhiteSpace(functionName))
-                {
-                    _functionName += "." + functionName;
-                }
+                SetFullFunctionName(functionName, keyPrefix);
+                return;
             }
+            
+            // Set the base function name (Lambda function name from environment)
+            var funcEnv = Environment.GetEnvironmentVariable(Constants.LambdaFunctionNameEnv);
+            _baseFunctionName = funcEnv ?? "testFunction";
 
             _idempotencyOptions = options;
 
@@ -137,6 +167,8 @@ internal void Configure(IdempotencyOptions options, string functionName, string
             _cache = cache;
             
             _isConfigured = true;
+            
+            SetFullFunctionName(functionName, keyPrefix);
         }
     }
 
@@ -356,7 +388,7 @@ private string GetHashedIdempotencyKey(JsonDocument data)
         }
 
         var hash = GenerateHash(node);
-        return _functionName + "#" + hash;
+        return _fullFunctionName.Value + "#" + hash;
     }
 
     /// <summary>
diff --git a/libraries/tests/AWS.Lambda.Powertools.Idempotency.Tests/Handlers/IdempotencyMultipleMethodsFunction.cs b/libraries/tests/AWS.Lambda.Powertools.Idempotency.Tests/Handlers/IdempotencyMultipleMethodsFunction.cs
@@ -0,0 +1,56 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * 
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ * 
+ *  http://aws.amazon.com/apache2.0
+ * 
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+using Amazon.Lambda.Core;
+using AWS.Lambda.Powertools.Idempotency.Tests.Model;
+
+namespace AWS.Lambda.Powertools.Idempotency.Tests.Handlers;
+
+/// <summary>
+/// Lambda function with multiple idempotent methods to test cross-method key isolation.
+/// This tests the scenario where two different methods are decorated with [Idempotent]
+/// and called with the same IdempotencyKey value - they should create separate entries
+/// in the persistence store.
+/// </summary>
+public class IdempotencyMultipleMethodsFunction
+{
+    public bool Method1Called { get; private set; }
+    public bool Method2Called { get; private set; }
+
+    public (Basket, Basket) HandleRequest([IdempotencyKey] string key, ILambdaContext context)
+    {
+        Idempotency.RegisterLambdaContext(context);
+
+        // Call both methods with the same key - they should each create their own idempotency record
+        var result1 = Method1(key);
+        var result2 = Method2(key);
+
+        return (result1, result2);
+    }
+
+    [Idempotent]
+    private Basket Method1([IdempotencyKey] string key)
+    {
+        Method1Called = true;
+        return new Basket(new Product(1, "Product from Method1", 10.0));
+    }
+
+    [Idempotent]
+    private Basket Method2([IdempotencyKey] string key)
+    {
+        Method2Called = true;
+        return new Basket(new Product(2, "Product from Method2", 20.0));
+    }
+}
diff --git a/libraries/tests/AWS.Lambda.Powertools.Idempotency.Tests/Internal/IdempotentAspectTests.cs b/libraries/tests/AWS.Lambda.Powertools.Idempotency.Tests/Internal/IdempotentAspectTests.cs
@@ -359,6 +359,45 @@ public async Task WhenIdempotency_Custom_Prefix_Key_Method()
         Assert.NotNull(record);
     }
 
+    [Fact]
+    public async Task Handle_WhenMultipleIdempotentMethods_WithSameKey_ShouldCreateSeparateRecords()
+    {
+        // Arrange
+        // This test validates the bug where two different methods decorated with [Idempotent]
+        // and called with the same IdempotencyKey value should create separate entries
+        // in the persistence store, but currently Method2 incorrectly uses Method1's cached result.
+        var store = new InMemoryPersistenceStore();
+        Idempotency.Configure(builder => builder.WithPersistenceStore(store));
+
+        var context = new TestLambdaContext
+        {
+            RemainingTime = TimeSpan.FromSeconds(30)
+        };
+
+        // Act
+        var function = new IdempotencyMultipleMethodsFunction();
+        var (result1, result2) = function.HandleRequest("same-key", context);
+
+        // Assert
+        // Both methods should have been called
+        function.Method1Called.Should().BeTrue("Method1 should be called on first invocation");
+        function.Method2Called.Should().BeTrue("Method2 should be called - it's a different method even with same key");
+
+        // Results should be different - each method returns its own product
+        result1.Products[0].Name.Should().Be("Product from Method1");
+        result2.Products[0].Name.Should().Be("Product from Method2");
+
+        // Debug: Show what keys were actually created
+        var allKeys = store.GetAllKeys().ToList();
+        
+        // Verify separate records exist in the store - should have 2 records with different method names
+        allKeys.Should().HaveCount(2, $"Expected 2 records but found: [{string.Join(", ", allKeys)}]");
+        
+        // Verify the keys contain the correct method names
+        allKeys.Should().Contain(k => k.StartsWith("testFunction.Method1#"), "Should have a record for Method1");
+        allKeys.Should().Contain(k => k.StartsWith("testFunction.Method2#"), "Should have a record for Method2");
+    }
+
     [Fact]
     public void Handle_WhenIdempotencyOnSubMethodNotAnnotated_ShouldThrowException()
     {
diff --git a/libraries/tests/AWS.Lambda.Powertools.Idempotency.Tests/Internal/InMemoryPersistenceStore.cs b/libraries/tests/AWS.Lambda.Powertools.Idempotency.Tests/Internal/InMemoryPersistenceStore.cs
@@ -15,6 +15,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using System.Threading.Tasks;
 using AWS.Lambda.Powertools.Idempotency.Persistence;
 
@@ -23,6 +24,12 @@ namespace AWS.Lambda.Powertools.Idempotency.Tests.Internal;
 public class InMemoryPersistenceStore: BasePersistenceStore
 {
     private readonly Dictionary<string, DataRecord> _records = new();
+    
+    /// <summary>
+    /// Gets all keys currently stored - useful for debugging tests
+    /// </summary>
+    public IEnumerable<string> GetAllKeys() => _records.Keys;
+    
     public override Task<DataRecord> GetRecord(string idempotencyKey)
     {
         return Task.FromResult(_records.ContainsKey(idempotencyKey) ? _records[idempotencyKey] : null);
