Multi-tracing support for Couchbase 3.2+

Author: ValentinZakharov

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;
+
+/** RequestSpan, which delegates all calls to two other RequestSpans */
+public class DelegatingRequestSpan implements RequestSpan {
+
+  private final RequestSpan ddSpan;
+  private final RequestSpan cncSpan;
+
+  public DelegatingRequestSpan(RequestSpan ddSpan, 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) {
+    // TODO: add null checks
+    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;
+  }
+}