fix(common): graceful fallback when LambdaTraceProvider is unavailable (#1176)

Author: hjgraca

libraries/src/AWS.Lambda.Powertools.Common/Core/PowertoolsConfigurations.cs

@@ -1,4 +1,7 @@
+using System;
 using System.Globalization;
+using System.Runtime.CompilerServices;
+using System.Threading;
 using Amazon.Lambda.Core;
 using AWS.Lambda.Powertools.Common.Core;
 
@@ -28,6 +31,13 @@ public class PowertoolsConfigurations : IPowertoolsConfigurations
     /// </summary>
     private static IPowertoolsConfigurations _instance;
 
+    /// <summary>
+    ///     Whether LambdaTraceProvider is available in the loaded Amazon.Lambda.Core assembly.
+    ///     0 = not yet checked, 1 = available, -1 = unavailable.
+    ///     Stored as int for atomic reads/writes via Volatile.
+    /// </summary>
+    private static int _traceProviderState; // 0 = unknown, 1 = available, -1 = unavailable
+
     /// <summary>
     ///     Initializes a new instance of the <see cref="PowertoolsConfigurations" /> class.
     /// </summary>
@@ -165,10 +175,12 @@ public bool GetEnvironmentVariableOrDefault(string variable, bool defaultValue)
 
     /// <summary>
     ///     Gets the X-Ray trace identifier.
+    ///     Uses LambdaTraceProvider.CurrentTraceId when available (Amazon.Lambda.Core >= 2.8.0)
+    ///     for correct trace ID isolation in concurrent Lambda executions (LMI).
+    ///     Falls back to the _X_AMZN_TRACE_ID environment variable for older runtimes.
     /// </summary>
     /// <value>The X-Ray trace identifier.</value>
-    public string XRayTraceId =>
-        LambdaTraceProvider.CurrentTraceId;
+    public string XRayTraceId => GetTraceId(() => GetEnvironmentVariable(Constants.XrayTraceIdEnv));
 
     /// <summary>
     ///     Gets a value indicating whether this instance is Lambda.
@@ -212,4 +224,39 @@ public bool GetEnvironmentVariableOrDefault(string variable, bool defaultValue)
     /// <inheritdoc />
     public string AwsInitializationType =>
         GetEnvironmentVariable(Constants.AWSInitializationTypeEnv);
+
+    private static string GetTraceId(Func<string> fallback)
+    {
+        var state = Volatile.Read(ref _traceProviderState);
+
+        if (state == 1)
+            return GetTraceIdFromProvider();
+
+        if (state == -1)
+            return fallback();
+
+        // First call — probe whether LambdaTraceProvider exists in the loaded runtime
+        try
+        {
+            var traceId = GetTraceIdFromProvider();
+            Volatile.Write(ref _traceProviderState, 1);
+            return traceId;
+        }
+        catch (TypeLoadException)
+        {
+            Volatile.Write(ref _traceProviderState, -1);
+            return fallback();
+        }
+    }
+
+    /// <summary>
+    ///     Isolated call to LambdaTraceProvider.CurrentTraceId.
+    ///     Must not be inlined so that the TypeLoadException is thrown only
+    ///     when this method is invoked, not when the caller is compiled.
+    /// </summary>
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static string GetTraceIdFromProvider()
+    {
+        return LambdaTraceProvider.CurrentTraceId;
+    }
 }

libraries/tests/AWS.Lambda.Powertools.Common.Tests/Core/PowertoolsConfigurationsTest.cs

@@ -534,5 +534,114 @@ public void IsLambdaEnvironment_WhenEnvironmentHasValue_ReturnsTrue()
         }
 
         #endregion
+
+        #region XRayTraceId Tests
+
+        [Fact]
+        public void XRayTraceId_WhenLambdaTraceProviderAvailable_ReturnsTraceId()
+        {
+            ResetTraceProviderState();
+
+            try
+            {
+                // Arrange
+                var environment = Substitute.For<IPowertoolsEnvironment>();
+                var configurations = new PowertoolsConfigurations(environment);
+
+                // Act - LambdaTraceProvider is available in test env (Amazon.Lambda.Core 2.8.0)
+                // Returns null/empty in test env (no active Lambda trace), but should not throw
+                var result = configurations.XRayTraceId;
+
+                // Assert - should not fall back to env var (provider was used instead)
+                environment.DidNotReceive()
+                    .GetEnvironmentVariable(Arg.Is<string>(i => i == Constants.XrayTraceIdEnv));
+            }
+            finally
+            {
+                ResetTraceProviderState();
+            }
+        }
+
+        [Fact]
+        public void XRayTraceId_WhenLambdaTraceProviderUnavailable_FallsBackToEnvironmentVariable()
+        {
+            ResetTraceProviderState();
+
+            try
+            {
+                // Arrange
+                var traceId = "Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1";
+                var environment = Substitute.For<IPowertoolsEnvironment>();
+                environment.GetEnvironmentVariable(Constants.XrayTraceIdEnv).Returns(traceId);
+                var configurations = new PowertoolsConfigurations(environment);
+
+                // Simulate LambdaTraceProvider not being available (older Amazon.Lambda.Core)
+                SetTraceProviderState(-1);
+
+                // Act
+                var result = configurations.XRayTraceId;
+
+                // Assert
+                environment.Received(1)
+                    .GetEnvironmentVariable(Arg.Is<string>(i => i == Constants.XrayTraceIdEnv));
+                Assert.Equal(traceId, result);
+            }
+            finally
+            {
+                ResetTraceProviderState();
+            }
+        }
+
+        [Fact]
+        public void XRayTraceId_WhenProviderCachedUnavailable_UsesEnvVarOnSubsequentCalls()
+        {
+            ResetTraceProviderState();
+
+            try
+            {
+                // Arrange
+                var environment = Substitute.For<IPowertoolsEnvironment>();
+                var configurations = new PowertoolsConfigurations(environment);
+
+                // Simulate LambdaTraceProvider not being available (cached)
+                SetTraceProviderState(-1);
+
+                var traceId1 = "Root=1-aaa;Parent=bbb;Sampled=1";
+                var traceId2 = "Root=1-ccc;Parent=ddd;Sampled=1";
+                environment.GetEnvironmentVariable(Constants.XrayTraceIdEnv).Returns(traceId1, traceId2);
+
+                // Act
+                var result1 = configurations.XRayTraceId;
+                var result2 = configurations.XRayTraceId;
+
+                // Assert - should go straight to env var on both calls (cached unavailable)
+                environment.Received(2)
+                    .GetEnvironmentVariable(Arg.Is<string>(i => i == Constants.XrayTraceIdEnv));
+                Assert.Equal(traceId1, result1);
+                Assert.Equal(traceId2, result2);
+            }
+            finally
+            {
+                ResetTraceProviderState();
+            }
+        }
+
+        private static void ResetTraceProviderState()
+        {
+            var field = typeof(PowertoolsConfigurations).GetField("_traceProviderState",
+                System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Static);
+            Assert.NotNull(field);
+            field.SetValue(null, 0);
+        }
+
+        private static void SetTraceProviderState(int state)
+        {
+            var field = typeof(PowertoolsConfigurations).GetField("_traceProviderState",
+                System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Static);
+            Assert.NotNull(field);
+            field.SetValue(null, state);
+        }
+
+        #endregion
     }
 }