Skip to content

Commit 3bc005f

Browse files
rithikanarayanrithikanarayan
authored
[AWS Lambda] Cleanup from PR#7835 (#7858)
## Summary of changes Cast object state to CallTargetState earlier in its usage to add type safety. Added type checks combined with variable declarations when using the state and scope attributes of CallTargetState. ## Reason for change Merged [earlier PR adding request id headers](#7835) to start and end invocation calls to the lambda extension, needed to improve code quality post-merge. ## Implementation details ## Test coverage ## Other details <!-- 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. MergeQueue is NOT enabled in this repository. If you have write access to the repo, the PR has 1-2 approvals (see above), and all of the required checks have passed, you can use the Squash and Merge button to merge the PR. If you don't have write access, or you need help, reach out in the #apm-dotnet channel in Slack. -->
1 parent a6063d8 commit 3bc005f

File tree

6 files changed

+59
-42
lines changed

6 files changed

+59
-42
lines changed

‎tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/HandlerWrapperSetHandlerIntegration.cs‎

Lines changed: 21 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ public object OnDelegateBegin<TArg1>(object sender, ref TArg1 arg)
7171
LambdaCommon.Log("DelegateWrapper Running OnDelegateBegin");
7272

7373
Scope scope;
74-
object requestid = null;
74+
object requestId = null;
7575
var proxyInstance = arg.DuckCast<IInvocationRequest>();
7676
if (proxyInstance == null)
7777
{
@@ -82,11 +82,11 @@ public object OnDelegateBegin<TArg1>(object sender, ref TArg1 arg)
8282
{
8383
var jsonString = ConvertPayloadStream(proxyInstance.InputStream);
8484
scope = LambdaCommon.SendStartInvocation(new LambdaRequestBuilder(), jsonString, proxyInstance.LambdaContext);
85-
requestid = proxyInstance.LambdaContext?.AwsRequestId;
85+
requestId = proxyInstance.LambdaContext?.AwsRequestId;
8686
}
8787

8888
LambdaCommon.Log("DelegateWrapper FINISHED Running OnDelegateBegin");
89-
return new CallTargetState(scope, requestid);
89+
return new CallTargetState(scope, requestId);
9090
}
9191

9292
public void OnException(object sender, Exception ex)
@@ -104,27 +104,31 @@ public TReturn OnDelegateEnd<TReturn>(object sender, TReturn returnValue, Except
104104
public async Task<TInnerReturn> OnDelegateEndAsync<TInnerReturn>(object sender, TInnerReturn returnValue, Exception exception, object state)
105105
{
106106
LambdaCommon.Log("DelegateWrapper Running OnDelegateEndAsync");
107-
try
107+
if (state is CallTargetState callTargetState)
108108
{
109-
var proxyInstance = returnValue.DuckCast<IInvocationResponse>();
110-
if (proxyInstance == null)
109+
try
111110
{
112-
LambdaCommon.Log("DuckCast.IInvocationResponse got null proxyInstance", debug: false);
113-
await LambdaCommon.EndInvocationAsync(string.Empty, exception, state, RequestBuilder).ConfigureAwait(false);
111+
var proxyInstance = returnValue.DuckCast<IInvocationResponse>();
112+
if (proxyInstance == null)
113+
{
114+
LambdaCommon.Log("DuckCast.IInvocationResponse got null proxyInstance", debug: false);
115+
await LambdaCommon.EndInvocationAsync(string.Empty, exception, callTargetState, RequestBuilder).ConfigureAwait(false);
116+
}
117+
else
118+
{
119+
var jsonString = ConvertPayloadStream(proxyInstance.OutputStream);
120+
await LambdaCommon.EndInvocationAsync(jsonString, exception, callTargetState, RequestBuilder).ConfigureAwait(false);
121+
}
114122
}
115-
else
123+
catch (Exception ex)
116124
{
117-
var jsonString = ConvertPayloadStream(proxyInstance.OutputStream);
118-
await LambdaCommon.EndInvocationAsync(jsonString, exception, state, RequestBuilder).ConfigureAwait(false);
125+
LambdaCommon.Log("OnDelegateEndAsync could not send payload to the extension", ex, false);
126+
await LambdaCommon.EndInvocationAsync(string.Empty, ex, callTargetState, RequestBuilder).ConfigureAwait(false);
119127
}
120-
}
121-
catch (Exception ex)
122-
{
123-
LambdaCommon.Log("OnDelegateEndAsync could not send payload to the extension", ex, false);
124-
await LambdaCommon.EndInvocationAsync(string.Empty, ex, state, RequestBuilder).ConfigureAwait(false);
128+
129+
LambdaCommon.Log("DelegateWrapper FINISHED Running OnDelegateEndAsync");
125130
}
126131

127-
LambdaCommon.Log("DelegateWrapper FINISHED Running OnDelegateEndAsync");
128132
return returnValue;
129133
}
130134
}

