Skip to content

Commit 73f194e

Browse files
amarzialiamarziali
authored
Reactor: early propagate span in context when subscribing (#8166)
* Reactor: early propagate span in context when subscribing * review
1 parent 9b219d4 commit 73f194e

File tree

5 files changed

+215
-28
lines changed

5 files changed

+215
-28
lines changed

‎dd-java-agent/instrumentation/reactor-core-3.1/src/latestDepTest/groovy/ReactorCoreTest.groovy

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,13 @@ import org.reactivestreams.Publisher
1515
import org.reactivestreams.Subscriber
1616
import org.reactivestreams.Subscription
1717
import reactor.core.publisher.Flux
18-
import reactor.core.publisher.Hooks
1918
import reactor.core.publisher.Mono
2019
import reactor.core.scheduler.Schedulers
2120
import reactor.util.context.Context
2221
import spock.lang.Shared
2322

2423
import java.time.Duration
24+
import java.util.concurrent.CompletableFuture
2525

2626
class ReactorCoreTest extends AgentTestRunner {
2727

@@ -443,12 +443,41 @@ class ReactorCoreTest extends AgentTestRunner {
443443

444444
def "test currentContext() calls on inner operator is not throwing a NPE on the advice"() {
445445
when:
446-
def mono = Flux.range(1, 100).windowUntil {it % 10 == 0}.count()
446+
def mono = Flux.range(1, 100).windowUntil { it % 10 == 0 }.count()
447447
then:
448448
// we are not interested into asserting a trace structure but only that the instrumentation error count is 0
449449
assert mono.block() == 10
450450
}
451451
452+
def "span in the context has to be activated when the publisher subscribes"() {
453+
when:
454+
// the mono is subscribed (block) when first is active.
455+
// However we expect that the span third will have second as parent and not first
456+
// because we set the parent explicitly in the reactor context (dd.span key)
457+
def result = runUnderTrace("first", {
458+
runUnderTrace("second", {
459+
def mono = Mono.defer {
460+
Mono.fromCompletionStage(CompletableFuture.supplyAsync {
461+
runUnderTrace("third", {
462+
"hello world"
463+
})
464+
})
465+
}.contextWrite(Context.of("dd.span", TEST_TRACER.activeSpan()))
466+
mono
467+
})
468+
.block()
469+
})
470+
then:
471+
assert result == "hello world"
472+
assertTraces(1, {
473+
trace(3, true) {
474+
basicSpan(it, "first")
475+
basicSpan(it, "second", span(0))
476+
basicSpan(it, "third", span(1))
477+
}
478+
})
479+
}
480+
452481
453482
@Trace(operationName = "trace-parent", resourceName = "trace-parent")
454483
def assemblePublisherUnderTrace(def publisherSupplier) {

‎dd-java-agent/instrumentation/reactor-core-3.1/src/main/java/datadog/trace/instrumentation/reactor/core/ContextSpanHelper.java

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
package datadog.trace.instrumentation.reactor.core;
2+
3+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
4+
import datadog.trace.bootstrap.instrumentation.api.WithAgentSpan;
5+
import javax.annotation.Nullable;
6+
import reactor.core.CoreSubscriber;
7+
import reactor.util.context.Context;
8+
9+
public class ContextSpanHelper {
10+
private static final String DD_SPAN_KEY = "dd.span";
11+
12+
private ContextSpanHelper() {}
13+
14+
@Nullable
15+
public static AgentSpan extractSpanFromSubscriberContext(final CoreSubscriber<?> subscriber) {
16+
if (subscriber == null) {
17+
return null;
18+
}
19+
Context context = null;
20+
try {
21+
context = subscriber.currentContext();
22+
} catch (Throwable ignored) {
23+
}
24+
if (context == null) {
25+
return null;
26+
}
27+
if (context.hasKey(DD_SPAN_KEY)) {
28+
Object maybeSpan = context.get(DD_SPAN_KEY);
29+
if (maybeSpan instanceof WithAgentSpan) {
30+
return ((WithAgentSpan) maybeSpan).asAgentSpan();
31+
}
32+
}
33+
return null;
34+
}
35+
}

‎dd-java-agent/instrumentation/reactor-core-3.1/src/main/java/datadog/trace/instrumentation/reactor/core/CorePublisherInstrumentation.java

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
package datadog.trace.instrumentation.reactor.core;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.hasSuperType;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
5+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
6+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
7+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
8+
import static datadog.trace.instrumentation.reactor.core.ContextSpanHelper.extractSpanFromSubscriberContext;
9+
import static net.bytebuddy.matcher.ElementMatchers.isStatic;
10+
import static net.bytebuddy.matcher.ElementMatchers.not;
11+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
12+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
13+
14+
import com.google.auto.service.AutoService;
15+
import datadog.trace.agent.tooling.Instrumenter;
16+
import datadog.trace.agent.tooling.InstrumenterModule;
17+
import datadog.trace.bootstrap.InstrumentationContext;
18+
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
19+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
20+
import java.util.HashMap;
21+
import java.util.Map;
22+
import net.bytebuddy.asm.Advice;
23+
import net.bytebuddy.description.type.TypeDescription;
24+
import net.bytebuddy.matcher.ElementMatcher;
25+
import org.reactivestreams.Publisher;
26+
import org.reactivestreams.Subscriber;
27+
import reactor.core.CoreSubscriber;
28+
29+
@AutoService(InstrumenterModule.class)
30+
public class CorePublisherInstrumentation extends InstrumenterModule.Tracing
31+
implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
32+
public CorePublisherInstrumentation() {
33+
super("reactor-core");
34+
}
35+
36+
@Override
37+
public String hierarchyMarkerType() {
38+
return "reactor.core.CoreSubscriber";
39+
}
40+
41+
@Override
42+
public ElementMatcher<TypeDescription> hierarchyMatcher() {
43+
return implementsInterface(named("reactor.core.CorePublisher")) // from 3.1.7
44+
.or(
45+
hasSuperType(
46+
namedOneOf(
47+
"reactor.core.publisher.Mono", "reactor.core.publisher.Flux"))); // < 3.1.7
48+
}
49+
50+
@Override
51+
public Map<String, String> contextStore() {
52+
final Map<String, String> ret = new HashMap<>();
53+
ret.put("org.reactivestreams.Subscriber", AgentSpan.class.getName());
54+
ret.put("org.reactivestreams.Publisher", AgentSpan.class.getName());
55+
return ret;
56+
}
57+
58+
@Override
59+
public String[] helperClassNames() {
60+
return new String[] {
61+
packageName + ".ContextSpanHelper",
62+
};
63+
}
64+
65+
@Override
66+
public void methodAdvice(MethodTransformer transformer) {
67+
transformer.applyAdvice(
68+
named("subscribe")
69+
.and(not(isStatic()))
70+
.and(takesArguments(1))
71+
.and(takesArgument(0, named("reactor.core.CoreSubscriber"))),
72+
getClass().getName() + "$PropagateContextSpanOnSubscribe");
73+
}
74+
75+
public static class PropagateContextSpanOnSubscribe {
76+
@Advice.OnMethodEnter(suppress = Throwable.class)
77+
public static AgentScope before(
78+
@Advice.This Publisher<?> self, @Advice.Argument(0) final CoreSubscriber<?> subscriber) {
79+
final AgentSpan span = extractSpanFromSubscriberContext(subscriber);
80+
81+
if (span != null) {
82+
// we force storing the span state linked to publisher and subscriber to the one explicitly
83+
// present in the context so that, if PublisherInstrumentation is kicking in after this
84+
// advice, it won't override that active span
85+
InstrumentationContext.get(Publisher.class, AgentSpan.class).put(self, span);
86+
InstrumentationContext.get(Subscriber.class, AgentSpan.class).put(subscriber, span);
87+
return activateSpan(span);
88+
}
89+
return null;
90+
}
91+
92+
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
93+
public static void after(@Advice.Enter final AgentScope scope) {
94+
if (scope != null) {
95+
scope.close();
96+
}
97+
}
98+
}
99+
}

‎dd-java-agent/instrumentation/reactor-core-3.1/src/main/java/datadog/trace/instrumentation/reactor/core/CoreSubscriberInstrumentation.java

Lines changed: 11 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -4,18 +4,17 @@
44
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
55
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
66
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
7+
import static datadog.trace.instrumentation.reactor.core.ContextSpanHelper.extractSpanFromSubscriberContext;
78

89
import com.google.auto.service.AutoService;
910
import datadog.trace.agent.tooling.Instrumenter;
1011
import datadog.trace.agent.tooling.InstrumenterModule;
1112
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1213
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
13-
import datadog.trace.bootstrap.instrumentation.api.WithAgentSpan;
1414
import net.bytebuddy.asm.Advice;
1515
import net.bytebuddy.description.type.TypeDescription;
1616
import net.bytebuddy.matcher.ElementMatcher;
1717
import reactor.core.CoreSubscriber;
18-
import reactor.util.context.Context;
1918

2019
@AutoService(InstrumenterModule.class)
2120
public class CoreSubscriberInstrumentation extends InstrumenterModule.Tracing
@@ -34,6 +33,13 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
3433
return implementsInterface(named(hierarchyMarkerType()));
3534
}
3635

36+
@Override
37+
public String[] helperClassNames() {
38+
return new String[] {
39+
packageName + ".ContextSpanHelper",
40+
};
41+
}
42+
3743
@Override
3844
public void methodAdvice(MethodTransformer transformer) {
3945
transformer.applyAdvice(
@@ -44,22 +50,9 @@ public void methodAdvice(MethodTransformer transformer) {
4450
public static class PropagateSpanInScopeAdvice {
4551
@Advice.OnMethodEnter(suppress = Throwable.class)
4652
public static AgentScope before(@Advice.This final CoreSubscriber<?> self) {
47-
Context context = null;
48-
try {
49-
context = self.currentContext();
50-
} catch (Throwable ignored) {
51-
}
52-
if (context == null) {
53-
return null;
54-
}
55-
if (context.hasKey("dd.span")) {
56-
Object maybeSpan = context.get("dd.span");
57-
if (maybeSpan instanceof WithAgentSpan) {
58-
AgentSpan span = ((WithAgentSpan) maybeSpan).asAgentSpan();
59-
if (span != null) {
60-
return activateSpan(span);
61-
}
62-
}
53+
final AgentSpan span = extractSpanFromSubscriberContext(self);
54+
if (span != null) {
55+
return activateSpan(span);
6356
}
6457
return null;
6558
}

‎dd-java-agent/instrumentation/reactor-core-3.1/src/test/groovy/ReactorCoreTest.groovy

Lines changed: 39 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import reactor.util.context.Context
2121
import spock.lang.Shared
2222

2323
import java.time.Duration
24+
import java.util.concurrent.CompletableFuture
2425

2526
class ReactorCoreTest extends AgentTestRunner {
2627

@@ -440,6 +441,44 @@ class ReactorCoreTest extends AgentTestRunner {
440441
})
441442
}
442443

444+
def "test currentContext() calls on inner operator is not throwing a NPE on the advice"() {
445+
when:
446+
def mono = Flux.range(1, 100).windowUntil { it % 10 == 0 }.count()
447+
then:
448+
// we are not interested into asserting a trace structure but only that the instrumentation error count is 0
449+
assert mono.block() == 11
450+
}
451+
452+
453+
def "span in the context has to be activated when the publisher subscribes"() {
454+
when:
455+
// the mono is subscribed (block) when first is active.
456+
// However we expect that the span third will have second as parent and not first
457+
// because we set the parent explicitly in the reactor context (dd.span key)
458+
def result = runUnderTrace("first", {
459+
runUnderTrace("second", {
460+
def mono = Mono.defer {
461+
Mono.fromCompletionStage(CompletableFuture.supplyAsync {
462+
runUnderTrace("third", {
463+
"hello world"
464+
})
465+
})
466+
}.subscriberContext(Context.of("dd.span", TEST_TRACER.activeSpan()))
467+
mono
468+
})
469+
.block()
470+
})
471+
then:
472+
assert result == "hello world"
473+
assertTraces(1, {
474+
trace(3, true) {
475+
basicSpan(it, "first")
476+
basicSpan(it, "second", span(0))
477+
basicSpan(it, "third", span(1))
478+
}
479+
})
480+
}
481+
443482
@Trace(operationName = "trace-parent", resourceName = "trace-parent")
444483
def assemblePublisherUnderTrace(def publisherSupplier) {
445484
def span = startSpan("publisher-parent")
@@ -490,14 +529,6 @@ class ReactorCoreTest extends AgentTestRunner {
490529
span.finish()
491530
}
492531
493-
def "test currentContext() calls on inner operator is not throwing a NPE on the advice"() {
494-
when:
495-
def mono = Flux.range(1, 100).windowUntil {it % 10 == 0}.count()
496-
then:
497-
// we are not interested into asserting a trace structure but only that the instrumentation error count is 0
498-
assert mono.block() == 11
499-
}
500-
501532
@Trace(operationName = "addOne", resourceName = "addOne")
502533
def static addOneFunc(int i) {
503534
return i + 1

0 commit comments

Comments
 (0)