Ensure tracer traceCounter is always shared with DDApi

Author: Andrew Kent

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

@@ -74,7 +74,7 @@ public int compare(final TraceInterceptor o1, final TraceInterceptor o2) {
           });
   private final CodecRegistry registry;
 
-  private final AtomicInteger traceCount = new AtomicInteger(0);
+  private final AtomicInteger traceCount;
 
   /** By default, report to local agent and collect all traces. */
   public DDTracer() {
@@ -178,10 +178,14 @@ public void run() {
     registry = new CodecRegistry();
     registry.register(Format.Builtin.HTTP_HEADERS, new HTTPCodec(taggedHeaders));
     registry.register(Format.Builtin.TEXT_MAP, new HTTPCodec(taggedHeaders));
-    if (this.writer instanceof DDAgentWriter && sampler instanceof DDApi.ResponseListener) {
+    if (this.writer instanceof DDAgentWriter) {
       final DDApi api = ((DDAgentWriter) this.writer).getApi();
-      api.addResponseListener((DDApi.ResponseListener) this.sampler);
-      api.addTraceCounter(traceCount);
+      traceCount = api.getTraceCounter();
+      if (sampler instanceof DDApi.ResponseListener) {
+        api.addResponseListener((DDApi.ResponseListener) this.sampler);
+      }
+    } else {
+      traceCount = new AtomicInteger(0);
     }
 
     registerClassLoader(ClassLoader.getSystemClassLoader());

dd-trace-ot/src/main/java/datadog/trace/common/writer/DDApi.java

@@ -35,7 +35,7 @@ public class DDApi {
   private final String tracesEndpoint;
   private final List<ResponseListener> responseListeners = new ArrayList<>();
 
-  private AtomicInteger traceCount;
+  private final AtomicInteger traceCount = new AtomicInteger(0);
   private volatile long nextAllowedLogTime = 0;
 
   private static final ObjectMapper objectMapper = new ObjectMapper(new MessagePackFactory());
@@ -59,8 +59,8 @@ public void addResponseListener(final ResponseListener listener) {
     }
   }
 
-  public void addTraceCounter(final AtomicInteger traceCount) {
-    this.traceCount = traceCount;
+  public AtomicInteger getTraceCounter() {
+    return traceCount;
   }
 
   /**

dd-trace-ot/src/test/groovy/datadog/trace/DDTracerTest.groovy

@@ -4,6 +4,8 @@ import datadog.opentracing.DDTracer
 import datadog.trace.api.Config
 import datadog.trace.common.sampling.AllSampler
 import datadog.trace.common.sampling.RateByServiceSampler
+import datadog.trace.common.writer.DDAgentWriter
+import datadog.trace.common.writer.DDApi
 import datadog.trace.common.writer.ListWriter
 import datadog.trace.common.writer.LoggingWriter
 import org.junit.Rule
@@ -112,4 +114,19 @@ class DDTracerTest extends Specification {
     tracer.writer == writer
     tracer.runtimeId.length() > 0
   }
+
+  def "Shares TraceCount with DDApi with #key = #value"() {
+    setup:
+    System.setProperty(PREFIX + key, value)
+    final DDTracer tracer = new DDTracer(new Config())
+
+    expect:
+    tracer.writer instanceof DDAgentWriter
+    tracer.traceCount.is(((DDAgentWriter) tracer.writer).getApi().traceCount)
+
+    where:
+    key                      | value
+    Config.PRIORITY_SAMPLING | "true"
+    Config.PRIORITY_SAMPLING | "false"
+  }
 }

dd-trace-ot/src/test/groovy/datadog/trace/api/writer/DDApiTest.groovy

@@ -9,7 +9,6 @@ import datadog.trace.common.writer.DDApi.ResponseListener
 import org.msgpack.jackson.dataformat.MessagePackFactory
 import spock.lang.Specification
 
-import java.util.concurrent.atomic.AtomicInteger
 import java.util.concurrent.atomic.AtomicReference
 
 import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer
@@ -69,6 +68,7 @@ class DDApiTest extends Specification {
 
     expect:
     client.tracesEndpoint == "http://localhost:${agent.address.port}/v0.4/traces"
+    client.getTraceCounter().addAndGet(traces.size()) >= 0
     client.sendTraces(traces)
     agent.lastRequest.contentType == "application/msgpack"
     agent.lastRequest.headers.get("Datadog-Meta-Lang") == "java"
@@ -130,8 +130,8 @@ class DDApiTest extends Specification {
     }
     def client = new DDApi("localhost", agent.address.port)
     client.addResponseListener(responseListener)
-    def traceCounter = new AtomicInteger(3)
-    client.addTraceCounter(traceCounter)
+    def traceCounter = client.getTraceCounter()
+    traceCounter.set(3)
 
     when:
     client.sendTraces([])