‎tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/ILambdaExtensionRequest.cs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
#if NET6_0_OR_GREATER
66

77
using System.Net;
8+
using Datadog.Trace.ClrProfiler.CallTarget;
89

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

@@ -20,7 +21,7 @@ internal interface ILambdaExtensionRequest
2021
/// Get the end invocation request
2122
/// </summary>
2223
/// <returns>The end invocation request</returns>
23-
WebRequest GetEndInvocationRequest(Scope scope, object state, bool isError);
24+
WebRequest GetEndInvocationRequest(CallTargetState stateObject, bool isError);
2425
}
2526

2627
#endif

‎tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaCommon.cs‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -63,9 +63,9 @@ internal static Scope SendStartInvocation(ILambdaExtensionRequest requestBuilder
6363
return CreatePlaceholderScope(tracer, headers);
6464
}
6565

66-
internal static void SendEndInvocation(ILambdaExtensionRequest requestBuilder, Scope scope, object state, bool isError, string data)
66+
internal static void SendEndInvocation(ILambdaExtensionRequest requestBuilder, CallTargetState stateObject, bool isError, string data)
6767
{
68-
var request = requestBuilder.GetEndInvocationRequest(scope, state, isError);
68+
var request = requestBuilder.GetEndInvocationRequest(stateObject, isError);
6969
WriteRequestPayload(request, data);
7070
using var response = (HttpWebResponse)request.GetResponse();
7171

@@ -75,10 +75,9 @@ internal static void SendEndInvocation(ILambdaExtensionRequest requestBuilder, S
7575
}
7676
}
7777

78-
internal static async Task EndInvocationAsync(string returnValue, Exception exception, object stateObject, ILambdaExtensionRequest requestBuilder)
78+
internal static async Task EndInvocationAsync(string returnValue, Exception exception, CallTargetState stateObject, ILambdaExtensionRequest requestBuilder)
7979
{
80-
var state = (CallTargetState)stateObject!;
81-
var scope = state.Scope;
80+
var scope = stateObject.Scope;
8281
try
8382
{
8483
await Task.WhenAll(
@@ -100,7 +99,7 @@ await Task.WhenAll(
10099
span.SetException(exception);
101100
}
102101

103-
SendEndInvocation(requestBuilder, scope, state.State, exception != null, returnValue);
102+
SendEndInvocation(requestBuilder, stateObject, exception != null, returnValue);
104103
}
105104
catch (Exception ex)
106105
{

‎tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/Lambda/LambdaRequestBuilder.cs‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
using System.Globalization;
88
using System.Net;
99
using Datadog.Trace.Agent.Transports;
10+
using Datadog.Trace.ClrProfiler.CallTarget;
1011
using Datadog.Trace.Util;
1112
#pragma warning disable CS0618 // WebRequest, HttpWebRequest, ServicePoint, and WebClient are obsolete. Use HttpClient instead.
1213

@@ -35,18 +36,18 @@ WebRequest ILambdaExtensionRequest.GetStartInvocationRequest()
3536
return request;
3637
}
3738

38-
WebRequest ILambdaExtensionRequest.GetEndInvocationRequest(Scope scope, object state, bool isError)
39+
WebRequest ILambdaExtensionRequest.GetEndInvocationRequest(CallTargetState stateObject, bool isError)
3940
{
4041
var request = WebRequest.Create(Uri + EndInvocationPath);
4142
request.Method = "POST";
4243
request.Headers.Set(HttpHeaderNames.TracingEnabled, "false");
4344

44-
if (state != null)
45+
if (stateObject.State is string state)
4546
{
46-
request.Headers.Set("lambda-runtime-aws-request-id", (string)state);
47+
request.Headers.Set("lambda-runtime-aws-request-id", state);
4748
}
4849

49-
if (scope is { Span: var span })
50+
if (stateObject.Scope is { Span: var span })
5051
{
5152
// TODO: add support for 128-bit trace ids in serverless
5253
request.Headers.Set(HttpHeaderNames.TraceId, span.TraceId128.Lower.ToString(CultureInfo.InvariantCulture));

‎tracer/test/Datadog.Trace.Tests/LambdaCommonTests.cs‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
using System.Net;
1010
using System.Threading.Tasks;
1111
using Datadog.Trace.ClrProfiler.AutoInstrumentation.AWS.Lambda;
12+
using Datadog.Trace.ClrProfiler.CallTarget;
1213
using Datadog.Trace.ExtensionMethods;
1314
using Datadog.Trace.TestHelpers;
1415
using Datadog.Trace.TestHelpers.TestTracer;
@@ -140,6 +141,7 @@ public async Task TestSendEndInvocationFailure()
140141
var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap();
141142
var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
142143
var state = "example-aws-request-id";
144+
var stateObject = new CallTargetState(scope, state);
143145

144146
var response = new Mock<HttpWebResponse>(MockBehavior.Loose);
145147
var responseStream = new Mock<Stream>(MockBehavior.Loose);
@@ -149,9 +151,9 @@ public async Task TestSendEndInvocationFailure()
149151
httpRequest.Setup(h => h.GetResponse()).Throws(new WebException());
150152
httpRequest.Setup(h => h.GetRequestStream()).Returns(responseStream.Object);
151153

152-
_lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(scope, state, true)).Returns(httpRequest.Object);
154+
_lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(stateObject, true)).Returns(httpRequest.Object);
153155

154-
Assert.Throws<WebException>(() => LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, state, true, "{}"));
156+
Assert.Throws<WebException>(() => LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, stateObject, true, "{}"));
155157
}
156158

