From b29d9843d3d1b8774a8350c9a67670ad65b689a4 Mon Sep 17 00:00:00 2001 From: Christopher Agocs Date: Tue, 22 Oct 2024 15:08:12 -0700 Subject: [PATCH] Backport 6401: Extract Upper64 bit trace ID from extension response (#6181) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary of changes I've updated the Lambda extension so it is capable of returning a 128 bit trace ID when a tracer calls the `/lambda/start-invocation` endpoint [in this PR](https://github.com/DataDog/datadog-agent/pull/27988) As per [the RFC](https://datadoghq.atlassian.net/wiki/spaces/RUMP/pages/3545630931/RFC+Support+128+bit+trace+IDs+in+RUM+SDKs#:~:text=For%20Datadog%20headers%2C%20the%20128%20bit%20trace%20id%20is%20sent%20in%20two%20parts%2C%20lower%2064%20bits%20in%20x%2Ddatadog%2Dtrace%2Did%20(decimal)%20and%20the%20higher%2064%20bits%20in%20x%2Ddatadog%2Dtags%20header%20under%20_dd.p.tid%20(hex)%20tag), the > lower 64 bits in `x-datadog-trace-id` (decimal) and the higher 64 bits in `x-datadog-tags` header under `_dd.p.tid` (hex) tag. This change modifies the function that calls `/lambda/start-invocation`, allowing it to pick out the upper 64 bits of the trace ID and set the resulting 128-bit trace ID in the extracted context. ## Reason for change The Lambda Extension may now return a 128 bit trace ID when a Step Function invokes a Lambda Function. ## Implementation details I rewrote LambdaCommon's `CreatePlaceholderScope` so it uses `SpanContextPropagator.Instance.Extract` rather than extracting trace context elements one by one. ## Test coverage Added a unit test for 128 bit trace IDs. Fixed existing unit tests so they pass a dictionary of headers to CreatePlaceholderScope. Removed a unit test that only passes SamplingPriority, since a distributed trace with only a sampling priority is hardly a distributed trace at all. ## Other details Backport to 2.x of #6041 --------- Co-authored-by: Lucas Pimentel Co-authored-by: Andrew Lock Co-authored-by: Daniel Romano <108014683+daniel-romano-DD@users.noreply.github.com> Co-authored-by: Steven Bouwkamp Co-authored-by: Anna Co-authored-by: NachoEchevarria <53266532+NachoEchevarria@users.noreply.github.com> Co-authored-by: William Conti <58711692+wconti27@users.noreply.github.com> Co-authored-by: Kevin Gosse Co-authored-by: Tony Redondo Co-authored-by: Gregory LEOCADIE --- .../AWS/Lambda/LambdaCommon.cs | 41 ++++++----------- .../Datadog.Trace.Tests/LambdaCommonTests.cs | 46 ++++++++----------- .../LambdaRequestBuilderTests.cs | 17 +++++-- 3 files changed, 44 insertions(+), 60 deletions(-) diff --git a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaCommon.cs b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaCommon.cs index 24f5c9bb23a0..26e5519be9c6 100644 --- a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaCommon.cs +++ b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaCommon.cs @@ -8,6 +8,9 @@ using System.Net; using System.Text; using System.Threading.Tasks; +using Datadog.Trace.ExtensionMethods; +using Datadog.Trace.Headers; +using Datadog.Trace.Propagators; using Datadog.Trace.Telemetry; using Datadog.Trace.Telemetry.Metrics; using Datadog.Trace.Util; @@ -21,31 +24,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 +41,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) diff --git a/tracer/test/Datadog.Trace.Tests/LambdaCommonTests.cs b/tracer/test/Datadog.Trace.Tests/LambdaCommonTests.cs index d7a1756c53be..d99a54eec600 100644 --- a/tracer/test/Datadog.Trace.Tests/LambdaCommonTests.cs +++ b/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; @@ -25,8 +26,9 @@ 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); @@ -34,27 +36,28 @@ public async Task TestCreatePlaceholderScopeSuccessWithTraceIdOnly() } [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,7 +67,7 @@ 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)); @@ -72,20 +75,6 @@ public async Task TestCreatePlaceholderScopeSuccessWithoutContext() scope.Span.SpanId.Should().BeGreaterThan(0); } - [Fact] - public async Task TestCreatePlaceholderScopeInvalidTraceId() - { - await using var tracer = TracerHelper.CreateWithFakeAgent(); - Assert.Throws(() => LambdaCommon.CreatePlaceholderScope(tracer, "invalid-trace-id", "-1")); - } - - [Fact] - public async Task TestCreatePlaceholderScopeInvalidSamplingPriority() - { - await using var tracer = TracerHelper.CreateWithFakeAgent(); - Assert.Throws(() => LambdaCommon.CreatePlaceholderScope(tracer, "1234", "invalid-sampling-priority")); - } - [Fact] [Trait("Category", "ArmUnsupported")] public void TestSendStartInvocationThrow() @@ -146,7 +135,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(MockBehavior.Loose); var responseStream = new Mock(MockBehavior.Loose); @@ -166,7 +156,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(MockBehavior.Loose); var responseStream = new Mock(MockBehavior.Loose); @@ -189,7 +180,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(MockBehavior.Loose); var responseStream = new Mock(MockBehavior.Loose); diff --git a/tracer/test/Datadog.Trace.Tests/LambdaRequestBuilderTests.cs b/tracer/test/Datadog.Trace.Tests/LambdaRequestBuilderTests.cs index 2797a0f82d89..dffadf7660cb 100644 --- a/tracer/test/Datadog.Trace.Tests/LambdaRequestBuilderTests.cs +++ b/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. // #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);