Merge pull request #345 from DataDog/ark/java-mdc

CorrelationIdentifier for MDC frameworks.

dd-java-agent/testing/src/test/groovy/TraceCorrelationTest.groovy

@@ -0,0 +1,25 @@
+import datadog.opentracing.DDSpan
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.api.CorrelationIdentifier
+import io.opentracing.Scope
+import io.opentracing.util.GlobalTracer
+
+class TraceCorrelationTest extends AgentTestRunner {
+
+  def "access trace correlation only under trace" () {
+    when:
+    Scope scope = GlobalTracer.get().buildSpan("myspan").startActive(true)
+    DDSpan span = (DDSpan) scope.span()
+
+    then:
+    CorrelationIdentifier.traceId == span.traceId
+    CorrelationIdentifier.spanId == span.spanId
+
+    when:
+    scope.close()
+
+    then:
+    CorrelationIdentifier.traceId == 0
+    CorrelationIdentifier.spanId == 0
+  }
+}

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

@@ -0,0 +1,45 @@
+package datadog.trace.api;
+
+import java.util.concurrent.atomic.AtomicReference;
+
+/**
+ * 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 long getTraceId() {
+    return provider.get().getTraceId();
+  }
+
+  public static long getSpanId() {
+    return provider.get().getSpanId();
+  }
+
+  public interface Provider {
+    long getTraceId();
+
+    long getSpanId();
+
+    Provider NO_OP =
+        new Provider() {
+          @Override
+          public long getTraceId() {
+            return 0;
+          }
+
+          @Override
+          public long getSpanId() {
+            return 0;
+          }
+        };
+  }
+}

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

@@ -7,6 +7,7 @@ import datadog.opentracing.propagation.ExtractedContext;
 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;
@@ -147,6 +148,8 @@ public class DDTracer implements io.opentracing.Tracer {
       addDecorator(decorator);
     }
 
+    CorrelationIdentifier.registerIfAbsent(OTTraceCorrelation.INSTANCE);
+
     log.info("New instance: {}", this);
   }
 

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

@@ -0,0 +1,29 @@
+package datadog.opentracing;
+
+import datadog.trace.api.CorrelationIdentifier;
+import io.opentracing.Span;
+import io.opentracing.util.GlobalTracer;
+
+public class OTTraceCorrelation implements CorrelationIdentifier.Provider {
+  public static final OTTraceCorrelation INSTANCE = new OTTraceCorrelation();
+
+  private OTTraceCorrelation() {}
+
+  @Override
+  public long getTraceId() {
+    final Span activeSpan = GlobalTracer.get().activeSpan();
+    if (activeSpan instanceof DDSpan) {
+      return ((DDSpan) activeSpan).getTraceId();
+    }
+    return 0;
+  }
+
+  @Override
+  public long getSpanId() {
+    final Span activeSpan = GlobalTracer.get().activeSpan();
+    if (activeSpan instanceof DDSpan) {
+      return ((DDSpan) activeSpan).getSpanId();
+    }
+    return 0;
+  }
+}