157159
[Fact]
@@ -162,6 +164,7 @@ public async Task TestSendEndInvocationSuccess()
162164
var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap();
163165
var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
164166
var state = "example-aws-request-id";
167+
var stateObject = new CallTargetState(scope, state);
165168

166169
var response = new Mock<HttpWebResponse>(MockBehavior.Loose);
167170
var responseStream = new Mock<Stream>(MockBehavior.Loose);
@@ -171,10 +174,10 @@ public async Task TestSendEndInvocationSuccess()
171174
httpRequest.Setup(h => h.GetResponse()).Returns(response.Object);
172175
httpRequest.Setup(h => h.GetRequestStream()).Returns(responseStream.Object);
173176

174-
_lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(scope, state, true)).Returns(httpRequest.Object);
177+
_lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(stateObject, true)).Returns(httpRequest.Object);
175178
var output = new StringWriter();
176179
Console.SetOut(output);
177-
LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, state, true, "{}");
180+
LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, stateObject, true, "{}");
178181
httpRequest.Verify(r => r.GetResponse(), Times.Once);
179182
Assert.Empty(output.ToString());
180183
}
@@ -187,6 +190,7 @@ public async Task TestSendEndInvocationFalse()
187190
var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap();
188191
var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
189192
var state = "example-aws-request-id";
193+
var stateObject = new CallTargetState(scope, state);
190194

191195
var response = new Mock<HttpWebResponse>(MockBehavior.Loose);
192196
var responseStream = new Mock<Stream>(MockBehavior.Loose);
@@ -196,10 +200,10 @@ public async Task TestSendEndInvocationFalse()
196200
httpRequest.Setup(h => h.GetResponse()).Returns(response.Object);
197201
httpRequest.Setup(h => h.GetRequestStream()).Returns(responseStream.Object);
198202

199-
_lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(scope, state, true)).Returns(httpRequest.Object);
203+
_lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(stateObject, true)).Returns(httpRequest.Object);
200204
var output = new StringWriter();
201205
Console.SetOut(output);
202-
LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, state, true, "{}");
206+
LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, stateObject, true, "{}");
203207
httpRequest.Verify(r => r.GetResponse(), Times.Once);
204208
Assert.Contains("Extension does not send a status 200 OK", output.ToString());
205209
}

‎tracer/test/Datadog.Trace.Tests/LambdaRequestBuilderTests.cs‎

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
using System.Net;
88
using System.Threading.Tasks;
99
using Datadog.Trace.ClrProfiler.AutoInstrumentation.AWS.Lambda;
10+
using Datadog.Trace.ClrProfiler.CallTarget;
1011
using Datadog.Trace.ExtensionMethods;
1112
using Datadog.Trace.TestHelpers.TestTracer;
1213
using FluentAssertions;
@@ -25,10 +26,10 @@ public async Task TestGetEndInvocationRequestWithError()
2526
var headers = new WebHeaderCollection().Wrap();
2627
var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
2728
var state = "example-aws-request-id";
29+
var stateObject = new CallTargetState(scope, state);
2830

2931
ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder();
30-
var request = requestBuilder.GetEndInvocationRequest(scope, state, isError: true);
31-
request.Headers.Get("x-datadog-invocation-error").Should().Be("true");
32+
var request = requestBuilder.GetEndInvocationRequest(stateObject, true);
3233
request.Headers.Get("x-datadog-tracing-enabled").Should().Be("false");
3334
request.Headers.Get("x-datadog-sampling-priority").Should().Be("1");
3435
request.Headers.Get("x-datadog-trace-id").Should().NotBeNull();
@@ -43,9 +44,10 @@ public async Task TestGetEndInvocationRequestWithoutError()
4344
var headers = new WebHeaderCollection().Wrap();
4445
var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
4546
var state = "example-aws-request-id";
47+
var stateObject = new CallTargetState(scope, state);
4648

