Refactor decorators and add service name mapping

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/JDBCMaps.java

@@ -17,13 +17,26 @@ public class JDBCMaps {
       Collections.synchronizedMap(new WeakHashMap<Connection, DBInfo>());
   public static final Map<PreparedStatement, String> preparedStatements =
       Collections.synchronizedMap(new WeakHashMap<PreparedStatement, String>());
-  public static final String UNKNOWN_QUERY = "Unknown Query";
+
+  private static final boolean RENAME_UNKNOWN =
+      Boolean.valueOf(getPropOrEnv("dd.trace.rename.unknown"));
+
+  public static final String DB_QUERY = RENAME_UNKNOWN ? "DB Query" : "Unknown Query";
 
   @Data
   public static class DBInfo {
-    public static DBInfo UNKNOWN = new DBInfo("null", "unknown", null);
+    public static DBInfo DEFAULT =
+        new DBInfo("null", RENAME_UNKNOWN ? "database" : "unknown", null);
     private final String url;
     private final String type;
     private final String user;
   }
+
+  private static String getPropOrEnv(final String name) {
+    return System.getProperty(name, System.getenv(propToEnvName(name)));
+  }
+
+  private static String propToEnvName(final String name) {
+    return name.toUpperCase().replace(".", "_");
+  }
 }

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/PreparedStatementInstrumentation.java

