TraceInterceptor registration api

Author: Andrew Kent

dd-trace-api/src/main/java/datadog/trace/api/CorrelationIdentifier.java

@@ -1,27 +1,19 @@
 package datadog.trace.api;
 
-import java.util.concurrent.atomic.AtomicReference;
+import datadog.trace.context.TracerBridge;
 
 /**
  * Utility class to access the active trace and span ids.
  *
  * <p>Intended for use with MDC frameworks.
  */
 public class CorrelationIdentifier {
-  private static final AtomicReference<Provider> provider = new AtomicReference<>(Provider.NO_OP);
-
-  public static void registerIfAbsent(Provider p) {
-    if (p != null && p != Provider.NO_OP) {
-      provider.compareAndSet(Provider.NO_OP, p);
-    }
-  }
-
   public static String getTraceId() {
-    return provider.get().getTraceId();
+    return TracerBridge.get().getTraceId();
   }
 
   public static String getSpanId() {
-    return provider.get().getSpanId();
+    return TracerBridge.get().getSpanId();
   }
 
   public interface Provider {

dd-trace-api/src/main/java/datadog/trace/context/TracerBridge.java

@@ -0,0 +1,63 @@
+package datadog.trace.context;
+
+import datadog.trace.api.interceptor.TraceInterceptor;
+import java.util.concurrent.atomic.AtomicReference;
+
+/** A global reference to the registered Datadog tracer. */
+public class TracerBridge {
+  private static final AtomicReference<Provider> provider = new AtomicReference<>(Provider.NO_OP);
+
+  public static void registerIfAbsent(Provider p) {
+    if (p != null && p != Provider.NO_OP) {
+      provider.compareAndSet(Provider.NO_OP, p);
+    }
+  }
+
+  public static Provider get() {
+    return provider.get();
+  }
+
+  /**
+   * Add a new interceptor to the tracer. Interceptors with duplicate priority to existing ones are
+   * ignored.
+   *
+   * @param traceInterceptor
+   * @return false if an interceptor with same priority exists.
+   */
+  public static boolean addTraceInterceptor(TraceInterceptor traceInterceptor) {
+    return get().addTraceInterceptor(traceInterceptor);
+  }
+
+  public interface Provider {
+    String getTraceId();
+
+    String getSpanId();
+
+    /**
+     * Add a new interceptor to the tracer. Interceptors with duplicate priority to existing ones
+     * are ignored.
+     *
+     * @param traceInterceptor
+     * @return false if an interceptor with same priority exists.
+     */
+    boolean addTraceInterceptor(TraceInterceptor traceInterceptor);
+
+    Provider NO_OP =
+        new Provider() {
+          @Override
+          public String getTraceId() {
+            return "0";
+          }
+
+          @Override
+          public String getSpanId() {
+            return "0";
+          }
+
+          @Override
+          public boolean addTraceInterceptor(TraceInterceptor traceInterceptor) {
+            return false;
+          }
+        };
+  }
+}

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

@@ -13,7 +13,6 @@
 import datadog.opentracing.propagation.HTTPCodec;
 import datadog.opentracing.scopemanager.ContextualScopeManager;
 import datadog.opentracing.scopemanager.ScopeContext;
-import datadog.trace.api.CorrelationIdentifier;
 import datadog.trace.api.interceptor.MutableSpan;
 import datadog.trace.api.interceptor.TraceInterceptor;
 import datadog.trace.api.sampling.PrioritySampling;
@@ -24,6 +23,7 @@
 import datadog.trace.common.writer.DDAgentWriter;
 import datadog.trace.common.writer.DDApi;
 import datadog.trace.common.writer.Writer;
+import datadog.trace.context.TracerBridge;
 import io.opentracing.Scope;
 import io.opentracing.ScopeManager;
 import io.opentracing.Span;
@@ -49,7 +49,7 @@
 
 /** DDTracer makes it easy to send traces and span to DD using the OpenTracing API. */
 @Slf4j
-public class DDTracer implements io.opentracing.Tracer, Closeable {
+public class DDTracer implements io.opentracing.Tracer, Closeable, TracerBridge.Provider {
 
   public static final String UNASSIGNED_DEFAULT_SERVICE_NAME = "unnamed-java-app";
 
@@ -162,7 +162,7 @@ public void run() {
       addDecorator(decorator);
     }
 
-    CorrelationIdentifier.registerIfAbsent(OTTraceCorrelation.INSTANCE);
+    TracerBridge.registerIfAbsent(this);
 
     log.info("New instance: {}", this);
   }
@@ -202,17 +202,6 @@ public void addDecorator(final AbstractDecorator decorator) {
     spanContextDecorators.put(decorator.getMatchingTag(), list);
   }
 
-  /**
-   * Add a new interceptor to the tracer. Interceptors with duplicate priority to existing ones are
-   * ignored.
-   *
-   * @param interceptor
-   * @return false if an interceptor with same priority exists.
-   */
-  public boolean addInterceptor(final TraceInterceptor interceptor) {
-    return interceptors.add(interceptor);
-  }
-
   public void addScopeContext(final ScopeContext context) {
     scopeManager.addScopeContext(context);
   }
@@ -227,7 +216,7 @@ public void registerClassLoader(final ClassLoader classLoader) {
     try {
       for (final TraceInterceptor interceptor :
           ServiceLoader.load(TraceInterceptor.class, classLoader)) {
-        addInterceptor(interceptor);
+        addTraceInterceptor(interceptor);
       }
     } catch (final ServiceConfigurationError e) {
       log.warn("Problem loading TraceInterceptor for classLoader: " + classLoader, e);
@@ -308,6 +297,29 @@ void incrementTraceCount() {
     traceCount.incrementAndGet();
   }
 
+  @Override
+  public String getTraceId() {
+    final Span activeSpan = this.activeSpan();
+    if (activeSpan instanceof DDSpan) {
+      return ((DDSpan) activeSpan).getTraceId();
+    }
+    return "0";
+  }
+
+  @Override
+  public String getSpanId() {
+    final Span activeSpan = this.activeSpan();
+    if (activeSpan instanceof DDSpan) {
+      return ((DDSpan) activeSpan).getSpanId();
+    }
+    return "0";
+  }
+
+  @Override
+  public boolean addTraceInterceptor(final TraceInterceptor interceptor) {
+    return interceptors.add(interceptor);
+  }
+
   @Override
   public void close() {
     PendingTrace.close();

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

@@ -4,14 +4,12 @@ import datadog.trace.common.writer.ListWriter
 import spock.lang.Shared
 import spock.lang.Specification
 
-class OTTraceCorrelationTest extends Specification {
+class TraceCorrelationTest extends Specification {
 
   static final WRITER = new ListWriter()
 
   @Shared
   DDTracer tracer = new DDTracer(WRITER)
-  @Shared
-  OTTraceCorrelation traceCorrelation = new OTTraceCorrelation(tracer)
 
   def scope = tracer.buildSpan("test").startActive(true)
 
@@ -24,24 +22,24 @@ class OTTraceCorrelationTest extends Specification {
     scope.close()
 
     expect:
-    "0" == traceCorrelation.getTraceId()
+    "0" == tracer.getTraceId()
   }
 
   def "get trace id with trace"() {
     expect:
-    ((DDSpan) scope.span()).traceId == traceCorrelation.getTraceId()
+    ((DDSpan) scope.span()).traceId == tracer.getTraceId()
   }
 
   def "get span id without span"() {
     setup:
     scope.close()
 
     expect:
-    "0" == traceCorrelation.getSpanId()
+    "0" == tracer.getSpanId()
   }
 
   def "get span id with trace"() {
     expect:
-    ((DDSpan) scope.span()).spanId == traceCorrelation.getSpanId()
+    ((DDSpan) scope.span()).spanId == tracer.getSpanId()
   }
 }

…pentracing/OTTraceCorrelationTest.groovy …/opentracing/TraceCorrelationTest.groovydd-trace-ot/src/test/groovy/datadog/opentracing/OTTraceCorrelationTest.groovy renamed to dd-trace-ot/src/test/groovy/datadog/opentracing/TraceCorrelationTest.groovy dd-trace-ot/src/test/groovy/datadog/opentracing/OTTraceCorrelationTest.groovy renamed to dd-trace-ot/src/test/groovy/datadog/opentracing/TraceCorrelationTest.groovy

@@ -3,6 +3,7 @@ package datadog.opentracing
 import datadog.trace.api.interceptor.MutableSpan
 import datadog.trace.api.interceptor.TraceInterceptor
 import datadog.trace.common.writer.ListWriter
+import datadog.trace.context.TracerBridge
 import spock.lang.Specification
 
 import java.util.concurrent.atomic.AtomicBoolean
@@ -143,4 +144,22 @@ class TraceInterceptorTest extends Specification {
     tags["thread.id"] != null
     tags.size() == 6
   }
+
+  def "register interceptor through bridge" () {
+    setup:
+    def interceptor = new TraceInterceptor() {
+      @Override
+      Collection<? extends MutableSpan> onTraceComplete(Collection<? extends MutableSpan> trace) {
+        return trace
+      }
+
+      @Override
+      int priority() {
+        return 38
+      }
+    }
+
+    expect:
+    TracerBridge.addTraceInterceptor(interceptor)
+  }
 }