Add decorator for couchbase instrumentation

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/decorator/DatabaseClientDecorator.java

@@ -0,0 +1,33 @@
+package datadog.trace.agent.decorator;
+
+import io.opentracing.Span;
+import io.opentracing.tag.Tags;
+
+public abstract class DatabaseClientDecorator<SESSION> extends ClientDecorator {
+
+  protected abstract String dbType();
+
+  protected abstract String dbUser(SESSION session);
+
+  protected abstract String dbInstance(SESSION session);
+
+  @Override
+  public Span afterStart(final Span span) {
+    assert span != null;
+    Tags.DB_TYPE.set(span, dbType());
+    return super.afterStart(span);
+  }
+
+  public Span onSession(final Span span, final SESSION statement) {
+    assert span != null;
+    Tags.DB_USER.set(span, dbUser(statement));
+    Tags.DB_INSTANCE.set(span, dbInstance(statement));
+    return span;
+  }
+
+  public Span onStatement(final Span span, final String statement) {
+    assert span != null;
+    Tags.DB_STATEMENT.set(span, statement);
+    return span;
+  }
+}

dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseBucketInstrumentation.java

@@ -1,6 +1,5 @@
 package datadog.trace.instrumentation.couchbase.client;
 
-import static io.opentracing.log.Fields.ERROR_OBJECT;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -12,12 +11,10 @@ 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.DDSpanTypes;
 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.Map;
@@ -48,6 +45,10 @@ public class CouchbaseBucketInstrumentation extends Instrumenter.Default {
   @Override
   public String[] helperClassNames() {
     return new String[] {
+      "datadog.trace.agent.decorator.BaseDecorator",
+      "datadog.trace.agent.decorator.ClientDecorator",
+      "datadog.trace.agent.decorator.DatabaseClientDecorator",
+      packageName + ".CouchbaseClientDecorator",
       getClass().getName() + "$TraceSpanStart",
       getClass().getName() + "$TraceSpanFinish",
       getClass().getName() + "$TraceSpanError",
@@ -113,14 +114,12 @@ public class CouchbaseBucketInstrumentation extends Instrumenter.Default {
 
       // just replace the no-op span.
       spanRef.set(
+          CouchbaseClientDecorator.INSTANCE.afterStart(
               GlobalTracer.get()
                   .buildSpan("couchbase.call")
-              .withTag(Tags.COMPONENT.getKey(), "couchbase-client")
-              .withTag(DDTags.SERVICE_NAME, "couchbase")
                   .withTag(DDTags.RESOURCE_NAME, resourceName)
-              .withTag(DDTags.SPAN_TYPE, DDSpanTypes.COUCHBASE)
                   .withTag("bucket", bucket)
-              .start());
+                  .start()));
     }
   }
 
@@ -136,6 +135,7 @@ public class CouchbaseBucketInstrumentation extends Instrumenter.Default {
       final Span span = spanRef.getAndSet(null);
 
       if (span != null) {
+        CouchbaseClientDecorator.INSTANCE.beforeFinish(span);
         span.finish();
       }
     }
@@ -152,8 +152,8 @@ public class CouchbaseBucketInstrumentation extends Instrumenter.Default {
     public void call(final Throwable throwable) {
       final Span span = spanRef.getAndSet(null);
       if (span != null) {
-        Tags.ERROR.set(span, true);
+        CouchbaseClientDecorator.INSTANCE.onError(span, throwable);
-        span.log(singletonMap(ERROR_OBJECT, throwable));
+        CouchbaseClientDecorator.INSTANCE.beforeFinish(span);
         span.finish();
       }
     }

dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseClientDecorator.java

@@ -0,0 +1,43 @@
+package datadog.trace.instrumentation.couchbase.client;
+
+import datadog.trace.agent.decorator.DatabaseClientDecorator;
+import datadog.trace.api.DDSpanTypes;
+
+class CouchbaseClientDecorator extends DatabaseClientDecorator {
+  public static final CouchbaseClientDecorator INSTANCE = new CouchbaseClientDecorator();
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"couchbase"};
+  }
+
+  @Override
+  protected String service() {
+    return "couchbase";
+  }
+
+  @Override
+  protected String component() {
+    return "couchbase-client";
+  }
+
+  @Override
+  protected String spanType() {
+    return DDSpanTypes.COUCHBASE;
+  }
+
+  @Override
+  protected String dbType() {
+    return "couchbase";
+  }
+
+  @Override
+  protected String dbUser(final Object o) {
+    return null;
+  }
+
+  @Override
+  protected String dbInstance(final Object o) {
+    return null;
+  }
+}

dd-java-agent/instrumentation/couchbase-2.0/src/main/java/datadog/trace/instrumentation/couchbase/client/CouchbaseClusterInstrumentation.java

@@ -1,6 +1,5 @@
 package datadog.trace.instrumentation.couchbase.client;
 
