Skip to content

Commit a1ce9ae

Browse files
ValentinZakharovValentinZakharov
authored
Multi-tracing support for Couchbase 3.2+ (#10147)
* Multi-tracing support for Couchbase 3.2+ * Multi tracing test
1 parent d989362 commit a1ce9ae

File tree

5 files changed

+325
-0
lines changed

5 files changed

+325
-0
lines changed

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package datadog.trace.instrumentation.couchbase_32.client;
22

3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
34
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
5+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
46

57
import com.google.auto.service.AutoService;
68
import datadog.trace.agent.tooling.Instrumenter;
@@ -23,6 +25,8 @@ public String[] helperClassNames() {
2325
packageName + ".DatadogRequestSpan",
2426
packageName + ".DatadogRequestSpan$1",
2527
packageName + ".DatadogRequestTracer",
28+
packageName + ".DelegatingRequestSpan",
29+
packageName + ".DelegatingRequestTracer"
2630
};
2731
}
2832

@@ -39,5 +43,8 @@ public String instrumentedType() {
3943
@Override
4044
public void methodAdvice(MethodTransformer transformer) {
4145
transformer.applyAdvice(isConstructor(), packageName + ".CoreEnvironmentBuilderAdvice");
46+
transformer.applyAdvice(
47+
isMethod().and(named("requestTracer")),
48+
packageName + ".CoreEnvironmentBuilderRequestTracerAdvice");
4249
}
4350
}

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

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package datadog.trace.instrumentation.couchbase_32.client;
2+
3+
import com.couchbase.client.core.Core;
4+
import com.couchbase.client.core.cnc.RequestTracer;
5+
import datadog.trace.bootstrap.ContextStore;
6+
import datadog.trace.bootstrap.InstrumentationContext;
7+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
8+
import net.bytebuddy.asm.Advice;
9+
10+
public class CoreEnvironmentBuilderRequestTracerAdvice {
11+
@Advice.OnMethodEnter(suppress = Throwable.class)
12+
public static void onEnter(
13+
@Advice.Argument(value = 0, readOnly = false) RequestTracer requestTracer) {
14+
15+
// already a delegating tracer
16+
if (requestTracer instanceof DelegatingRequestTracer) {
17+
return;
18+
}
19+
20+
// already a datadog tracer
21+
if (requestTracer instanceof DatadogRequestTracer) {
22+
return;
23+
}
24+
25+
ContextStore<Core, String> coreContext = InstrumentationContext.get(Core.class, String.class);
26+
27+
DatadogRequestTracer datadogTracer = new DatadogRequestTracer(AgentTracer.get(), coreContext);
28+
29+
// if the app didn't set a custom tracer, use only datadog tracer
30+
if (requestTracer == null) {
31+
requestTracer = datadogTracer;
32+
return;
33+
}
34+
35+
// Wrap custom datadog and cnc tracers into a delegating
36+
requestTracer = new DelegatingRequestTracer(datadogTracer, requestTracer);
37+
}
38+
}

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

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
package datadog.trace.instrumentation.couchbase_32.client;
2+
3+
import com.couchbase.client.core.cnc.RequestSpan;
4+
import com.couchbase.client.core.msg.RequestContext;
5+
import java.time.Instant;
6+
import javax.annotation.Nonnull;
7+
8+
/** RequestSpan, which delegates all calls to two other RequestSpans */
9+
public class DelegatingRequestSpan implements RequestSpan {
10+
11+
private final RequestSpan ddSpan;
12+
private final RequestSpan cncSpan;
13+
14+
public DelegatingRequestSpan(@Nonnull RequestSpan ddSpan, @Nonnull RequestSpan cncSpan) {
15+
this.ddSpan = ddSpan;
16+
this.cncSpan = cncSpan;
17+
}
18+
19+
public RequestSpan getDatadogSpan() {
20+
return ddSpan;
21+
}
22+
23+
public RequestSpan getCncSpan() {
24+
return cncSpan;
25+
}
26+
27+
@Override
28+
public void attribute(String key, String value) {
29+
ddSpan.attribute(key, value);
30+
cncSpan.attribute(key, value);
31+
}
32+
33+
@Override
34+
public void attribute(String key, boolean value) {
35+
ddSpan.attribute(key, value);
36+
cncSpan.attribute(key, value);
37+
}
38+
39+
@Override
40+
public void attribute(String key, long value) {
41+
ddSpan.attribute(key, value);
42+
cncSpan.attribute(key, value);
43+
}
44+
45+
@Override
46+
public void event(String name, Instant timestamp) {
47+
ddSpan.event(name, timestamp);
48+
cncSpan.event(name, timestamp);
49+
}
50+
51+
@Override
52+
public void status(StatusCode status) {
53+
ddSpan.status(status);
54+
cncSpan.status(status);
55+
}
56+
57+
@Override
58+
public void end() {
59+
try {
60+
ddSpan.end();
61+
} finally {
62+
// guarantee cnc spans get ended even if ddSpan.end() throws exception
63+
cncSpan.end();
64+
}
65+
}
66+
67+
@Override
68+
public void requestContext(RequestContext requestContext) {
69+
ddSpan.requestContext(requestContext);
70+
cncSpan.requestContext(requestContext);
71+
}
72+
}

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

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
package datadog.trace.instrumentation.couchbase_32.client;
2+
3+
import com.couchbase.client.core.cnc.RequestSpan;
4+
import com.couchbase.client.core.cnc.RequestTracer;
5+
import com.couchbase.client.core.cnc.tracing.NoopRequestSpan;
6+
import java.time.Duration;
7+
import reactor.core.publisher.Mono;
8+
9+
public class DelegatingRequestTracer implements RequestTracer {
10+
11+
private final DatadogRequestTracer ddTracer;
12+
private final RequestTracer cncTracer;
13+
14+
public DelegatingRequestTracer(DatadogRequestTracer ddTracer, RequestTracer cncTracer) {
15+
this.ddTracer = ddTracer;
16+
this.cncTracer = cncTracer;
17+
}
18+
19+
@Override
20+
public RequestSpan requestSpan(String name, RequestSpan parent) {
21+
RequestSpan ddParentSpan = unwrapDatadogParentSpan(parent);
22+
RequestSpan cncParentSpan = unwrapCncParentSpan(parent);
23+
24+
RequestSpan ddSpan = ddTracer != null ? ddTracer.requestSpan(name, ddParentSpan) : null;
25+
RequestSpan cncSpan = cncTracer != null ? cncTracer.requestSpan(name, cncParentSpan) : null;
26+
27+
// no tracers are present - return noop span
28+
if (ddSpan == null && cncSpan == null) {
29+
return NoopRequestSpan.INSTANCE;
30+
}
31+
32+
// only one tracer is present - no need to delegate
33+
if (ddSpan == null) {
34+
return cncSpan;
35+
}
36+
if (cncSpan == null) {
37+
return ddSpan;
38+
}
39+
40+
return new DelegatingRequestSpan(ddSpan, cncSpan);
41+
}
42+
43+
@Override
44+
public Mono<Void> start() {
45+
Mono<Void> primary = ddTracer != null ? ddTracer.start() : Mono.empty();
46+
Mono<Void> secondary = cncTracer != null ? cncTracer.start() : Mono.empty();
47+
return Mono.when(primary, secondary);
48+
}
49+
50+
@Override
51+
public Mono<Void> stop(Duration timeout) {
52+
Mono<Void> primary = ddTracer != null ? ddTracer.stop(timeout) : Mono.empty();
53+
Mono<Void> secondary = cncTracer != null ? cncTracer.stop(timeout) : Mono.empty();
54+
return Mono.when(primary, secondary);
55+
}
56+
57+
private static RequestSpan unwrapDatadogParentSpan(RequestSpan parent) {
58+
if (parent instanceof DelegatingRequestSpan) {
59+
return ((DelegatingRequestSpan) parent).getDatadogSpan();
60+
}
61+
return parent;
62+
}
63+
64+
private static RequestSpan unwrapCncParentSpan(RequestSpan parent) {
65+
if (parent instanceof DelegatingRequestSpan) {
66+
return ((DelegatingRequestSpan) parent).getCncSpan();
67+
}
68+
return parent;
69+
}
70+
}

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

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,13 @@
11
import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
22
import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
33

