Skip to content

Commit

Permalink
cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
amarziali committed Nov 28, 2024
1 parent 1d21fec commit 50b9acd
Show file tree
Hide file tree
Showing 6 changed files with 17 additions and 32 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;

Expand All @@ -10,7 +11,6 @@
import datadog.trace.agent.tooling.InstrumenterModule;
import datadog.trace.bootstrap.InstrumentationContext;
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
import net.bytebuddy.asm.Advice;
import org.mule.runtime.api.event.EventContext;
import org.mule.runtime.core.api.event.CoreEvent;
Expand Down Expand Up @@ -47,7 +47,7 @@ public static AgentScope before(@Advice.Argument(0) final CoreEvent event) {
SpanState spanState =
InstrumentationContext.get(EventContext.class, SpanState.class).get(event.getContext());
if (spanState != null && spanState.getEventContextSpan() != null) {
return AgentTracer.activateSpan(spanState.getSpanContextSpan());
return activateSpan(spanState.getSpanContextSpan());
}
return null;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -71,19 +71,19 @@ private Component findComponent(final InitialSpanInfo initialSpanInfo) {
return null;
}

private void activateOnContext(@Nonnull final EventContext eventContext, final AgentSpan span) {
private void linkToContext(@Nonnull final EventContext eventContext, final AgentSpan span) {
final SpanState previousState = eventContextStore.get(eventContext);
final AgentSpan spanToActivate;
final AgentSpan spanToLink;
if (span != null) {
spanToActivate = span;
spanToLink = span;
} else if (previousState != null) {
spanToActivate = previousState.getEventContextSpan();
spanToLink = previousState.getEventContextSpan();
} else {
spanToActivate = null;
spanToLink = null;
}

eventContextStore.put(
eventContext, new SpanState(spanToActivate, previousState).withSpanContextSpan(span));
eventContext, new SpanState(spanToLink, previousState).withSpanContextSpan(span));
}

private void handleNewSpan(CoreEvent event, InitialSpanInfo spanInfo) {
Expand All @@ -96,7 +96,7 @@ private void handleNewSpan(CoreEvent event, InitialSpanInfo spanInfo) {

final AgentSpan span =
DECORATE.onMuleSpan(findParent(eventContext), spanInfo, event, findComponent(spanInfo));
activateOnContext(eventContext, span);
linkToContext(eventContext, span);
}

private void handleEndOfSpan(CoreEvent event) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,37 +8,20 @@
import datadog.trace.agent.tooling.Instrumenter;
import datadog.trace.agent.tooling.InstrumenterModule;
import datadog.trace.bootstrap.InstrumentationContext;
import java.util.Collections;
import java.util.Map;
import net.bytebuddy.asm.Advice;
import org.mule.runtime.api.component.Component;
import org.mule.runtime.tracer.api.span.info.InitialSpanInfo;
import org.mule.runtime.tracer.customization.impl.info.ExecutionInitialSpanInfo;

@AutoService(InstrumenterModule.class)
public class ExecutionInitialSpanInfoInstrumentation extends InstrumenterModule.Tracing
public class ExecutionInitialSpanInfoInstrumentation extends AbstractMuleInstrumentation
implements Instrumenter.ForSingleType {
public ExecutionInitialSpanInfoInstrumentation() {
super("mule");
}

@Override
protected boolean defaultEnabled() {
return false;
}

@Override
public String instrumentedType() {
return "org.mule.runtime.tracer.customization.impl.info.ExecutionInitialSpanInfo";
}

@Override
public Map<String, String> contextStore() {
return Collections.singletonMap(
"org.mule.runtime.tracer.api.span.info.InitialSpanInfo",
"org.mule.runtime.api.component.Component");
}

@Override
public void methodAdvice(MethodTransformer transformer) {
transformer.applyAdvice(
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
package datadog.trace.instrumentation.mule4;

import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.MULE_CORRELATION_ID;
import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.MULE_LOCATION;

import datadog.trace.api.Functions;
import datadog.trace.api.cache.DDCache;
Expand Down Expand Up @@ -73,14 +75,14 @@ public AgentSpan onMuleSpan(
// here we have to use the forEachAttribute since each specialized InitialSpanInfo class can add
// different things through this method. Using the map version is not the same.
spanInfo.forEachAttribute((s, s2) -> span.setTag(TAG_CACHE.computeIfAbsent(s, TAG_ADDER), s2));
span.setTag("mule.correlation_id", event.getCorrelationId());
span.setTag(MULE_CORRELATION_ID, event.getCorrelationId());
// cache the resource name might be complex since it depends on a couple of keys
String extraDetail = null;
if (component != null) {
extraDetail = COMPONENT_DOC_CACHE.computeIfAbsent(component, COMPONENT_DOC_ADDER);
}
if (extraDetail == null) {
extraDetail = (String) span.getTag("mule.location");
extraDetail = (String) span.getTag(MULE_LOCATION);
}
if (extraDetail != null) {
span.setResourceName(spanInfo.getName() + " " + extraDetail);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -154,9 +154,6 @@ class MuleForkedTest extends WithHttpServer<MuleTestContainer> {

assertTraces(1) {
trace(4 + 3 * names.size(), new TreeComparator(trace(0))) { traceAssert ->
for (int i = 0; i < (4 + 3 * names.size()); i++) {
System.err.println(span(i))
}
span {
operationName operation()
resourceName "PUT /pfe-request"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -116,4 +116,7 @@ public class InstrumentationTags {

public static final String TIBCO_NODE = "tibco.node";
public static final String TIBCO_VERSION = "tibco.version";

public static final String MULE_CORRELATION_ID = "mule.correlation_id";
public static final String MULE_LOCATION = "mule.location";
}

0 comments on commit 50b9acd

Please sign in to comment.