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 50f370f976e0..2a55a521716a 100644 --- a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/HandlerWrapperSetHandlerIntegration.cs +++ b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/HandlerWrapperSetHandlerIntegration.cs @@ -71,7 +71,7 @@ public object OnDelegateBegin(object sender, ref TArg1 arg) LambdaCommon.Log("DelegateWrapper Running OnDelegateBegin"); Scope scope; - object requestid = null; + object requestId = null; var proxyInstance = arg.DuckCast(); if (proxyInstance == null) { @@ -82,11 +82,11 @@ public object OnDelegateBegin(object sender, ref TArg1 arg) { var jsonString = ConvertPayloadStream(proxyInstance.InputStream); scope = LambdaCommon.SendStartInvocation(new LambdaRequestBuilder(), jsonString, proxyInstance.LambdaContext); - requestid = proxyInstance.LambdaContext?.AwsRequestId; + requestId = proxyInstance.LambdaContext?.AwsRequestId; } LambdaCommon.Log("DelegateWrapper FINISHED Running OnDelegateBegin"); - return new CallTargetState(scope, requestid); + return new CallTargetState(scope, requestId); } public void OnException(object sender, Exception ex) @@ -104,27 +104,31 @@ public TReturn OnDelegateEnd(object sender, TReturn returnValue, Except public async Task OnDelegateEndAsync(object sender, TInnerReturn returnValue, Exception exception, object state) { LambdaCommon.Log("DelegateWrapper Running OnDelegateEndAsync"); - try + if (state is CallTargetState callTargetState) { - var proxyInstance = returnValue.DuckCast(); - if (proxyInstance == null) + try { - LambdaCommon.Log("DuckCast.IInvocationResponse got null proxyInstance", debug: false); - await LambdaCommon.EndInvocationAsync(string.Empty, exception, state, RequestBuilder).ConfigureAwait(false); + var proxyInstance = returnValue.DuckCast(); + if (proxyInstance == null) + { + LambdaCommon.Log("DuckCast.IInvocationResponse got null proxyInstance", debug: false); + await LambdaCommon.EndInvocationAsync(string.Empty, exception, callTargetState, RequestBuilder).ConfigureAwait(false); + } + else + { + var jsonString = ConvertPayloadStream(proxyInstance.OutputStream); + await LambdaCommon.EndInvocationAsync(jsonString, exception, callTargetState, RequestBuilder).ConfigureAwait(false); + } } - else + catch (Exception ex) { - var jsonString = ConvertPayloadStream(proxyInstance.OutputStream); - await LambdaCommon.EndInvocationAsync(jsonString, exception, state, RequestBuilder).ConfigureAwait(false); + LambdaCommon.Log("OnDelegateEndAsync could not send payload to the extension", ex, false); + await LambdaCommon.EndInvocationAsync(string.Empty, ex, callTargetState, RequestBuilder).ConfigureAwait(false); } - } - catch (Exception ex) - { - LambdaCommon.Log("OnDelegateEndAsync could not send payload to the extension", ex, false); - await LambdaCommon.EndInvocationAsync(string.Empty, ex, state, RequestBuilder).ConfigureAwait(false); + + LambdaCommon.Log("DelegateWrapper FINISHED Running OnDelegateEndAsync"); } - LambdaCommon.Log("DelegateWrapper FINISHED Running OnDelegateEndAsync"); return returnValue; } } 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 c2d78cca1b06..4bbabffdbeda 100644 --- a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/ILambdaExtensionRequest.cs +++ b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/ILambdaExtensionRequest.cs @@ -5,6 +5,7 @@ #if NET6_0_OR_GREATER using System.Net; +using Datadog.Trace.ClrProfiler.CallTarget; namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.AWS.Lambda; @@ -20,7 +21,7 @@ internal interface ILambdaExtensionRequest /// Get the end invocation request /// /// The end invocation request - WebRequest GetEndInvocationRequest(Scope scope, object state, bool isError); + WebRequest GetEndInvocationRequest(CallTargetState stateObject, 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 e9c2203de200..b472003bc700 100644 --- a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaCommon.cs +++ b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaCommon.cs @@ -63,9 +63,9 @@ internal static Scope SendStartInvocation(ILambdaExtensionRequest requestBuilder return CreatePlaceholderScope(tracer, headers); } - internal static void SendEndInvocation(ILambdaExtensionRequest requestBuilder, Scope scope, object state, bool isError, string data) + internal static void SendEndInvocation(ILambdaExtensionRequest requestBuilder, CallTargetState stateObject, bool isError, string data) { - var request = requestBuilder.GetEndInvocationRequest(scope, state, isError); + var request = requestBuilder.GetEndInvocationRequest(stateObject, isError); WriteRequestPayload(request, data); using var response = (HttpWebResponse)request.GetResponse(); @@ -75,10 +75,9 @@ internal static void SendEndInvocation(ILambdaExtensionRequest requestBuilder, S } } - internal static async Task EndInvocationAsync(string returnValue, Exception exception, object stateObject, ILambdaExtensionRequest requestBuilder) + internal static async Task EndInvocationAsync(string returnValue, Exception exception, CallTargetState stateObject, ILambdaExtensionRequest requestBuilder) { - var state = (CallTargetState)stateObject!; - var scope = state.Scope; + var scope = stateObject.Scope; try { await Task.WhenAll( @@ -100,7 +99,7 @@ await Task.WhenAll( span.SetException(exception); } - SendEndInvocation(requestBuilder, scope, state.State, exception != null, returnValue); + SendEndInvocation(requestBuilder, stateObject, 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 0ef8441dcdfc..56fe65abeb2c 100644 --- a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaRequestBuilder.cs +++ b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaRequestBuilder.cs @@ -7,6 +7,7 @@ using System.Globalization; using System.Net; using Datadog.Trace.Agent.Transports; +using Datadog.Trace.ClrProfiler.CallTarget; using Datadog.Trace.Util; #pragma warning disable CS0618 // WebRequest, HttpWebRequest, ServicePoint, and WebClient are obsolete. Use HttpClient instead. @@ -35,18 +36,18 @@ WebRequest ILambdaExtensionRequest.GetStartInvocationRequest() return request; } - WebRequest ILambdaExtensionRequest.GetEndInvocationRequest(Scope scope, object state, bool isError) + WebRequest ILambdaExtensionRequest.GetEndInvocationRequest(CallTargetState stateObject, bool isError) { var request = WebRequest.Create(Uri + EndInvocationPath); request.Method = "POST"; request.Headers.Set(HttpHeaderNames.TracingEnabled, "false"); - if (state != null) + if (stateObject.State is string state) { - request.Headers.Set("lambda-runtime-aws-request-id", (string)state); + request.Headers.Set("lambda-runtime-aws-request-id", state); } - if (scope is { Span: var span }) + if (stateObject.Scope is { Span: var span }) { // TODO: add support for 128-bit trace ids in serverless request.Headers.Set(HttpHeaderNames.TraceId, span.TraceId128.Lower.ToString(CultureInfo.InvariantCulture)); diff --git a/tracer/test/Datadog.Trace.Tests/LambdaCommonTests.cs b/tracer/test/Datadog.Trace.Tests/LambdaCommonTests.cs index 055915868279..833359a3fac7 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.ClrProfiler.CallTarget; using Datadog.Trace.ExtensionMethods; using Datadog.Trace.TestHelpers; using Datadog.Trace.TestHelpers.TestTracer; @@ -140,6 +141,7 @@ public async Task TestSendEndInvocationFailure() var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); var state = "example-aws-request-id"; + var stateObject = new CallTargetState(scope, state); var response = new Mock(MockBehavior.Loose); var responseStream = new Mock(MockBehavior.Loose); @@ -149,9 +151,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, state, true)).Returns(httpRequest.Object); + _lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(stateObject, true)).Returns(httpRequest.Object); - Assert.Throws(() => LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, state, true, "{}")); + Assert.Throws(() => LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, stateObject, true, "{}")); } [Fact] @@ -162,6 +164,7 @@ public async Task TestSendEndInvocationSuccess() var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); var state = "example-aws-request-id"; + var stateObject = new CallTargetState(scope, state); var response = new Mock(MockBehavior.Loose); var responseStream = new Mock(MockBehavior.Loose); @@ -171,10 +174,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, state, true)).Returns(httpRequest.Object); + _lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(stateObject, true)).Returns(httpRequest.Object); var output = new StringWriter(); Console.SetOut(output); - LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, state, true, "{}"); + LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, stateObject, true, "{}"); httpRequest.Verify(r => r.GetResponse(), Times.Once); Assert.Empty(output.ToString()); } @@ -187,6 +190,7 @@ public async Task TestSendEndInvocationFalse() var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); var state = "example-aws-request-id"; + var stateObject = new CallTargetState(scope, state); var response = new Mock(MockBehavior.Loose); var responseStream = new Mock(MockBehavior.Loose); @@ -196,10 +200,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, state, true)).Returns(httpRequest.Object); + _lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(stateObject, true)).Returns(httpRequest.Object); var output = new StringWriter(); Console.SetOut(output); - LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, state, true, "{}"); + LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, stateObject, 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 064566aa5d85..3b2118e9733c 100644 --- a/tracer/test/Datadog.Trace.Tests/LambdaRequestBuilderTests.cs +++ b/tracer/test/Datadog.Trace.Tests/LambdaRequestBuilderTests.cs @@ -7,6 +7,7 @@ using System.Net; using System.Threading.Tasks; using Datadog.Trace.ClrProfiler.AutoInstrumentation.AWS.Lambda; +using Datadog.Trace.ClrProfiler.CallTarget; using Datadog.Trace.ExtensionMethods; using Datadog.Trace.TestHelpers.TestTracer; using FluentAssertions; @@ -25,10 +26,10 @@ public async Task TestGetEndInvocationRequestWithError() var headers = new WebHeaderCollection().Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); var state = "example-aws-request-id"; + var stateObject = new CallTargetState(scope, state); ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder(); - var request = requestBuilder.GetEndInvocationRequest(scope, state, isError: true); - request.Headers.Get("x-datadog-invocation-error").Should().Be("true"); + var request = requestBuilder.GetEndInvocationRequest(stateObject, 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(); @@ -43,9 +44,10 @@ public async Task TestGetEndInvocationRequestWithoutError() var headers = new WebHeaderCollection().Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); var state = "example-aws-request-id"; + var stateObject = new CallTargetState(scope, state); ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder(); - var request = requestBuilder.GetEndInvocationRequest(scope, state, isError: false); + var request = requestBuilder.GetEndInvocationRequest(stateObject, 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"); @@ -61,9 +63,10 @@ public async Task TestGetEndInvocationRequestWithScope() var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" } }.Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); var state = "example-aws-request-id"; + var stateObject = new CallTargetState(scope, state); ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder(); - var request = requestBuilder.GetEndInvocationRequest(scope, state, isError: false); + var request = requestBuilder.GetEndInvocationRequest(stateObject, 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"); @@ -77,7 +80,9 @@ public void TestGetEndInvocationRequestWithoutScope() { ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder(); var state = "example-aws-request-id"; - var request = requestBuilder.GetEndInvocationRequest(scope: null, state, isError: false); + var stateObject = new CallTargetState(scope: null, state); + + var request = requestBuilder.GetEndInvocationRequest(stateObject, 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(); @@ -92,9 +97,10 @@ public async Task TestGetEndInvocationRequestWithoutState() await using var tracer = TracerHelper.CreateWithFakeAgent(); var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" } }.Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); + var stateObject = new CallTargetState(scope, state: null); ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder(); - var request = requestBuilder.GetEndInvocationRequest(scope, state: null, isError: false); + var request = requestBuilder.GetEndInvocationRequest(stateObject, 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"); @@ -110,6 +116,7 @@ public async Task TestGetEndInvocationRequestWithErrorTags() var headers = new WebHeaderCollection().Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); var state = "example-aws-request-id"; + var stateObject = new CallTargetState(scope, state); var errorMsg = "Exception"; var errorType = "Exception"; @@ -123,7 +130,7 @@ public async Task TestGetEndInvocationRequestWithErrorTags() var expectedErrorStack = Convert.ToBase64String(System.Text.Encoding.UTF8.GetBytes(errorStack)); ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder(); - var request = requestBuilder.GetEndInvocationRequest(scope, state, true); + var request = requestBuilder.GetEndInvocationRequest(stateObject, 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); @@ -142,9 +149,10 @@ public async Task TestGetEndInvocationRequestWithoutErrorTags() var headers = new WebHeaderCollection().Wrap(); var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers); var state = "example-aws-request-id"; + var stateObject = new CallTargetState(scope, state); ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder(); - var request = requestBuilder.GetEndInvocationRequest(scope, state, true); + var request = requestBuilder.GetEndInvocationRequest(stateObject, 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();