4+
import com.couchbase.client.core.cnc.RequestSpan
5+
import com.couchbase.client.core.cnc.RequestTracer
46
import com.couchbase.client.core.env.TimeoutConfig
57
import com.couchbase.client.core.error.CouchbaseException
68
import com.couchbase.client.core.error.DocumentNotFoundException
79
import com.couchbase.client.core.error.ParsingFailureException
10+
import com.couchbase.client.core.msg.RequestContext
811
import com.couchbase.client.java.Bucket
912
import com.couchbase.client.java.Cluster
1013
import com.couchbase.client.java.ClusterOptions
@@ -19,10 +22,13 @@ import datadog.trace.api.DDTags
1922
import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags
2023
import datadog.trace.bootstrap.instrumentation.api.Tags
2124
import datadog.trace.core.DDSpan
25+
import java.time.Instant
26+
import java.util.concurrent.CopyOnWriteArrayList
2227
import org.slf4j.Logger
2328
import org.slf4j.LoggerFactory
2429
import org.testcontainers.couchbase.BucketDefinition
2530
import org.testcontainers.couchbase.CouchbaseContainer
31+
import reactor.core.publisher.Mono
2632
import spock.lang.Shared
2733

2834
import java.time.Duration
@@ -394,6 +400,69 @@ abstract class CouchbaseClient32Test extends VersionedNamingTestBase {
394400
}
395401
}
396402

