adding a list of traces to rebuild the entire trace for DD

src/main/java/com/datadoghq/trace/impl/DDSpan.java

@@ -1,6 +1,9 @@
 package com.datadoghq.trace.impl;
 
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
 import java.util.Map;
+import java.util.Optional;
 
 import com.fasterxml.jackson.annotation.JsonGetter;
 
@@ -17,19 +20,27 @@ public class DDSpan implements io.opentracing.Span {
     protected long startTimeNano; // Only used to measure nano time durations
     protected long durationNano;
     protected final DDSpanContext context;
+    protected final LinkedHashSet<Span> traces;
 
     DDSpan(
             Tracer tracer,
             String operationName,
+            LinkedHashSet<Span> traces,
             Map<String, Object> tags,
             Long timestamp,
             DDSpanContext context) {
+
         this.tracer = tracer;
         this.operationName = operationName;
+        this.traces = Optional.ofNullable(traces).orElse(new LinkedHashSet<>());
         this.tags = tags;
-        this.startTime = System.currentTimeMillis()*1000000;
+        this.startTime = System.currentTimeMillis() * 1000000;
         this.startTimeNano = System.nanoTime();
         this.context = context;
+
+        // track each span of the trace
+        this.traces.add(this);
+
     }
 
     public SpanContext context() {
@@ -37,17 +48,26 @@ public class DDSpan implements io.opentracing.Span {
     }
 
     public void finish() {
-        this.durationNano = System.nanoTime() - startTimeNano;
+        finish(System.nanoTime());
     }
 
     public void finish(long stopTimeMicro) {
         this.durationNano = stopTimeMicro * 1000L - startTime;
+        if (this.isRootSpan()) {
+            this.traces.stream()
+                    .filter(s -> ((DDSpanContext) s.context()).getSpanId() != ((DDSpanContext) this.context()).getSpanId())
+                    .forEach(s -> s.finish());
+        }
     }
 
     public void close() {
         this.finish();
     }
 
+    private boolean isRootSpan() {
+        return context.getTraceId() == context.getSpanId();
+    }
+
     public io.opentracing.Span setTag(String tag, String value) {
         return this.setTag(tag, value);
     }
@@ -91,9 +111,11 @@ public class DDSpan implements io.opentracing.Span {
     }
 
     public Span setOperationName(String operationName) {
-    	if(this.operationName!=null)
-    		throw new IllegalArgumentException("The operationName is already assigned.");
-    		
+        // FIXME: @renaud, the operationName (mandatory) is always set by the constructor
+        // FIXME: should be an UnsupportedOperation if we don't want to update the operationName + final
+        if (this.operationName != null) {
+            throw new IllegalArgumentException("The operationName is already assigned.");
+        }
         this.operationName = operationName;
         return this;
     }
@@ -159,4 +181,8 @@ public class DDSpan implements io.opentracing.Span {
     public int getError() {
         return context.getErrorFlag() ? 1 : 0;
     }
+
+    public LinkedHashSet<Span> getTraces() {
+        return traces;
+    }
 }

src/main/java/com/datadoghq/trace/impl/Tracer.java

@@ -1,9 +1,6 @@
 package com.datadoghq.trace.impl;
 
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Optional;
+import java.util.*;
 
 import com.datadoghq.trace.Utils.TracerLogger;
 
@@ -34,7 +31,7 @@ public class Tracer implements io.opentracing.Tracer {
         private final String operationName;
         private Map<String, Object> tags = new HashMap<>();
         private Long timestamp;
-        private SpanContext parent;
+        private DDSpan parent;
         private String serviceName;
         private String resourceName;
         private boolean errorFlag;
@@ -45,12 +42,14 @@ public class Tracer implements io.opentracing.Tracer {
         }
 
         public Tracer.SpanBuilder asChildOf(SpanContext spanContext) {
-            this.parent = spanContext;
-            return this;
+            throw new UnsupportedOperationException("Should be a complete span");
+            //this.parent = spanContext;
+            //return this;
         }
 
         public Tracer.SpanBuilder asChildOf(Span span) {
-            return asChildOf(span.context());
+            this.parent = (DDSpan) span;
+            return this;
         }
 
         public Tracer.SpanBuilder addReference(String referenceType, SpanContext spanContext) {
@@ -106,9 +105,15 @@ public class Tracer implements io.opentracing.Tracer {
             DDSpanContext context = buildTheSpanContext();
             logger.startNewSpan(this.operationName, context.getSpanId());
 
+            LinkedHashSet traces = null;
+            if (this.parent != null) {
+                traces = parent.getTraces();
+            }
+
             return new DDSpan(
                     Tracer.this,
                     this.operationName,
+                    traces,
                     this.tags,
                     this.timestamp,
                     context);
@@ -120,7 +125,7 @@ public class Tracer implements io.opentracing.Tracer {
 
             long generatedId = generateNewId();
             if (this.parent != null) {
-                DDSpanContext p = (DDSpanContext) this.parent;
+                DDSpanContext p = (DDSpanContext) this.parent.context();
                 context = new DDSpanContext(
                         p.getTraceId(),
                         generatedId,
@@ -154,7 +159,7 @@ public class Tracer implements io.opentracing.Tracer {
             if (parent == null) {
                 return Collections.emptyList();
             }
-            return parent.baggageItems();
+            return parent.context().baggageItems();
         }
     }
 

src/test/java/Example.java

@@ -27,10 +27,8 @@ public class Example {
                 .withServiceName("service-name")
                 .start();
 
-        child.finish();
 
-        writer.write(parent);
-        writer.write(child);
+
 
         writer.close();
 

src/test/java/com/datadoghq/trace/impl/DDSpanBuilderTest.java

@@ -5,11 +5,13 @@ import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Map;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 public class DDSpanBuilderTest {
@@ -168,4 +170,28 @@ public class DDSpanBuilderTest {
 
     }
 
+    @Test
+    public void shouldTrackAllSpanInTrace() {
+
+        ArrayList<DDSpan> spans = new ArrayList<DDSpan>();
+        ArrayList<DDSpan> spansClone = new ArrayList<DDSpan>();
+        final int nbSamples = 10;
+
+        DDSpan root = (DDSpan) tracer.buildSpan("fake_O").start();
+        spans.add(root);
+
+        for (int i = 1; i <= 10; i++) {
+            spans.add((DDSpan) tracer.buildSpan("fake_" + i).asChildOf(spans.get(i - 1)).start());
+        }
+
+        assertThat(root.getTraces()).hasSize(nbSamples + 1);
+        assertThat(root.getTraces()).containsAll(spans);
+        assertThat(spans.get((int) (Math.random() * nbSamples)).getTraces()).containsAll(spans);
+
+        root.finish();
+        spans.forEach(span -> assertThat(span.getDurationNano()).isNotNull());
+
+
+    }
+
 }

src/test/java/com/datadoghq/trace/impl/DDSpanTest.java

@@ -28,6 +28,7 @@ public class DDSpanTest {
                 "fakeName",
                 null,
                 null,
+                null,
                 context
         );
 
@@ -50,13 +51,13 @@ public class DDSpanTest {
                 expectedOperationName1,
                 null,
                 null,
+                null,
                 null
         );
 
         assertThat(span.getOperationName()).isEqualTo(expectedOperationName1);
 
         span.setOperationName(expectedOperationName2);
-
         assertThat(span.getOperationName()).isEqualTo(expectedOperationName1);
     }