-import static io.opentracing.log.Fields.ERROR_OBJECT;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -12,12 +11,10 @@ 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.DDSpanTypes;
 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.Map;
@@ -48,6 +45,10 @@ public class CouchbaseClusterInstrumentation extends Instrumenter.Default {
   @Override
   public String[] helperClassNames() {
     return new String[] {
+      "datadog.trace.agent.decorator.BaseDecorator",
+      "datadog.trace.agent.decorator.ClientDecorator",
+      "datadog.trace.agent.decorator.DatabaseClientDecorator",
+      packageName + ".CouchbaseClientDecorator",
       getClass().getName() + "$TraceSpanStart",
       getClass().getName() + "$TraceSpanFinish",
       getClass().getName() + "$TraceSpanError",
@@ -108,13 +109,11 @@ public class CouchbaseClusterInstrumentation extends Instrumenter.Default {
 
       // just replace the no-op span.
       spanRef.set(
+          CouchbaseClientDecorator.INSTANCE.afterStart(
               GlobalTracer.get()
                   .buildSpan("couchbase.call")
-              .withTag(Tags.COMPONENT.getKey(), "couchbase-client")
-              .withTag(DDTags.SERVICE_NAME, "couchbase")
                   .withTag(DDTags.RESOURCE_NAME, resourceName)
-              .withTag(DDTags.SPAN_TYPE, DDSpanTypes.COUCHBASE)
+                  .start()));
-              .start());
     }
   }
 
@@ -130,6 +129,7 @@ public class CouchbaseClusterInstrumentation extends Instrumenter.Default {
       final Span span = spanRef.getAndSet(null);
 
       if (span != null) {
+        CouchbaseClientDecorator.INSTANCE.beforeFinish(span);
         span.finish();
       }
     }
@@ -146,8 +146,8 @@ public class CouchbaseClusterInstrumentation extends Instrumenter.Default {
     public void call(final Throwable throwable) {
       final Span span = spanRef.getAndSet(null);
       if (span != null) {
-        Tags.ERROR.set(span, true);
+        CouchbaseClientDecorator.INSTANCE.onError(span, throwable);
-        span.log(singletonMap(ERROR_OBJECT, throwable));
+        CouchbaseClientDecorator.INSTANCE.beforeFinish(span);
         span.finish();
       }
     }

dd-java-agent/instrumentation/couchbase-2.0/src/test/groovy/CouchbaseClientTest.groovy

@@ -25,6 +25,8 @@ class CouchbaseClientTest extends AbstractCouchbaseTest {
           tags {
             "$DDTags.SPAN_TYPE" DDSpanTypes.COUCHBASE
             "$Tags.COMPONENT.key" "couchbase-client"
+            "$Tags.DB_TYPE.key" "couchbase"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
             defaultTags()
           }
         }
@@ -52,6 +54,8 @@ class CouchbaseClientTest extends AbstractCouchbaseTest {
           tags {
             "$DDTags.SPAN_TYPE" DDSpanTypes.COUCHBASE
             "$Tags.COMPONENT.key" "couchbase-client"
+            "$Tags.DB_TYPE.key" "couchbase"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
             "bucket" bkt.name()
             defaultTags()
           }
@@ -79,6 +83,8 @@ class CouchbaseClientTest extends AbstractCouchbaseTest {
           tags {
             "$DDTags.SPAN_TYPE" DDSpanTypes.COUCHBASE
             "$Tags.COMPONENT.key" "couchbase-client"
+            "$Tags.DB_TYPE.key" "couchbase"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
             "bucket" bkt.name()
             defaultTags()
           }
@@ -122,6 +128,8 @@ class CouchbaseClientTest extends AbstractCouchbaseTest {
           tags {
             "$DDTags.SPAN_TYPE" DDSpanTypes.COUCHBASE
             "$Tags.COMPONENT.key" "couchbase-client"
+            "$Tags.DB_TYPE.key" "couchbase"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
             "bucket" bkt.name()
             defaultTags()
           }

dd-java-agent/instrumentation/couchbase-2.0/src/test/groovy/springdata/CouchbaseSpringRepositoryTest.groovy

@@ -85,6 +85,8 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
           tags {
             "$DDTags.SPAN_TYPE" DDSpanTypes.COUCHBASE
             "$Tags.COMPONENT.key" "couchbase-client"
+            "$Tags.DB_TYPE.key" "couchbase"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
             "bucket" bucketCouchbase.name()
             defaultTags()
           }
@@ -115,6 +117,8 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
           tags {
             "$DDTags.SPAN_TYPE" DDSpanTypes.COUCHBASE
             "$Tags.COMPONENT.key" "couchbase-client"
+            "$Tags.DB_TYPE.key" "couchbase"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
             "bucket" bucketCouchbase.name()
             defaultTags()
           }
@@ -138,6 +142,8 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
           tags {
             "$DDTags.SPAN_TYPE" DDSpanTypes.COUCHBASE
             "$Tags.COMPONENT.key" "couchbase-client"
+            "$Tags.DB_TYPE.key" "couchbase"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
             "bucket" bucketCouchbase.name()
             defaultTags()
           }
@@ -164,6 +170,8 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
           tags {
             "$DDTags.SPAN_TYPE" DDSpanTypes.COUCHBASE
             "$Tags.COMPONENT.key" "couchbase-client"
+            "$Tags.DB_TYPE.key" "couchbase"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
             "bucket" bucketCouchbase.name()
             defaultTags()
           }
@@ -179,6 +187,8 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
           tags {
             "$DDTags.SPAN_TYPE" DDSpanTypes.COUCHBASE
             "$Tags.COMPONENT.key" "couchbase-client"
+            "$Tags.DB_TYPE.key" "couchbase"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
             "bucket" bucketCouchbase.name()
             defaultTags()
           }
@@ -194,6 +204,8 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
           tags {
             "$DDTags.SPAN_TYPE" DDSpanTypes.COUCHBASE
             "$Tags.COMPONENT.key" "couchbase-client"
+            "$Tags.DB_TYPE.key" "couchbase"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
             "bucket" bucketCouchbase.name()
             defaultTags()
           }
@@ -220,6 +232,8 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
           tags {
             "$DDTags.SPAN_TYPE" DDSpanTypes.COUCHBASE
             "$Tags.COMPONENT.key" "couchbase-client"
+            "$Tags.DB_TYPE.key" "couchbase"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
             "bucket" bucketCouchbase.name()
             defaultTags()
           }
@@ -235,6 +249,8 @@ class CouchbaseSpringRepositoryTest extends AbstractCouchbaseTest {
           tags {
             "$DDTags.SPAN_TYPE" DDSpanTypes.COUCHBASE
             "$Tags.COMPONENT.key" "couchbase-client"
+            "$Tags.DB_TYPE.key" "couchbase"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
             "bucket" bucketCouchbase.name()
             defaultTags()
           }

dd-java-agent/instrumentation/couchbase-2.0/src/test/groovy/springdata/CouchbaseSpringTemplateTest.groovy

@@ -51,6 +51,8 @@ class CouchbaseSpringTemplateTest extends AbstractCouchbaseTest {
           tags {
             "$DDTags.SPAN_TYPE" DDSpanTypes.COUCHBASE
             "$Tags.COMPONENT.key" "couchbase-client"
+            "$Tags.DB_TYPE.key" "couchbase"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
             "bucket" name
             defaultTags()
           }
@@ -66,6 +68,8 @@ class CouchbaseSpringTemplateTest extends AbstractCouchbaseTest {
           tags {
             "$DDTags.SPAN_TYPE" DDSpanTypes.COUCHBASE
             "$Tags.COMPONENT.key" "couchbase-client"
+            "$Tags.DB_TYPE.key" "couchbase"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
             "bucket" name
             defaultTags()
           }
@@ -81,6 +85,8 @@ class CouchbaseSpringTemplateTest extends AbstractCouchbaseTest {
           tags {
             "$DDTags.SPAN_TYPE" DDSpanTypes.COUCHBASE
             "$Tags.COMPONENT.key" "couchbase-client"
+            "$Tags.DB_TYPE.key" "couchbase"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
             "bucket" name
             defaultTags()
           }
@@ -96,6 +102,8 @@ class CouchbaseSpringTemplateTest extends AbstractCouchbaseTest {
           tags {
             "$DDTags.SPAN_TYPE" DDSpanTypes.COUCHBASE
             "$Tags.COMPONENT.key" "couchbase-client"
+            "$Tags.DB_TYPE.key" "couchbase"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
             "bucket" name
             defaultTags()
           }

dd-trace-ot/src/main/java/datadog/opentracing/decorators/DBTypeDecorator.java

@@ -13,8 +13,8 @@ public class DBTypeDecorator extends AbstractDecorator {
 
   public DBTypeDecorator() {
     super();
-    this.setMatchingTag(Tags.DB_TYPE.getKey());
+    setMatchingTag(Tags.DB_TYPE.getKey());
-    this.setReplacementTag(DDTags.SERVICE_NAME);
+    setReplacementTag(DDTags.SERVICE_NAME);
   }
 
   @Override
@@ -22,6 +22,10 @@ public class DBTypeDecorator extends AbstractDecorator {
 
     // Assign service name
     if (!super.shouldSetTag(context, tag, value)) {
+      if ("couchbase".equals(value)) {
+        // Couchbase instrumentation has different behavior.
+        return true;
+      }
       // Assign span type to DB
       // Special case: Mongo, set to mongodb
       if ("mongo".equals(value)) {