Skip to content

Commit 5738b43

Browse files
rithikanarayanlucaspimentel
rithikanarayan
and
lucaspimentel
authored
Add request-id as header to lambda start/end invocation (#7835)
## Summary of changes Adds` lambda-runtime-aws-request-id` header in start invocation and end invocation requests from tracer to extension. Assigns header to the aws request id of the span. ## Reason for change Allows the lambda extension to identify which start/end invocation requests belong to which invocation if multiple Lambda requests are being served concurrently. ## Implementation details Uses the LambdaContext from the `IInvocationRequest proxyinstance` to get the aws-request-id in `SendStartInvocation` and adds it to request headers in `WriteRequestHeaders`. Saves the aws-request-id as part of the state passed from `OnDelegateBegin `to `EndInvocationAsync`. `EndInvocationAsync` passes the state down to `GetEndInvocationRequest`, which adds it to the request headers. Updates the LambdaCommonTests and LambdaRequestBuilderTests to match new function signatures with state passed down and to test for the assignment to the lambda-runtime-aws-request-id. ## Test coverage Added unit test in LambdaRequestBuilderTests to check for ` lambda-runtime-aws-request-id` header and value in end invocation request. Manually tested for headers in start and end invocation using custom dd-trace-dotnet layer `arn:aws:lambda:sa-east-1:425362996713:layer:dotnet-request-id-header-rithika:2` and custom lambda extension `arn:aws:lambda:sa-east-1:425362996713:layer:Debug-extension-rithika:1` which prints the start and end invocation headers. Sample trace with logs of the headers [here](https://support-admin.us1.prod.dog/admin/switch_handle_get/org_id/1543931?next_url=%2Fapm%2Ftrace%2F219302768028734365%3FgraphType%3Dservice_map%26shouldShowLegend%3Dtrue%26spanID%3D4621724130062486128%26timeHint%3D1763567215307.0002%26trace%3D2193027680287343654621724130062486128%26traceQuery%3D) ## 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. --> --------- Co-authored-by: Lucas Pimentel <lucas.pimentel@datadoghq.com>
1 parent c068213 commit 5738b43

File tree

7 files changed

+83
-34
lines changed

7 files changed

+83
-34
lines changed

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

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +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;
7475
var proxyInstance = arg.DuckCast<IInvocationRequest>();
7576
if (proxyInstance == null)
7677
{
@@ -80,11 +81,12 @@ public object OnDelegateBegin<TArg1>(object sender, ref TArg1 arg)
8081
else
8182
{
8283
var jsonString = ConvertPayloadStream(proxyInstance.InputStream);
83-
scope = LambdaCommon.SendStartInvocation(new LambdaRequestBuilder(), jsonString, proxyInstance.LambdaContext?.ClientContext?.Custom);
84+
scope = LambdaCommon.SendStartInvocation(new LambdaRequestBuilder(), jsonString, proxyInstance.LambdaContext);
85+
requestid = proxyInstance.LambdaContext?.AwsRequestId;
8486
}
8587

8688
LambdaCommon.Log("DelegateWrapper FINISHED Running OnDelegateBegin");
87-
return new CallTargetState(scope);
89+
return new CallTargetState(scope, requestid);
8890
}
8991

9092
public void OnException(object sender, Exception ex)
@@ -108,18 +110,18 @@ public async Task<TInnerReturn> OnDelegateEndAsync<TInnerReturn>(object sender,
108110
if (proxyInstance == null)
109111
{
110112
LambdaCommon.Log("DuckCast.IInvocationResponse got null proxyInstance", debug: false);
111-
await LambdaCommon.EndInvocationAsync(string.Empty, exception, ((CallTargetState)state!).Scope, RequestBuilder).ConfigureAwait(false);
113+
await LambdaCommon.EndInvocationAsync(string.Empty, exception, state, RequestBuilder).ConfigureAwait(false);
112114
}
113115
else
114116
{
115117
var jsonString = ConvertPayloadStream(proxyInstance.OutputStream);
116-
await LambdaCommon.EndInvocationAsync(jsonString, exception, ((CallTargetState)state!).Scope, RequestBuilder).ConfigureAwait(false);
118+
await LambdaCommon.EndInvocationAsync(jsonString, exception, state, RequestBuilder).ConfigureAwait(false);
117119
}
118120
}
119121
catch (Exception ex)
120122
{
121123
LambdaCommon.Log("OnDelegateEndAsync could not send payload to the extension", ex, false);
122-
await LambdaCommon.EndInvocationAsync(string.Empty, ex, ((CallTargetState)state!).Scope, RequestBuilder).ConfigureAwait(false);
124+
await LambdaCommon.EndInvocationAsync(string.Empty, ex, state, RequestBuilder).ConfigureAwait(false);
123125
}
124126

125127
LambdaCommon.Log("DelegateWrapper FINISHED Running OnDelegateEndAsync");

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ internal interface ILambdaExtensionRequest
2020
/// Get the end invocation request
2121
/// </summary>
2222
/// <returns>The end invocation request</returns>
23-
WebRequest GetEndInvocationRequest(Scope scope, bool isError);
23+
WebRequest GetEndInvocationRequest(Scope scope, object state, bool isError);
2424
}
2525