@@ -66,7 +66,7 @@ public static Scope startSpan(@Advice.This final PreparedStatement statement) {
 
       JDBCMaps.DBInfo dbInfo = JDBCMaps.connectionInfo.get(connection);
       if (dbInfo == null) {
-        dbInfo = JDBCMaps.DBInfo.UNKNOWN;
+        dbInfo = JDBCMaps.DBInfo.DEFAULT;
       }
       final Scope scope =
           GlobalTracer.get().buildSpan(dbInfo.getType() + ".query").startActive(true);
@@ -77,7 +77,7 @@ public static Scope startSpan(@Advice.This final PreparedStatement statement) {
       Tags.COMPONENT.set(span, "java-jdbc-prepared_statement");
 
       span.setTag(DDTags.SERVICE_NAME, dbInfo.getType());
-      span.setTag(DDTags.RESOURCE_NAME, sql == null ? JDBCMaps.UNKNOWN_QUERY : sql);
+      span.setTag(DDTags.RESOURCE_NAME, sql == null ? JDBCMaps.DB_QUERY : sql);
       span.setTag(DDTags.SPAN_TYPE, "sql");
       span.setTag("span.origin.type", statement.getClass().getName());
       span.setTag("db.jdbc.url", dbInfo.getUrl());

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/StatementInstrumentation.java

@@ -66,7 +66,7 @@ public static Scope startSpan(
 
       JDBCMaps.DBInfo dbInfo = JDBCMaps.connectionInfo.get(connection);
       if (dbInfo == null) {
-        dbInfo = JDBCMaps.DBInfo.UNKNOWN;
+        dbInfo = JDBCMaps.DBInfo.DEFAULT;
       }
 
       final Scope scope =

dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaConsumerInstrumentation.java

@@ -95,7 +95,7 @@ public static class ConsumeScopeAction
 
     @Override
     public void decorate(final Tracer.SpanBuilder spanBuilder, final ConsumerRecord record) {
-      final String topic = record.topic() == null ? "unknown" : record.topic();
+      final String topic = record.topic() == null ? "kafka" : record.topic();
       final SpanContext spanContext =
           GlobalTracer.get()
               .extract(Format.Builtin.TEXT_MAP, new TextMapExtractAdapter(record.headers()));

dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java

@@ -73,7 +73,7 @@ public static Scope startSpan(
       callback = new ProducerCallback(callback, scope);
 
       final Span span = scope.span();
-      final String topic = record.topic() == null ? "unknown" : record.topic();
+      final String topic = record.topic() == null ? "kafka" : record.topic();
       if (record.partition() != null) {
         span.setTag("kafka.partition", record.partition());
       }

dd-trace-ot/dd-trace-ot.gradle

@@ -10,7 +10,6 @@ apply from: "${rootDir}/gradle/publish.gradle"
 minimumBranchCoverage = 0.5
 minimumInstructionCoverage = 0.6
 whitelistedInstructionClasses += whitelistedBranchClasses += [
-  'datadog.opentracing.decorators.*',
   'datadog.trace.common.writer.ListWriter',
   'datadog.trace.common.sampling.PrioritySampling'
 ]

dd-trace-ot/src/main/java/datadog/opentracing/DDSpanContext.java

@@ -4,7 +4,6 @@
 import datadog.opentracing.decorators.AbstractDecorator;
 import datadog.trace.api.DDTags;
 import datadog.trace.api.sampling.PrioritySampling;
-import io.opentracing.tag.Tags;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -238,25 +237,14 @@ public synchronized void setTag(final String tag, final Object value) {
       return;
     }
 
-    if (tag.equals(DDTags.SERVICE_NAME)) {
-      setServiceName(value.toString());
-      return;
-    } else if (tag.equals(DDTags.RESOURCE_NAME)) {
-      setResourceName(value.toString());
-      return;
-    } else if (tag.equals(DDTags.SPAN_TYPE)) {
-      setSpanType(value.toString());
-      return;
-    }
-
-    this.tags.put(tag, value);
+    boolean addTag = true;
 
     // Call decorators
     final List<AbstractDecorator> decorators = tracer.getSpanContextDecorators(tag);
     if (decorators != null && value != null) {
       for (final AbstractDecorator decorator : decorators) {
         try {
-          decorator.afterSetTag(this, tag, value);
+          addTag &= decorator.shouldSetTag(this, tag, value);
         } catch (final Throwable ex) {
           log.debug(
               "Could not decorate the span decorator={}: {}",
@@ -265,10 +253,9 @@ public synchronized void setTag(final String tag, final Object value) {
         }
       }
     }
-    // Error management
-    if (Tags.ERROR.getKey().equals(tag)
-        && Boolean.TRUE.equals(value instanceof String ? Boolean.valueOf((String) value) : value)) {
-      this.errorFlag = true;
+
+    if (addTag) {
+      this.tags.put(tag, value);
     }
   }
 

dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java

@@ -56,7 +56,7 @@ public class DDTracer implements io.opentracing.Tracer {
   final ContextualScopeManager scopeManager = new ContextualScopeManager();
 
   /** A set of tags that are added to every span */
-  private final Map<String, Object> spanTags;
+  private final Map<String, String> spanTags;
 
   /** Span context decorators */
   private final Map<String, List<AbstractDecorator>> spanContextDecorators =
@@ -88,24 +88,31 @@ public DDTracer(final Properties config) {
         config.getProperty(DDTraceConfig.SERVICE_NAME),
         Writer.Builder.forConfig(config),
         Sampler.Builder.forConfig(config),
-        DDTraceConfig.parseMap(config.getProperty(DDTraceConfig.SPAN_TAGS)));
+        DDTraceConfig.parseMap(config.getProperty(DDTraceConfig.SPAN_TAGS)),
+        DDTraceConfig.parseMap(config.getProperty(DDTraceConfig.SERVICE_MAPPING)));
     log.debug("Using config: {}", config);
   }
 
   public DDTracer(final String serviceName, final Writer writer, final Sampler sampler) {
-    this(serviceName, writer, sampler, Collections.<String, Object>emptyMap());
+    this(
+        serviceName,
+        writer,
+        sampler,
+        Collections.<String, String>emptyMap(),
+        Collections.<String, String>emptyMap());
   }
 
   public DDTracer(
       final String serviceName,
       final Writer writer,
       final Sampler sampler,
-      final Map<String, Object> spanTags) {
+      final Map<String, String> defaultSpanTags,
+      final Map<String, String> serviceNameMappings) {
     this.serviceName = serviceName;
     this.writer = writer;
     this.writer.start();
     this.sampler = sampler;
-    this.spanTags = spanTags;
+    this.spanTags = defaultSpanTags;
 
     registry = new CodecRegistry();
     registry.register(Format.Builtin.HTTP_HEADERS, new HTTPCodec());
@@ -118,7 +125,8 @@ public DDTracer(
 
     registerClassLoader(ClassLoader.getSystemClassLoader());
 
-    final List<AbstractDecorator> decorators = DDDecoratorsFactory.createBuiltinDecorators();
+    final List<AbstractDecorator> decorators =
+        DDDecoratorsFactory.createBuiltinDecorators(serviceNameMappings);
     for (final AbstractDecorator decorator : decorators) {
       log.debug("Loading decorator: {}", decorator.getClass().getSimpleName());
       addDecorator(decorator);
@@ -132,7 +140,8 @@ public DDTracer(final Writer writer) {
         UNASSIGNED_DEFAULT_SERVICE_NAME,
         writer,
         new AllSampler(),
-        DDTraceConfig.parseMap(new DDTraceConfig().getProperty(DDTraceConfig.SPAN_TAGS)));
+        DDTraceConfig.parseMap(new DDTraceConfig().getProperty(DDTraceConfig.SPAN_TAGS)),
+        DDTraceConfig.parseMap(new DDTraceConfig().getProperty(DDTraceConfig.SERVICE_MAPPING)));
   }
 
   /**
@@ -302,7 +311,9 @@ public class DDSpanBuilder implements SpanBuilder {
 
     // Builder attributes
     private Map<String, Object> tags =
-        spanTags.isEmpty() ? Collections.<String, Object>emptyMap() : new HashMap<>(spanTags);
+        spanTags.isEmpty()
+            ? Collections.<String, Object>emptyMap()
+            : new HashMap<String, Object>(spanTags);
     private long timestampMicro;
     private SpanContext parent;
     private String serviceName;

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

@@ -1,7 +1,6 @@
 package datadog.opentracing.decorators;
 
 import datadog.opentracing.DDSpanContext;
-import datadog.trace.api.DDTags;
 
 /**
  * Span decorators are called when new tags are written and proceed to various remappings and
@@ -13,27 +12,20 @@ public abstract class AbstractDecorator {
 
   private Object matchingValue;
 
-  private String setTag;
+  private String replacementTag;
 
-  private String setValue;
+  private String replacementValue;
 
-  public boolean afterSetTag(final DDSpanContext context, final String tag, final Object value) {
+  public boolean shouldSetTag(final DDSpanContext context, final String tag, final Object value) {
     if (this.getMatchingValue() == null || this.getMatchingValue().equals(value)) {
-      final String targetTag = getSetTag() == null ? tag : getSetTag();
-      final String targetValue = getSetValue() == null ? String.valueOf(value) : getSetValue();
-
-      if (targetTag.equals(DDTags.SERVICE_NAME)) {
-        context.setServiceName(targetValue);
-      } else if (targetTag.equals(DDTags.RESOURCE_NAME)) {
-        context.setResourceName(targetValue);
-      } else if (targetTag.equals(DDTags.SPAN_TYPE)) {
-        context.setSpanType(targetValue);
-      } else {
-        context.setTag(targetTag, targetValue);
-      }
-      return true;
-    } else {
+      final String targetTag = getReplacementTag() == null ? tag : getReplacementTag();
+      final String targetValue =
+          getReplacementValue() == null ? String.valueOf(value) : getReplacementValue();
+
+      context.setTag(targetTag, targetValue);
       return false;
+    } else {
+      return true;
     }
   }
 
@@ -53,19 +45,19 @@ public void setMatchingValue(final Object value) {
     this.matchingValue = value;
   }
 
-  public String getSetTag() {
-    return setTag;
+  public String getReplacementTag() {
+    return replacementTag;
   }
 
-  public void setSetTag(final String targetTag) {
-    this.setTag = targetTag;
+  public void setReplacementTag(final String targetTag) {
+    this.replacementTag = targetTag;
   }
 
-  public String getSetValue() {
-    return setValue;
+  public String getReplacementValue() {
+    return replacementValue;
   }
 
-  public void setSetValue(final String targetValue) {
-    this.setValue = targetValue;
+  public void setReplacementValue(final String targetValue) {
+    this.replacementValue = targetValue;
   }
 }

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

@@ -9,11 +9,11 @@ public class DBStatementAsResourceName extends AbstractDecorator {
   public DBStatementAsResourceName() {
     super();
     this.setMatchingTag(Tags.DB_STATEMENT.getKey());
-    this.setSetTag(DDTags.RESOURCE_NAME);
+    this.setReplacementTag(DDTags.RESOURCE_NAME);
   }
 
   @Override
-  public boolean afterSetTag(final DDSpanContext context, final String tag, final Object value) {
+  public boolean shouldSetTag(final DDSpanContext context, final String tag, final Object value) {
 
     // Special case: Mongo
     // Skip the decorators
@@ -23,15 +23,15 @@ public boolean afterSetTag(final DDSpanContext context, final String tag, final
     }
 
     // Assign service name
-    if (super.afterSetTag(context, tag, value)) {
+    if (!super.shouldSetTag(context, tag, value)) {
       // TODO: remove properly the tag (immutable at this time)
       // the `db.statement` tag must be removed because it will be set
       // by the Datadog Trace Agent as `sql.query`; here we're removing
       // a duplicate that will not be obfuscated with the current Datadog
       // Trace Agent version.
       context.setTag(Tags.DB_STATEMENT.getKey(), null);
-      return true;
+      return false;
     }
-    return false;
+    return true;
   }
 }

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

@@ -13,14 +13,14 @@ public class DBTypeDecorator extends AbstractDecorator {
   public DBTypeDecorator() {
     super();
     this.setMatchingTag(Tags.DB_TYPE.getKey());
-    this.setSetTag(DDTags.SERVICE_NAME);
+    this.setReplacementTag(DDTags.SERVICE_NAME);
   }
 
   @Override
-  public boolean afterSetTag(final DDSpanContext context, final String tag, final Object value) {
+  public boolean shouldSetTag(final DDSpanContext context, final String tag, final Object value) {
 
     // Assign service name
-    if (super.afterSetTag(context, tag, value)) {
+    if (!super.shouldSetTag(context, tag, value)) {
       // Assign span type to DB
       // Special case: Mongo, set to mongodb
       if ("mongo".equals(value)) {
@@ -33,8 +33,7 @@ public boolean afterSetTag(final DDSpanContext context, final String tag, final
       }
       // Works for: mongo, cassandra, jdbc
       context.setOperationName(String.valueOf(value) + ".query");
-      return true;
     }
-    return false;
+    return true;
   }
 }

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

@@ -1,32 +1,33 @@
 package datadog.opentracing.decorators;
 
-import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 
 /** Create DDSpanDecorators */
 public class DDDecoratorsFactory {
-  public static List<AbstractDecorator> createBuiltinDecorators() {
-    List<AbstractDecorator> builtin = new ArrayList<AbstractDecorator>(8);
-    {
-      final HTTPComponent httpDecorator1 = new HTTPComponent();
-      httpDecorator1.setMatchingTag("component");
-      httpDecorator1.setMatchingValue("okhttp");
-      builtin.add(httpDecorator1);
-    }
-    {
-      final HTTPComponent httpDecorator2 = new HTTPComponent();
-      httpDecorator2.setMatchingTag("component");
-      httpDecorator2.setMatchingValue("java-aws-sdk");
-      builtin.add(httpDecorator2);
-    }
-    builtin.add(new ErrorFlag());
-    builtin.add(new DBTypeDecorator());
-    builtin.add(new DBStatementAsResourceName());
-    builtin.add(new OperationDecorator());
-    builtin.add(new Status404Decorator());
-    builtin.add(new URLAsResourceName());
-    builtin.add(new Status5XXDecorator());
+  public static List<AbstractDecorator> createBuiltinDecorators(
+      final Map<String, String> mappings) {
+    final HTTPComponent httpDecorator1 = new HTTPComponent();
+    httpDecorator1.setMatchingTag("component");
+    httpDecorator1.setMatchingValue("okhttp");
 
-    return builtin;
+    final HTTPComponent httpDecorator2 = new HTTPComponent();
+    httpDecorator2.setMatchingTag("component");
+    httpDecorator2.setMatchingValue("java-aws-sdk");
+
+    return Arrays.asList(
+        new DBStatementAsResourceName(),
+        new DBTypeDecorator(),
+        new ErrorFlag(),
+        httpDecorator1,
+        httpDecorator2,
+        new OperationDecorator(),
+        new ResourceNameDecorator(),
+        new ServiceNameDecorator(mappings),
+        new SpanTypeDecorator(),
+        new Status5XXDecorator(),
+        new Status404Decorator(),
+        new URLAsResourceName());
   }
 }