diff --git a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/HandlerWrapperSetHandlerIntegration.cs b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/HandlerWrapperSetHandlerIntegration.cs index 8f3bee451276..50f370f976e0 100644 --- a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/HandlerWrapperSetHandlerIntegration.cs +++ b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/HandlerWrapperSetHandlerIntegration.cs @@ -71,6 +71,7 @@ public object OnDelegateBegin(object sender, ref TArg1 arg) LambdaCommon.Log("DelegateWrapper Running OnDelegateBegin"); Scope scope; + object requestid = null; var proxyInstance = arg.DuckCast(); if (proxyInstance == null) { @@ -80,11 +81,12 @@ public object OnDelegateBegin(object sender, ref TArg1 arg) else { var jsonString = ConvertPayloadStream(proxyInstance.InputStream); - scope = LambdaCommon.SendStartInvocation(new LambdaRequestBuilder(), jsonString, proxyInstance.LambdaContext?.ClientContext?.Custom); + scope = LambdaCommon.SendStartInvocation(new LambdaRequestBuilder(), jsonString, proxyInstance.LambdaContext); + requestid = proxyInstance.LambdaContext?.AwsRequestId; } LambdaCommon.Log("DelegateWrapper FINISHED Running OnDelegateBegin"); - return new CallTargetState(scope); + return new CallTargetState(scope, requestid); } public void OnException(object sender, Exception ex) @@ -108,18 +110,18 @@ public async Task OnDelegateEndAsync(object sender, if (proxyInstance == null) { LambdaCommon.Log("DuckCast.IInvocationResponse got null proxyInstance", debug: false); - await LambdaCommon.EndInvocationAsync(string.Empty, exception, ((CallTargetState)state!).Scope, RequestBuilder).ConfigureAwait(false); + await LambdaCommon.EndInvocationAsync(string.Empty, exception, state, RequestBuilder).ConfigureAwait(false); } else { var jsonString = ConvertPayloadStream(proxyInstance.OutputStream); - await LambdaCommon.EndInvocationAsync(jsonString, exception, ((CallTargetState)state!).Scope, RequestBuilder).ConfigureAwait(false); + await LambdaCommon.EndInvocationAsync(jsonString, exception, state, RequestBuilder).ConfigureAwait(false); } } catch (Exception ex) { LambdaCommon.Log("OnDelegateEndAsync could not send payload to the extension", ex, false); - await LambdaCommon.EndInvocationAsync(string.Empty, ex, ((CallTargetState)state!).Scope, RequestBuilder).ConfigureAwait(false); + await LambdaCommon.EndInvocationAsync(string.Empty, ex, state, RequestBuilder).ConfigureAwait(false); } LambdaCommon.Log("DelegateWrapper FINISHED Running OnDelegateEndAsync"); diff --git a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/ILambdaExtensionRequest.cs b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/ILambdaExtensionRequest.cs index 86bdb1bf91de..c2d78cca1b06 100644 --- a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/ILambdaExtensionRequest.cs +++ b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/ILambdaExtensionRequest.cs @@ -20,7 +20,7 @@ internal interface ILambdaExtensionRequest /// Get the end invocation request /// /// The end invocation request - WebRequest GetEndInvocationRequest(Scope scope, bool isError); + WebRequest GetEndInvocationRequest(Scope scope, object state, bool isError); } #endif 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 7579f10c5314..e9c2203de200 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,8 @@ using System.Net; using System.Text; using System.Threading.Tasks; +using Datadog.Trace.ClrProfiler.AutoInstrumentation.AWS.SDK; +using Datadog.Trace.ClrProfiler.CallTarget; using Datadog.Trace.Headers; using Datadog.Trace.Propagators; using Datadog.Trace.Telemetry; @@ -22,6 +24,7 @@ internal abstract class LambdaCommon private const string PlaceholderOperationName = "placeholder-operation"; private const double ServerlessMaxWaitingFlushTime = 3; private const string LogLevelEnvName = "DD_LOG_LEVEL"; + private const string LambdaRuntimeAwsRequestIdHeader = "lambda-runtime-aws-request-id"; internal static Scope CreatePlaceholderScope(Tracer tracer, NameValueHeadersCollection headers) { @@ -38,11 +41,16 @@ internal static Scope CreatePlaceholderScope(Tracer tracer, NameValueHeadersColl return tracer.TracerManager.ScopeManager.Activate(span, false); } - internal static Scope SendStartInvocation(ILambdaExtensionRequest requestBuilder, string data, IDictionary context) + internal static Scope SendStartInvocation(ILambdaExtensionRequest requestBuilder, string data, ILambdaContext context) { var request = requestBuilder.GetStartInvocationRequest(); WriteRequestPayload(request, data); - WriteRequestHeaders(request, context); + WriteRequestHeaders(request, context?.ClientContext?.Custom); + if (context?.AwsRequestId != null) + { + request.Headers.Add(LambdaRuntimeAwsRequestIdHeader, context.AwsRequestId); + } + using var response = (HttpWebResponse)request.GetResponse(); var headers = response.Headers.Wrap(); @@ -55,9 +63,9 @@ internal static Scope SendStartInvocation(ILambdaExtensionRequest requestBuilder return CreatePlaceholderScope(tracer, headers); } - internal static void SendEndInvocation(ILambdaExtensionRequest requestBuilder, Scope scope, bool isError, string data) + internal static void SendEndInvocation(ILambdaExtensionRequest requestBuilder, Scope scope, object state, bool isError, string data) { - var request = requestBuilder.GetEndInvocationRequest(scope, isError); + var request = requestBuilder.GetEndInvocationRequest(scope, state, isError); WriteRequestPayload(request, data); using var response = (HttpWebResponse)request.GetResponse(); @@ -67,8 +75,10 @@ internal static void SendEndInvocation(ILambdaExtensionRequest requestBuilder, S } } - internal static async Task EndInvocationAsync(string returnValue, Exception exception, Scope scope, ILambdaExtensionRequest requestBuilder) + internal static async Task EndInvocationAsync(string returnValue, Exception exception, object stateObject, ILambdaExtensionRequest requestBuilder) { + var state = (CallTargetState)stateObject!; + var scope = state.Scope; try { await Task.WhenAll( @@ -90,7 +100,7 @@ await Task.WhenAll( span.SetException(exception); } - SendEndInvocation(requestBuilder, scope, exception != null, returnValue); + SendEndInvocation(requestBuilder, scope, state.State, exception != null, returnValue); } catch (Exception ex) { diff --git a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaRequestBuilder.cs b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaRequestBuilder.cs index cef67f436911..0ef8441dcdfc 100644 --- a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaRequestBuilder.cs +++ b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaRequestBuilder.cs @@ -35,12 +35,17 @@ WebRequest ILambdaExtensionRequest.GetStartInvocationRequest() return request; } - WebRequest ILambdaExtensionRequest.GetEndInvocationRequest(Scope scope, bool isError) + WebRequest ILambdaExtensionRequest.GetEndInvocationRequest(Scope scope, object state, bool isError) { var request = WebRequest.Create(Uri + EndInvocationPath); request.Method = "POST"; request.Headers.Set(HttpHeaderNames.TracingEnabled, "false"); + if (state != null) + { + request.Headers.Set("lambda-runtime-aws-request-id", (string)state); + } + if (scope is { Span: var span }) { // TODO: add support for 128-bit trace ids in serverless diff --git a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/SDK/ILambdaContext.cs b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/SDK/ILambdaContext.cs index af2137c43a85..9aa5ddb2dbc3 100644 --- a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/SDK/ILambdaContext.cs +++ b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/SDK/ILambdaContext.cs @@ -12,12 +12,12 @@ namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.AWS.SDK; /// internal interface ILambdaContext { - // /// - // /// Gets the AWS request ID associated with the request. - // /// This is the same ID returned to the client that called invoke(). - // /// This ID is reused for retries on the same request. - // /// - // string AwsRequestId { get; } + /// + /// Gets the AWS request ID associated with the request. + /// This is the same ID returned to the client that called invoke(). + /// This ID is reused for retries on the same request. + /// + string AwsRequestId { get; } /// /// Gets information about the client application and device when invoked diff --git a/tracer/test/Datadog.Trace.Tests/LambdaCommonTests.cs b/tracer/test/Datadog.Trace.Tests/LambdaCommonTests.cs index 2bf9e7ffb93d..055915868279 100644 --- a/tracer/test/Datadog.Trace.Tests/LambdaCommonTests.cs +++ b/tracer/test/Datadog.Trace.Tests/LambdaCommonTests.cs @@ -91,7 +91,7 @@ public void TestSendStartInvocationThrow() _lambdaRequestMock.Setup(lr => lr.GetStartInvocationRequest()).Returns(httpRequest.Object); - Assert.Throws(() => LambdaCommon.SendStartInvocation(_lambdaRequestMock.Object, "{}", new Dictionary())); + Assert.Throws(() => LambdaCommon.SendStartInvocation(_lambdaRequestMock.Object, "{}", null)); } [Fact] @@ -110,7 +110,7 @@ public void TestSendStartInvocationNull() _lambdaRequestMock.Setup(lr => lr.GetStartInvocationRequest()).Returns(httpRequest.Object); - LambdaCommon.SendStartInvocation(_lambdaRequestMock.Object, "{}", new Dictionary()).Should().BeNull(); + LambdaCommon.SendStartInvocation(_lambdaRequestMock.Object, "{}", null).Should().BeNull(); } [Fact] @@ -129,7 +129,7 @@ public void TestSendStartInvocationSuccess() _lambdaRequestMock.Setup(lr => lr.GetStartInvocationRequest()).Returns(httpRequest.Object); - LambdaCommon.SendStartInvocation(_lambdaRequestMock.Object, "{}", new Dictionary()).Should().NotBeNull(); + LambdaCommon.SendStartInvocation(_lambdaRequestMock.Object, "{}", null).Should().NotBeNull(); } [Fact] @@ -139,6 +139,7 @@ public async Task TestSendEndInvocationFailure() await using var tracer = TracerHelper.CreateWithFakeAgent(); var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); + var state = "example-aws-request-id"; var response = new Mock(MockBehavior.Loose); var responseStream = new Mock(MockBehavior.Loose); @@ -148,9 +149,9 @@ public async Task TestSendEndInvocationFailure() httpRequest.Setup(h => h.GetResponse()).Throws(new WebException()); httpRequest.Setup(h => h.GetRequestStream()).Returns(responseStream.Object); - _lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(scope, true)).Returns(httpRequest.Object); + _lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(scope, state, true)).Returns(httpRequest.Object); - Assert.Throws(() => LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, true, "{}")); + Assert.Throws(() => LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, state, true, "{}")); } [Fact] @@ -160,6 +161,7 @@ public async Task TestSendEndInvocationSuccess() await using var tracer = TracerHelper.CreateWithFakeAgent(); var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); + var state = "example-aws-request-id"; var response = new Mock(MockBehavior.Loose); var responseStream = new Mock(MockBehavior.Loose); @@ -169,10 +171,10 @@ public async Task TestSendEndInvocationSuccess() httpRequest.Setup(h => h.GetResponse()).Returns(response.Object); httpRequest.Setup(h => h.GetRequestStream()).Returns(responseStream.Object); - _lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(scope, true)).Returns(httpRequest.Object); + _lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(scope, state, true)).Returns(httpRequest.Object); var output = new StringWriter(); Console.SetOut(output); - LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, true, "{}"); + LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, state, true, "{}"); httpRequest.Verify(r => r.GetResponse(), Times.Once); Assert.Empty(output.ToString()); } @@ -184,6 +186,7 @@ public async Task TestSendEndInvocationFalse() await using var tracer = TracerHelper.CreateWithFakeAgent(); var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); + var state = "example-aws-request-id"; var response = new Mock(MockBehavior.Loose); var responseStream = new Mock(MockBehavior.Loose); @@ -193,10 +196,10 @@ public async Task TestSendEndInvocationFalse() httpRequest.Setup(h => h.GetResponse()).Returns(response.Object); httpRequest.Setup(h => h.GetRequestStream()).Returns(responseStream.Object); - _lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(scope, true)).Returns(httpRequest.Object); + _lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(scope, state, true)).Returns(httpRequest.Object); var output = new StringWriter(); Console.SetOut(output); - LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, true, "{}"); + LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, state, true, "{}"); httpRequest.Verify(r => r.GetResponse(), Times.Once); Assert.Contains("Extension does not send a status 200 OK", output.ToString()); } diff --git a/tracer/test/Datadog.Trace.Tests/LambdaRequestBuilderTests.cs b/tracer/test/Datadog.Trace.Tests/LambdaRequestBuilderTests.cs index f71b5580f442..064566aa5d85 100644 --- a/tracer/test/Datadog.Trace.Tests/LambdaRequestBuilderTests.cs +++ b/tracer/test/Datadog.Trace.Tests/LambdaRequestBuilderTests.cs @@ -24,14 +24,16 @@ public async Task TestGetEndInvocationRequestWithError() await using var tracer = TracerHelper.CreateWithFakeAgent(); var headers = new WebHeaderCollection().Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); + var state = "example-aws-request-id"; ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder(); - var request = requestBuilder.GetEndInvocationRequest(scope, isError: true); + var request = requestBuilder.GetEndInvocationRequest(scope, state, isError: true); request.Headers.Get("x-datadog-invocation-error").Should().Be("true"); request.Headers.Get("x-datadog-tracing-enabled").Should().Be("false"); request.Headers.Get("x-datadog-sampling-priority").Should().Be("1"); request.Headers.Get("x-datadog-trace-id").Should().NotBeNull(); request.Headers.Get("x-datadog-span-id").Should().NotBeNull(); + request.Headers.Get("lambda-runtime-aws-request-id").Should().Be("example-aws-request-id"); } [Fact] @@ -40,14 +42,16 @@ public async Task TestGetEndInvocationRequestWithoutError() await using var tracer = TracerHelper.CreateWithFakeAgent(); var headers = new WebHeaderCollection().Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); + var state = "example-aws-request-id"; ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder(); - var request = requestBuilder.GetEndInvocationRequest(scope, isError: false); + var request = requestBuilder.GetEndInvocationRequest(scope, state, isError: false); request.Headers.Get("x-datadog-invocation-error").Should().BeNull(); request.Headers.Get("x-datadog-tracing-enabled").Should().Be("false"); request.Headers.Get("x-datadog-sampling-priority").Should().Be("1"); request.Headers.Get("x-datadog-trace-id").Should().NotBeNull(); request.Headers.Get("x-datadog-span-id").Should().NotBeNull(); + request.Headers.Get("lambda-runtime-aws-request-id").Should().Be("example-aws-request-id"); } [Fact] @@ -56,26 +60,47 @@ public async Task TestGetEndInvocationRequestWithScope() await using var tracer = TracerHelper.CreateWithFakeAgent(); var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" } }.Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); + var state = "example-aws-request-id"; ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder(); - var request = requestBuilder.GetEndInvocationRequest(scope, isError: false); + var request = requestBuilder.GetEndInvocationRequest(scope, state, isError: false); request.Headers.Get("x-datadog-invocation-error").Should().BeNull(); request.Headers.Get("x-datadog-tracing-enabled").Should().Be("false"); request.Headers.Get("x-datadog-sampling-priority").Should().Be("1"); request.Headers.Get("x-datadog-trace-id").Should().Be("1234"); request.Headers.Get("x-datadog-span-id").Should().NotBeNull(); + request.Headers.Get("lambda-runtime-aws-request-id").Should().Be("example-aws-request-id"); } [Fact] public void TestGetEndInvocationRequestWithoutScope() { ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder(); - var request = requestBuilder.GetEndInvocationRequest(scope: null, isError: false); + var state = "example-aws-request-id"; + var request = requestBuilder.GetEndInvocationRequest(scope: null, state, isError: false); request.Headers.Get("x-datadog-invocation-error").Should().BeNull(); request.Headers.Get("x-datadog-tracing-enabled").Should().Be("false"); request.Headers.Get("x-datadog-sampling-priority").Should().BeNull(); request.Headers.Get("x-datadog-trace-id").Should().BeNull(); request.Headers.Get("x-datadog-span-id").Should().BeNull(); + request.Headers.Get("lambda-runtime-aws-request-id").Should().Be("example-aws-request-id"); + } + + [Fact] + public async Task TestGetEndInvocationRequestWithoutState() + { + await using var tracer = TracerHelper.CreateWithFakeAgent(); + var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" } }.Wrap(); + var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); + + ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder(); + var request = requestBuilder.GetEndInvocationRequest(scope, state: null, isError: false); + request.Headers.Get("x-datadog-invocation-error").Should().BeNull(); + request.Headers.Get("x-datadog-tracing-enabled").Should().Be("false"); + request.Headers.Get("x-datadog-sampling-priority").Should().Be("1"); + request.Headers.Get("x-datadog-trace-id").Should().Be("1234"); + request.Headers.Get("x-datadog-span-id").Should().NotBeNull(); + request.Headers.Get("lambda-runtime-aws-request-id").Should().BeNull(); } [Fact] @@ -84,6 +109,7 @@ public async Task TestGetEndInvocationRequestWithErrorTags() await using var tracer = TracerHelper.CreateWithFakeAgent(); var headers = new WebHeaderCollection().Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); + var state = "example-aws-request-id"; var errorMsg = "Exception"; var errorType = "Exception"; @@ -97,7 +123,7 @@ public async Task TestGetEndInvocationRequestWithErrorTags() var expectedErrorStack = Convert.ToBase64String(System.Text.Encoding.UTF8.GetBytes(errorStack)); ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder(); - var request = requestBuilder.GetEndInvocationRequest(scope, true); + var request = requestBuilder.GetEndInvocationRequest(scope, state, true); request.Headers.Get("x-datadog-invocation-error").Should().NotBeNull(); request.Headers.Get("x-datadog-invocation-error-msg").Should().Be(expectedErrorMsg); request.Headers.Get("x-datadog-invocation-error-type").Should().Be(expectedErrorType); @@ -106,6 +132,7 @@ public async Task TestGetEndInvocationRequestWithErrorTags() request.Headers.Get("x-datadog-sampling-priority").Should().Be("1"); request.Headers.Get("x-datadog-trace-id").Should().NotBeNull(); request.Headers.Get("x-datadog-span-id").Should().NotBeNull(); + request.Headers.Get("lambda-runtime-aws-request-id").Should().Be("example-aws-request-id"); } [Fact] @@ -114,9 +141,10 @@ public async Task TestGetEndInvocationRequestWithoutErrorTags() await using var tracer = TracerHelper.CreateWithFakeAgent(); var headers = new WebHeaderCollection().Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); + var state = "example-aws-request-id"; ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder(); - var request = requestBuilder.GetEndInvocationRequest(scope, true); + var request = requestBuilder.GetEndInvocationRequest(scope, state, true); request.Headers.Get("x-datadog-invocation-error").Should().NotBeNull(); request.Headers.Get("x-datadog-invocation-error-msg").Should().BeNull(); request.Headers.Get("x-datadog-invocation-error-type").Should().BeNull(); @@ -125,6 +153,7 @@ public async Task TestGetEndInvocationRequestWithoutErrorTags() request.Headers.Get("x-datadog-sampling-priority").Should().Be("1"); request.Headers.Get("x-datadog-trace-id").Should().NotBeNull(); request.Headers.Get("x-datadog-span-id").Should().NotBeNull(); + request.Headers.Get("lambda-runtime-aws-request-id").Should().Be("example-aws-request-id"); } } }