2626
#endif

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

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
using System.Net;
99
using System.Text;
1010
using System.Threading.Tasks;
11+
using Datadog.Trace.ClrProfiler.AutoInstrumentation.AWS.SDK;
12+
using Datadog.Trace.ClrProfiler.CallTarget;
1113
using Datadog.Trace.Headers;
1214
using Datadog.Trace.Propagators;
1315
using Datadog.Trace.Telemetry;
@@ -22,6 +24,7 @@ internal abstract class LambdaCommon
2224
private const string PlaceholderOperationName = "placeholder-operation";
2325
private const double ServerlessMaxWaitingFlushTime = 3;
2426
private const string LogLevelEnvName = "DD_LOG_LEVEL";
27+
private const string LambdaRuntimeAwsRequestIdHeader = "lambda-runtime-aws-request-id";
2528

2629
internal static Scope CreatePlaceholderScope(Tracer tracer, NameValueHeadersCollection headers)
2730
{
@@ -38,11 +41,16 @@ internal static Scope CreatePlaceholderScope(Tracer tracer, NameValueHeadersColl
3841
return tracer.TracerManager.ScopeManager.Activate(span, false);
3942
}
4043

41-
internal static Scope SendStartInvocation(ILambdaExtensionRequest requestBuilder, string data, IDictionary<string, string> context)
44+
internal static Scope SendStartInvocation(ILambdaExtensionRequest requestBuilder, string data, ILambdaContext context)
4245
{
4346
var request = requestBuilder.GetStartInvocationRequest();
4447
WriteRequestPayload(request, data);
45-
WriteRequestHeaders(request, context);
48+
WriteRequestHeaders(request, context?.ClientContext?.Custom);
49+
if (context?.AwsRequestId != null)
50+
{
51+
request.Headers.Add(LambdaRuntimeAwsRequestIdHeader, context.AwsRequestId);
52+
}
53+
4654
using var response = (HttpWebResponse)request.GetResponse();
4755

4856
var headers = response.Headers.Wrap();
@@ -55,9 +63,9 @@ internal static Scope SendStartInvocation(ILambdaExtensionRequest requestBuilder
5563
return CreatePlaceholderScope(tracer, headers);
5664
}
5765

58-
internal static void SendEndInvocation(ILambdaExtensionRequest requestBuilder, Scope scope, bool isError, string data)
66+
internal static void SendEndInvocation(ILambdaExtensionRequest requestBuilder, Scope scope, object state, bool isError, string data)
5967
{
60-
var request = requestBuilder.GetEndInvocationRequest(scope, isError);
68+
var request = requestBuilder.GetEndInvocationRequest(scope, state, isError);
6169
WriteRequestPayload(request, data);
6270
using var response = (HttpWebResponse)request.GetResponse();
6371

@@ -67,8 +75,10 @@ internal static void SendEndInvocation(ILambdaExtensionRequest requestBuilder, S
6775
}
6876
}
6977

70-
internal static async Task EndInvocationAsync(string returnValue, Exception exception, Scope scope, ILambdaExtensionRequest requestBuilder)
78+
internal static async Task EndInvocationAsync(string returnValue, Exception exception, object stateObject, ILambdaExtensionRequest requestBuilder)
7179
{
80+
var state = (CallTargetState)stateObject!;
81+
var scope = state.Scope;
7282
try
7383
{
7484
await Task.WhenAll(
@@ -90,7 +100,7 @@ await Task.WhenAll(
90100
span.SetException(exception);
91101
}
92102

93-
SendEndInvocation(requestBuilder, scope, exception != null, returnValue);
103+
SendEndInvocation(requestBuilder, scope, state.State, exception != null, returnValue);
94104
}
95105
catch (Exception ex)
96106
{

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,12 +35,17 @@ WebRequest ILambdaExtensionRequest.GetStartInvocationRequest()
3535
return request;
3636
}
3737

38-
WebRequest ILambdaExtensionRequest.GetEndInvocationRequest(Scope scope, bool isError)
38+
WebRequest ILambdaExtensionRequest.GetEndInvocationRequest(Scope scope, object state, bool isError)
3939
{
4040
var request = WebRequest.Create(Uri + EndInvocationPath);
4141
request.Method = "POST";
4242
request.Headers.Set(HttpHeaderNames.TracingEnabled, "false");
4343

44+
if (state != null)
45+
{
46+
request.Headers.Set("lambda-runtime-aws-request-id", (string)state);
47+
}
48+
4449
if (scope is { Span: var span })
4550
{
4651
// TODO: add support for 128-bit trace ids in serverless

‎tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AWS/SDK/ILambdaContext.cs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,12 @@ namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.AWS.SDK;
1212
/// </summary>
1313
internal interface ILambdaContext
1414
{
15-
// /// <summary>
16-
// /// Gets the AWS request ID associated with the request.
17-
// /// This is the same ID returned to the client that called invoke().
18-
// /// This ID is reused for retries on the same request.
19-
// /// </summary>
20-
// string AwsRequestId { get; }
15+
/// <summary>
16+
/// Gets the AWS request ID associated with the request.
17+
/// This is the same ID returned to the client that called invoke().
18+
/// This ID is reused for retries on the same request.
19+
/// </summary>
20+
string AwsRequestId { get; }
2121

2222
/// <summary>
2323
/// Gets information about the client application and device when invoked

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

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ public void TestSendStartInvocationThrow()
9191

9292
_lambdaRequestMock.Setup(lr => lr.GetStartInvocationRequest()).Returns(httpRequest.Object);
9393

94-
Assert.Throws<WebException>(() => LambdaCommon.SendStartInvocation(_lambdaRequestMock.Object, "{}", new Dictionary<string, string>()));
94+
Assert.Throws<WebException>(() => LambdaCommon.SendStartInvocation(_lambdaRequestMock.Object, "{}", null));
9595
}
9696

9797
[Fact]
@@ -110,7 +110,7 @@ public void TestSendStartInvocationNull()
110110

111111
_lambdaRequestMock.Setup(lr => lr.GetStartInvocationRequest()).Returns(httpRequest.Object);
112112

113-
LambdaCommon.SendStartInvocation(_lambdaRequestMock.Object, "{}", new Dictionary<string, string>()).Should().BeNull();
113+
LambdaCommon.SendStartInvocation(_lambdaRequestMock.Object, "{}", null).Should().BeNull();
114114
}
115115

116116
[Fact]
@@ -129,7 +129,7 @@ public void TestSendStartInvocationSuccess()
129129

130130
_lambdaRequestMock.Setup(lr => lr.GetStartInvocationRequest()).Returns(httpRequest.Object);
131131

132-
LambdaCommon.SendStartInvocation(_lambdaRequestMock.Object, "{}", new Dictionary<string, string>()).Should().NotBeNull();
132+
LambdaCommon.SendStartInvocation(_lambdaRequestMock.Object, "{}", null).Should().NotBeNull();
133133
}
134134

135135
[Fact]
@@ -139,6 +139,7 @@ public async Task TestSendEndInvocationFailure()
139139
await using var tracer = TracerHelper.CreateWithFakeAgent();
140140
var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap();
141141
var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
142+
var state = "example-aws-request-id";
142143

143144
var response = new Mock<HttpWebResponse>(MockBehavior.Loose);
144145
var responseStream = new Mock<Stream>(MockBehavior.Loose);
@@ -148,9 +149,9 @@ public async Task TestSendEndInvocationFailure()
148149
httpRequest.Setup(h => h.GetResponse()).Throws(new WebException());
149150
httpRequest.Setup(h => h.GetRequestStream()).Returns(responseStream.Object);
150151

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

153-
Assert.Throws<WebException>(() => LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, true, "{}"));
154+
Assert.Throws<WebException>(() => LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, state, true, "{}"));
154155
}
155156

156157
[Fact]
@@ -160,6 +161,7 @@ public async Task TestSendEndInvocationSuccess()
160161
await using var tracer = TracerHelper.CreateWithFakeAgent();
161162
var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap();
162163
var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
164+
var state = "example-aws-request-id";
163165

164166
var response = new Mock<HttpWebResponse>(MockBehavior.Loose);
165167
var responseStream = new Mock<Stream>(MockBehavior.Loose);
@@ -169,10 +171,10 @@ public async Task TestSendEndInvocationSuccess()
169171
httpRequest.Setup(h => h.GetResponse()).Returns(response.Object);
170172
httpRequest.Setup(h => h.GetRequestStream()).Returns(responseStream.Object);
171173

172-
_lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(scope, true)).Returns(httpRequest.Object);
174+
_lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(scope, state, true)).Returns(httpRequest.Object);
173175
var output = new StringWriter();
174176
Console.SetOut(output);
175-
LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, true, "{}");
177+
LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, state, true, "{}");
176178
httpRequest.Verify(r => r.GetResponse(), Times.Once);
177179
Assert.Empty(output.ToString());
178180
}
@@ -184,6 +186,7 @@ public async Task TestSendEndInvocationFalse()
184186
await using var tracer = TracerHelper.CreateWithFakeAgent();
185187
var headers = new WebHeaderCollection { { HttpHeaderNames.TraceId, "1234" }, { HttpHeaderNames.SamplingPriority, "-1" } }.Wrap();
186188
var scope = LambdaCommon.CreatePlaceholderScope(tracer, headers);
189+
var state = "example-aws-request-id";
187190

188191
var response = new Mock<HttpWebResponse>(MockBehavior.Loose);
189192
var responseStream = new Mock<Stream>(MockBehavior.Loose);
@@ -193,10 +196,10 @@ public async Task TestSendEndInvocationFalse()
193196
httpRequest.Setup(h => h.GetResponse()).Returns(response.Object);
194197
httpRequest.Setup(h => h.GetRequestStream()).Returns(responseStream.Object);
195198

196-
_lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(scope, true)).Returns(httpRequest.Object);
199+
_lambdaRequestMock.Setup(lr => lr.GetEndInvocationRequest(scope, state, true)).Returns(httpRequest.Object);
197200
var output = new StringWriter();
198201
Console.SetOut(output);
199-
LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, true, "{}");
202+
LambdaCommon.SendEndInvocation(_lambdaRequestMock.Object, scope, state, true, "{}");
200203
httpRequest.Verify(r => r.GetResponse(), Times.Once);
201204
Assert.Contains("Extension does not send a status 200 OK", output.ToString());
202205
}

0 commit comments

Comments
 (0)