Skip to content

Commit

Permalink
Backport 6401: Extract Upper64 bit trace ID from extension response (#…
Browse files Browse the repository at this point in the history
…6181)

## 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](DataDog/datadog-agent#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
<!-- Fixes #{issue} -->

<!-- ⚠️ Note: where possible, please obtain 2 approvals prior to
merging. Unless CODEOWNERS specifies otherwise, for external teams it is
typically best to have one review from a team member, and one review
from apm-dotnet. Trivial changes do not require 2 reviews. -->

---------

Co-authored-by: Lucas Pimentel <[email protected]>
Co-authored-by: Andrew Lock <[email protected]>
Co-authored-by: Daniel Romano <[email protected]>
Co-authored-by: Steven Bouwkamp <[email protected]>
Co-authored-by: Anna <[email protected]>
Co-authored-by: NachoEchevarria <[email protected]>
Co-authored-by: William Conti <[email protected]>
Co-authored-by: Kevin Gosse <[email protected]>
Co-authored-by: Tony Redondo <[email protected]>
Co-authored-by: Gregory LEOCADIE <[email protected]>
  • Loading branch information
11 people authored Oct 22, 2024
1 parent 4397347 commit b29d984
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 60 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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);
Expand All @@ -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)
Expand Down
46 changes: 19 additions & 27 deletions tracer/test/Datadog.Trace.Tests/LambdaCommonTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -25,36 +26,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);
}
Expand All @@ -64,28 +67,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()
Expand Down Expand Up @@ -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<HttpWebResponse>(MockBehavior.Loose);
var responseStream = new Mock<Stream>(MockBehavior.Loose);
Expand All @@ -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<HttpWebResponse>(MockBehavior.Loose);
var responseStream = new Mock<Stream>(MockBehavior.Loose);
Expand All @@ -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<HttpWebResponse>(MockBehavior.Loose);
var responseStream = new Mock<Stream>(MockBehavior.Loose);
Expand Down
17 changes: 12 additions & 5 deletions tracer/test/Datadog.Trace.Tests/LambdaRequestBuilderTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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");
Expand All @@ -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);
Expand Down

0 comments on commit b29d984

Please sign in to comment.