Skip to content

Commit 8be300e

Browse files
nhulstonnhulston
authored
Inject trace context into EventBridge detail (#7613)
Inject trace context into EventBridge payloads
1 parent e5e5427 commit 8be300e

File tree

9 files changed

+761
-2
lines changed

9 files changed

+761
-2
lines changed

‎dd-java-agent/instrumentation/aws-java-eventbridge-2.0/build.gradle

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
muzzle {
2+
pass {
3+
group = "software.amazon.awssdk"
4+
module = "eventbridge"
5+
versions = "[2.7,3)"
6+
assertInverse = true
7+
}
8+
}
9+
10+
apply from: "$rootDir/gradle/java.gradle"
11+
12+
addTestSuiteForDir('latestDepTest', 'test')
13+
addTestSuiteExtendingForDir('latestDepForkedTest', 'latestDepTest', 'test')
14+
15+
dependencies {
16+
compileOnly group: 'software.amazon.awssdk', name: 'eventbridge', version: '2.27.19'
17+
18+
// Include httpclient instrumentation for testing because it is a dependency for aws-sdk.
19+
testImplementation project(':dd-java-agent:instrumentation:apache-httpclient-4')
20+
testImplementation project(':dd-java-agent:instrumentation:aws-java-sdk-2.2')
21+
testImplementation 'software.amazon.awssdk:eventbridge:2.27.23'
22+
// SQS and SNS are used to act as the "targets" of the EB bus.
23+
testImplementation 'software.amazon.awssdk:sqs:2.27.23'
24+
testImplementation 'software.amazon.awssdk:sns:2.27.23'
25+
testImplementation 'org.testcontainers:localstack:1.20.1'
26+
27+
latestDepTestImplementation group: 'software.amazon.awssdk', name: 'eventbridge', version: '+'
28+
}
29+
30+
tasks.withType(Test).configureEach {
31+
usesService(testcontainersLimit)
32+
}

‎dd-java-agent/instrumentation/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeClientInstrumentation.java

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
package datadog.trace.instrumentation.aws.v2.eventbridge;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
5+
6+
import com.google.auto.service.AutoService;
7+
import datadog.trace.agent.tooling.Instrumenter;
8+
import datadog.trace.agent.tooling.InstrumenterModule;
9+
import java.util.List;
10+
import net.bytebuddy.asm.Advice;
11+
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
12+
13+
@AutoService(InstrumenterModule.class)
14+
public final class EventBridgeClientInstrumentation extends InstrumenterModule.Tracing
15+
implements Instrumenter.ForSingleType {
16+
public EventBridgeClientInstrumentation() {
17+
super("eventbridge");
18+
}
19+
20+
@Override
21+
public String instrumentedType() {
22+
return "software.amazon.awssdk.core.client.builder.SdkDefaultClientBuilder";
23+
}
24+
25+
@Override
26+
public void methodAdvice(MethodTransformer transformer) {
27+
transformer.applyAdvice(
28+
isMethod().and(named("resolveExecutionInterceptors")),
29+
EventBridgeClientInstrumentation.class.getName() + "$AwsEventBridgeBuilderAdvice");
30+
}
31+
32+
@Override
33+
public String[] helperClassNames() {
34+
return new String[] {
35+
packageName + ".EventBridgeInterceptor", packageName + ".TextMapInjectAdapter"
36+
};
37+
}
38+
39+
public static class AwsEventBridgeBuilderAdvice {
40+
@Advice.OnMethodExit(suppress = Throwable.class)
41+
public static void addHandler(@Advice.Return final List<ExecutionInterceptor> interceptors) {
42+
for (ExecutionInterceptor interceptor : interceptors) {
43+
if (interceptor instanceof EventBridgeInterceptor) {
44+
return; // list already has our interceptor, return to builder
45+
}
46+
}
47+
interceptors.add(new EventBridgeInterceptor());
48+
}
49+
}
50+
}

‎dd-java-agent/instrumentation/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
package datadog.trace.instrumentation.aws.v2.eventbridge;
2+
3+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
4+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig;
5+
import static datadog.trace.core.datastreams.TagsProcessor.BUS_TAG;
6+
import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_OUT;
7+
import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_TAG;
8+
import static datadog.trace.core.datastreams.TagsProcessor.TYPE_TAG;
9+
import static datadog.trace.instrumentation.aws.v2.eventbridge.TextMapInjectAdapter.SETTER;
10+
11+
import datadog.trace.bootstrap.InstanceStore;
12+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
13+
import datadog.trace.bootstrap.instrumentation.api.PathwayContext;
14+
import java.util.ArrayList;
15+
import java.util.LinkedHashMap;
16+
import java.util.List;
17+
import org.slf4j.Logger;
18+
import org.slf4j.LoggerFactory;
19+
import software.amazon.awssdk.core.SdkRequest;
20+
import software.amazon.awssdk.core.interceptor.Context;
21+
import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
22+
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
23+
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
24+
import software.amazon.awssdk.services.eventbridge.model.PutEventsRequest;
25+
import software.amazon.awssdk.services.eventbridge.model.PutEventsRequestEntry;
26+
27+
public class EventBridgeInterceptor implements ExecutionInterceptor {
28+
private static final Logger log = LoggerFactory.getLogger(EventBridgeInterceptor.class);
29+
30+
public static final ExecutionAttribute<AgentSpan> SPAN_ATTRIBUTE =
31+
InstanceStore.of(ExecutionAttribute.class)
32+
.putIfAbsent("DatadogSpan", () -> new ExecutionAttribute<>("DatadogSpan"));
33+
34+
private static final String START_TIME_KEY = "x-datadog-start-time";
35+
private static final String RESOURCE_NAME_KEY = "x-datadog-resource-name";
36+
37+
@Override
38+
public SdkRequest modifyRequest(
39+
Context.ModifyRequest context, ExecutionAttributes executionAttributes) {
40+
if (!(context.request() instanceof PutEventsRequest)) {
41+
return context.request();
42+
}
43+
44+
PutEventsRequest request = (PutEventsRequest) context.request();
45+
List<PutEventsRequestEntry> modifiedEntries = new ArrayList<>(request.entries().size());
46+
long startTime = System.currentTimeMillis();
47+
48+
for (PutEventsRequestEntry entry : request.entries()) {
49+
StringBuilder detailBuilder = new StringBuilder(entry.detail().trim());
50+
if (detailBuilder.length() == 0) {
51+
detailBuilder.append("{}");
52+
}
53+
if (detailBuilder.charAt(detailBuilder.length() - 1) != '}') {
54+
log.debug(
55+
"Unable to parse detail JSON. Not injecting trace context into EventBridge payload.");
56+
modifiedEntries.add(entry); // Add the original entry without modification
57+
continue;
58+
}
59+
60+
String traceContext =
61+
getTraceContextToInject(executionAttributes, entry.eventBusName(), startTime);
62+
detailBuilder.setLength(detailBuilder.length() - 1); // Remove the last bracket
63+
if (detailBuilder.length() > 1) {
64+
detailBuilder.append(", "); // Only add a comma if detail is not empty.
65+
}
66+
67+
detailBuilder
68+
.append("\"")
69+
.append(PathwayContext.DATADOG_KEY)
70+
.append("\": ")
71+
.append(traceContext)
72+
.append('}');
73+
74+
String modifiedDetail = detailBuilder.toString();
75+
PutEventsRequestEntry modifiedEntry = entry.toBuilder().detail(modifiedDetail).build();
76+
modifiedEntries.add(modifiedEntry);
77+
}
78+
79+
return request.toBuilder().entries(modifiedEntries).build();
80+
}
81+
82+
private String getTraceContextToInject(
83+
ExecutionAttributes executionAttributes, String eventBusName, long startTime) {
84+
final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
85+
StringBuilder jsonBuilder = new StringBuilder();
86+
jsonBuilder.append('{');
87+
88+
// Inject trace context
89+
propagate().inject(span, jsonBuilder, SETTER);
90+
91+
if (traceConfig().isDataStreamsEnabled()) {
92+
propagate().injectPathwayContext(span, jsonBuilder, SETTER, getTags(eventBusName));
93+
}
94+
95+
// Add bus name and start time
96+
jsonBuilder
97+
.append(" \"")
98+
.append(START_TIME_KEY)
99+
.append("\": \"")
100+
.append(startTime)
101+
.append("\", ");
102+
jsonBuilder
103+
.append(" \"")
104+
.append(RESOURCE_NAME_KEY)
105+
.append("\": \"")
106+
.append(eventBusName)
107+
.append("\"");
108+
109+
jsonBuilder.append('}');
110+
return jsonBuilder.toString();
111+
}
112+
113+
private LinkedHashMap<String, String> getTags(String eventBusName) {
114+
LinkedHashMap<String, String> sortedTags = new LinkedHashMap<>();
115+
sortedTags.put(DIRECTION_TAG, DIRECTION_OUT);
116+
sortedTags.put(BUS_TAG, eventBusName);
117+
sortedTags.put(TYPE_TAG, "bus");
118+
119+
return sortedTags;
120+
}
121+
}

‎dd-java-agent/instrumentation/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/TextMapInjectAdapter.java

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
package datadog.trace.instrumentation.aws.v2.eventbridge;
2+
3+
import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
4+
5+
public class TextMapInjectAdapter implements AgentPropagation.Setter<StringBuilder> {
6+
7+
public static final TextMapInjectAdapter SETTER = new TextMapInjectAdapter();
8+
9+
@Override
10+
public void set(final StringBuilder builder, final String key, final String value) {
11+
builder.append('"').append(key).append("\":\"").append(value).append("\",");
12+
}
13+
}

0 commit comments

Comments
 (0)