Extract Upper64 bit trace ID from extension response

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaCommon.cs

@@ -8,6 +8,8 @@
 using System.Net;
 using System.Text;
 using System.Threading.Tasks;
+using Datadog.Trace.Headers;
+using Datadog.Trace.Propagators;
 using Datadog.Trace.Telemetry;
 using Datadog.Trace.Telemetry.Metrics;
 using Datadog.Trace.Util;
@@ -21,31 +23,12 @@ internal abstract class LambdaCommon
     private const double ServerlessMaxWaitingFlushTime = 3;
     private const string LogLevelEnvName = "DD_LOG_LEVEL";
 
-    internal static Scope CreatePlaceholderScope(Tracer tracer, string traceId, string samplingPriority)
+    internal static Scope CreatePlaceholderScope(Tracer tracer, NameValueHeadersCollection headers)
     {
-        Span span;
+        var spanContext = SpanContextPropagator.Instance.Extract(headers);
+        TelemetryFactory.Metrics.RecordCountSpanCreated(MetricTags.IntegrationName.AwsLambda);
 
-        if (traceId == null)
-        {
-            Log("traceId not found");
-            span = tracer.StartSpan(PlaceholderOperationName, tags: null, serviceName: PlaceholderServiceName, addToTraceContext: false);
-        }
-        else
-        {
-            Log($"creating the placeholder traceId = {traceId}");
-            span = tracer.StartSpan(PlaceholderOperationName, tags: null, serviceName: PlaceholderServiceName, traceId: (TraceId)Convert.ToUInt64(traceId), addToTraceContext: false);
-        }
-
-        if (samplingPriority == null)
-        {
-            Log("samplingPriority not found");
-            _ = span.Context.TraceContext?.GetOrMakeSamplingDecision();
-        }
-        else
-        {
-            Log($"setting the placeholder sampling priority to = {samplingPriority}");
-            span.Context.TraceContext?.SetSamplingPriority(Convert.ToInt32(samplingPriority), notifyDistributedTracer: false);
-        }
+        var span = spanContext != null ? tracer.StartSpan(PlaceholderOperationName, tags: null, parent: spanContext, serviceName: PlaceholderServiceName, addToTraceContext: false) : tracer.StartSpan(PlaceholderOperationName, tags: null, serviceName: PlaceholderServiceName, addToTraceContext: false);
 
         TelemetryFactory.Metrics.RecordCountSpanCreated(MetricTags.IntegrationName.AwsLambda);
         return tracer.TracerManager.ScopeManager.Activate(span, false);
@@ -57,14 +40,15 @@ internal static Scope SendStartInvocation(ILambdaExtensionRequest requestBuilder
         WriteRequestPayload(request, data);
         WriteRequestHeaders(request, context);
         var response = (HttpWebResponse)request.GetResponse();
-        var traceId = response.Headers.Get(HttpHeaderNames.TraceId);
-        var samplingPriority = response.Headers.Get(HttpHeaderNames.SamplingPriority);
-        if (ValidateOkStatus(response))
+
+        var headers = response.Headers.Wrap();
+        if (!ValidateOkStatus(response))
         {
-            return CreatePlaceholderScope(Tracer.Instance, traceId, samplingPriority);
+            return null;
         }
 
-        return null;
+        var tracer = Tracer.Instance;
+        return CreatePlaceholderScope(tracer, headers);
     }
 
     internal static void SendEndInvocation(ILambdaExtensionRequest requestBuilder, Scope scope, bool isError, string data)

tracer/test/Datadog.Trace.Tests/LambdaCommonTests.cs

@@ -9,6 +9,7 @@
 using System.Net;
 using System.Threading.Tasks;
 using Datadog.Trace.ClrProfiler.AutoInstrumentation.AWS.Lambda;
+using Datadog.Trace.ExtensionMethods;
 using Datadog.Trace.TestHelpers;
 
 using FluentAssertions;
@@ -17,6 +18,8 @@
 
 namespace Datadog.Trace.Tests
 {
+    extern alias DatadogTraceManual;
+
     public class LambdaCommonTests
     {
         private readonly Mock<ILambdaExtensionRequest> _lambdaRequestMock = new();
@@ -25,36 +28,38 @@ public class LambdaCommonTests
         public async Task TestCreatePlaceholderScopeSuccessWithTraceIdOnly()
         {
             await using var tracer = TracerHelper.CreateWithFakeAgent();
-            var scope = LambdaCommon.CreatePlaceholderScope(tracer, "1234", null);
 
+            var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" } }.Wrap();
+            var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
             scope.Should().NotBeNull();
             scope.Span.TraceId128.Should().Be((TraceId)1234);
             ((ISpan)scope.Span).TraceId.Should().Be(1234);
             scope.Span.SpanId.Should().BeGreaterThan(0);
         }
 
         [Fact]
-        public async Task TestCreatePlaceholderScopeSuccessWithSamplingPriorityOnly()
+        public async Task TestCreatePlaceholderScopeSuccessWith64BitTraceIdContext()
         {
             await using var tracer = TracerHelper.CreateWithFakeAgent();
-            var scope = LambdaCommon.CreatePlaceholderScope(tracer, null, "-1");
+            var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap();
+            var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
 
             scope.Should().NotBeNull();
-            scope.Span.TraceId128.Should().BeGreaterThan(TraceId.Zero);
-            ((ISpan)scope.Span).TraceId.Should().BeGreaterThan(0);
+            scope.Span.TraceId128.Should().Be((TraceId)1234);
+            ((ISpan)scope.Span).TraceId.Should().Be(1234);
             scope.Span.SpanId.Should().BeGreaterThan(0);
             scope.Span.Context.TraceContext.SamplingPriority.Should().Be(-1);
         }
 
         [Fact]
-        public async Task TestCreatePlaceholderScopeSuccessWithFullContext()
+        public async Task TestCreatePlaceholderScopeSuccessWith128BitTraceIdContext()
         {
             await using var tracer = TracerHelper.CreateWithFakeAgent();
-            var scope = LambdaCommon.CreatePlaceholderScope(tracer, "1234", "-1");
+            var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "5744042798732701615" }, { HttpHeaderNames.SamplingPriority, "-1" }, { HttpHeaderNames.PropagatedTags, "_dd.p.tid=1914fe7789eb32be" } }.Wrap();
+            var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
 
             scope.Should().NotBeNull();
-            scope.Span.TraceId128.Should().Be((TraceId)1234);
-            ((ISpan)scope.Span).TraceId.Should().Be(1234);
+            scope.Span.TraceId128.ToString().Should().Be("1914fe7789eb32be4fb6f07e011a6faf");
             scope.Span.SpanId.Should().BeGreaterThan(0);
             scope.Span.Context.TraceContext.SamplingPriority.Should().Be(-1);
         }
@@ -64,28 +69,14 @@ public async Task TestCreatePlaceholderScopeSuccessWithFullContext()
         public async Task TestCreatePlaceholderScopeSuccessWithoutContext()
         {
             await using var tracer = TracerHelper.CreateWithFakeAgent();
-            var scope = LambdaCommon.CreatePlaceholderScope(tracer, null, null);
+            var scope = LambdaCommon.CreatePlaceholderScope(tracer, new WebHeaderCollection().Wrap());
 
             scope.Should().NotBeNull();
             scope.Span.TraceId128.Should().BeGreaterThan((TraceId.Zero));
             ((ISpan)scope.Span).TraceId.Should().BeGreaterThan(0);
             scope.Span.SpanId.Should().BeGreaterThan(0);
         }
 
-        [Fact]
-        public async Task TestCreatePlaceholderScopeInvalidTraceId()
-        {
-            await using var tracer = TracerHelper.CreateWithFakeAgent();
-            Assert.Throws<FormatException>(() => LambdaCommon.CreatePlaceholderScope(tracer, "invalid-trace-id", "-1"));
-        }
-
-        [Fact]
-        public async Task TestCreatePlaceholderScopeInvalidSamplingPriority()
-        {
-            await using var tracer = TracerHelper.CreateWithFakeAgent();
-            Assert.Throws<FormatException>(() => LambdaCommon.CreatePlaceholderScope(tracer, "1234", "invalid-sampling-priority"));
-        }
-
         [Fact]
         [Trait("Category", "ArmUnsupported")]
         public void TestSendStartInvocationThrow()
@@ -146,7 +137,8 @@ public void TestSendStartInvocationSuccess()
         public async Task TestSendEndInvocationFailure()
         {
             await using var tracer = TracerHelper.CreateWithFakeAgent();
-            var scope = LambdaCommon.CreatePlaceholderScope(tracer, "1234", "-1");
+            var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap();
+            var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
 
             var response = new Mock<HttpWebResponse>(MockBehavior.Loose);
             var responseStream = new Mock<Stream>(MockBehavior.Loose);
@@ -166,7 +158,8 @@ public async Task TestSendEndInvocationFailure()
         public async Task TestSendEndInvocationSuccess()
         {
             await using var tracer = TracerHelper.CreateWithFakeAgent();
-            var scope = LambdaCommon.CreatePlaceholderScope(tracer, "1234", "-1");
+            var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap();
+            var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
 
             var response = new Mock<HttpWebResponse>(MockBehavior.Loose);
             var responseStream = new Mock<Stream>(MockBehavior.Loose);
@@ -189,7 +182,8 @@ public async Task TestSendEndInvocationSuccess()
         public async Task TestSendEndInvocationFalse()
         {
             await using var tracer = TracerHelper.CreateWithFakeAgent();
-            var scope = LambdaCommon.CreatePlaceholderScope(tracer, "1234", "-1");
+            var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap();
+            var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
 
             var response = new Mock<HttpWebResponse>(MockBehavior.Loose);
             var responseStream = new Mock<Stream>(MockBehavior.Loose);

tracer/test/Datadog.Trace.Tests/LambdaRequestBuilderTests.cs

@@ -3,8 +3,10 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
 // </copyright>
 #if NET6_0_OR_GREATER
+using System.Net;
 using System.Threading.Tasks;
 using Datadog.Trace.ClrProfiler.AutoInstrumentation.AWS.Lambda;
+using Datadog.Trace.ExtensionMethods;
 using Datadog.Trace.TestHelpers;
 
 using FluentAssertions;
@@ -20,7 +22,8 @@ public class LambdaRequestBuilderTests
         public async Task TestGetEndInvocationRequestWithError()
         {
             await using var tracer = TracerHelper.CreateWithFakeAgent();
-            var scope = LambdaCommon.CreatePlaceholderScope(tracer, traceId: null, samplingPriority: null);
+            var headers = new WebHeaderCollection().Wrap();
+            var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
 
             ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder();
             var request = requestBuilder.GetEndInvocationRequest(scope, isError: true);
@@ -35,7 +38,8 @@ public async Task TestGetEndInvocationRequestWithError()
         public async Task TestGetEndInvocationRequestWithoutError()
         {
             await using var tracer = TracerHelper.CreateWithFakeAgent();
-            var scope = LambdaCommon.CreatePlaceholderScope(tracer, traceId: null, samplingPriority: null);
+            var headers = new WebHeaderCollection().Wrap();
+            var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
 
             ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder();
             var request = requestBuilder.GetEndInvocationRequest(scope, isError: false);
@@ -50,7 +54,8 @@ public async Task TestGetEndInvocationRequestWithoutError()
         public async Task TestGetEndInvocationRequestWithScope()
         {
             await using var tracer = TracerHelper.CreateWithFakeAgent();
-            var scope = LambdaCommon.CreatePlaceholderScope(tracer, traceId: "1234", samplingPriority: "-1");
+            var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap();
+            var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
 
             ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder();
             var request = requestBuilder.GetEndInvocationRequest(scope, isError: false);
@@ -77,7 +82,8 @@ public void TestGetEndInvocationRequestWithoutScope()
         public async Task TestGetEndInvocationRequestWithErrorTags()
         {
             await using var tracer = TracerHelper.CreateWithFakeAgent();
-            var scope = LambdaCommon.CreatePlaceholderScope(tracer, traceId: null, samplingPriority: null);
+            var headers = new WebHeaderCollection().Wrap();
+            var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
             scope.Span.SetTag("error.msg", "Exception");
             scope.Span.SetTag("error.type", "Exception");
             scope.Span.SetTag("error.stack", "everything is " + System.Environment.NewLine + "fine");
@@ -98,7 +104,8 @@ public async Task TestGetEndInvocationRequestWithErrorTags()
         public async Task TestGetEndInvocationRequestWithoutErrorTags()
         {
             await using var tracer = TracerHelper.CreateWithFakeAgent();
-            var scope = LambdaCommon.CreatePlaceholderScope(tracer, null, null);
+            var headers = new WebHeaderCollection().Wrap();
+            var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
 
             ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder();
             var request = requestBuilder.GetEndInvocationRequest(scope, true);