Separate out bucket and cluster instrumentation
Include bucket name in bucket spans
ADD
dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseBucketInstrumentation.java
+package datadog.trace.instrumentation.couchbase.client;
+
+import static io.opentracing.log.Fields.ERROR_OBJECT;
+import static net.bytebuddy.matcher.ElementMatchers.isInterface;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+
+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.Span;
+import io.opentracing.noop.NoopSpan;
+import io.opentracing.tag.Tags;
+import io.opentracing.util.GlobalTracer;
+import java.lang.reflect.Method;
+import java.util.Collections;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
+import net.bytebuddy.asm.Advice;
+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 {
+
+ public CouchbaseBucketInstrumentation() {
+ super("couchbase");
+ }
+
+ @Override
+ public ElementMatcher<TypeDescription> typeMatcher() {
+ return not(isInterface())
+ .and(
+ named("com.couchbase.client.java.bucket.DefaultAsyncBucketManager")
+ .or(named("com.couchbase.client.java.CouchbaseAsyncBucket")));
+ }
+
+ @Override
+ public String[] helperClassNames() {
+ return new String[] {
+ getClass().getName() + "$TraceSpanStart",
+ getClass().getName() + "$TraceSpanFinish",
+ getClass().getName() + "$TraceSpanError",
+ };
+ }
+
+ @Override
+ public Map<ElementMatcher, String> transformers() {
+ return Collections.<ElementMatcher, String>singletonMap(
+ isMethod().and(isPublic()).and(returns(named("rx.Observable"))),
+ CouchbaseClientAdvice.class.getName());
+ }
+
+ public static class CouchbaseClientAdvice {
+
+ @Advice.OnMethodEnter
+ public static int trackCallDepth() {
+ return CallDepthThreadLocalMap.incrementCallDepth(CouchbaseCluster.class);
+ }
+
+ @Advice.OnMethodExit
+ 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() + "(" + bucket + ")";
+
+ // just replace the no-op span.
+ spanRef.set(
+ GlobalTracer.get()
+ .buildSpan("couchbase.call")
+ .withTag(DDTags.SERVICE_NAME, "couchbase")
+ .withTag(DDTags.RESOURCE_NAME, resourceName)
+ .withTag("bucket", bucket)
+ .start());
+ }
+ }
+
+ 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) {
+ 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) {
+ Tags.ERROR.set(span, true);
+ span.log(Collections.singletonMap(ERROR_OBJECT, throwable));
+ span.finish();
+ }
+ }
+ }
+}

