Skip to content

Commit b83178b

Browse files
amarzialiamarziali
authored
Move rxjava2 advices to context tracking (#10482)
Move rxjava2 advices to context tracking beautify Co-authored-by: andrea.marziali <andrea.marziali@datadoghq.com>
1 parent d8f2599 commit b83178b

12 files changed

Lines changed: 170 additions & 262 deletions

‎dd-java-agent/instrumentation/rxjava/rxjava-2.0/src/main/java/datadog/trace/instrumentation/rxjava2/CompletableInstrumentation.java‎

Lines changed: 16 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,49 +1,28 @@
11
package datadog.trace.instrumentation.rxjava2;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
6-
import static java.util.Collections.singletonMap;
74
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
85
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
96
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
107
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
118

12-
import com.google.auto.service.AutoService;
9+
import datadog.context.Context;
10+
import datadog.context.ContextScope;
1311
import datadog.trace.agent.tooling.Instrumenter;
14-
import datadog.trace.agent.tooling.InstrumenterModule;
1512
import datadog.trace.bootstrap.InstrumentationContext;
16-
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
17-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
13+
import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge;
1814
import io.reactivex.Completable;
1915
import io.reactivex.CompletableObserver;
20-
import java.util.Map;
2116
import net.bytebuddy.asm.Advice;
2217

23-
@AutoService(InstrumenterModule.class)
24-
public final class CompletableInstrumentation extends InstrumenterModule.Tracing
18+
public final class CompletableInstrumentation
2519
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
26-
public CompletableInstrumentation() {
27-
super("rxjava");
28-
}
2920

3021
@Override
3122
public String instrumentedType() {
3223
return "io.reactivex.Completable";
3324
}
3425

35-
@Override
36-
public String[] helperClassNames() {
37-
return new String[] {
38-
packageName + ".TracingCompletableObserver",
39-
};
40-
}
41-
42-
@Override
43-
public Map<String, String> contextStore() {
44-
return singletonMap("io.reactivex.Completable", AgentSpan.class.getName());
45-
}
46-
4726
@Override
4827
public void methodAdvice(MethodTransformer transformer) {
4928
transformer.applyAdvice(isConstructor(), getClass().getName() + "$CaptureParentSpanAdvice");
@@ -58,33 +37,34 @@ public void methodAdvice(MethodTransformer transformer) {
5837
public static class CaptureParentSpanAdvice {
5938
@Advice.OnMethodExit(suppress = Throwable.class)
6039
public static void onConstruct(@Advice.This final Completable completable) {
61-
AgentSpan parentSpan = activeSpan();
62-
if (parentSpan != null) {
63-
InstrumentationContext.get(Completable.class, AgentSpan.class).put(completable, parentSpan);
40+
Context parentContext = Java8BytecodeBridge.getCurrentContext();
41+
if (parentContext != null && parentContext != Java8BytecodeBridge.getRootContext()) {
42+
InstrumentationContext.get(Completable.class, Context.class)
43+
.put(completable, parentContext);
6444
}
6545
}
6646
}
6747

6848
public static class PropagateParentSpanAdvice {
6949
@Advice.OnMethodEnter(suppress = Throwable.class)
70-
public static AgentScope onSubscribe(
50+
public static ContextScope onSubscribe(
7151
@Advice.This final Completable completable,
7252
@Advice.Argument(value = 0, readOnly = false) CompletableObserver observer) {
7353
if (observer != null) {
74-
AgentSpan parentSpan =
75-
InstrumentationContext.get(Completable.class, AgentSpan.class).get(completable);
76-
if (parentSpan != null) {
54+
Context parentContext =
55+
InstrumentationContext.get(Completable.class, Context.class).get(completable);
56+
if (parentContext != null) {
7757
// wrap the observer so spans from its events treat the captured span as their parent
78-
observer = new TracingCompletableObserver(observer, parentSpan);
79-
// activate the span here in case additional observers are created during subscribe
80-
return activateSpan(parentSpan);
58+
observer = new TracingCompletableObserver(observer, parentContext);
59+
// attach the context here in case additional observers are created during subscribe
60+
return parentContext.attach();
8161
}
8262
}
8363
return null;
8464
}
8565

8666
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
87-
public static void closeScope(@Advice.Enter final AgentScope scope) {
67+
public static void closeScope(@Advice.Enter final ContextScope scope) {
8868
if (scope != null) {
8969
scope.close();
9070
}

‎dd-java-agent/instrumentation/rxjava/rxjava-2.0/src/main/java/datadog/trace/instrumentation/rxjava2/FlowableInstrumentation.java‎

Lines changed: 15 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,49 +1,28 @@
11
package datadog.trace.instrumentation.rxjava2;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
6-
import static java.util.Collections.singletonMap;
74
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
85
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
96
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
107
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
118

12-
import com.google.auto.service.AutoService;
9+
import datadog.context.Context;
10+
import datadog.context.ContextScope;
1311
import datadog.trace.agent.tooling.Instrumenter;
14-
import datadog.trace.agent.tooling.InstrumenterModule;
1512
import datadog.trace.bootstrap.InstrumentationContext;
16-
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
17-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
13+
import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge;
1814
import io.reactivex.Flowable;
19-
import java.util.Map;
2015
import net.bytebuddy.asm.Advice;
2116
import org.reactivestreams.Subscriber;
2217

23-
@AutoService(InstrumenterModule.class)
24-
public final class FlowableInstrumentation extends InstrumenterModule.Tracing
18+
public final class FlowableInstrumentation
2519
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
26-
public FlowableInstrumentation() {
27-
super("rxjava");
28-
}
2920

3021
@Override
3122
public String instrumentedType() {
3223
return "io.reactivex.Flowable";
3324
}
3425

35-
@Override
36-
public String[] helperClassNames() {
37-
return new String[] {
38-
packageName + ".TracingSubscriber",
39-
};
40-
}
41-
42-
@Override
43-
public Map<String, String> contextStore() {
44-
return singletonMap("io.reactivex.Flowable", AgentSpan.class.getName());
45-
}
46-
4726
@Override
4827
public void methodAdvice(MethodTransformer transformer) {
4928
transformer.applyAdvice(isConstructor(), getClass().getName() + "$CaptureParentSpanAdvice");
@@ -58,33 +37,33 @@ public void methodAdvice(MethodTransformer transformer) {
5837
public static class CaptureParentSpanAdvice {
5938
@Advice.OnMethodExit(suppress = Throwable.class)
6039
public static void onConstruct(@Advice.This final Flowable<?> flowable) {
61-
AgentSpan parentSpan = activeSpan();
62-
if (parentSpan != null) {
63-
InstrumentationContext.get(Flowable.class, AgentSpan.class).put(flowable, parentSpan);
40+
Context parentContext = Java8BytecodeBridge.getCurrentContext();
41+
if (parentContext != null && parentContext != Java8BytecodeBridge.getRootContext()) {
42+
InstrumentationContext.get(Flowable.class, Context.class).put(flowable, parentContext);
6443
}
6544
}
6645
}
6746

6847
public static class PropagateParentSpanAdvice {
6948
@Advice.OnMethodEnter(suppress = Throwable.class)
70-
public static AgentScope onSubscribe(
49+
public static ContextScope onSubscribe(
7150
@Advice.This final Flowable<?> flowable,
7251
@Advice.Argument(value = 0, readOnly = false) Subscriber<?> subscriber) {
7352
if (subscriber != null) {
74-
AgentSpan parentSpan =
75-
InstrumentationContext.get(Flowable.class, AgentSpan.class).get(flowable);
76-
if (parentSpan != null) {
53+
Context parentContext =
54+
InstrumentationContext.get(Flowable.class, Context.class).get(flowable);
55+
if (parentContext != null) {
7756
// wrap the subscriber so spans from its events treat the captured span as their parent
78-
subscriber = new TracingSubscriber<>(subscriber, parentSpan);
79-
// activate the span here in case additional subscribers are created during subscribe
80-
return activateSpan(parentSpan);
57+
subscriber = new TracingSubscriber<>(subscriber, parentContext);
58+
// attach the context here in case additional subscribers are created during subscribe
59+
return parentContext.attach();
8160
}
8261
}
8362
return null;
8463
}
8564

8665
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
87-
public static void closeScope(@Advice.Enter final AgentScope scope) {
66+
public static void closeScope(@Advice.Enter final ContextScope scope) {
8867
if (scope != null) {
8968
scope.close();
9069
}

‎dd-java-agent/instrumentation/rxjava/rxjava-2.0/src/main/java/datadog/trace/instrumentation/rxjava2/MaybeInstrumentation.java‎

Lines changed: 14 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,49 +1,27 @@
11
package datadog.trace.instrumentation.rxjava2;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
6-
import static java.util.Collections.singletonMap;
74
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
85
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
96
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
107
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
118

12-
import com.google.auto.service.AutoService;
9+
import datadog.context.Context;
10+
import datadog.context.ContextScope;
1311
import datadog.trace.agent.tooling.Instrumenter;
14-
import datadog.trace.agent.tooling.InstrumenterModule;
1512
import datadog.trace.bootstrap.InstrumentationContext;
16-
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
17-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
13+
import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge;
1814
import io.reactivex.Maybe;
1915
import io.reactivex.MaybeObserver;
20-
import java.util.Map;
2116
import net.bytebuddy.asm.Advice;
2217

23-
@AutoService(InstrumenterModule.class)
24-
public final class MaybeInstrumentation extends InstrumenterModule.Tracing
18+
public final class MaybeInstrumentation
2519
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
26-
public MaybeInstrumentation() {
27-
super("rxjava");
28-
}
29-
3020
@Override
3121
public String instrumentedType() {
3222
return "io.reactivex.Maybe";
3323
}
3424

35-
@Override
36-
public String[] helperClassNames() {
37-
return new String[] {
38-
packageName + ".TracingMaybeObserver",
39-
};
40-
}
41-
42-
@Override
43-
public Map<String, String> contextStore() {
44-
return singletonMap("io.reactivex.Maybe", AgentSpan.class.getName());
45-
}
46-
4725
@Override
4826
public void methodAdvice(MethodTransformer transformer) {
4927
transformer.applyAdvice(isConstructor(), getClass().getName() + "$CaptureParentSpanAdvice");
@@ -58,32 +36,32 @@ public void methodAdvice(MethodTransformer transformer) {
5836
public static class CaptureParentSpanAdvice {
5937
@Advice.OnMethodExit(suppress = Throwable.class)
6038
public static void onConstruct(@Advice.This final Maybe<?> maybe) {
61-
AgentSpan parentSpan = activeSpan();
62-
if (parentSpan != null) {
63-
InstrumentationContext.get(Maybe.class, AgentSpan.class).put(maybe, parentSpan);
39+
Context parentContext = Java8BytecodeBridge.getCurrentContext();
40+
if (parentContext != null && parentContext != Java8BytecodeBridge.getRootContext()) {
41+
InstrumentationContext.get(Maybe.class, Context.class).put(maybe, parentContext);
6442
}
6543
}
6644
}
6745

6846
public static class PropagateParentSpanAdvice {
6947
@Advice.OnMethodEnter(suppress = Throwable.class)
70-
public static AgentScope onSubscribe(
48+
public static ContextScope onSubscribe(
7149
@Advice.This final Maybe<?> maybe,
7250
@Advice.Argument(value = 0, readOnly = false) MaybeObserver<?> observer) {
7351
if (observer != null) {
74-
AgentSpan parentSpan = InstrumentationContext.get(Maybe.class, AgentSpan.class).get(maybe);
75-
if (parentSpan != null) {
52+
Context parentContext = InstrumentationContext.get(Maybe.class, Context.class).get(maybe);
53+
if (parentContext != null) {
7654
// wrap the observer so spans from its events treat the captured span as their parent
77-
observer = new TracingMaybeObserver<>(observer, parentSpan);
78-
// activate the span here in case additional observers are created during subscribe
79-
return activateSpan(parentSpan);
55+
observer = new TracingMaybeObserver<>(observer, parentContext);
56+
// attach the context here in case additional observers are created during subscribe
57+
return parentContext.attach();
8058
}
8159
}
8260
return null;
8361
}
8462

8563
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
86-
public static void closeScope(@Advice.Enter final AgentScope scope) {
64+
public static void closeScope(@Advice.Enter final ContextScope scope) {
8765
if (scope != null) {
8866
scope.close();
8967
}

‎dd-java-agent/instrumentation/rxjava/rxjava-2.0/src/main/java/datadog/trace/instrumentation/rxjava2/ObservableInstrumentation.java‎

Lines changed: 15 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,49 +1,27 @@
11
package datadog.trace.instrumentation.rxjava2;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
6-
import static java.util.Collections.singletonMap;
74
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
85
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
96
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
107
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
118

12-
import com.google.auto.service.AutoService;
9+
import datadog.context.Context;
10+
import datadog.context.ContextScope;
1311
import datadog.trace.agent.tooling.Instrumenter;
14-
import datadog.trace.agent.tooling.InstrumenterModule;
1512
import datadog.trace.bootstrap.InstrumentationContext;
16-
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
17-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
13+
import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge;
1814
import io.reactivex.Observable;
1915
import io.reactivex.Observer;
20-
import java.util.Map;
2116
import net.bytebuddy.asm.Advice;
2217

23-
@AutoService(InstrumenterModule.class)
24-
public final class ObservableInstrumentation extends InstrumenterModule.Tracing
18+
public final class ObservableInstrumentation
2519
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
26-
public ObservableInstrumentation() {
27-
super("rxjava");
28-
}
29-
3020
@Override
3121
public String instrumentedType() {
3222
return "io.reactivex.Observable";
3323
}
3424

35-
@Override
36-
public String[] helperClassNames() {
37-
return new String[] {
38-
packageName + ".TracingObserver",
39-
};
40-
}
41-
42-
@Override
43-
public Map<String, String> contextStore() {
44-
return singletonMap("io.reactivex.Observable", AgentSpan.class.getName());
45-
}
46-
4725
@Override
4826
public void methodAdvice(MethodTransformer transformer) {
4927
transformer.applyAdvice(isConstructor(), getClass().getName() + "$CaptureParentSpanAdvice");
@@ -58,33 +36,33 @@ public void methodAdvice(MethodTransformer transformer) {
5836
public static class CaptureParentSpanAdvice {
5937
@Advice.OnMethodExit(suppress = Throwable.class)
6038
public static void onConstruct(@Advice.This final Observable<?> observable) {
61-
AgentSpan parentSpan = activeSpan();
62-
if (parentSpan != null) {
63-
InstrumentationContext.get(Observable.class, AgentSpan.class).put(observable, parentSpan);
39+
Context parentContext = Java8BytecodeBridge.getCurrentContext();
40+
if (parentContext != null) {
41+
InstrumentationContext.get(Observable.class, Context.class).put(observable, parentContext);
6442
}
6543
}
6644
}
6745

6846
public static class PropagateParentSpanAdvice {
6947
@Advice.OnMethodEnter(suppress = Throwable.class)
70-
public static AgentScope onSubscribe(
48+
public static ContextScope onSubscribe(
7149
@Advice.This final Observable<?> observable,
7250
@Advice.Argument(value = 0, readOnly = false) Observer<?> observer) {
7351
if (observer != null) {
74-
AgentSpan parentSpan =
75-
InstrumentationContext.get(Observable.class, AgentSpan.class).get(observable);
76-
if (parentSpan != null) {
52+
Context parentContext =
53+
InstrumentationContext.get(Observable.class, Context.class).get(observable);
54+
if (parentContext != null && parentContext != Java8BytecodeBridge.getRootContext()) {
7755
// wrap the observer so spans from its events treat the captured span as their parent
78-
observer = new TracingObserver<>(observer, parentSpan);
79-
// activate the span here in case additional observers are created during subscribe
80-
return activateSpan(parentSpan);
56+
observer = new TracingObserver<>(observer, parentContext);
57+
// attach the context here in case additional observers are created during subscribe
58+
return parentContext.attach();
8159
}
8260
}
8361
return null;
8462
}
8563

8664
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
87-
public static void closeScope(@Advice.Enter final AgentScope scope) {
65+
public static void closeScope(@Advice.Enter final ContextScope scope) {
8866
if (scope != null) {
8967
scope.close();
9068
}

0 commit comments

Comments
 (0)