diff --git a/libraries/src/AWS.Lambda.Powertools.Common/Core/PowertoolsConfigurations.cs b/libraries/src/AWS.Lambda.Powertools.Common/Core/PowertoolsConfigurations.cs index c1371013..d5bb32a3 100644 --- a/libraries/src/AWS.Lambda.Powertools.Common/Core/PowertoolsConfigurations.cs +++ b/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 /// private static IPowertoolsConfigurations _instance; + /// + /// 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. + /// + private static int _traceProviderState; // 0 = unknown, 1 = available, -1 = unavailable + /// /// Initializes a new instance of the class. /// @@ -165,10 +175,12 @@ public bool GetEnvironmentVariableOrDefault(string variable, bool defaultValue) /// /// 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. /// /// The X-Ray trace identifier. - public string XRayTraceId => - LambdaTraceProvider.CurrentTraceId; + public string XRayTraceId => GetTraceId(() => GetEnvironmentVariable(Constants.XrayTraceIdEnv)); /// /// Gets a value indicating whether this instance is Lambda. @@ -212,4 +224,39 @@ public bool GetEnvironmentVariableOrDefault(string variable, bool defaultValue) /// public string AwsInitializationType => GetEnvironmentVariable(Constants.AWSInitializationTypeEnv); + + private static string GetTraceId(Func 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(); + } + } + + /// + /// 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. + /// + [MethodImpl(MethodImplOptions.NoInlining)] + private static string GetTraceIdFromProvider() + { + return LambdaTraceProvider.CurrentTraceId; + } } \ No newline at end of file diff --git a/libraries/tests/AWS.Lambda.Powertools.Common.Tests/Core/PowertoolsConfigurationsTest.cs b/libraries/tests/AWS.Lambda.Powertools.Common.Tests/Core/PowertoolsConfigurationsTest.cs index 5e450307..784cda2f 100644 --- a/libraries/tests/AWS.Lambda.Powertools.Common.Tests/Core/PowertoolsConfigurationsTest.cs +++ b/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(); + 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(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(); + 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(i => i == Constants.XrayTraceIdEnv)); + Assert.Equal(traceId, result); + } + finally + { + ResetTraceProviderState(); + } + } + + [Fact] + public void XRayTraceId_WhenProviderCachedUnavailable_UsesEnvVarOnSubsequentCalls() + { + ResetTraceProviderState(); + + try + { + // Arrange + var environment = Substitute.For(); + 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(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 } } \ No newline at end of file