Skip to content

Commit 038534c

Browse files
RanVakninRanVaknin
authored
Add support for lambda traces via threadLocal (#6295)
* Add support for lambda traces for concurrent environments * Add async support * fix access modifier * Add afterExecution/onExecutionFailure hooks to restore MDC trace ID * Fix minor issues * Add test for pre-execution failures
1 parent fcf7d4d commit 038534c

File tree

3 files changed

+297
-5
lines changed

3 files changed

+297
-5
lines changed

‎core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java

Lines changed: 37 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,11 @@
1616
package software.amazon.awssdk.awscore.interceptor;
1717

1818
import java.util.Optional;
19+
import org.slf4j.MDC;
1920
import software.amazon.awssdk.annotations.SdkProtectedApi;
2021
import software.amazon.awssdk.awscore.internal.interceptor.TracingSystemSetting;
2122
import software.amazon.awssdk.core.interceptor.Context;
23+
import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
2224
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
2325
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
2426
import software.amazon.awssdk.http.SdkHttpRequest;
@@ -32,27 +34,57 @@
3234
public class TraceIdExecutionInterceptor implements ExecutionInterceptor {
3335
private static final String TRACE_ID_HEADER = "X-Amzn-Trace-Id";
3436
private static final String LAMBDA_FUNCTION_NAME_ENVIRONMENT_VARIABLE = "AWS_LAMBDA_FUNCTION_NAME";
37+
private static final String CONCURRENT_TRACE_ID_KEY = "AWS_LAMBDA_X_TRACE_ID";
38+
private static final ExecutionAttribute<String> TRACE_ID = new ExecutionAttribute<>("TraceId");
39+
40+
@Override
41+
public void beforeExecution(Context.BeforeExecution context, ExecutionAttributes executionAttributes) {
42+
String traceId = MDC.get(CONCURRENT_TRACE_ID_KEY);
43+
if (traceId != null) {
44+
executionAttributes.putAttribute(TRACE_ID, traceId);
45+
}
46+
}
3547

3648
@Override
3749
public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) {
3850
Optional<String> traceIdHeader = traceIdHeader(context);
3951
if (!traceIdHeader.isPresent()) {
40-
Optional<String> lambdafunctionName = lambdaFunctionNameEnvironmentVariable();
41-
Optional<String> traceId = traceId();
52+
Optional<String> lambdaFunctionName = lambdaFunctionNameEnvironmentVariable();
53+
Optional<String> traceId = traceId(executionAttributes);
4254

43-
if (lambdafunctionName.isPresent() && traceId.isPresent()) {
55+
if (lambdaFunctionName.isPresent() && traceId.isPresent()) {
4456
return context.httpRequest().copy(r -> r.putHeader(TRACE_ID_HEADER, traceId.get()));
4557
}
4658
}
47-
4859
return context.httpRequest();
4960
}
5061

62+
@Override
63+
public void afterExecution(Context.AfterExecution context, ExecutionAttributes executionAttributes) {
64+
saveTraceId(executionAttributes);
65+
}
66+
67+
@Override
68+
public void onExecutionFailure(Context.FailedExecution context, ExecutionAttributes executionAttributes) {
69+
saveTraceId(executionAttributes);
70+
}
71+
72+
private static void saveTraceId(ExecutionAttributes executionAttributes) {
73+
String traceId = executionAttributes.getAttribute(TRACE_ID);
74+
if (traceId != null) {
75+
MDC.put(CONCURRENT_TRACE_ID_KEY, executionAttributes.getAttribute(TRACE_ID));
76+
}
77+
}
78+
5179
private Optional<String> traceIdHeader(Context.ModifyHttpRequest context) {
5280
return context.httpRequest().firstMatchingHeader(TRACE_ID_HEADER);
5381
}
5482

55-
private Optional<String> traceId() {
83+
private Optional<String> traceId(ExecutionAttributes executionAttributes) {
84+
Optional<String> traceId = Optional.ofNullable(executionAttributes.getAttribute(TRACE_ID));
85+
if (traceId.isPresent()) {
86+
return traceId;
87+
}
5688
return TracingSystemSetting._X_AMZN_TRACE_ID.getStringValue();
5789
}
5890

‎core/aws-core/src/test/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptorTest.java

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import java.util.Properties;
2222
import org.junit.jupiter.api.Test;
2323
import org.mockito.Mockito;
24+
import org.slf4j.MDC;
2425
import software.amazon.awssdk.core.SdkRequest;
2526
import software.amazon.awssdk.core.interceptor.Context;
2627
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
@@ -111,6 +112,78 @@ public void headerNotAddedIfNoTraceIdEnvVar() {
111112
});
112113
}
113114

115+
@Test
116+
public void modifyHttpRequest_whenMultiConcurrencyModeWithMdc_shouldAddTraceIdHeader() {
117+
EnvironmentVariableHelper.run(env -> {
118+
resetRelevantEnvVars(env);
119+
env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
120+
MDC.put("AWS_LAMBDA_X_TRACE_ID", "mdc-trace-123");
121+
122+
try {
123+
TraceIdExecutionInterceptor interceptor = new TraceIdExecutionInterceptor();
124+
ExecutionAttributes executionAttributes = new ExecutionAttributes();
125+
126+
interceptor.beforeExecution(null, executionAttributes);
127+
Context.ModifyHttpRequest context = context();
128+
129+
SdkHttpRequest request = interceptor.modifyHttpRequest(context, executionAttributes);
130+
assertThat(request.firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123");
131+
} finally {
132+
MDC.remove("AWS_LAMBDA_X_TRACE_ID");
133+
}
134+
});
135+
}
136+
137+
@Test
138+
public void modifyHttpRequest_whenMultiConcurrencyModeWithBothMdcAndSystemProperty_shouldUseMdcValue() {
139+
EnvironmentVariableHelper.run(env -> {
140+
resetRelevantEnvVars(env);
141+
env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
142+
143+
MDC.put("AWS_LAMBDA_X_TRACE_ID", "mdc-trace-123");
144+
Properties props = System.getProperties();
145+
props.setProperty("com.amazonaws.xray.traceHeader", "sys-prop-345");
146+
147+
try {
148+
TraceIdExecutionInterceptor interceptor = new TraceIdExecutionInterceptor();
149+
ExecutionAttributes executionAttributes = new ExecutionAttributes();
150+
151+
interceptor.beforeExecution(null, executionAttributes);
152+
153+
Context.ModifyHttpRequest context = context();
154+
SdkHttpRequest request = interceptor.modifyHttpRequest(context, executionAttributes);
155+
156+
assertThat(request.firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123");
157+
} finally {
158+
MDC.remove("AWS_LAMBDA_X_TRACE_ID");
159+
props.remove("com.amazonaws.xray.traceHeader");
160+
}
161+
});
162+
}
163+
164+
@Test
165+
public void modifyHttpRequest_whenNotInLambdaEnvironmentWithMdc_shouldNotAddHeader() {
166+
EnvironmentVariableHelper.run(env -> {
167+
resetRelevantEnvVars(env);
168+
169+
MDC.put("AWS_LAMBDA_X_TRACE_ID", "should-be-ignored");
170+
171+
try {
172+
TraceIdExecutionInterceptor interceptor = new TraceIdExecutionInterceptor();
173+
ExecutionAttributes executionAttributes = new ExecutionAttributes();
174+
175+
interceptor.beforeExecution(null, executionAttributes);
176+
177+
Context.ModifyHttpRequest context = context();
178+
SdkHttpRequest request = interceptor.modifyHttpRequest(context, executionAttributes);
179+
180+
assertThat(request.firstMatchingHeader("X-Amzn-Trace-Id")).isEmpty();
181+
} finally {
182+
MDC.remove("AWS_LAMBDA_X_TRACE_ID");
183+
}
184+
});
185+
}
186+
114187
private Context.ModifyHttpRequest context() {
115188
return context(SdkHttpRequest.builder()
116189
.uri(URI.create("https://localhost"))

‎test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java

Lines changed: 187 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,24 @@
1717

1818
import static org.assertj.core.api.Assertions.assertThat;
1919

20+
import java.util.List;
21+
import java.util.concurrent.atomic.AtomicReference;
2022
import org.junit.jupiter.api.Test;
23+
import org.slf4j.MDC;
2124
import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider;
2225
import software.amazon.awssdk.awscore.interceptor.TraceIdExecutionInterceptor;
26+
import software.amazon.awssdk.core.interceptor.Context;
27+
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
28+
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
2329
import software.amazon.awssdk.http.AbortableInputStream;
2430
import software.amazon.awssdk.http.HttpExecuteResponse;
31+
import software.amazon.awssdk.http.SdkHttpRequest;
2532
import software.amazon.awssdk.http.SdkHttpResponse;
2633
import software.amazon.awssdk.regions.Region;
34+
import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonAsyncClient;
2735
import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonClient;
2836
import software.amazon.awssdk.testutils.EnvironmentVariableHelper;
37+
import software.amazon.awssdk.testutils.service.http.MockAsyncHttpClient;
2938
import software.amazon.awssdk.testutils.service.http.MockSyncHttpClient;
3039
import software.amazon.awssdk.utils.StringInputStream;
3140

@@ -56,4 +65,182 @@ public void traceIdInterceptorIsEnabled() {
5665
}
5766
});
5867
}
68+
69+
@Test
70+
public void traceIdInterceptorPreservesTraceIdAcrossRetries() {
71+
EnvironmentVariableHelper.run(env -> {
72+
env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
73+
MDC.put("AWS_LAMBDA_X_TRACE_ID", "mdc-trace-123");
74+
75+
try (MockAsyncHttpClient mockHttpClient = new MockAsyncHttpClient();
76+
ProtocolRestJsonAsyncClient client = ProtocolRestJsonAsyncClient.builder()
77+
.region(Region.US_WEST_2)
78+
.credentialsProvider(AnonymousCredentialsProvider.create())
79+
.httpClient(mockHttpClient)
80+
.build()) {
81+
82+
mockHttpClient.stubResponses(
83+
HttpExecuteResponse.builder()
84+
.response(SdkHttpResponse.builder().statusCode(500).build())
85+
.responseBody(AbortableInputStream.create(new StringInputStream("{}")))
86+
.build(),
87+
HttpExecuteResponse.builder()
88+
.response(SdkHttpResponse.builder().statusCode(500).build())
89+
.responseBody(AbortableInputStream.create(new StringInputStream("{}")))
90+
.build(),
91+
HttpExecuteResponse.builder().response(SdkHttpResponse.builder().statusCode(200).build())
92+
.responseBody(AbortableInputStream.create(new StringInputStream("{}")))
93+
.build());
94+
95+
client.allTypes().join();
96+
97+
List<SdkHttpRequest> requests = mockHttpClient.getRequests();
98+
assertThat(requests).hasSize(3);
99+
100+
assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123");
101+
assertThat(requests.get(1).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123");
102+
assertThat(requests.get(2).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123");
103+
104+
} finally {
105+
MDC.clear();
106+
}
107+
});
108+
}
109+
110+
@Test
111+
public void traceIdInterceptorPreservesTraceIdAcrossChainedFutures() {
112+
EnvironmentVariableHelper.run(env -> {
113+
env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
114+
MDC.put("AWS_LAMBDA_X_TRACE_ID", "mdc-trace-123");
115+
116+
try (MockAsyncHttpClient mockHttpClient = new MockAsyncHttpClient();
117+
ProtocolRestJsonAsyncClient client = ProtocolRestJsonAsyncClient.builder()
118+
.region(Region.US_WEST_2)
119+
.credentialsProvider(AnonymousCredentialsProvider.create())
120+
.httpClient(mockHttpClient)
121+
.build()) {
122+
123+
mockHttpClient.stubResponses(
124+
HttpExecuteResponse.builder()
125+
.response(SdkHttpResponse.builder().statusCode(200).build())
126+
.responseBody(AbortableInputStream.create(new StringInputStream("{}")))
127+
.build(),
128+
HttpExecuteResponse.builder()
129+
.response(SdkHttpResponse.builder().statusCode(200).build())
130+
.responseBody(AbortableInputStream.create(new StringInputStream("{}")))
131+
.build()
132+
);
133+
134+
client.allTypes()
135+
.thenRun(() -> {
136+
client.allTypes().join();
137+
})
138+
.join();
139+
140+
List<SdkHttpRequest> requests = mockHttpClient.getRequests();
141+
142+
assertThat(requests).hasSize(2);
143+
144+
assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123");
145+
assertThat(requests.get(1).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123");
146+
147+
} finally {
148+
MDC.clear();
149+
}
150+
});
151+
}
152+
153+
@Test
154+
public void traceIdInterceptorPreservesTraceIdAcrossExceptionallyCompletedFutures() {
155+
EnvironmentVariableHelper.run(env -> {
156+
env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
157+
MDC.put("AWS_LAMBDA_X_TRACE_ID", "mdc-trace-123");
158+
159+
try (MockAsyncHttpClient mockHttpClient = new MockAsyncHttpClient();
160+
ProtocolRestJsonAsyncClient client = ProtocolRestJsonAsyncClient.builder()
161+
.region(Region.US_WEST_2)
162+
.credentialsProvider(AnonymousCredentialsProvider.create())
163+
.httpClient(mockHttpClient)
164+
.build()) {
165+
166+
mockHttpClient.stubResponses(
167+
HttpExecuteResponse.builder()
168+
.response(SdkHttpResponse.builder().statusCode(400).build())
169+
.responseBody(AbortableInputStream.create(new StringInputStream("{}")))
170+
.build(),
171+
HttpExecuteResponse.builder()
172+
.response(SdkHttpResponse.builder().statusCode(200).build())
173+
.responseBody(AbortableInputStream.create(new StringInputStream("{}")))
174+
.build()
175+
);
176+
177+
client.allTypes()
178+
.exceptionally(throwable -> {
179+
client.allTypes().join();
180+
return null;
181+
}).join();
182+
183+
List<SdkHttpRequest> requests = mockHttpClient.getRequests();
184+
185+
assertThat(requests).hasSize(2);
186+
187+
assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123");
188+
assertThat(requests.get(1).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123");
189+
190+
} finally {
191+
MDC.clear();
192+
}
193+
});
194+
}
195+
196+
@Test
197+
public void traceIdInterceptorPreservesTraceIdAcrossExceptionallyCompletedFuturesThrownInPreExecution() {
198+
EnvironmentVariableHelper.run(env -> {
199+
env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
200+
MDC.put("AWS_LAMBDA_X_TRACE_ID", "mdc-trace-123");
201+
202+
ExecutionInterceptor throwingInterceptor = new ExecutionInterceptor() {
203+
private boolean hasThrown = false;
204+
205+
@Override
206+
public void beforeMarshalling(Context.BeforeMarshalling context, ExecutionAttributes executionAttributes) {
207+
if (!hasThrown) {
208+
hasThrown = true;
209+
throw new RuntimeException("failing in pre execution");
210+
}
211+
}
212+
};
213+
214+
try (MockAsyncHttpClient mockHttpClient = new MockAsyncHttpClient();
215+
ProtocolRestJsonAsyncClient client = ProtocolRestJsonAsyncClient.builder()
216+
.region(Region.US_WEST_2)
217+
.credentialsProvider(AnonymousCredentialsProvider.create())
218+
.overrideConfiguration(o -> o.addExecutionInterceptor(throwingInterceptor))
219+
.httpClient(mockHttpClient)
220+
.build()) {
221+
222+
mockHttpClient.stubResponses(
223+
HttpExecuteResponse.builder()
224+
.response(SdkHttpResponse.builder().statusCode(200).build())
225+
.responseBody(AbortableInputStream.create(new StringInputStream("{}")))
226+
.build()
227+
);
228+
229+
client.allTypes()
230+
.exceptionally(throwable -> {
231+
client.allTypes().join();
232+
return null;
233+
}).join();
234+
235+
List<SdkHttpRequest> requests = mockHttpClient.getRequests();
236+
237+
assertThat(requests).hasSize(1);
238+
assertThat(requests.get(0).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("mdc-trace-123");
239+
240+
} finally {
241+
MDC.clear();
242+
}
243+
});
244+
}
59245
}
246+

0 commit comments

Comments
 (0)