Skip to content

Move rxjava2 advices to context tracking #10482

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
amarziali wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Move rxjava2 advices to context tracking #10482

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6fe6d04
Move rxjava2 advices to context tracking
amarziali Jan 30, 2026
fd432d2
beautify
amarziali Jan 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 16 additions & 36 deletions ...a-2.0/src/main/java/datadog/trace/instrumentation/rxjava2/CompletableInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,49 +1,28 @@
package datadog.trace.instrumentation.rxjava2;

import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
import static java.util.Collections.singletonMap;
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;

import com.google.auto.service.AutoService;
import datadog.context.Context;
import datadog.context.ContextScope;
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 datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge;
import io.reactivex.Completable;
import io.reactivex.CompletableObserver;
import java.util.Map;
import net.bytebuddy.asm.Advice;

@AutoService(InstrumenterModule.class)
public final class CompletableInstrumentation extends InstrumenterModule.Tracing
public final class CompletableInstrumentation
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
public CompletableInstrumentation() {
super("rxjava");
}

@Override
public String instrumentedType() {
return "io.reactivex.Completable";
}

@Override
public String[] helperClassNames() {
return new String[] {
packageName + ".TracingCompletableObserver",
};
}

@Override
public Map<String, String> contextStore() {
return singletonMap("io.reactivex.Completable", AgentSpan.class.getName());
}

