From d2a8b7fb28bfa3ff17abe4dbf3a824b852d9b0a1 Mon Sep 17 00:00:00 2001
From: Andrew Kent <andrew.kent@datadoghq.com>
Date: Thu, 8 Nov 2018 10:48:19 -0800
Subject: [PATCH] Ensure tracer traceCounter is always shared with DDApi

---
 .../main/java/datadog/opentracing/DDTracer.java | 12 ++++++++----
 .../java/datadog/trace/common/writer/DDApi.java |  6 +++---
 .../groovy/datadog/trace/DDTracerTest.groovy    | 17 +++++++++++++++++
 .../datadog/trace/api/writer/DDApiTest.groovy   |  6 +++---
 4 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java b/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java
index fb8e99f4fa..43f6efceae 100644
--- a/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java
+++ b/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java
@@ -74,7 +74,7 @@ public class DDTracer implements io.opentracing.Tracer, Closeable, datadog.trace
           });
   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 class DDTracer implements io.opentracing.Tracer, Closeable, datadog.trace
     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());
diff --git a/dd-trace-ot/src/main/java/datadog/trace/common/writer/DDApi.java b/dd-trace-ot/src/main/java/datadog/trace/common/writer/DDApi.java
index 69d8f7d152..1b0c65b963 100644
--- a/dd-trace-ot/src/main/java/datadog/trace/common/writer/DDApi.java
+++ b/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 class DDApi {
     }
   }
 
-  public void addTraceCounter(final AtomicInteger traceCount) {
-    this.traceCount = traceCount;
+  public AtomicInteger getTraceCounter() {
+    return traceCount;
   }
 
   /**
diff --git a/dd-trace-ot/src/test/groovy/datadog/trace/DDTracerTest.groovy b/dd-trace-ot/src/test/groovy/datadog/trace/DDTracerTest.groovy
index bf8f1fdbde..0845d1fa04 100644
--- a/dd-trace-ot/src/test/groovy/datadog/trace/DDTracerTest.groovy
+++ b/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"
+  }
 }
diff --git a/dd-trace-ot/src/test/groovy/datadog/trace/api/writer/DDApiTest.groovy b/dd-trace-ot/src/test/groovy/datadog/trace/api/writer/DDApiTest.groovy
index 7bb7daa50b..45de8cf82f 100644
--- a/dd-trace-ot/src/test/groovy/datadog/trace/api/writer/DDApiTest.groovy
+++ b/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([])