From 40f8c132192c3c2e534c8f6bb42445aa8c2a8f7e Mon Sep 17 00:00:00 2001 From: Valentin Zakharov Date: Wed, 10 Dec 2025 10:04:20 +0100 Subject: [PATCH 1/2] Multi-tracing support for Couchbase 3.2+ --- ...CoreEnvironmentBuilderInstrumentation.java | 7 ++ ...EnvironmentBuilderRequestTracerAdvice.java | 38 ++++++++++ .../client/DelegatingRequestSpan.java | 72 +++++++++++++++++++ .../client/DelegatingRequestTracer.java | 70 ++++++++++++++++++ 4 files changed, 187 insertions(+) create mode 100644 dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/CoreEnvironmentBuilderRequestTracerAdvice.java create mode 100644 dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/DelegatingRequestSpan.java create mode 100644 dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/DelegatingRequestTracer.java diff --git a/dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/CoreEnvironmentBuilderInstrumentation.java b/dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/CoreEnvironmentBuilderInstrumentation.java index 9d73dbd64df..db9928c9cf9 100644 --- a/dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/CoreEnvironmentBuilderInstrumentation.java +++ b/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"); } } diff --git a/dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/CoreEnvironmentBuilderRequestTracerAdvice.java b/dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/CoreEnvironmentBuilderRequestTracerAdvice.java new file mode 100644 index 00000000000..a2e7c3473f1 --- /dev/null +++ b/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 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); + } +} diff --git a/dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/DelegatingRequestSpan.java b/dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/DelegatingRequestSpan.java new file mode 100644 index 00000000000..bb4cfb46d55 --- /dev/null +++ b/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); + } +} diff --git a/dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/DelegatingRequestTracer.java b/dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/DelegatingRequestTracer.java new file mode 100644 index 00000000000..932ffca045a --- /dev/null +++ b/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 start() { + Mono primary = ddTracer != null ? ddTracer.start() : Mono.empty(); + Mono secondary = cncTracer != null ? cncTracer.start() : Mono.empty(); + return Mono.when(primary, secondary); + } + + @Override + public Mono stop(Duration timeout) { + Mono primary = ddTracer != null ? ddTracer.stop(timeout) : Mono.empty(); + Mono 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; + } +} From ba9e98f4948b1234ea590004ded79a726785e695 Mon Sep 17 00:00:00 2001 From: Valentin Zakharov Date: Wed, 10 Dec 2025 11:25:39 +0100 Subject: [PATCH 2/2] Multi tracing test --- .../client/DelegatingRequestSpan.java | 4 +- .../test/groovy/CouchbaseClient32Test.groovy | 138 ++++++++++++++++++ 2 files changed, 140 insertions(+), 2 deletions(-) diff --git a/dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/DelegatingRequestSpan.java b/dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/DelegatingRequestSpan.java index bb4cfb46d55..eb0cfbc68d1 100644 --- a/dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/DelegatingRequestSpan.java +++ b/dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/main/java/datadog/trace/instrumentation/couchbase_32/client/DelegatingRequestSpan.java @@ -3,6 +3,7 @@ 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 { @@ -10,7 +11,7 @@ public class DelegatingRequestSpan implements RequestSpan { private final RequestSpan ddSpan; private final RequestSpan cncSpan; - public DelegatingRequestSpan(RequestSpan ddSpan, RequestSpan cncSpan) { + public DelegatingRequestSpan(@Nonnull RequestSpan ddSpan, @Nonnull RequestSpan cncSpan) { this.ddSpan = ddSpan; this.cncSpan = cncSpan; } @@ -25,7 +26,6 @@ public RequestSpan getCncSpan() { @Override public void attribute(String key, String value) { - // TODO: add null checks ddSpan.attribute(key, value); cncSpan.attribute(key, value); } diff --git a/dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/test/groovy/CouchbaseClient32Test.groovy b/dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/test/groovy/CouchbaseClient32Test.groovy index 10150d91ef3..ff7e9968cc8 100644 --- a/dd-java-agent/instrumentation/couchbase/couchbase-3.2/src/test/groovy/CouchbaseClient32Test.groovy +++ b/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 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 spans = new CopyOnWriteArrayList<>() + + @Override + RequestSpan requestSpan(String requestName, RequestSpan parent) { + def span = new TestRequestSpan(requestName, parent) + spans.add(span) + return span + } + + @Override + Mono start() { + return Mono.empty() + } + + @Override + Mono stop(Duration timeout) { + return Mono.empty() + } + } + + static class TestRequestSpan implements RequestSpan { + + final String name + final RequestSpan parent + final Map attributes = new LinkedHashMap<>() + final List 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 {