@Override
public void methodAdvice(MethodTransformer transformer) {
transformer.applyAdvice(isConstructor(), getClass().getName() + "$CaptureParentSpanAdvice");
Expand All @@ -58,33 +37,34 @@ public void methodAdvice(MethodTransformer transformer) {
public static class CaptureParentSpanAdvice {
@Advice.OnMethodExit(suppress = Throwable.class)
public static void onConstruct(@Advice.This final Completable completable) {
AgentSpan parentSpan = activeSpan();
if (parentSpan != null) {
InstrumentationContext.get(Completable.class, AgentSpan.class).put(completable, parentSpan);
Context parentContext = Java8BytecodeBridge.getCurrentContext();
if (parentContext != null && parentContext != Java8BytecodeBridge.getRootContext()) {
InstrumentationContext.get(Completable.class, Context.class)
.put(completable, parentContext);
}
}
}

public static class PropagateParentSpanAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static AgentScope onSubscribe(
public static ContextScope onSubscribe(
@Advice.This final Completable completable,
@Advice.Argument(value = 0, readOnly = false) CompletableObserver observer) {
if (observer != null) {
AgentSpan parentSpan =
InstrumentationContext.get(Completable.class, AgentSpan.class).get(completable);
if (parentSpan != null) {
Context parentContext =
InstrumentationContext.get(Completable.class, Context.class).get(completable);
if (parentContext != null) {
// wrap the observer so spans from its events treat the captured span as their parent
observer = new TracingCompletableObserver(observer, parentSpan);
// activate the span here in case additional observers are created during subscribe
return activateSpan(parentSpan);
observer = new TracingCompletableObserver(observer, parentContext);
// attach the context here in case additional observers are created during subscribe
return parentContext.attach();
}
}
return null;
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void closeScope(@Advice.Enter final AgentScope scope) {
public static void closeScope(@Advice.Enter final ContextScope scope) {
if (scope != null) {
scope.close();
}
Expand Down
51 changes: 15 additions & 36 deletions ...java-2.0/src/main/java/datadog/trace/instrumentation/rxjava2/FlowableInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,49 +1,28 @@
package datadog.trace.instrumentation.rxjava2;

import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
import static java.util.Collections.singletonMap;
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;

import com.google.auto.service.AutoService;
import datadog.context.Context;
import datadog.context.ContextScope;
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 datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge;
import io.reactivex.Flowable;
import java.util.Map;
import net.bytebuddy.asm.Advice;
import org.reactivestreams.Subscriber;

@AutoService(InstrumenterModule.class)
public final class FlowableInstrumentation extends InstrumenterModule.Tracing
public final class FlowableInstrumentation
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
public FlowableInstrumentation() {
super("rxjava");
}

@Override
public String instrumentedType() {
return "io.reactivex.Flowable";
}

@Override
public String[] helperClassNames() {
return new String[] {
packageName + ".TracingSubscriber",
};
}

@Override
public Map<String, String> contextStore() {
return singletonMap("io.reactivex.Flowable", AgentSpan.class.getName());
}

@Override
public void methodAdvice(MethodTransformer transformer) {
transformer.applyAdvice(isConstructor(), getClass().getName() + "$CaptureParentSpanAdvice");
Expand All @@ -58,33 +37,33 @@ public void methodAdvice(MethodTransformer transformer) {
public static class CaptureParentSpanAdvice {
@Advice.OnMethodExit(suppress = Throwable.class)
public static void onConstruct(@Advice.This final Flowable<?> flowable) {
AgentSpan parentSpan = activeSpan();
if (parentSpan != null) {
InstrumentationContext.get(Flowable.class, AgentSpan.class).put(flowable, parentSpan);
Context parentContext = Java8BytecodeBridge.getCurrentContext();
if (parentContext != null && parentContext != Java8BytecodeBridge.getRootContext()) {
InstrumentationContext.get(Flowable.class, Context.class).put(flowable, parentContext);
}
}
}

public static class PropagateParentSpanAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static AgentScope onSubscribe(
public static ContextScope onSubscribe(
@Advice.This final Flowable<?> flowable,
@Advice.Argument(value = 0, readOnly = false) Subscriber<?> subscriber) {
if (subscriber != null) {
AgentSpan parentSpan =
InstrumentationContext.get(Flowable.class, AgentSpan.class).get(flowable);
if (parentSpan != null) {
Context parentContext =
InstrumentationContext.get(Flowable.class, Context.class).get(flowable);
if (parentContext != null) {
// wrap the subscriber so spans from its events treat the captured span as their parent
subscriber = new TracingSubscriber<>(subscriber, parentSpan);
// activate the span here in case additional subscribers are created during subscribe
return activateSpan(parentSpan);
subscriber = new TracingSubscriber<>(subscriber, parentContext);
// attach the context here in case additional subscribers are created during subscribe
return parentContext.attach();
}
}
return null;
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void closeScope(@Advice.Enter final AgentScope scope) {
public static void closeScope(@Advice.Enter final ContextScope scope) {
if (scope != null) {
scope.close();
}
Expand Down
50 changes: 14 additions & 36 deletions .../rxjava-2.0/src/main/java/datadog/trace/instrumentation/rxjava2/MaybeInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,49 +1,27 @@
package datadog.trace.instrumentation.rxjava2;

import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
import static java.util.Collections.singletonMap;
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;

import com.google.auto.service.AutoService;
import datadog.context.Context;
import datadog.context.ContextScope;
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 datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge;
import io.reactivex.Maybe;
import io.reactivex.MaybeObserver;
import java.util.Map;
import net.bytebuddy.asm.Advice;

@AutoService(InstrumenterModule.class)
public final class MaybeInstrumentation extends InstrumenterModule.Tracing
public final class MaybeInstrumentation
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
public MaybeInstrumentation() {
super("rxjava");
}

@Override
public String instrumentedType() {
return "io.reactivex.Maybe";
}

@Override
public String[] helperClassNames() {
return new String[] {
packageName + ".TracingMaybeObserver",
};
}

@Override
public Map<String, String> contextStore() {
return singletonMap("io.reactivex.Maybe", AgentSpan.class.getName());
}

@Override
public void methodAdvice(MethodTransformer transformer) {
transformer.applyAdvice(isConstructor(), getClass().getName() + "$CaptureParentSpanAdvice");
Expand All @@ -58,32 +36,32 @@ public void methodAdvice(MethodTransformer transformer) {
public static class CaptureParentSpanAdvice {
@Advice.OnMethodExit(suppress = Throwable.class)
public static void onConstruct(@Advice.This final Maybe<?> maybe) {
AgentSpan parentSpan = activeSpan();
if (parentSpan != null) {
InstrumentationContext.get(Maybe.class, AgentSpan.class).put(maybe, parentSpan);
Context parentContext = Java8BytecodeBridge.getCurrentContext();
if (parentContext != null && parentContext != Java8BytecodeBridge.getRootContext()) {
InstrumentationContext.get(Maybe.class, Context.class).put(maybe, parentContext);
}
}
}

public static class PropagateParentSpanAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static AgentScope onSubscribe(
public static ContextScope onSubscribe(
@Advice.This final Maybe<?> maybe,
@Advice.Argument(value = 0, readOnly = false) MaybeObserver<?> observer) {
if (observer != null) {
AgentSpan parentSpan = InstrumentationContext.get(Maybe.class, AgentSpan.class).get(maybe);
if (parentSpan != null) {
Context parentContext = InstrumentationContext.get(Maybe.class, Context.class).get(maybe);
if (parentContext != null) {
// wrap the observer so spans from its events treat the captured span as their parent
observer = new TracingMaybeObserver<>(observer, parentSpan);
// activate the span here in case additional observers are created during subscribe
return activateSpan(parentSpan);
observer = new TracingMaybeObserver<>(observer, parentContext);
// attach the context here in case additional observers are created during subscribe
return parentContext.attach();
}
}
return null;
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void closeScope(@Advice.Enter final AgentScope scope) {
public static void closeScope(@Advice.Enter final ContextScope scope) {
if (scope != null) {
scope.close();
}
Expand Down
52 changes: 15 additions & 37 deletions ...va-2.0/src/main/java/datadog/trace/instrumentation/rxjava2/ObservableInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,49 +1,27 @@
package datadog.trace.instrumentation.rxjava2;

import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
import static java.util.Collections.singletonMap;
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;

import com.google.auto.service.AutoService;
import datadog.context.Context;
import datadog.context.ContextScope;
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 datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge;
import io.reactivex.Observable;
import io.reactivex.Observer;
import java.util.Map;
import net.bytebuddy.asm.Advice;

@AutoService(InstrumenterModule.class)
public final class ObservableInstrumentation extends InstrumenterModule.Tracing
public final class ObservableInstrumentation
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
public ObservableInstrumentation() {
super("rxjava");
}

@Override
public String instrumentedType() {
return "io.reactivex.Observable";
}

@Override
public String[] helperClassNames() {
return new String[] {
packageName + ".TracingObserver",
};
}

@Override
public Map<String, String> contextStore() {
return singletonMap("io.reactivex.Observable", AgentSpan.class.getName());
}

@Override
public void methodAdvice(MethodTransformer transformer) {
transformer.applyAdvice(isConstructor(), getClass().getName() + "$CaptureParentSpanAdvice");
Expand All @@ -58,33 +36,33 @@ public void methodAdvice(MethodTransformer transformer) {
public static class CaptureParentSpanAdvice {
@Advice.OnMethodExit(suppress = Throwable.class)
public static void onConstruct(@Advice.This final Observable<?> observable) {
AgentSpan parentSpan = activeSpan();
if (parentSpan != null) {
InstrumentationContext.get(Observable.class, AgentSpan.class).put(observable, parentSpan);
Context parentContext = Java8BytecodeBridge.getCurrentContext();
if (parentContext != null) {
InstrumentationContext.get(Observable.class, Context.class).put(observable, parentContext);
}
}
}

public static class PropagateParentSpanAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static AgentScope onSubscribe(
public static ContextScope onSubscribe(
@Advice.This final Observable<?> observable,
@Advice.Argument(value = 0, readOnly = false) Observer<?> observer) {
if (observer != null) {
AgentSpan parentSpan =
InstrumentationContext.get(Observable.class, AgentSpan.class).get(observable);
if (parentSpan != null) {
Context parentContext =
InstrumentationContext.get(Observable.class, Context.class).get(observable);
if (parentContext != null && parentContext != Java8BytecodeBridge.getRootContext()) {
// wrap the observer so spans from its events treat the captured span as their parent
observer = new TracingObserver<>(observer, parentSpan);
// activate the span here in case additional observers are created during subscribe
return activateSpan(parentSpan);
observer = new TracingObserver<>(observer, parentContext);
// attach the context here in case additional observers are created during subscribe
return parentContext.attach();
}
}
return null;
}

@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void closeScope(@Advice.Enter final AgentScope scope) {
public static void closeScope(@Advice.Enter final ContextScope scope) {
if (scope != null) {
scope.close();
}
Expand Down
Loading