updating flare to implement reporter and store only root spans

Author: mhlidd

dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java

@@ -250,8 +250,8 @@ PropagationTags.Factory getPropagationTagsFactory() {
   public void onRootSpanFinished(AgentSpan root, EndpointTracker tracker) {
     if (!root.isOutbound()) {
       profilingContextIntegration.onRootSpanFinished(root, tracker);
-      TracerDump.suspendSpansFromRootSpan(root);
     }
+    TracerDump.suspendRootSpan(root);
   }
 
   /**
@@ -276,6 +276,7 @@ public EndpointTracker onRootSpanStarted(AgentSpan root) {
     if (!root.isOutbound()) {
       return profilingContextIntegration.onRootSpanStarted(root);
     }
+    TracerDump.addActiveSpan(root);
     return null;
   }
 
@@ -1300,7 +1301,6 @@ public CoreSpanBuilder ignoreActiveSpan() {
     private DDSpan buildSpan() {
       addTerminatedContextAsLinks();
       DDSpan span = DDSpan.create(instrumentationName, timestampMicro, buildSpanContext(), links);
-      TracerDump.addActiveSpan(span, span.isLocalRootSpan());
       if (span.isLocalRootSpan()) {
         EndpointTracker tracker = tracer.onRootSpanStarted(span);
         if (tracker != null) {

dd-trace-core/src/main/java/datadog/trace/core/PendingTrace.java

@@ -448,4 +448,8 @@ public static long getDurationNano(CoreSpan<?> span) {
     PendingTrace trace = (PendingTrace) traceCollector;
     return trace.getLastWriteTime() - span.getStartTime();
   }
+
+  public ConcurrentLinkedDeque<DDSpan> getSpans() {
+    return spans;
+  }
 }

dd-trace-core/src/main/java/datadog/trace/core/flare/TracerFlareService.java

@@ -9,7 +9,6 @@
 import datadog.trace.api.time.TimeUtils;
 import datadog.trace.core.CoreTracer;
 import datadog.trace.core.DDTraceCoreInfo;
-import datadog.trace.core.util.TracerDump;
 import datadog.trace.logging.GlobalLogLevelSwitcher;
 import datadog.trace.logging.LogLevel;
 import datadog.trace.util.AgentTaskScheduler;
@@ -216,7 +215,6 @@ private byte[] buildFlareZip(long startMillis, long endMillis, boolean dumpThrea
       addRuntime(zip);
       tracer.addTracerReportToFlare(zip);
       TracerFlare.addReportsToFlare(zip);
-      TracerDump.dumpTrace(zip);
       if (dumpThreads) {
         addThreadDump(zip);
       }

dd-trace-core/src/main/java/datadog/trace/core/util/TracerDump.java

@@ -3,41 +3,45 @@
 import datadog.trace.api.flare.TracerFlare;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.core.DDSpan;
+import datadog.trace.core.PendingTrace;
 import java.io.IOException;
 import java.lang.ref.WeakReference;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
+import java.util.HashSet;
+import java.util.Set;
 import java.util.zip.ZipOutputStream;
 
-public final class TracerDump {
-  private static final Map<WeakReference<DDSpan>, List<WeakReference<DDSpan>>> activeSpans =
-      new HashMap<>();
+public final class TracerDump implements TracerFlare.Reporter {
 
-  public static void addActiveSpan(final DDSpan span, boolean isRootSpan) {
-    WeakReference<DDSpan> weakRootSpan =
-        isRootSpan ? new WeakReference<>(span) : new WeakReference<>(span.getLocalRootSpan());
-    if (isRootSpan) {
-      List<WeakReference<DDSpan>> childSpans = new ArrayList<>();
-      activeSpans.put(weakRootSpan, childSpans);
-    }
-    WeakReference<DDSpan> childSpan = new WeakReference<>(span);
-    activeSpans.get(weakRootSpan).add(childSpan);
+  private static final TracerDump INSTANCE = new TracerDump();
+
+  public static void register() {
+    TracerFlare.addReporter(INSTANCE);
+  }
+
+  private static final Set<WeakReference<DDSpan>> rootSpans = new HashSet<>();
+
+  public static void addActiveSpan(final AgentSpan rootSpan) {
+    DDSpan rootDDSpan = (DDSpan) rootSpan;
+    WeakReference<DDSpan> weakRootSpan = new WeakReference<>(rootDDSpan);
+    rootSpans.add(weakRootSpan);
   }
 
-  public static void suspendSpansFromRootSpan(AgentSpan rootSpan) {
+  public static void suspendRootSpan(final AgentSpan rootSpan) {
     DDSpan rootDDSpan = (DDSpan) rootSpan;
     WeakReference<DDSpan> weakRootSpan = new WeakReference<>(rootDDSpan);
-    activeSpans.remove(weakRootSpan);
+    rootSpans.remove(weakRootSpan);
   }
 
-  public static void dumpTrace(ZipOutputStream zip) throws IOException {
-    for (List<WeakReference<DDSpan>> spans : activeSpans.values()) {
-      for (WeakReference<DDSpan> weakSpan : spans) {
-        DDSpan span = weakSpan.get();
-        if (span != null) {
-          TracerFlare.addText(zip, "trace_dump.txt", span.toString());
+  @Override
+  public void addReportToFlare(ZipOutputStream zip) throws IOException {
+    for (WeakReference<DDSpan> weakRootSpan : rootSpans) {
+      DDSpan rootSpan = weakRootSpan.get();
+      if (rootSpan != null) {
+        PendingTrace trace = (PendingTrace) rootSpan.context().getTraceCollector();
+        if (trace != null) {
+          for (DDSpan span : trace.getSpans()) {
+            TracerFlare.addText(zip, "trace_dump.txt", span.toString());
+          }
         }
       }
     }