Skip to content

Commit 7dee697

Browse files
ojproductionsojproductions
authored
add request-id to lambda start/end invocation calls (#9991)
* add request-id to lambda start/end invocation calls * tests * commenting and naming * remove unused coreTracer param
1 parent 1c2d839 commit 7dee697

File tree

5 files changed

+41
-27
lines changed

5 files changed

+41
-27
lines changed

‎dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -88,15 +88,15 @@ static AgentScope enter(
8888
if (CallDepthThreadLocalMap.incrementCallDepth(RequestHandler.class) > 0) {
8989
return null;
9090
}
91-
92-
AgentSpanContext lambdaContext = AgentTracer.get().notifyExtensionStart(in);
91+
String lambdaRequestId = awsContext.getAwsRequestId();
92+
AgentSpanContext lambdaContext = AgentTracer.get().notifyExtensionStart(in, lambdaRequestId);
9393
final AgentSpan span;
9494
if (null == lambdaContext) {
9595
span = startSpan(INVOCATION_SPAN_NAME);
9696
} else {
9797
span = startSpan(INVOCATION_SPAN_NAME, lambdaContext);
9898
}
99-
span.setTag("request_id", awsContext.getAwsRequestId());
99+
span.setTag("request_id", lambdaRequestId);
100100

101101
final AgentScope scope = activateSpan(span);
102102
return scope;
@@ -107,6 +107,7 @@ static void exit(
107107
@Origin String method,
108108
@Enter final AgentScope scope,
109109
@Advice.Argument(1) final Object result,
110+
@Advice.Argument(2) final Context awsContext,
110111
@Advice.Thrown final Throwable throwable) {
111112

112113
if (scope == null) {
@@ -120,8 +121,10 @@ static void exit(
120121
if (throwable != null) {
121122
span.addThrowable(throwable);
122123
}
124+
String lambdaRequestId = awsContext.getAwsRequestId();
125+
123126
span.finish();
124-
AgentTracer.get().notifyExtensionEnd(span, result, null != throwable);
127+
AgentTracer.get().notifyExtensionEnd(span, result, null != throwable, lambdaRequestId);
125128
} finally {
126129
scope.close();
127130
}

‎dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1199,13 +1199,14 @@ public void closeActive() {
11991199
}
12001200

12011201
@Override
1202-
public AgentSpanContext notifyExtensionStart(Object event) {
1203-
return LambdaHandler.notifyStartInvocation(this, event);
1202+
public AgentSpanContext notifyExtensionStart(Object event, String lambdaRequestId) {
1203+
return LambdaHandler.notifyStartInvocation(event, lambdaRequestId);
12041204
}
12051205

12061206
@Override
1207-
public void notifyExtensionEnd(AgentSpan span, Object result, boolean isError) {
1208-
LambdaHandler.notifyEndInvocation(span, result, isError);
1207+
public void notifyExtensionEnd(
1208+
AgentSpan span, Object result, boolean isError, String lambdaRequestId) {
1209+
LambdaHandler.notifyEndInvocation(span, result, isError, lambdaRequestId);
12091210
}
12101211

12111212
@Override

‎dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
import datadog.trace.api.DDTags;
1010
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
1111
import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
12-
import datadog.trace.core.CoreTracer;
1312
import java.io.ByteArrayInputStream;
1413
import java.io.ByteArrayOutputStream;
1514
import java.nio.charset.StandardCharsets;
@@ -43,6 +42,7 @@ public class LambdaHandler {
4342
private static final String DATADOG_INVOCATION_ERROR_MSG = "x-datadog-invocation-error-msg";
4443
private static final String DATADOG_INVOCATION_ERROR_TYPE = "x-datadog-invocation-error-type";
4544
private static final String DATADOG_INVOCATION_ERROR_STACK = "x-datadog-invocation-error-stack";
45+
private static final String LAMBDA_RUNTIME_AWS_REQUEST_ID = "lambda-runtime-aws-request-id";
4646

4747
private static final String START_INVOCATION = "/lambda/start-invocation";
4848
private static final String END_INVOCATION = "/lambda/end-invocation";
@@ -72,14 +72,15 @@ public class LambdaHandler {
7272

7373
private static String EXTENSION_BASE_URL = "http://127.0.0.1:8124";
7474

75-
public static AgentSpanContext notifyStartInvocation(CoreTracer tracer, Object event) {
75+
public static AgentSpanContext notifyStartInvocation(Object event, String lambdaRequestId) {
7676
RequestBody body = RequestBody.create(jsonMediaType, writeValueAsString(event));
7777
try (Response response =
7878
HTTP_CLIENT
7979
.newCall(
8080
new Request.Builder()
8181
.url(EXTENSION_BASE_URL + START_INVOCATION)
8282
.addHeader(DATADOG_META_LANG, "java")
83+
.addHeader(LAMBDA_RUNTIME_AWS_REQUEST_ID, lambdaRequestId)
8384
.post(body)
8485
.build())
8586
.execute()) {
@@ -99,7 +100,8 @@ public static AgentSpanContext notifyStartInvocation(CoreTracer tracer, Object e
99100
return null;
100101
}
101102

102-
public static boolean notifyEndInvocation(AgentSpan span, Object result, boolean isError) {
103+
public static boolean notifyEndInvocation(
104+
AgentSpan span, Object result, boolean isError, String lambdaRequestId) {
103105
if (null == span || null == span.getSamplingPriority()) {
104106
log.error(
105107
"could not notify the extension as the lambda span is null or no sampling priority has been found");
@@ -113,6 +115,7 @@ public static boolean notifyEndInvocation(AgentSpan span, Object result, boolean
113115
.addHeader(DATADOG_SPAN_ID, DDSpanId.toString(span.getSpanId()))
114116
.addHeader(DATADOG_SAMPLING_PRIORITY, span.getSamplingPriority().toString())
115117
.addHeader(DATADOG_META_LANG, "java")
118+
.addHeader(LAMBDA_RUNTIME_AWS_REQUEST_ID, lambdaRequestId)
116119
.post(body);
117120

118121
Object errorMessage = span.getTag(DDTags.ERROR_MSG);

‎dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy‎

Lines changed: 18 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -50,11 +50,12 @@ class LambdaHandlerTest extends DDCoreSpecification {
5050
LambdaHandler.setExtensionBaseUrl(server.address.toString())
5151

5252
when:
53-
def objTest = LambdaHandler.notifyStartInvocation(ct, obj)
53+
def objTest = LambdaHandler.notifyStartInvocation(obj, "lambda-request-123")
5454

5555
then:
5656
objTest.getTraceId().toString() == traceId
5757
objTest.getSamplingPriority() == samplingPriority
58+
server.lastRequest.headers.get("lambda-runtime-aws-request-id") == "lambda-request-123"
5859

5960
cleanup:
6061
server.close()
@@ -84,11 +85,12 @@ class LambdaHandlerTest extends DDCoreSpecification {
8485
LambdaHandler.setExtensionBaseUrl(server.address.toString())
8586

8687
when:
87-
def objTest = LambdaHandler.notifyStartInvocation(ct, obj)
88+
def objTest = LambdaHandler.notifyStartInvocation(obj, "lambda-request-123")
8889

8990
then:
9091
objTest.getTraceId().toHexString() == traceId
9192
objTest.getSamplingPriority() == samplingPriority
93+
server.lastRequest.headers.get("lambda-runtime-aws-request-id") == "lambda-request-123"
9294

9395
cleanup:
9496
server.close()
@@ -115,10 +117,11 @@ class LambdaHandlerTest extends DDCoreSpecification {
115117
LambdaHandler.setExtensionBaseUrl(server.address.toString())
116118

117119
when:
118-
def objTest = LambdaHandler.notifyStartInvocation(ct, obj)
120+
def objTest = LambdaHandler.notifyStartInvocation(obj, "my-lambda-request")
119121

120122
then:
121123
objTest == expected
124+
server.lastRequest.headers.get("lambda-runtime-aws-request-id") == "my-lambda-request"
122125

123126
cleanup:
124127
server.close()
@@ -148,22 +151,23 @@ class LambdaHandlerTest extends DDCoreSpecification {
148151
}
149152

150153
when:
151-
def result = LambdaHandler.notifyEndInvocation(span, lambdaResult, boolValue)
154+
def result = LambdaHandler.notifyEndInvocation(span, lambdaResult, boolValue, lambdaReqIdHeaderValue)
152155

153156
then:
154157
server.lastRequest.headers.get("x-datadog-invocation-error") == eHeaderValue
155158
server.lastRequest.headers.get("x-datadog-trace-id") == tIdHeaderValue
156159
server.lastRequest.headers.get("x-datadog-span-id") == sIdHeaderValue
157160
server.lastRequest.headers.get("x-datadog-sampling-priority") == sPIdHeaderValue
161+
server.lastRequest.headers.get("lambda-runtime-aws-request-id") == lambdaReqIdHeaderValue
158162
result == expected
159163

160164
cleanup:
161165
server.close()
162166

163167
where:
164-
expected | eHeaderValue | tIdHeaderValue | sIdHeaderValue | sPIdHeaderValue | lambdaResult | boolValue
165-
true | "true" | "1234" | "5678" | "2" | {} | true
166-
true | null | "1234" | "5678" | "2" | "12345" | false
168+
expected | eHeaderValue | tIdHeaderValue | sIdHeaderValue | sPIdHeaderValue | lambdaResult | boolValue | lambdaReqIdHeaderValue
169+
true | "true" | "1234" | "5678" | "2" | {} | true | "request123"
170+
true | null | "1234" | "5678" | "2" | "12345" | false | "request456"
167171
}
168172

169173
def "test end invocation failure"() {
@@ -185,19 +189,20 @@ class LambdaHandlerTest extends DDCoreSpecification {
185189
}
186190

187191
when:
188-
def result = LambdaHandler.notifyEndInvocation(span, lambdaResult, boolValue)
192+
def result = LambdaHandler.notifyEndInvocation(span, lambdaResult, boolValue, lambdaReqIdHeaderValue)
189193

190194
then:
191195
result == expected
192196
server.lastRequest.headers.get("x-datadog-invocation-error") == headerValue
197+
server.lastRequest.headers.get("lambda-runtime-aws-request-id") == lambdaReqIdHeaderValue
193198

194199
cleanup:
195200
server.close()
196201

197202
where:
198-
expected | headerValue | lambdaResult | boolValue
199-
false | "true" | {} | true
200-
false | null | "12345" | false
203+
expected | headerValue | lambdaResult | boolValue | lambdaReqIdHeaderValue
204+
false | "true" | {} | true | "request123"
205+
false | null | "12345" | false | "request456"
201206
}
202207

203208
def "test end invocation success with error metadata"() {
@@ -222,13 +227,14 @@ class LambdaHandlerTest extends DDCoreSpecification {
222227
}
223228

224229
when:
225-
LambdaHandler.notifyEndInvocation(span, {}, true)
230+
LambdaHandler.notifyEndInvocation(span, {}, true, "lambda-request-123")
226231

227232
then:
228233
server.lastRequest.headers.get("x-datadog-invocation-error") == "true"
229234
server.lastRequest.headers.get("x-datadog-invocation-error-msg") == "custom error message"
230235
server.lastRequest.headers.get("x-datadog-invocation-error-type") == "java.lang.Throwable"
231236
server.lastRequest.headers.get("x-datadog-invocation-error-stack") == "ZXJyb3JTdGFjawogCXRlc3Q="
237+
server.lastRequest.headers.get("lambda-runtime-aws-request-id") == "lambda-request-123"
232238

233239
cleanup:
234240
server.close()

‎internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -406,9 +406,9 @@ default SpanBuilder singleSpanBuilder(CharSequence spanName) {
406406

407407
CallbackProvider getUniversalCallbackProvider();
408408

409-
AgentSpanContext notifyExtensionStart(Object event);
409+
AgentSpanContext notifyExtensionStart(Object event, String lambdaRequestId);
410410

411-
void notifyExtensionEnd(AgentSpan span, Object result, boolean isError);
411+
void notifyExtensionEnd(AgentSpan span, Object result, boolean isError, String lambdaRequestId);
412412

413413
AgentDataStreamsMonitoring getDataStreamsMonitoring();
414414

@@ -653,12 +653,13 @@ public EndpointTracker onRootSpanStarted(AgentSpan root) {
653653
}
654654

655655
@Override
656-
public AgentSpanContext notifyExtensionStart(Object event) {
656+
public AgentSpanContext notifyExtensionStart(Object event, String lambdaRequestId) {
657657
return null;
658658
}
659659

660660
@Override
661-
public void notifyExtensionEnd(AgentSpan span, Object result, boolean isError) {}
661+
public void notifyExtensionEnd(
662+
AgentSpan span, Object result, boolean isError, String lambdaRequestId) {}
662663

663664
@Override
664665
public AgentDataStreamsMonitoring getDataStreamsMonitoring() {

0 commit comments

Comments
 (0)