feat(tracing): thread safety and async context preservation in Powertools Tracing (#1082)

Signed-off-by: Henrique Graca <[email protected]>

Author: hjgraca

libraries/src/AWS.Lambda.Powertools.Tracing/Internal/TracingAspect.cs

@@ -2,6 +2,7 @@
 using System.Linq;
 using System.Runtime.ExceptionServices;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using AspectInjector.Broker;
 using AWS.Lambda.Powertools.Common;
@@ -28,14 +29,15 @@ public class TracingAspect
     private readonly IXRayRecorder _xRayRecorder;
 
     /// <summary>
-    ///     If true, capture annotations
+    ///     Thread-safe flag for capturing annotations. Uses int for Interlocked operations.
+    ///     1 = should capture, 0 = already captured
     /// </summary>
-    private static bool _captureAnnotations = true;
+    private static int _captureAnnotations = 1;
 
     /// <summary>
-    ///     If true, annotations have been captured
+    ///     If true, annotations have been captured by this invocation's execution context
     /// </summary>
-    private bool _isAnnotationsCaptured;
+    private static readonly AsyncLocal<bool> _isAnnotationsCaptured = new();
 
     /// <summary>
     /// Aspect constructor
@@ -117,8 +119,12 @@ public object Around(
         }
         finally
         {
-            if (_isAnnotationsCaptured)
-                _captureAnnotations = true;
+            // Reset the capture flag if this execution context captured annotations
+            if (_isAnnotationsCaptured.Value)
+            {
+                Interlocked.Exchange(ref _captureAnnotations, 1);
+                _isAnnotationsCaptured.Value = false;
+            }
         }
     }
 
@@ -127,12 +133,12 @@ private void BeginSegment(string segmentName, string @namespace)
         _xRayRecorder.BeginSubsegment(segmentName);
         _xRayRecorder.SetNamespace(@namespace);
 
-        if (_captureAnnotations)
+        // Use Interlocked.CompareExchange for thread-safe check-and-set
+        // Only one thread will successfully change from 1 to 0
+        if (Interlocked.CompareExchange(ref _captureAnnotations, 0, 1) == 1)
         {
             _xRayRecorder.AddAnnotation("ColdStart", LambdaLifecycleTracker.IsColdStart);
-
-            _captureAnnotations = false;
-            _isAnnotationsCaptured = true;
+            _isAnnotationsCaptured.Value = true;
 
             if (_powertoolsConfigurations.IsServiceDefined)
                 _xRayRecorder.AddAnnotation("Service", _powertoolsConfigurations.Service);
@@ -231,6 +237,7 @@ private bool CaptureError(TracingCaptureMode captureMode)
     internal static void ResetForTest()
     {
         LambdaLifecycleTracker.Reset();
-        _captureAnnotations = true;
+        Interlocked.Exchange(ref _captureAnnotations, 1);
+        _isAnnotationsCaptured.Value = false;
     }
 }

libraries/src/Directory.Packages.props

@@ -21,7 +21,7 @@
     <PackageVersion Include="AWSSDK.SecretsManager" Version="4.0.1.2" />
     <PackageVersion Include="AWSSDK.SimpleSystemsManagement" Version="4.0.3.2" />
     <PackageVersion Include="Microsoft.Extensions.Configuration" Version="8.0.0" />
-    <PackageVersion Include="AWSSDK.XRay" Version="4.0.2.1" />
+    <PackageVersion Include="AWSSDK.XRay" Version="4.0.3.7" />
     <PackageVersion Include="AWSXRayRecorder.Core" Version="2.15.0" />
     <PackageVersion Include="Amazon.Lambda.DynamoDBEvents" Version="3.1.1" />
     <PackageVersion Include="Amazon.Lambda.KinesisEvents" Version="3.0.0" />

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

@@ -26,6 +26,7 @@
     <ItemGroup>
         <ProjectReference Include="..\..\src\AWS.Lambda.Powertools.Logging\AWS.Lambda.Powertools.Logging.csproj" />
         <ProjectReference Include="..\..\src\AWS.Lambda.Powertools.Metrics\AWS.Lambda.Powertools.Metrics.csproj" />
+        <ProjectReference Include="..\..\src\AWS.Lambda.Powertools.Tracing\AWS.Lambda.Powertools.Tracing.csproj" />
     </ItemGroup>
 
 </Project>