Skip to content

Commit 01633c9

Browse files
amarzialiamarziali
committed
proto
1 parent 62b2aa0 commit 01633c9

File tree

8 files changed

+268
-17
lines changed

8 files changed

+268
-17
lines changed

‎dd-java-agent/instrumentation/mule-4/build.gradle‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ muzzle {
2727

2828
apply from: "$rootDir/gradle/java.gradle"
2929

30-
def muleVersion = '4.2.2'
30+
def muleVersion = '4.5.0'
3131
def muleBaseDir = "$buildDir/mule"
3232
def appDir = "$projectDir/application"
3333
def generatedResourcesDir = "$buildDir/generated-resources/test"

‎dd-java-agent/instrumentation/mule-4/src/main/java/datadog/trace/instrumentation/mule4/CurrentEventHelper.java‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
44

5+
import datadog.trace.api.Pair;
56
import datadog.trace.bootstrap.ContextStore;
67
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
78
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -14,19 +15,21 @@ public class CurrentEventHelper {
1415
private static final ThreadLocal<AgentScope> currentEventScope = new ThreadLocal<>();
1516

1617
public static void handleEventChange(
17-
final PrivilegedEvent event, ContextStore<EventContext, AgentSpan> contextStore) {
18+
final PrivilegedEvent event, ContextStore<EventContext, Pair> contextStore) {
19+
attachSpanToEventContext(event == null ? null : event.getContext(), contextStore);
20+
}
21+
22+
public static void attachSpanToEventContext(
23+
final EventContext eventContext, ContextStore<EventContext, Pair> contextStore) {
1824
final AgentScope currentScope = currentEventScope.get();
1925
if (null != currentScope) {
2026
currentScope.close();
2127
}
2228
AgentScope newScope = null;
23-
if (null != event) {
24-
EventContext eventContext = event.getContext();
25-
if (null != eventContext) {
26-
AgentSpan span = contextStore.get(eventContext);
27-
if (null != span) {
28-
newScope = activateSpan(span);
29-
}
29+
if (null != eventContext) {
30+
Pair<AgentSpan, Object> pair = contextStore.get(eventContext);
31+
if (null != pair && pair.hasLeft()) {
32+
newScope = activateSpan(pair.getLeft());
3033
}
3134
}
3235
currentEventScope.set(newScope);

‎dd-java-agent/instrumentation/mule-4/src/main/java/datadog/trace/instrumentation/mule4/DDEventTracer.java‎

Lines changed: 172 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,172 @@
1+
package datadog.trace.instrumentation.mule4;
2+
3+
import datadog.trace.api.Pair;
4+
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
5+
import datadog.trace.bootstrap.ContextStore;
6+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
7+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
8+
import java.util.Map;
9+
import java.util.Optional;
10+
import java.util.function.Supplier;
11+
import org.mule.runtime.api.event.EventContext;
12+
import org.mule.runtime.api.message.Error;
13+
import org.mule.runtime.api.profiling.tracing.Span;
14+
import org.mule.runtime.api.profiling.tracing.SpanError;
15+
import org.mule.runtime.core.api.event.CoreEvent;
16+
import org.mule.runtime.tracer.api.EventTracer;
17+
import org.mule.runtime.tracer.api.context.getter.DistributedTraceContextGetter;
18+
import org.mule.runtime.tracer.api.sniffer.SpanSnifferManager;
19+
import org.mule.runtime.tracer.api.span.info.InitialSpanInfo;
20+
import org.mule.runtime.tracer.api.span.validation.Assertion;
21+
22+
public class DDEventTracer implements EventTracer<CoreEvent> {
23+
24+
private final ContextStore<Span, AgentSpan> spanStore;
25+
private final ContextStore<EventContext, Pair> contextStore;
26+
private final EventTracer<CoreEvent> delegate;
27+
28+
public DDEventTracer(
29+
ContextStore<Span, AgentSpan> spanStore,
30+
ContextStore<EventContext, Pair> contextStore,
31+
EventTracer<CoreEvent> delegate) {
32+
this.spanStore = spanStore;
33+
this.contextStore = contextStore;
34+
this.delegate = delegate;
35+
}
36+
37+
private AgentSpan fromMuleSpan(Span muleSpan, InitialSpanInfo spanInfo) {
38+
final AgentSpan parent =
39+
muleSpan.getParent() != null ? spanStore.get(muleSpan.getParent()) : null;
40+
41+
final AgentSpan span;
42+
if (parent == null) {
43+
span = AgentTracer.startSpan("mule", muleSpan.getName());
44+
} else {
45+
span = AgentTracer.startSpan("mule", muleSpan.getName(), parent.context());
46+
}
47+
spanInfo.forEachAttribute(span::setTag);
48+
return span;
49+
}
50+
51+
private void activateOnContext(EventContext eventContext, AgentSpan span, Span muleSpan) {
52+
contextStore.put(eventContext, Pair.of(span, muleSpan));
53+
CurrentEventHelper.attachSpanToEventContext(eventContext, contextStore);
54+
}
55+
56+
private Optional<Span> handleNewSpan(
57+
CoreEvent event, Optional<Span> maybeSpan, InitialSpanInfo spanInfo) {
58+
if (!maybeSpan.isPresent() || CallDepthThreadLocalMap.incrementCallDepth(Span.class) > 0) {
59+
return maybeSpan;
60+
}
61+
try {
62+
final Span muleSpan = maybeSpan.get();
63+
final AgentSpan span = fromMuleSpan(muleSpan, spanInfo);
64+
spanStore.put(muleSpan, span);
65+
activateOnContext(event.getContext(), span, muleSpan);
66+
} finally {
67+
CallDepthThreadLocalMap.reset(Span.class);
68+
}
69+
return maybeSpan;
70+
}
71+
72+
private void handleEndOfSpan(CoreEvent event) {
73+
if (CallDepthThreadLocalMap.incrementCallDepth(Span.class) > 0) {
74+
return;
75+
}
76+
try {
77+
final Pair<AgentSpan, Span> pair = contextStore.get(event.getContext());
78+
if (pair != null && pair.hasLeft() && pair.hasRight()) {
79+
final AgentSpan span = pair.getLeft();
80+
final Span muleSpan = pair.getRight();
81+
spanStore.remove(muleSpan);
82+
if (muleSpan.hasErrors()) {
83+
span.setError(true);
84+
for (final SpanError spanError : muleSpan.getErrors()) {
85+
if (spanError.getError() != null && spanError.getError().getCause() != null) {
86+
span.addThrowable(spanError.getError().getCause());
87+
}
88+
}
89+
}
90+
span.finish();
91+
final Span muleParent = muleSpan.getParent();
92+
if (muleParent != null) {
93+
final AgentSpan parent = spanStore.get(muleParent);
94+
if (parent != null) {
95+
activateOnContext(event.getContext(), parent, muleParent);
96+
}
97+
}
98+
}
99+
} finally {
100+
CallDepthThreadLocalMap.reset(Span.class);
101+
}
102+
}
103+
104+
@Override
105+
public Optional<Span> startSpan(CoreEvent event, InitialSpanInfo spanInfo) {
106+
return handleNewSpan(event, delegate.startSpan(event, spanInfo), spanInfo);
107+
}
108+
109+
@Override
110+
public Optional<Span> startSpan(CoreEvent event, InitialSpanInfo spanInfo, Assertion assertion) {
111+
return handleNewSpan(event, delegate.startSpan(event, spanInfo, assertion), spanInfo);
112+
}
113+
114+
@Override
115+
public void endCurrentSpan(CoreEvent event) {
116+
delegate.endCurrentSpan(event);
117+
handleEndOfSpan(event);
118+
}
119+
120+
@Override
121+
public void endCurrentSpan(CoreEvent event, Assertion condition) {
122+
delegate.endCurrentSpan(event, condition);
123+
handleEndOfSpan(event);
124+
}
125+
126+
@Override
127+
public void injectDistributedTraceContext(
128+
EventContext eventContext, DistributedTraceContextGetter distributedTraceContextGetter) {
129+
delegate.injectDistributedTraceContext(eventContext, distributedTraceContextGetter);
130+
}
131+
132+
@Override
133+
public void recordErrorAtCurrentSpan(
134+
CoreEvent event, Supplier<Error> errorSupplier, boolean isErrorEscapingCurrentSpan) {
135+
delegate.recordErrorAtCurrentSpan(event, errorSupplier, isErrorEscapingCurrentSpan);
136+
}
137+
138+
@Override
139+
public void setCurrentSpanName(CoreEvent event, String name) {
140+
delegate.setCurrentSpanName(event, name);
141+
Pair<AgentSpan, Span> pair = contextStore.get(event.getContext());
142+
if (pair != null && pair.hasLeft()) {
143+
pair.getLeft().setOperationName(name);
144+
}
145+
}
146+
147+
@Override
148+
public void addCurrentSpanAttribute(CoreEvent event, String key, String value) {
149+
delegate.addCurrentSpanAttribute(event, key, value);
150+
Pair<AgentSpan, Span> pair = contextStore.get(event.getContext());
151+
if (pair != null && pair.hasLeft()) {
152+
pair.getLeft().setTag(key, value);
153+
}
154+
}
155+
156+
@Override
157+
public void addCurrentSpanAttributes(CoreEvent event, Map<String, String> attributes) {
158+
delegate.addCurrentSpanAttributes(event, attributes);
159+
Pair<AgentSpan, Span> pair = contextStore.get(event.getContext());
160+
if (pair != null && pair.hasLeft()) {
161+
final AgentSpan span = pair.getLeft();
162+
for (Map.Entry<String, String> entry : attributes.entrySet()) {
163+
span.setTag(entry.getKey(), entry.getValue());
164+
}
165+
}
166+
}
167+
168+
@Override
169+
public SpanSnifferManager getSpanSnifferManager() {
170+
return delegate.getSpanSnifferManager();
171+
}
172+
}

‎dd-java-agent/instrumentation/mule-4/src/main/java/datadog/trace/instrumentation/mule4/EventContextCreationAdvice.java‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,14 @@
22

33
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
44

5+
import datadog.trace.api.Pair;
56
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
67
import datadog.trace.bootstrap.ContextStore;
78
import datadog.trace.bootstrap.InstrumentationContext;
89
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
910
import net.bytebuddy.asm.Advice;
1011
import org.mule.runtime.api.event.EventContext;
12+
import org.mule.runtime.api.profiling.tracing.Span;
1113
import org.mule.runtime.core.internal.event.DefaultEventContext;
1214

1315
public class EventContextCreationAdvice {
@@ -27,8 +29,8 @@ public static void onExit(
2729
}
2830
CallDepthThreadLocalMap.reset(EventContext.class);
2931

30-
final ContextStore<EventContext, AgentSpan> contextStore =
31-
InstrumentationContext.get(EventContext.class, AgentSpan.class);
32+
final ContextStore<EventContext, Pair> contextStore =
33+
InstrumentationContext.get(EventContext.class, Pair.class);
3234
AgentSpan span = null;
3335

3436
// This is a roundabout way to know if we are in the constructor for DefaultEventContext or
@@ -38,10 +40,13 @@ public static void onExit(
3840
} else if (arg instanceof EventContext) {
3941
// This means that we are in the constructor for ChildContext and we should copy the span
4042
// from the parent EventContext which is the first argument.
41-
span = contextStore.get((EventContext) arg);
43+
Pair<AgentSpan, Object> pair = contextStore.get((EventContext) arg);
44+
if (pair != null && pair.hasLeft()) {
45+
span = pair.getLeft();
46+
}
4247
}
4348
if (null != span) {
44-
contextStore.put(zis, span);
49+
contextStore.put(zis, Pair.of(span, null));
4550
}
4651
}
4752
}

‎dd-java-agent/instrumentation/mule-4/src/main/java/datadog/trace/instrumentation/mule4/EventContextInstrumentation.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
import com.google.auto.service.AutoService;
77
import datadog.trace.agent.tooling.Instrumenter;
88
import datadog.trace.agent.tooling.InstrumenterModule;
9+
import datadog.trace.api.Pair;
10+
911
import java.util.Map;
1012

1113
/**
@@ -38,7 +40,7 @@ public String[] knownMatchingTypes() {
3840
public Map<String, String> contextStore() {
3941
return singletonMap(
4042
"org.mule.runtime.api.event.EventContext",
41-
"datadog.trace.bootstrap.instrumentation.api.AgentSpan");
43+
Pair.class.getName());
4244
}
4345

4446
@Override

‎dd-java-agent/instrumentation/mule-4/src/main/java/datadog/trace/instrumentation/mule4/EventTracerInstrumentation.java‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
package datadog.trace.instrumentation.mule4;
2+
3+
import com.google.auto.service.AutoService;
4+
import datadog.trace.agent.tooling.Instrumenter;
5+
import datadog.trace.agent.tooling.InstrumenterModule;
6+
import datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers;
7+
import datadog.trace.api.Pair;
8+
import datadog.trace.bootstrap.InstrumentationContext;
9+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
10+
import java.util.HashMap;
11+
import java.util.Map;
12+
import net.bytebuddy.asm.Advice;
13+
import org.mule.runtime.api.event.EventContext;
14+
import org.mule.runtime.api.profiling.tracing.Span;
15+
import org.mule.runtime.core.api.event.CoreEvent;
16+
import org.mule.runtime.tracer.api.EventTracer;
17+
18+
@AutoService(InstrumenterModule.class)
19+
public class EventTracerInstrumentation extends InstrumenterModule.Tracing
20+
implements Instrumenter.ForSingleType {
21+
public EventTracerInstrumentation() {
22+
super("mule-otel");
23+
}
24+
25+
@Override
26+
public Map<String, String> contextStore() {
27+
final Map<String, String> ret = new HashMap<>();
28+
ret.put("org.mule.runtime.api.event.EventContext", Pair.class.getName());
29+
ret.put(
30+
"org.mule.runtime.api.profiling.tracing.Span",
31+
"datadog.trace.bootstrap.instrumentation.api.AgentSpan");
32+
return ret;
33+
}
34+
35+
@Override
36+
public String[] helperClassNames() {
37+
return new String[] {
38+
packageName + ".CurrentEventHelper", packageName + ".DDEventTracer",
39+
};
40+
}
41+
42+
@Override
43+
public void methodAdvice(MethodTransformer transformer) {
44+
transformer.applyAdvice(
45+
NameMatchers.named("updateSelectedCoreEventTracer"),
46+
getClass().getName() + "$SwapCoreTracerAdvice");
47+
}
48+
49+
@Override
50+
public String instrumentedType() {
51+
return "org.mule.runtime.tracer.impl.SelectableCoreEventTracer";
52+
}
53+
54+
public static class SwapCoreTracerAdvice {
55+
@Advice.OnMethodExit(onThrowable = Throwable.class)
56+
public static void afterInit(
57+
@Advice.FieldValue(value = "selectedCoreEventTracer", readOnly = false)
58+
EventTracer<CoreEvent> eventTracer) {
59+
System.err.println("SWAPPING " + eventTracer);
60+
eventTracer =
61+
new DDEventTracer(
62+
InstrumentationContext.get(Span.class, AgentSpan.class),
63+
InstrumentationContext.get(EventContext.class, Pair.class),
64+
eventTracer);
65+
}
66+
}
67+
}

‎dd-java-agent/instrumentation/mule-4/src/main/java/datadog/trace/instrumentation/mule4/PrivilegedEventInstrumentation.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
import com.google.auto.service.AutoService;
77
import datadog.trace.agent.tooling.Instrumenter;
88
import datadog.trace.agent.tooling.InstrumenterModule;
9+
import datadog.trace.api.Pair;
10+
911
import java.util.Map;
1012

1113
/**
@@ -35,7 +37,7 @@ public String instrumentedType() {
3537
public Map<String, String> contextStore() {
3638
return singletonMap(
3739
"org.mule.runtime.api.event.EventContext",
38-
"datadog.trace.bootstrap.instrumentation.api.AgentSpan");
40+
Pair.class.getName());
3941
}
4042

4143
@Override

‎dd-java-agent/instrumentation/mule-4/src/main/java/datadog/trace/instrumentation/mule4/PrivilegedEventSetCurrentAdvice.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
package datadog.trace.instrumentation.mule4;
22

3+
import datadog.trace.api.Pair;
34
import datadog.trace.bootstrap.InstrumentationContext;
4-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
55
import net.bytebuddy.asm.Advice;
66
import org.mule.runtime.api.event.EventContext;
77
import org.mule.runtime.core.privileged.event.PrivilegedEvent;
@@ -10,6 +10,6 @@ public class PrivilegedEventSetCurrentAdvice {
1010
@Advice.OnMethodEnter(suppress = Throwable.class)
1111
public static void onEnter(@Advice.Argument(0) final PrivilegedEvent event) {
1212
CurrentEventHelper.handleEventChange(
13-
event, InstrumentationContext.get(EventContext.class, AgentSpan.class));
13+
event, InstrumentationContext.get(EventContext.class, Pair.class));
1414
}
1515
}

0 commit comments

Comments
 (0)