-
Notifications
You must be signed in to change notification settings - Fork 293
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Reactor: early propagate span in context when subscribing (#8166)
* Reactor: early propagate span in context when subscribing * review
- Loading branch information
Showing
5 changed files
with
215 additions
and
28 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
35 changes: 35 additions & 0 deletions
35
...-core-3.1/src/main/java/datadog/trace/instrumentation/reactor/core/ContextSpanHelper.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,35 @@ | ||
| package datadog.trace.instrumentation.reactor.core; | ||
|
|
||
| import datadog.trace.bootstrap.instrumentation.api.AgentSpan; | ||
| import datadog.trace.bootstrap.instrumentation.api.WithAgentSpan; | ||
| import javax.annotation.Nullable; | ||
| import reactor.core.CoreSubscriber; | ||
| import reactor.util.context.Context; | ||
|
|
||
| public class ContextSpanHelper { | ||
| private static final String DD_SPAN_KEY = "dd.span"; | ||
|
|
||
| private ContextSpanHelper() {} | ||
|
|
||
| @Nullable | ||
| public static AgentSpan extractSpanFromSubscriberContext(final CoreSubscriber<?> subscriber) { | ||
| if (subscriber == null) { | ||
| return null; | ||
| } | ||
| Context context = null; | ||
| try { | ||
| context = subscriber.currentContext(); | ||
| } catch (Throwable ignored) { | ||
| } | ||
| if (context == null) { | ||
| return null; | ||
| } | ||
| if (context.hasKey(DD_SPAN_KEY)) { | ||
| Object maybeSpan = context.get(DD_SPAN_KEY); | ||
| if (maybeSpan instanceof WithAgentSpan) { | ||
| return ((WithAgentSpan) maybeSpan).asAgentSpan(); | ||
| } | ||
| } | ||
| return null; | ||
| } | ||
| } |
99 changes: 99 additions & 0 deletions
99
...rc/main/java/datadog/trace/instrumentation/reactor/core/CorePublisherInstrumentation.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,99 @@ | ||
| package datadog.trace.instrumentation.reactor.core; | ||
|
|
||
| import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.hasSuperType; | ||
| import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; | ||
| 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 datadog.trace.instrumentation.reactor.core.ContextSpanHelper.extractSpanFromSubscriberContext; | ||
| import static net.bytebuddy.matcher.ElementMatchers.isStatic; | ||
| import static net.bytebuddy.matcher.ElementMatchers.not; | ||
| import static net.bytebuddy.matcher.ElementMatchers.takesArgument; | ||
| import static net.bytebuddy.matcher.ElementMatchers.takesArguments; | ||
|
|
||
| import com.google.auto.service.AutoService; | ||
| import datadog.trace.agent.tooling.Instrumenter; | ||
| import datadog.trace.agent.tooling.InstrumenterModule; | ||
| import datadog.trace.bootstrap.InstrumentationContext; | ||
| import datadog.trace.bootstrap.instrumentation.api.AgentScope; | ||
| import datadog.trace.bootstrap.instrumentation.api.AgentSpan; | ||
| import java.util.HashMap; | ||
| import java.util.Map; | ||
| import net.bytebuddy.asm.Advice; | ||
| import net.bytebuddy.description.type.TypeDescription; | ||
| import net.bytebuddy.matcher.ElementMatcher; | ||
| import org.reactivestreams.Publisher; | ||
| import org.reactivestreams.Subscriber; | ||
| import reactor.core.CoreSubscriber; | ||
|
|
||
| @AutoService(InstrumenterModule.class) | ||
| public class CorePublisherInstrumentation extends InstrumenterModule.Tracing | ||
| implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { | ||
| public CorePublisherInstrumentation() { | ||
| super("reactor-core"); | ||
| } | ||
|
|
||
| @Override | ||
| public String hierarchyMarkerType() { | ||
| return "reactor.core.CoreSubscriber"; | ||
| } | ||
|
|
||
| @Override | ||
| public ElementMatcher<TypeDescription> hierarchyMatcher() { | ||
| return implementsInterface(named("reactor.core.CorePublisher")) // from 3.1.7 | ||
| .or( | ||
| hasSuperType( | ||
| namedOneOf( | ||
| "reactor.core.publisher.Mono", "reactor.core.publisher.Flux"))); // < 3.1.7 | ||
| } | ||
|
|
||
| @Override | ||
| public Map<String, String> contextStore() { | ||
| final Map<String, String> ret = new HashMap<>(); | ||
| ret.put("org.reactivestreams.Subscriber", AgentSpan.class.getName()); | ||
| ret.put("org.reactivestreams.Publisher", AgentSpan.class.getName()); | ||
| return ret; | ||
| } | ||
|
|
||
| @Override | ||
| public String[] helperClassNames() { | ||
| return new String[] { | ||
| packageName + ".ContextSpanHelper", | ||
| }; | ||
| } | ||
|
|
||
| @Override | ||
| public void methodAdvice(MethodTransformer transformer) { | ||
| transformer.applyAdvice( | ||
| named("subscribe") | ||
| .and(not(isStatic())) | ||
| .and(takesArguments(1)) | ||
| .and(takesArgument(0, named("reactor.core.CoreSubscriber"))), | ||
| getClass().getName() + "$PropagateContextSpanOnSubscribe"); | ||
| } | ||
|
|
||
| public static class PropagateContextSpanOnSubscribe { | ||
| @Advice.OnMethodEnter(suppress = Throwable.class) | ||
| public static AgentScope before( | ||
| @Advice.This Publisher<?> self, @Advice.Argument(0) final CoreSubscriber<?> subscriber) { | ||
| final AgentSpan span = extractSpanFromSubscriberContext(subscriber); | ||
|
|
||
| if (span != null) { | ||
| // we force storing the span state linked to publisher and subscriber to the one explicitly | ||
| // present in the context so that, if PublisherInstrumentation is kicking in after this | ||
| // advice, it won't override that active span | ||
| InstrumentationContext.get(Publisher.class, AgentSpan.class).put(self, span); | ||
| InstrumentationContext.get(Subscriber.class, AgentSpan.class).put(subscriber, span); | ||
| return activateSpan(span); | ||
| } | ||
| return null; | ||
| } | ||
|
|
||
| @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) | ||
| public static void after(@Advice.Enter final AgentScope scope) { | ||
| if (scope != null) { | ||
| scope.close(); | ||
| } | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters