rxJava 2 instrumentation

dd-java-agent/instrumentation/rxjava-2/rxjava-2.gradle

@@ -0,0 +1,32 @@
+// Set properties before any plugins get loaded
+ext {
+  minJavaVersionForTests = JavaVersion.VERSION_1_8
+}
+
+muzzle {
+  pass {
+    group = "io.reactivex.rxjava2"
+    module = "rxjava"
+    versions = "[2.0.0,)"
+  }
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+
+apply plugin: 'org.unbroken-dome.test-sets'
+
+testSets {
+  latestDepTest {
+    dirName = 'test'
+  }
+}
+
+dependencies {
+  compileOnly group: 'org.reactivestreams', name: 'reactive-streams', version: '1.0.0'
+  compileOnly group: 'io.reactivex.rxjava2', name: 'rxjava', version: '2.0.0'
+
+  testCompile project(':dd-java-agent:instrumentation:trace-annotation')
+
+  testCompile group: 'io.reactivex.rxjava2', name: 'rxjava', version: '2.0.5'
+  latestDepTestCompile group: 'io.reactivex.rxjava2', name: 'rxjava', version: '+'
+}

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

@@ -0,0 +1,97 @@
+package datadog.trace.instrumentation.rxjava2;
+
+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.named;
+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.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import io.reactivex.Completable;
+import io.reactivex.CompletableObserver;
+import java.util.HashMap;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public final class CompletableInstrumentation extends Instrumenter.Default {
+  public CompletableInstrumentation() {
+    super("rxjava");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("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 Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
+    transformers.put(isConstructor(), getClass().getName() + "$CaptureParentSpanAdvice");
+    transformers.put(
+        isMethod()
+            .and(named("subscribe"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, named("io.reactivex.CompletableObserver"))),
+        getClass().getName() + "$PropagateParentSpanAdvice");
+    return transformers;
+  }
+
+  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);
+      }
+    }
+  }
+
+  public static class PropagateParentSpanAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope 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) {
+          // 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);
+        }
+      }
+      return null;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void closeScope(@Advice.Enter final AgentScope scope) {
+      if (scope != null) {
+        scope.close();
+      }
+    }
+  }
+}

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

@@ -0,0 +1,97 @@
+package datadog.trace.instrumentation.rxjava2;
+
+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.named;
+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.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import io.reactivex.Flowable;
+import java.util.HashMap;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.reactivestreams.Subscriber;
+
+@AutoService(Instrumenter.class)
+public final class FlowableInstrumentation extends Instrumenter.Default {
+  public FlowableInstrumentation() {
+    super("rxjava");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("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 Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
+    transformers.put(isConstructor(), getClass().getName() + "$CaptureParentSpanAdvice");
+    transformers.put(
+        isMethod()
+            .and(named("subscribe"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, named("org.reactivestreams.Subscriber"))),
+        getClass().getName() + "$PropagateParentSpanAdvice");
+    return transformers;
+  }
+
+  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);
+      }
+    }
+  }
+
+  public static class PropagateParentSpanAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope 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) {
+          // 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);
+        }
+      }
+      return null;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void closeScope(@Advice.Enter final AgentScope scope) {
+      if (scope != null) {
+        scope.close();
+      }
+    }
+  }
+}

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

@@ -0,0 +1,96 @@
+package datadog.trace.instrumentation.rxjava2;
+
+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.named;
+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.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import io.reactivex.Maybe;
+import io.reactivex.MaybeObserver;
+import java.util.HashMap;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public final class MaybeInstrumentation extends Instrumenter.Default {
+  public MaybeInstrumentation() {
+    super("rxjava");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("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 Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
+    transformers.put(isConstructor(), getClass().getName() + "$CaptureParentSpanAdvice");
+    transformers.put(
+        isMethod()
+            .and(named("subscribe"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, named("io.reactivex.MaybeObserver"))),
+        getClass().getName() + "$PropagateParentSpanAdvice");
+    return transformers;
+  }
+
+  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);
+      }
+    }
+  }
+
+  public static class PropagateParentSpanAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope 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) {
+          // 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);
+        }
+      }
+      return null;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void closeScope(@Advice.Enter final AgentScope scope) {
+      if (scope != null) {
+        scope.close();
+      }
+    }
+  }
+}