Multi-tracing support for Couchbase 3.2+

dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/CoreEnvironmentBuilderInstrumentation.java

@@ -1,6 +1,8 @@
 package datadog.trace.instrumentation.couchbase_32.client;
 
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
@@ -23,6 +25,8 @@ public String[] helperClassNames() {
       packageName + ".DatadogRequestSpan",
       packageName + ".DatadogRequestSpan$1",
       packageName + ".DatadogRequestTracer",
+      packageName + ".DelegatingRequestSpan",
+      packageName + ".DelegatingRequestTracer"
     };
   }
 
@@ -39,5 +43,8 @@ public String instrumentedType() {
   @Override
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(isConstructor(), packageName + ".CoreEnvironmentBuilderAdvice");
+    transformer.applyAdvice(
+        isMethod().and(named("requestTracer")),
+        packageName + ".CoreEnvironmentBuilderRequestTracerAdvice");
   }
 }

dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/CoreEnvironmentBuilderRequestTracerAdvice.java

@@ -0,0 +1,38 @@
+package datadog.trace.instrumentation.couchbase_32.client;
+
+import com.couchbase.client.core.Core;
+import com.couchbase.client.core.cnc.RequestTracer;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import net.bytebuddy.asm.Advice;
+
+public class CoreEnvironmentBuilderRequestTracerAdvice {
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static void onEnter(
+      @Advice.Argument(value = 0, readOnly = false) RequestTracer requestTracer) {
+
+    // already a delegating tracer
+    if (requestTracer instanceof DelegatingRequestTracer) {
+      return;
+    }
+
+    // already a datadog tracer
+    if (requestTracer instanceof DatadogRequestTracer) {
+      return;
+    }
+
+    ContextStore<Core, String> coreContext = InstrumentationContext.get(Core.class, String.class);
+
+    DatadogRequestTracer datadogTracer = new DatadogRequestTracer(AgentTracer.get(), coreContext);
+
+    // if the app didn't set a custom tracer, use only datadog tracer
+    if (requestTracer == null) {
+      requestTracer = datadogTracer;
+      return;
+    }
+
+    // Wrap custom datadog and cnc tracers into a delegating
+    requestTracer = new DelegatingRequestTracer(datadogTracer, requestTracer);
+  }
+}

dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/DelegatingRequestSpan.java

@@ -0,0 +1,72 @@
+package datadog.trace.instrumentation.couchbase_32.client;
+
+import com.couchbase.client.core.cnc.RequestSpan;
+import com.couchbase.client.core.msg.RequestContext;
+import java.time.Instant;
+import javax.annotation.Nonnull;
+
+/** RequestSpan, which delegates all calls to two other RequestSpans */
+public class DelegatingRequestSpan implements RequestSpan {
+
+  private final RequestSpan ddSpan;
+  private final RequestSpan cncSpan;
+
+  public DelegatingRequestSpan(@Nonnull RequestSpan ddSpan, @Nonnull RequestSpan cncSpan) {
+    this.ddSpan = ddSpan;
+    this.cncSpan = cncSpan;
+  }
+
+  public RequestSpan getDatadogSpan() {
+    return ddSpan;
+  }
+
+  public RequestSpan getCncSpan() {
+    return cncSpan;
+  }
+
+  @Override
+  public void attribute(String key, String value) {
+    ddSpan.attribute(key, value);
+    cncSpan.attribute(key, value);
+  }
+
+  @Override
+  public void attribute(String key, boolean value) {
+    ddSpan.attribute(key, value);
+    cncSpan.attribute(key, value);
+  }
+
+  @Override
+  public void attribute(String key, long value) {
+    ddSpan.attribute(key, value);
+    cncSpan.attribute(key, value);
+  }
+
+  @Override
+  public void event(String name, Instant timestamp) {
+    ddSpan.event(name, timestamp);
+    cncSpan.event(name, timestamp);
+  }
+
+  @Override
+  public void status(StatusCode status) {
+    ddSpan.status(status);
+    cncSpan.status(status);
+  }
+
+  @Override
+  public void end() {
+    try {
+      ddSpan.end();
+    } finally {
+      // guarantee cnc spans get ended even if ddSpan.end() throws exception
+      cncSpan.end();
+    }
+  }
+
+  @Override
+  public void requestContext(RequestContext requestContext) {
+    ddSpan.requestContext(requestContext);
+    cncSpan.requestContext(requestContext);
+  }
+}

dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/DelegatingRequestTracer.java

@@ -0,0 +1,70 @@
+package datadog.trace.instrumentation.couchbase_32.client;
+
+import com.couchbase.client.core.cnc.RequestSpan;
+import com.couchbase.client.core.cnc.RequestTracer;
+import com.couchbase.client.core.cnc.tracing.NoopRequestSpan;
+import java.time.Duration;
+import reactor.core.publisher.Mono;
+
+public class DelegatingRequestTracer implements RequestTracer {
+
+  private final DatadogRequestTracer ddTracer;
+  private final RequestTracer cncTracer;
+
+  public DelegatingRequestTracer(DatadogRequestTracer ddTracer, RequestTracer cncTracer) {
+    this.ddTracer = ddTracer;
+    this.cncTracer = cncTracer;
+  }
+
+  @Override
+  public RequestSpan requestSpan(String name, RequestSpan parent) {
+    RequestSpan ddParentSpan = unwrapDatadogParentSpan(parent);
+    RequestSpan cncParentSpan = unwrapCncParentSpan(parent);
+
+    RequestSpan ddSpan = ddTracer != null ? ddTracer.requestSpan(name, ddParentSpan) : null;
+    RequestSpan cncSpan = cncTracer != null ? cncTracer.requestSpan(name, cncParentSpan) : null;
+
+    // no tracers are present - return noop span
+    if (ddSpan == null && cncSpan == null) {
+      return NoopRequestSpan.INSTANCE;
+    }
+
+    // only one tracer is present - no need to delegate
+    if (ddSpan == null) {
+      return cncSpan;
+    }
+    if (cncSpan == null) {
+      return ddSpan;
+    }
+
+    return new DelegatingRequestSpan(ddSpan, cncSpan);
+  }
+
+  @Override
+  public Mono<Void> start() {
+    Mono<Void> primary = ddTracer != null ? ddTracer.start() : Mono.empty();
+    Mono<Void> secondary = cncTracer != null ? cncTracer.start() : Mono.empty();
+    return Mono.when(primary, secondary);
+  }
+
+  @Override
+  public Mono<Void> stop(Duration timeout) {
+    Mono<Void> primary = ddTracer != null ? ddTracer.stop(timeout) : Mono.empty();
+    Mono<Void> secondary = cncTracer != null ? cncTracer.stop(timeout) : Mono.empty();
+    return Mono.when(primary, secondary);
+  }
+
+  private static RequestSpan unwrapDatadogParentSpan(RequestSpan parent) {
+    if (parent instanceof DelegatingRequestSpan) {
+      return ((DelegatingRequestSpan) parent).getDatadogSpan();
+    }
+    return parent;
+  }
+
+  private static RequestSpan unwrapCncParentSpan(RequestSpan parent) {
+    if (parent instanceof DelegatingRequestSpan) {
+      return ((DelegatingRequestSpan) parent).getCncSpan();
+    }
+    return parent;
+  }
+}

dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/test/groovy/CouchbaseClient32Test.groovy

@@ -1,10 +1,13 @@
 import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
 import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
 
+import com.couchbase.client.core.cnc.RequestSpan
+import com.couchbase.client.core.cnc.RequestTracer
 import com.couchbase.client.core.env.TimeoutConfig
 import com.couchbase.client.core.error.CouchbaseException
 import com.couchbase.client.core.error.DocumentNotFoundException
 import com.couchbase.client.core.error.ParsingFailureException
+import com.couchbase.client.core.msg.RequestContext
 import com.couchbase.client.java.Bucket
 import com.couchbase.client.java.Cluster
 import com.couchbase.client.java.ClusterOptions
@@ -19,10 +22,13 @@ import datadog.trace.api.DDTags
 import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags
 import datadog.trace.bootstrap.instrumentation.api.Tags
 import datadog.trace.core.DDSpan
+import java.time.Instant
+import java.util.concurrent.CopyOnWriteArrayList
 import org.slf4j.Logger
 import org.slf4j.LoggerFactory
 import org.testcontainers.couchbase.BucketDefinition
 import org.testcontainers.couchbase.CouchbaseContainer
+import reactor.core.publisher.Mono
 import spock.lang.Shared
 
 import java.time.Duration
@@ -394,6 +400,69 @@ abstract class CouchbaseClient32Test extends VersionedNamingTestBase {
     }
   }
 