4749
ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder();
48-
var request = requestBuilder.GetEndInvocationRequest(scope, state, isError: false);
50+
var request = requestBuilder.GetEndInvocationRequest(stateObject, isError: false);
4951
request.Headers.Get("x-datadog-invocation-error").Should().BeNull();
5052
request.Headers.Get("x-datadog-tracing-enabled").Should().Be("false");
5153
request.Headers.Get("x-datadog-sampling-priority").Should().Be("1");
@@ -61,9 +63,10 @@ public async Task TestGetEndInvocationRequestWithScope()
6163
var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" } }.Wrap();
6264
var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
6365
var state = "example-aws-request-id";
66+
var stateObject = new CallTargetState(scope, state);
6467

6568
ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder();
66-
var request = requestBuilder.GetEndInvocationRequest(scope, state, isError: false);
69+
var request = requestBuilder.GetEndInvocationRequest(stateObject, isError: false);
6770
request.Headers.Get("x-datadog-invocation-error").Should().BeNull();
6871
request.Headers.Get("x-datadog-tracing-enabled").Should().Be("false");
6972
request.Headers.Get("x-datadog-sampling-priority").Should().Be("1");
@@ -77,7 +80,9 @@ public void TestGetEndInvocationRequestWithoutScope()
7780
{
7881
ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder();
7982
var state = "example-aws-request-id";
80-
var request = requestBuilder.GetEndInvocationRequest(scope: null, state, isError: false);
83+
var stateObject = new CallTargetState(scope: null, state);
84+
85+
var request = requestBuilder.GetEndInvocationRequest(stateObject, isError: false);
8186
request.Headers.Get("x-datadog-invocation-error").Should().BeNull();
8287
request.Headers.Get("x-datadog-tracing-enabled").Should().Be("false");
8388
request.Headers.Get("x-datadog-sampling-priority").Should().BeNull();
@@ -92,9 +97,10 @@ public async Task TestGetEndInvocationRequestWithoutState()
9297
await using var tracer = TracerHelper.CreateWithFakeAgent();
9398
var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" } }.Wrap();
9499
var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
100+
var stateObject = new CallTargetState(scope, state: null);
95101

96102
ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder();
97-
var request = requestBuilder.GetEndInvocationRequest(scope, state: null, isError: false);
103+
var request = requestBuilder.GetEndInvocationRequest(stateObject, isError: false);
98104
request.Headers.Get("x-datadog-invocation-error").Should().BeNull();
99105
request.Headers.Get("x-datadog-tracing-enabled").Should().Be("false");
100106
request.Headers.Get("x-datadog-sampling-priority").Should().Be("1");
@@ -110,6 +116,7 @@ public async Task TestGetEndInvocationRequestWithErrorTags()
110116
var headers = new WebHeaderCollection().Wrap();
111117
var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
112118
var state = "example-aws-request-id";
119+
var stateObject = new CallTargetState(scope, state);
113120

114121
var errorMsg = "Exception";
115122
var errorType = "Exception";
@@ -123,7 +130,7 @@ public async Task TestGetEndInvocationRequestWithErrorTags()
123130
var expectedErrorStack = Convert.ToBase64String(System.Text.Encoding.UTF8.GetBytes(errorStack));
124131

125132
ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder();
126-
var request = requestBuilder.GetEndInvocationRequest(scope, state, true);
133+
var request = requestBuilder.GetEndInvocationRequest(stateObject, true);
127134
request.Headers.Get("x-datadog-invocation-error").Should().NotBeNull();
128135
request.Headers.Get("x-datadog-invocation-error-msg").Should().Be(expectedErrorMsg);
129136
request.Headers.Get("x-datadog-invocation-error-type").Should().Be(expectedErrorType);
@@ -142,9 +149,10 @@ public async Task TestGetEndInvocationRequestWithoutErrorTags()
142149
var headers = new WebHeaderCollection().Wrap();
143150
var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
144151
var state = "example-aws-request-id";
152+
var stateObject = new CallTargetState(scope, state);
145153

146154
ILambdaExtensionRequest requestBuilder = new LambdaRequestBuilder();
147-
var request = requestBuilder.GetEndInvocationRequest(scope, state, true);
155+
var request = requestBuilder.GetEndInvocationRequest(stateObject, true);
148156
request.Headers.Get("x-datadog-invocation-error").Should().NotBeNull();
149157
request.Headers.Get("x-datadog-invocation-error-msg").Should().BeNull();
150158
request.Headers.Get("x-datadog-invocation-error-type").Should().BeNull();

0 commit comments

Comments
 (0)