Revert "Couchbase Async Subscriptions"

dd-java-agent/instrumentation/couchbase-2.0/couchbase-2.0.gradle

@@ -44,10 +44,8 @@ muzzle {
 }
 
 dependencies {
-  compile project(':dd-java-agent:instrumentation:rxjava-1')
-
   compileOnly group: 'com.couchbase.client', name: 'java-client', version: '2.0.0'
-  
+
   testCompile group: 'com.couchbase.mock', name: 'CouchbaseMock', version: '1.5.19'
 
   testCompile group: 'org.springframework.data', name: 'spring-data-couchbase', version: '2.0.0.RELEASE'

dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseBucketInstrumentation.java

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.couchbase.client;
 
+import static datadog.trace.instrumentation.couchbase.client.CouchbaseClientDecorator.DECORATE;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -11,14 +12,22 @@
 import com.couchbase.client.java.CouchbaseCluster;
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.api.DDTags;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import io.opentracing.Scope;
+import io.opentracing.Span;
+import io.opentracing.noop.NoopSpan;
+import io.opentracing.util.GlobalTracer;
 import java.lang.reflect.Method;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import rx.Observable;
+import rx.functions.Action0;
+import rx.functions.Action1;
 
 @AutoService(Instrumenter.class)
 public class CouchbaseBucketInstrumentation extends Instrumenter.Default {
@@ -38,15 +47,13 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      "rx.DDTracingUtil",
       "datadog.trace.agent.decorator.BaseDecorator",
       "datadog.trace.agent.decorator.ClientDecorator",
       "datadog.trace.agent.decorator.DatabaseClientDecorator",
-      "datadog.trace.instrumentation.rxjava.SpanFinishingSubscription",
-      "datadog.trace.instrumentation.rxjava.TracedSubscriber",
-      "datadog.trace.instrumentation.rxjava.TracedOnSubscribe",
       packageName + ".CouchbaseClientDecorator",
-      packageName + ".CouchbaseOnSubscribe",
+      getClass().getName() + "$TraceSpanStart",
+      getClass().getName() + "$TraceSpanFinish",
+      getClass().getName() + "$TraceSpanError",
     };
   }
 
@@ -69,13 +76,94 @@ public static void subscribeResult(
         @Advice.Enter final int callDepth,
         @Advice.Origin final Method method,
         @Advice.FieldValue("bucket") final String bucket,
+        @Advice.AllArguments final Object[] args,
         @Advice.Return(readOnly = false) Observable result) {
       if (callDepth > 0) {
         return;
       }
       CallDepthThreadLocalMap.reset(CouchbaseCluster.class);
+      final AtomicReference<Span> spanRef = new AtomicReference<>();
+      result =
+          result
+              .doOnSubscribe(new TraceSpanStart(method, bucket, spanRef))
+              .doOnCompleted(new TraceSpanFinish(spanRef))
+              .doOnError(new TraceSpanError(spanRef));
+    }
+  }
+
+  public static class TraceSpanStart implements Action0 {
+    private final Method method;
+    private final String bucket;
+    private final AtomicReference<Span> spanRef;
+
+    public TraceSpanStart(
+        final Method method, final String bucket, final AtomicReference<Span> spanRef) {
+      this.method = method;
+      this.bucket = bucket;
+      this.spanRef = spanRef;
+    }
+
+    @Override
+    public void call() {
+      // This is called each time an observer has a new subscriber, but we should only time it once.
+      if (!spanRef.compareAndSet(null, NoopSpan.INSTANCE)) {
+        return;
+      }
+      final Class<?> declaringClass = method.getDeclaringClass();
+      final String className =
+          declaringClass.getSimpleName().replace("CouchbaseAsync", "").replace("DefaultAsync", "");
+      final String resourceName = className + "." + method.getName();
+
+      final Span span =
+          GlobalTracer.get()
+              .buildSpan("couchbase.call")
+              .withTag(DDTags.RESOURCE_NAME, resourceName)
+              .withTag("bucket", bucket)
+              .start();
+      try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, false)) {
+        // just replace the no-op span.
+        spanRef.set(DECORATE.afterStart(span));
+      }
+    }
+  }
 
-      result = Observable.create(new CouchbaseOnSubscribe(result, method, bucket));
+  public static class TraceSpanFinish implements Action0 {
+    private final AtomicReference<Span> spanRef;
+
+    public TraceSpanFinish(final AtomicReference<Span> spanRef) {
+      this.spanRef = spanRef;
+    }
+
+    @Override
+    public void call() {
+      final Span span = spanRef.getAndSet(null);
+
+      if (span != null) {
+        try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, false)) {
+          DECORATE.beforeFinish(span);
+          span.finish();
+        }
+      }
+    }
+  }
+
+  public static class TraceSpanError implements Action1<Throwable> {
+    private final AtomicReference<Span> spanRef;
+
+    public TraceSpanError(final AtomicReference<Span> spanRef) {
+      this.spanRef = spanRef;
+    }
+
+    @Override
+    public void call(final Throwable throwable) {
+      final Span span = spanRef.getAndSet(null);
+      if (span != null) {
+        try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, false)) {
+          DECORATE.onError(span, throwable);
+          DECORATE.beforeFinish(span);
+          span.finish();
+        }
+      }
     }
   }
 }

dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseClusterInstrumentation.java

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.couchbase.client;
 
+import static datadog.trace.instrumentation.couchbase.client.CouchbaseClientDecorator.DECORATE;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -11,14 +12,22 @@
 import com.couchbase.client.java.CouchbaseCluster;
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.api.DDTags;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import io.opentracing.Scope;
+import io.opentracing.Span;
+import io.opentracing.noop.NoopSpan;
+import io.opentracing.util.GlobalTracer;
 import java.lang.reflect.Method;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import rx.Observable;
+import rx.functions.Action0;
+import rx.functions.Action1;
 
 @AutoService(Instrumenter.class)
 public class CouchbaseClusterInstrumentation extends Instrumenter.Default {
@@ -38,22 +47,20 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      "rx.DDTracingUtil",
       "datadog.trace.agent.decorator.BaseDecorator",
       "datadog.trace.agent.decorator.ClientDecorator",
       "datadog.trace.agent.decorator.DatabaseClientDecorator",
-      "datadog.trace.instrumentation.rxjava.SpanFinishingSubscription",
-      "datadog.trace.instrumentation.rxjava.TracedSubscriber",
-      "datadog.trace.instrumentation.rxjava.TracedOnSubscribe",
       packageName + ".CouchbaseClientDecorator",
-      packageName + ".CouchbaseOnSubscribe",
+      getClass().getName() + "$TraceSpanStart",
+      getClass().getName() + "$TraceSpanFinish",
+      getClass().getName() + "$TraceSpanError",
     };
   }
 
   @Override
   public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
     return singletonMap(
-        isMethod().and(isPublic()).and(returns(named("rx.Observable"))).and(not(named("core"))),
+        isMethod().and(isPublic()).and(returns(named("rx.Observable"))),
         CouchbaseClientAdvice.class.getName());
   }
 
@@ -73,7 +80,84 @@ public static void subscribeResult(
         return;
       }
       CallDepthThreadLocalMap.reset(CouchbaseCluster.class);
-      result = Observable.create(new CouchbaseOnSubscribe(result, method, null));
+      final AtomicReference<Span> spanRef = new AtomicReference<>();
+      result =
+          result
+              .doOnSubscribe(new TraceSpanStart(method, spanRef))
+              .doOnCompleted(new TraceSpanFinish(spanRef))
+              .doOnError(new TraceSpanError(spanRef));
+    }
+  }
+
+  public static class TraceSpanStart implements Action0 {
+    private final Method method;
+    private final AtomicReference<Span> spanRef;
+
+    public TraceSpanStart(final Method method, final AtomicReference<Span> spanRef) {
+      this.method = method;
+      this.spanRef = spanRef;
+    }
+
+    @Override
+    public void call() {
+      // This is called each time an observer has a new subscriber, but we should only time it once.
+      if (!spanRef.compareAndSet(null, NoopSpan.INSTANCE)) {
+        return;
+      }
+      final Class<?> declaringClass = method.getDeclaringClass();
+      final String className =
+          declaringClass.getSimpleName().replace("CouchbaseAsync", "").replace("DefaultAsync", "");
+      final String resourceName = className + "." + method.getName();
+
+      final Span span =
+          GlobalTracer.get()
+              .buildSpan("couchbase.call")
+              .withTag(DDTags.RESOURCE_NAME, resourceName)
+              .start();
+      try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, false)) {
+        // just replace the no-op span.
+        spanRef.set(DECORATE.afterStart(scope.span()));
+      }
+    }
+  }
+
+  public static class TraceSpanFinish implements Action0 {
+    private final AtomicReference<Span> spanRef;
+
+    public TraceSpanFinish(final AtomicReference<Span> spanRef) {
+      this.spanRef = spanRef;
+    }
+
+    @Override
+    public void call() {
+      final Span span = spanRef.getAndSet(null);
+
+      if (span != null) {
+        try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, false)) {
+          DECORATE.beforeFinish(span);
+          span.finish();
+        }
+      }
+    }
+  }
+
+  public static class TraceSpanError implements Action1<Throwable> {
+    private final AtomicReference<Span> spanRef;
+
+    public TraceSpanError(final AtomicReference<Span> spanRef) {
+      this.spanRef = spanRef;
+    }
+
+    @Override
+    public void call(final Throwable throwable) {
+      final Span span = spanRef.getAndSet(null);
+      if (span != null) {
+        try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, false)) {
+          DECORATE.onError(span, throwable);
+          DECORATE.beforeFinish(span);
+          span.finish();
+        }
+      }
     }
   }
 }