+  def "check basic spans with custom request tracer"() {
+    setup:
+    def customTracer = new TestRequestTracer()
+
+    ClusterEnvironment environmentWithCustomTracer = ClusterEnvironment.builder()
+      .timeoutConfig(TimeoutConfig.kvTimeout(Duration.ofSeconds(10)))
+      .requestTracer(customTracer)
+      .build()
+
+    def connectionString = "couchbase://${couchbase.host}:${couchbase.bootstrapCarrierDirectPort},${couchbase.host}:${couchbase.bootstrapHttpDirectPort}=manager"
+
+    Cluster localCluster = Cluster.connect(
+      connectionString,
+      ClusterOptions
+      .clusterOptions(couchbase.username, couchbase.password)
+      .environment(environmentWithCustomTracer)
+      )
+    Bucket localBucket = localCluster.bucket(BUCKET)
+    localBucket.waitUntilReady(Duration.ofSeconds(30))
+    def collection = localBucket.defaultCollection()
+
+    when:
+    collection.get("data 0")
+
+    then:
+    assertTraces(1) {
+      sortSpansByStart()
+      trace(2) {
+        assertCouchbaseCall(it, "get", [
+          'db.couchbase.collection' : '_default',
+          'db.couchbase.document_id': { String },
+          'db.couchbase.retries'    : { Long },
+          'db.couchbase.scope'      : '_default',
+          'db.couchbase.service'    : 'kv',
+          'db.name'                 : BUCKET,
+          'db.operation'            : 'get'
+        ])
+        assertCouchbaseDispatchCall(it, span(0), [
+          'db.couchbase.collection'     : '_default',
+          'db.couchbase.document_id'    : { String },
+          'db.couchbase.scope'          : '_default',
+          'db.name'                     : BUCKET
+        ])
+      }
+    }
+
+    and: "custom tracer also saw spans"
+    customTracer.spans.size() > 0
+    customTracer.spans*.ended.every { it == true }
+
+    cleanup:
+    try {
+      localCluster?.disconnect()
+    } catch (Throwable t) {
+      LOGGER.debug("Unable to properly disconnect localCluster in custom tracer test", t)
+    }
+    try {
+      environmentWithCustomTracer?.shutdown()
+    } catch (Throwable t) {
+      LOGGER.debug("Unable to properly shutdown environmentWithCustomTracer", t)
+    }
+  }
+
   void assertCouchbaseCall(TraceAssert trace, String name, Map<String, Serializable> extraTags, boolean internal = false, Throwable ex = null) {
     assertCouchbaseCall(trace, name, extraTags, null, internal, ex)
   }
@@ -453,6 +522,75 @@ abstract class CouchbaseClient32Test extends VersionedNamingTestBase {
     }
     assertCouchbaseCall(trace, 'dispatch_to_server', allExtraTags, parentSpan, true, null)
   }
+
+  static class TestRequestTracer implements RequestTracer {
+
+    final List<TestRequestSpan> spans = new CopyOnWriteArrayList<>()
+
+    @Override
+    RequestSpan requestSpan(String requestName, RequestSpan parent) {
+      def span = new TestRequestSpan(requestName, parent)
+      spans.add(span)
+      return span
+    }
+
+    @Override
+    Mono<Void> start() {
+      return Mono.empty()
+    }
+
+    @Override
+    Mono<Void> stop(Duration timeout) {
+      return Mono.empty()
+    }
+  }
+
+  static class TestRequestSpan implements RequestSpan {
+
+    final String name
+    final RequestSpan parent
+    final Map<String, Object> attributes = new LinkedHashMap<>()
+    final List<String> events = []
+    volatile boolean ended = false
+
+    TestRequestSpan(String name, RequestSpan parent) {
+      this.name = name
+      this.parent = parent
+    }
+
+    @Override
+    void end() {
+      ended = true
+    }
+
+    @Override
+    void attribute(String key, String value) {
+      attributes.put(key, value)
+    }
+
+    @Override
+    void attribute(String key, boolean value) {
+      attributes.put(key, value)
+    }
+
+    @Override
+    void attribute(String key, long value) {
+      attributes.put(key, value)
+    }
+
+    @Override
+    void event(String name, Instant timestamp) {
+      events.add(name)
+    }
+
+    @Override
+    void status(StatusCode status) {
+    }
+
+    @Override
+    void requestContext(RequestContext requestContext) {
+    }
+  }
 }
 
 class CouchbaseClient32V0Test extends CouchbaseClient32Test {