Skip to content

Commit 002d909

Browse files
hjgracahjgraca
authored
feat(idempotency): enhance thread safety with AsyncLocal context and thread-safe configuration (#1084)
1 parent 3565891 commit 002d909

14 files changed

+3576
-32
lines changed

‎libraries/src/AWS.Lambda.Powertools.Idempotency/Idempotency.cs‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using System;
22
using System.Text.Json.Serialization;
3+
using System.Threading;
34
using Amazon.Lambda.Core;
45
using AWS.Lambda.Powertools.Common;
56
using AWS.Lambda.Powertools.Idempotency.Internal.Serializers;
@@ -16,6 +17,12 @@ namespace AWS.Lambda.Powertools.Idempotency;
1617
/// </summary>
1718
public sealed class Idempotency
1819
{
20+
/// <summary>
21+
/// AsyncLocal storage for per-invocation LambdaContext.
22+
/// This ensures each concurrent Lambda invocation has its own isolated context.
23+
/// </summary>
24+
private static readonly AsyncLocal<ILambdaContext> _lambdaContext = new();
25+
1926
/// <summary>
2027
/// The general configurations for the idempotency
2128
/// </summary>
@@ -66,18 +73,25 @@ public static void Configure(Action<IdempotencyBuilder> configurationAction)
6673
}
6774

6875
/// <summary>
69-
/// Holds ILambdaContext
76+
/// Holds ILambdaContext using AsyncLocal for per-invocation isolation.
77+
/// Each concurrent Lambda invocation will have its own isolated context.
7078
/// </summary>
71-
public ILambdaContext LambdaContext { get; private set; }
79+
public ILambdaContext LambdaContext
80+
{
81+
get => _lambdaContext.Value;
82+
private set => _lambdaContext.Value = value;
83+
}
7284

7385
/// <summary>
7486
/// Can be used in a method which is not the handler to capture the Lambda context,
7587
/// to calculate the remaining time before the invocation times out.
88+
/// This method is thread-safe and stores the context in AsyncLocal storage,
89+
/// ensuring isolation between concurrent Lambda invocations.
7690
/// </summary>
77-
/// <param name="context"></param>
91+
/// <param name="context">The Lambda context for the current invocation</param>
7892
public static void RegisterLambdaContext(ILambdaContext context)
7993
{
80-
Instance.LambdaContext = context;
94+
_lambdaContext.Value = context;
8195
}
8296

8397
/// <summary>

‎libraries/src/AWS.Lambda.Powertools.Idempotency/InternalsVisibleTo.cs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,4 +15,5 @@
1515

1616
using System.Runtime.CompilerServices;
1717

18-
[assembly: InternalsVisibleTo("AWS.Lambda.Powertools.Idempotency.Tests")]
18+
[assembly: InternalsVisibleTo("AWS.Lambda.Powertools.Idempotency.Tests")]
19+
[assembly: InternalsVisibleTo("AWS.Lambda.Powertools.ConcurrencyTests")]

‎libraries/src/AWS.Lambda.Powertools.Idempotency/Persistence/BasePersistenceStore.cs‎

Lines changed: 79 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,16 @@ namespace AWS.Lambda.Powertools.Idempotency.Persistence;
1818
/// </summary>
1919
public abstract class BasePersistenceStore : IPersistenceStore
2020
{
21+
/// <summary>
22+
/// Lock object for thread-safe configuration
23+
/// </summary>
24+
private readonly object _configureLock = new object();
25+
26+
/// <summary>
27+
/// Flag indicating whether the store has been configured
28+
/// </summary>
29+
private volatile bool _isConfigured;
30+
2131
/// <summary>
2232
/// Idempotency Options
2333
/// </summary>
@@ -39,50 +49,95 @@ public abstract class BasePersistenceStore : IPersistenceStore
3949
private LRUCache<string, DataRecord> _cache = null!;
4050

4151
/// <summary>
42-
/// Initialize the base persistence layer from the configuration settings
52+
/// Initialize the base persistence layer from the configuration settings.
53+
/// This method is thread-safe and idempotent - multiple calls with the same parameters are safe.
4354
/// </summary>
4455
/// <param name="idempotencyOptions">Idempotency configuration settings</param>
4556
/// <param name="functionName">The name of the function being decorated</param>
4657
/// <param name="keyPrefix"></param>
4758
public void Configure(IdempotencyOptions idempotencyOptions, string functionName, string keyPrefix)
4859
{
49-
if (!string.IsNullOrEmpty(keyPrefix))
60+
// Fast path - already configured
61+
if (_isConfigured) return;
62+
63+
lock (_configureLock)
5064
{
51-
_functionName = keyPrefix;
52-
}
53-
else
54-
{
55-
var funcEnv = Environment.GetEnvironmentVariable(Constants.LambdaFunctionNameEnv);
56-
57-
_functionName = funcEnv ?? "testFunction";
58-
if (!string.IsNullOrWhiteSpace(functionName))
65+
// Double-check pattern
66+
if (_isConfigured) return;
67+
68+
if (!string.IsNullOrEmpty(keyPrefix))
5969
{
60-
_functionName += "." + functionName;
70+
_functionName = keyPrefix;
6171
}
62-
}
72+
else
73+
{
74+
var funcEnv = Environment.GetEnvironmentVariable(Constants.LambdaFunctionNameEnv);
6375

64-
_idempotencyOptions = idempotencyOptions;
76+
_functionName = funcEnv ?? "testFunction";
77+
if (!string.IsNullOrWhiteSpace(functionName))
78+
{
79+
_functionName += "." + functionName;
80+
}
81+
}
6582

66-
if (!string.IsNullOrWhiteSpace(_idempotencyOptions.PayloadValidationJmesPath))
67-
{
68-
PayloadValidationEnabled = true;
69-
}
83+
_idempotencyOptions = idempotencyOptions;
7084

71-
var useLocalCache = _idempotencyOptions.UseLocalCache;
72-
if (useLocalCache)
73-
{
74-
_cache = new LRUCache<string, DataRecord>(_idempotencyOptions.LocalCacheMaxItems);
85+
if (!string.IsNullOrWhiteSpace(_idempotencyOptions.PayloadValidationJmesPath))
86+
{
87+
PayloadValidationEnabled = true;
88+
}
89+
90+
var useLocalCache = _idempotencyOptions.UseLocalCache;
91+
if (useLocalCache)
92+
{
93+
_cache = new LRUCache<string, DataRecord>(_idempotencyOptions.LocalCacheMaxItems);
94+
}
95+
96+
_isConfigured = true;
7597
}
7698
}
7799

78100
/// <summary>
79-
/// For test purpose only (adding a cache to mock)
101+
/// For test purpose only (adding a cache to mock).
102+
/// This method is thread-safe and idempotent.
80103
/// </summary>
81104
internal void Configure(IdempotencyOptions options, string functionName, string keyPrefix,
82105
LRUCache<string, DataRecord> cache)
83106
{
84-
Configure(options, functionName, keyPrefix);
85-
_cache = cache;
107+
// Fast path - already configured
108+
if (_isConfigured) return;
109+
110+
lock (_configureLock)
111+
{
112+
// Double-check pattern
113+
if (_isConfigured) return;
114+
115+
if (!string.IsNullOrEmpty(keyPrefix))
116+
{
117+
_functionName = keyPrefix;
118+
}
119+
else
120+
{
121+
var funcEnv = Environment.GetEnvironmentVariable(Constants.LambdaFunctionNameEnv);
122+
123+
_functionName = funcEnv ?? "testFunction";
124+
if (!string.IsNullOrWhiteSpace(functionName))
125+
{
126+
_functionName += "." + functionName;
127+
}
128+
}
129+
130+
_idempotencyOptions = options;
131+
132+
if (!string.IsNullOrWhiteSpace(_idempotencyOptions.PayloadValidationJmesPath))
133+
{
134+
PayloadValidationEnabled = true;
135+
}
136+
137+
_cache = cache;
138+
139+
_isConfigured = true;
140+
}
86141
}
87142

88143
/// <summary>

‎libraries/tests/AWS.Lambda.Powertools.ConcurrencyTests/AWS.Lambda.Powertools.ConcurrencyTests.csproj‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,11 @@
2121
<PrivateAssets>all</PrivateAssets>
2222
<IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
2323
</PackageReference>
24+
<PackageReference Include="Amazon.Lambda.TestUtilities" />
2425
</ItemGroup>
2526

2627
<ItemGroup>
28+
<ProjectReference Include="..\..\src\AWS.Lambda.Powertools.Idempotency\AWS.Lambda.Powertools.Idempotency.csproj" />
2729
<ProjectReference Include="..\..\src\AWS.Lambda.Powertools.Logging\AWS.Lambda.Powertools.Logging.csproj" />
2830
<ProjectReference Include="..\..\src\AWS.Lambda.Powertools.Metrics\AWS.Lambda.Powertools.Metrics.csproj" />
2931
<ProjectReference Include="..\..\src\AWS.Lambda.Powertools.Tracing\AWS.Lambda.Powertools.Tracing.csproj" />

0 commit comments

Comments
 (0)