Skip to content

Cleanup from PR#7835 #7858

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 3 commits into
base: master
Choose a base branch
from
+59 −42
Draft

Cleanup from PR#7835 #7858

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
02b9ef2
Style improvements, casting state to CallTargetState sooner
rithikanarayan Nov 24, 2025
6528bbe
Updated unit tests
rithikanarayan Nov 24, 2025
2a752c2
nit: naming convention:
rithikanarayan Nov 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 21 additions & 17 deletions ...g.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/HandlerWrapperSetHandlerIntegration.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ public object OnDelegateBegin<TArg1>(object sender, ref TArg1 arg)
LambdaCommon.Log("DelegateWrapper Running OnDelegateBegin");

Scope scope;
object requestid = null;
object requestId = null;
var proxyInstance = arg.DuckCast<IInvocationRequest>();
if (proxyInstance == null)
{
Expand All @@ -82,11 +82,11 @@ public object OnDelegateBegin<TArg1>(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)
Expand All @@ -104,27 +104,31 @@ public TReturn OnDelegateEnd<TReturn>(object sender, TReturn returnValue, Except
public async Task<TInnerReturn> OnDelegateEndAsync<TInnerReturn>(object sender, TInnerReturn returnValue, Exception exception, object state)
{
LambdaCommon.Log("DelegateWrapper Running OnDelegateEndAsync");
try
if (state is CallTargetState callTargetState)
Comment on lines 104 to +107
Copy link
Member

@andrewlock andrewlock Nov 24, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should never pass CallTargetState to an object parameter, because it will cause boxing and extra allocation 🙂 Unfortunately I'm not very familiar with this integration, so I can't suggest a great alternative, but if you need to pass two different types of state object, it's probably best to just create two different methods with the two possible parameter types

Copy link
Contributor Author

@rithikanarayan rithikanarayan Nov 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think it's possible to pass in CallTargetState directly (or avoid passing in an object) because this function comes from DelegateInstrumentation.cs, which is not a lambda-specific file and specifies that an object is passed in to this function. The state is already boxed when it comes to us.

{
var proxyInstance = returnValue.DuckCast<IInvocationResponse>();
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<IInvocationResponse>();
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;
}
}
Expand Down
3 changes: 2 additions & 1 deletion ...r/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/ILambdaExtensionRequest.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
#if NET6_0_OR_GREATER

using System.Net;
using Datadog.Trace.ClrProfiler.CallTarget;

namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.AWS.Lambda;

Expand All @@ -20,7 +21,7 @@ internal interface ILambdaExtensionRequest
/// Get the end invocation request
/// </summary>
/// <returns>The end invocation request</returns>
WebRequest GetEndInvocationRequest(Scope scope, object state, bool isError);
WebRequest GetEndInvocationRequest(CallTargetState stateObject, bool isError);
}

#endif
11 changes: 5 additions & 6 deletions tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaCommon.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();

Expand All @@ -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(
Expand All @@ -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)
{
Expand Down
9 changes: 5 additions & 4 deletions tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaRequestBuilder.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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));
Expand Down
16 changes: 10 additions & 6 deletions tracer/test/Datadog.Trace.Tests/LambdaCommonTests.cs
View file
Open in desktop
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.ClrProfiler.CallTarget;
using Datadog.Trace.ExtensionMethods;
using Datadog.Trace.TestHelpers;
using Datadog.Trace.TestHelpers.TestTracer;
Expand Down Expand Up @@ -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<HttpWebResponse>(MockBehavior.Loose);
var responseStream = new Mock<Stream>(MockBehavior.Loose);
Expand All @@ -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<WebException>(() => LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, state, true, "{}"));
Assert.Throws<WebException>(() => LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, stateObject, true, "{}"));
}

[Fact]
Expand All @@ -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<HttpWebResponse>(MockBehavior.Loose);
var responseStream = new Mock<Stream>(MockBehavior.Loose);
Expand All @@ -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());
}
Expand All @@ -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<HttpWebResponse>(MockBehavior.Loose);
var responseStream = new Mock<Stream>(MockBehavior.Loose);
Expand All @@ -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());
}
Expand Down
24 changes: 16 additions & 8 deletions tracer/test/Datadog.Trace.Tests/LambdaRequestBuilderTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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();
Expand All @@ -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");
Expand All @@ -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");
Expand All @@ -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();
Expand All @@ -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");
Expand All @@ -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";
Expand All @@ -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);
Expand All @@ -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();
Expand Down
Loading