RENAME
dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseClientInstrumentation.java
dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseClusterInstrumentation.java
package datadog.trace.instrumentation.couchbase.client;
-import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
import static io.opentracing.log.Fields.ERROR_OBJECT;
import static net.bytebuddy.matcher.ElementMatchers.isInterface;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -30,9 +29,9 @@ import rx.functions.Action0;
import rx.functions.Action1;
@AutoService(Instrumenter.class)
-public class CouchbaseClientInstrumentation extends Instrumenter.Default {
+public class CouchbaseClusterInstrumentation extends Instrumenter.Default {
- public CouchbaseClientInstrumentation() {
+ public CouchbaseClusterInstrumentation() {
super("couchbase");
}
@@ -40,11 +39,8 @@ public class CouchbaseClientInstrumentation extends Instrumenter.Default {
public ElementMatcher<TypeDescription> typeMatcher() {
return not(isInterface())
.and(
- safeHasSuperType(
- named("com.couchbase.client.java.cluster.AsyncClusterManager")
- .or(named("com.couchbase.client.java.bucket.AsyncBucketManager"))
- .or(named("com.couchbase.client.java.AsyncCluster"))
- .or(named("com.couchbase.client.java.AsyncBucket"))));
+ named("com.couchbase.client.java.cluster.DefaultAsyncClusterManager")
+ .or(named("com.couchbase.client.java.CouchbaseAsyncCluster")));
}
@Override

MODIFY
dd-java-agent/instrumentation/couchbase-2.0/src/test/groovy/CouchbaseClientTest.groovy
dd-java-agent/instrumentation/couchbase-2.0/src/test/groovy/CouchbaseClientTest.groovy
import com.couchbase.client.java.Bucket
import com.couchbase.client.java.document.JsonDocument
import com.couchbase.client.java.document.json.JsonObject
+import com.couchbase.client.java.query.N1qlQuery
+import com.couchbase.client.java.query.N1qlQueryRow
import util.AbstractCouchbaseTest
import static datadog.trace.agent.test.asserts.ListWriterAssert.assertTraces
@@ -8,15 +10,53 @@ import static datadog.trace.agent.test.asserts.ListWriterAssert.assertTraces
class CouchbaseClientTest extends AbstractCouchbaseTest {
def "test client #type"() {
- setup:
+ when:
+ manager.hasBucket(bucketSettings.name())
+
+ then:
+ assertTraces(TEST_WRITER, 1) {
+ trace(0, 1) {
+ span(0) {
+ serviceName "couchbase"
+ resourceName "ClusterManager.hasBucket"
+ operationName "couchbase.call"
+ errored false
+ parent()
+ tags {
+ defaultTags()
+ }
+ }
+ }
+ }
+ TEST_WRITER.clear()
+
+ when:
// Connect to the bucket and open it
Bucket bkt = cluster.openBucket(bucketSettings.name())
// Create a JSON document and store it with the ID "helloworld"
JsonObject content = JsonObject.create().put("hello", "world")
+ def inserted = bkt.upsert(JsonDocument.create("helloworld", content))
+
+ then:
+ assertTraces(TEST_WRITER, 1) {
+ trace(0, 1) {
+ span(0) {
+ serviceName "couchbase"
+ resourceName "Bucket.upsert(${bkt.name()})"
+ operationName "couchbase.call"
+ errored false
+ parent()
+ tags {
+ "bucket" bkt.name()
+ defaultTags()
+ }
+ }
+ }
+ }
+ TEST_WRITER.clear()
when:
- def inserted = bkt.upsert(JsonDocument.create("helloworld", content))
def found = bkt.get("helloworld")
then:
@@ -24,27 +64,83 @@ class CouchbaseClientTest extends AbstractCouchbaseTest {
found.content().getString("hello") == "world"
and:
- assertTraces(TEST_WRITER, 2) {
+ assertTraces(TEST_WRITER, 1) {
+ trace(0, 1) {
+ span(0) {
+ serviceName "couchbase"
+ resourceName "Bucket.get(${bkt.name()})"
+ operationName "couchbase.call"
+ errored false
+ parent()
+ tags {
+ "bucket" bkt.name()
+ defaultTags()
+ }
+ }
+ }
+ }
+ TEST_WRITER.clear()
+
+
+ where:
+ bucketSettings << [bucketCouchbase, bucketMemcache, bucketEphemeral]
+ type = bucketSettings.type().name()
+ }
+
+ def "test query"() {
+ when:
+ // Connect to the bucket and open it
+ Bucket bkt = cluster.openBucket(bucketSettings.name())
+
+ // Create a JSON document and store it with the ID "helloworld"
+ JsonObject content = JsonObject.create().put("hello", "world")
+ def inserted = bkt.upsert(JsonDocument.create("helloworld", content))
+
+ then:
+ inserted != null
+ assertTraces(TEST_WRITER, 1) {
trace(0, 1) {
span(0) {
serviceName "couchbase"
- resourceName "Bucket.upsert"
+ resourceName "Bucket.upsert(${bkt.name()})"
operationName "couchbase.call"
errored false
parent()
tags {
+ "bucket" bkt.name()
defaultTags()
}
}
}
- trace(1, 1) {
+ }
+ TEST_WRITER.clear()
+
+ when:
+ def result = bkt.query(
+ N1qlQuery.parameterized(
+ "SELECT * FROM `${bkt.name()}` WHERE hello = \$hello",
+ JsonObject.create()
+ .put("hello", "world")
+ )
+ )
+
+ then:
+ result.parseSuccess()
+ result.finalSuccess()
+ N1qlQueryRow row = result.first()
+ row.value().get(bkt.name()).get("hello") == "world"
+
+ and:
+ assertTraces(TEST_WRITER, 1) {
+ trace(0, 1) {
span(0) {
serviceName "couchbase"
- resourceName "Bucket.get"
+ resourceName "Bucket.query(${bkt.name()})"
operationName "couchbase.call"
errored false
parent()
tags {
+ "bucket" bkt.name()
defaultTags()
}
}
@@ -52,7 +148,7 @@ class CouchbaseClientTest extends AbstractCouchbaseTest {
}
where:
- bucketSettings << [bucketCouchbase, bucketMemcache, bucketEphemeral]
+ bucketSettings << [bucketCouchbase] // Only couchbase buckets support queries.
type = bucketSettings.type().name()
}
}

MODIFY
dd-java-agent/instrumentation/couchbase-2.0/src/test/groovy/springdata/CouchbaseSpringRepositoryTest.groovy
dd-java-agent/instrumentation/couchbase-2.0/src/test/groovy/springdata/CouchbaseSpringRepositoryTest.groovy
@@ -54,11 +54,12 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
trace(0, 1) {
span(0) {
serviceName "couchbase"
- resourceName "Bucket.query"
+ resourceName "Bucket.query(${bucketCouchbase.name()})"
operationName "couchbase.call"
errored false
parent()
tags {
+ "bucket" bucketCouchbase.name()
defaultTags()
}
}
@@ -73,7 +74,7 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
when:
def doc = new Doc()
- then:
+ then: // CREATE
repo.save(doc) == doc
and:
@@ -81,11 +82,12 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
trace(0, 1) {
span(0) {
serviceName "couchbase"
- resourceName "Bucket.upsert"
+ resourceName "Bucket.upsert(${bucketCouchbase.name()})"
operationName "couchbase.call"
errored false
parent()
tags {
+ "bucket" bucketCouchbase.name()
defaultTags()
}
}
@@ -93,7 +95,7 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
}
TEST_WRITER.clear()
- and:
+ and: // RETRIEVE
FIND(repo, "1") == doc
and:
@@ -101,11 +103,12 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
trace(0, 1) {
span(0) {
serviceName "couchbase"
- resourceName "Bucket.get"
+ resourceName "Bucket.get(${bucketCouchbase.name()})"
operationName "couchbase.call"
errored false
parent()
tags {
+ "bucket" bucketCouchbase.name()
defaultTags()
}
}
@@ -114,22 +117,25 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
TEST_WRITER.clear()
when:
- doc.data == "other data"
+ doc.data = "other data"
- then:
+ then: // UPDATE
repo.save(doc) == doc
FIND(repo, "1") == doc
+ // findAll doesn't seem to be working.
+// repo.findAll().asList() == [doc]
and:
assertTraces(TEST_WRITER, 2) {
trace(0, 1) {
span(0) {
serviceName "couchbase"
- resourceName "Bucket.upsert"
+ resourceName "Bucket.upsert(${bucketCouchbase.name()})"
operationName "couchbase.call"
errored false
parent()
tags {
+ "bucket" bucketCouchbase.name()
defaultTags()
}
}
@@ -137,11 +143,12 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
trace(1, 1) {
span(0) {
serviceName "couchbase"
- resourceName "Bucket.get"
+ resourceName "Bucket.get(${bucketCouchbase.name()})"
operationName "couchbase.call"
errored false
parent()
tags {
+ "bucket" bucketCouchbase.name()
defaultTags()
}
}
@@ -149,7 +156,7 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
}
TEST_WRITER.clear()
- when:
+ when: // DELETE
repo.delete("1")
then:
@@ -160,11 +167,12 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
trace(0, 1) {
span(0) {
serviceName "couchbase"
- resourceName "Bucket.remove"
+ resourceName "Bucket.remove(${bucketCouchbase.name()})"
operationName "couchbase.call"
errored false
parent()
tags {
+ "bucket" bucketCouchbase.name()
defaultTags()
}
}
@@ -172,11 +180,12 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
trace(1, 1) {
span(0) {
serviceName "couchbase"
- resourceName "Bucket.query"
+ resourceName "Bucket.query(${bucketCouchbase.name()})"
operationName "couchbase.call"
errored false
parent()
tags {
+ "bucket" bucketCouchbase.name()
defaultTags()
}
}

MODIFY
dd-java-agent/instrumentation/couchbase-2.0/src/test/groovy/springdata/CouchbaseSpringTemplateTest.groovy
dd-java-agent/instrumentation/couchbase-2.0/src/test/groovy/springdata/CouchbaseSpringTemplateTest.groovy
@@ -44,11 +44,12 @@ class CouchbaseSpringTemplateTest extends AbstractCouchbaseTest {
trace(0, 1) {
span(0) {
serviceName "couchbase"
- resourceName "Bucket.upsert"
+ resourceName "Bucket.upsert($name)"
operationName "couchbase.call"
errored false
parent()
tags {
+ "bucket" name
defaultTags()
}
}
@@ -56,11 +57,12 @@ class CouchbaseSpringTemplateTest extends AbstractCouchbaseTest {
trace(1, 1) {
span(0) {
serviceName "couchbase"
- resourceName "Bucket.get"
+ resourceName "Bucket.get($name)"
operationName "couchbase.call"
errored false
parent()
tags {
+ "bucket" name
defaultTags()
}
}
@@ -68,11 +70,12 @@ class CouchbaseSpringTemplateTest extends AbstractCouchbaseTest {
trace(2, 1) {
span(0) {
serviceName "couchbase"
- resourceName "Bucket.remove"
+ resourceName "Bucket.remove($name)"
operationName "couchbase.call"
errored false
parent()
tags {
+ "bucket" name
defaultTags()
}
}
@@ -80,11 +83,12 @@ class CouchbaseSpringTemplateTest extends AbstractCouchbaseTest {
trace(3, 1) {
span(0) {
serviceName "couchbase"
- resourceName "Bucket.get"
+ resourceName "Bucket.get($name)"
operationName "couchbase.call"
errored false
parent()
tags {
+ "bucket" name
defaultTags()
}
}

MODIFY
dd-java-agent/instrumentation/couchbase-2.0/src/test/groovy/util/AbstractCouchbaseTest.groovy
dd-java-agent/instrumentation/couchbase-2.0/src/test/groovy/util/AbstractCouchbaseTest.groovy
@@ -79,9 +79,9 @@ class AbstractCouchbaseTest extends AgentTestRunner {
@Shared
def couchbaseContainer
@Shared
- CouchbaseCluster cluster
+ protected CouchbaseCluster cluster
@Shared
- ClusterManager manager
+ protected ClusterManager manager
def setupSpec() {
@@ -185,15 +185,16 @@ class AbstractCouchbaseTest extends AgentTestRunner {
Bucket bucket = cluster.openBucket(bucketSettings.name(), bucketSettings.password())
- bucket.query(Index.createPrimaryIndex().on(bucketSetting.name()))
-
- // We don't have a good way to tell that all traces are reported
- // since we don't know how many there will be.
- Thread.sleep(150)
- TEST_WRITER.clear() // remove traces generated by insertBucket
+ if (BucketType.COUCHBASE.equals(bucketSettings.type())) {
+ bucket.query(Index.dropPrimaryIndex(bucketSetting.name()))
+ def result = bucket.query(Index.createPrimaryIndex().on(bucketSetting.name()))
+ if (!result.finalSuccess()) {
+ // Lets try again.
+ result = bucket.query(Index.createPrimaryIndex().on(bucketSetting.name()))
+ }
+ assert result.finalSuccess()
// Create view for SpringRepository's findAll()
- if (BucketType.COUCHBASE.equals(bucketSettings.type())) {
bucket.bucketManager().insertDesignDocument(
DesignDocument.create("doc", Collections.singletonList(DefaultView.create("all",
'''
@@ -206,5 +207,10 @@ class AbstractCouchbaseTest extends AgentTestRunner {
)))
)
}
+
+ // We don't have a good way to tell that all traces are reported
+ // since we don't know how many there will be.
+ Thread.sleep(150)
+ TEST_WRITER.clear() // remove traces generated by insertBucket
}
}

MODIFY
dd-java-agent/instrumentation/elasticsearch-transport-2/src/test/groovy/springdata/Elasticsearch2SpringRepositoryTest.groovy
dd-java-agent/instrumentation/elasticsearch-transport-2/src/test/groovy/springdata/Elasticsearch2SpringRepositoryTest.groovy
@@ -156,7 +156,7 @@ class Elasticsearch2SpringRepositoryTest extends AgentTestRunner {
TEST_WRITER.clear()
when:
- doc.data == "other data"
+ doc.data = "other data"
then:
repo.index(doc) == doc

Separate out bucket and cluster instrumentation

Include bucket name in bucket spans