403+
def "check basic spans with custom request tracer"() {
404+
setup:
405+
def customTracer = new TestRequestTracer()
406+
407+
ClusterEnvironment environmentWithCustomTracer = ClusterEnvironment.builder()
408+
.timeoutConfig(TimeoutConfig.kvTimeout(Duration.ofSeconds(10)))
409+
.requestTracer(customTracer)
410+
.build()
411+
412+
def connectionString = "couchbase://${couchbase.host}:${couchbase.bootstrapCarrierDirectPort},${couchbase.host}:${couchbase.bootstrapHttpDirectPort}=manager"
413+
414+
Cluster localCluster = Cluster.connect(
415+
connectionString,
416+
ClusterOptions
417+
.clusterOptions(couchbase.username, couchbase.password)
418+
.environment(environmentWithCustomTracer)
419+
)
420+
Bucket localBucket = localCluster.bucket(BUCKET)
421+
localBucket.waitUntilReady(Duration.ofSeconds(30))
422+
def collection = localBucket.defaultCollection()
423+
424+
when:
425+
collection.get("data 0")
426+
427+
then:
428+
assertTraces(1) {
429+
sortSpansByStart()
430+
trace(2) {
431+
assertCouchbaseCall(it, "get", [
432+
'db.couchbase.collection' : '_default',
433+
'db.couchbase.document_id': { String },
434+
'db.couchbase.retries' : { Long },
435+
'db.couchbase.scope' : '_default',
436+
'db.couchbase.service' : 'kv',
437+
'db.name' : BUCKET,
438+
'db.operation' : 'get'
439+
])
440+
assertCouchbaseDispatchCall(it, span(0), [
441+
'db.couchbase.collection' : '_default',
442+
'db.couchbase.document_id' : { String },
443+
'db.couchbase.scope' : '_default',
444+
'db.name' : BUCKET
445+
])
446+
}
447+
}
448+
449+
and: "custom tracer also saw spans"
450+
customTracer.spans.size() > 0
451+
customTracer.spans*.ended.every { it == true }
452+
453+
cleanup:
454+
try {
455+
localCluster?.disconnect()
456+
} catch (Throwable t) {
457+
LOGGER.debug("Unable to properly disconnect localCluster in custom tracer test", t)
458+
}
459+
try {
460+
environmentWithCustomTracer?.shutdown()
461+
} catch (Throwable t) {
462+
LOGGER.debug("Unable to properly shutdown environmentWithCustomTracer", t)
463+
}
464+
}
465+
397466
void assertCouchbaseCall(TraceAssert trace, String name, Map<String, Serializable> extraTags, boolean internal = false, Throwable ex = null) {
398467
assertCouchbaseCall(trace, name, extraTags, null, internal, ex)
399468
}
@@ -453,6 +522,75 @@ abstract class CouchbaseClient32Test extends VersionedNamingTestBase {
453522
}
454523
assertCouchbaseCall(trace, 'dispatch_to_server', allExtraTags, parentSpan, true, null)
455524
}
525+
526+
static class TestRequestTracer implements RequestTracer {
527+
528+
final List<TestRequestSpan> spans = new CopyOnWriteArrayList<>()
529+
530+
@Override
531+
RequestSpan requestSpan(String requestName, RequestSpan parent) {
532+
def span = new TestRequestSpan(requestName, parent)
533+
spans.add(span)
534+
return span
535+
}
536+
537+
@Override
538+
Mono<Void> start() {
539+
return Mono.empty()
540+
}
541+
542+
@Override
543+
Mono<Void> stop(Duration timeout) {
544+
return Mono.empty()
545+
}
546+
}
547+
548+
static class TestRequestSpan implements RequestSpan {
549+
550+
final String name
551+
final RequestSpan parent
552+
final Map<String, Object> attributes = new LinkedHashMap<>()
553+
final List<String> events = []
554+
volatile boolean ended = false
555+
556+
TestRequestSpan(String name, RequestSpan parent) {
557+
this.name = name
558+
this.parent = parent
559+
}
560+
561+
@Override
562+
void end() {
563+
ended = true
564+
}
565+
566+
@Override
567+
void attribute(String key, String value) {
568+
attributes.put(key, value)
569+
}
570+
571+
@Override
572+
void attribute(String key, boolean value) {
573+
attributes.put(key, value)
574+
}
575+
576+
@Override
577+
void attribute(String key, long value) {
578+
attributes.put(key, value)
579+
}
580+
581+
@Override
582+
void event(String name, Instant timestamp) {
583+
events.add(name)
584+
}
585+
586+
@Override
587+
void status(StatusCode status) {
588+
}
589+
590+
@Override
591+
void requestContext(RequestContext requestContext) {
592+
}
593+
}
456594
}
457595

458596
class CouchbaseClient32V0Test extends CouchbaseClient32Test {

0 commit comments

Comments
 (0)