diff --git a/src/main/java/com/datadoghq/trace/impl/DDSpan.java b/src/main/java/com/datadoghq/trace/impl/DDSpan.java index 2ddd7c3df8..963ba680bb 100644 --- a/src/main/java/com/datadoghq/trace/impl/DDSpan.java +++ b/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 traces; DDSpan( Tracer tracer, String operationName, + LinkedHashSet traces, Map 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 getTraces() { + return traces; + } } diff --git a/src/main/java/com/datadoghq/trace/impl/Tracer.java b/src/main/java/com/datadoghq/trace/impl/Tracer.java index 212bef9eae..07b4fa15bb 100644 --- a/src/main/java/com/datadoghq/trace/impl/Tracer.java +++ b/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 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(); } } diff --git a/src/test/java/Example.java b/src/test/java/Example.java index 9ef723ad60..8553b8c4ff 100644 --- a/src/test/java/Example.java +++ b/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(); diff --git a/src/test/java/com/datadoghq/trace/impl/DDSpanBuilderTest.java b/src/test/java/com/datadoghq/trace/impl/DDSpanBuilderTest.java index bf2a456c7c..f6b5acd438 100644 --- a/src/test/java/com/datadoghq/trace/impl/DDSpanBuilderTest.java +++ b/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 spans = new ArrayList(); + ArrayList spansClone = new ArrayList(); + 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()); + + + } + } \ No newline at end of file diff --git a/src/test/java/com/datadoghq/trace/impl/DDSpanTest.java b/src/test/java/com/datadoghq/trace/impl/DDSpanTest.java index 71d2188658..7dbdd244f9 100644 --- a/src/test/java/com/datadoghq/trace/impl/DDSpanTest.java +++ b/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); }