diff --git a/dd-trace-ot/src/main/java/datadog/opentracing/PendingTrace.java b/dd-trace-ot/src/main/java/datadog/opentracing/PendingTrace.java index 9d39b2e41f..52bb6cd5c5 100644 --- a/dd-trace-ot/src/main/java/datadog/opentracing/PendingTrace.java +++ b/dd-trace-ot/src/main/java/datadog/opentracing/PendingTrace.java @@ -16,6 +16,8 @@ import java.util.concurrent.ThreadFactory; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + import lombok.extern.slf4j.Slf4j; @Slf4j @@ -41,6 +43,7 @@ public class PendingTrace extends ConcurrentLinkedDeque { Collections.newSetFromMap(new ConcurrentHashMap, Boolean>()); private final AtomicInteger pendingReferenceCount = new AtomicInteger(0); + private final AtomicReference> rootSpan = new AtomicReference<>(); /** Ensure a trace is never written multiple times */ private final AtomicBoolean isWritten = new AtomicBoolean(false); @@ -74,6 +77,7 @@ public class PendingTrace extends ConcurrentLinkedDeque { log.debug("{} - span registered for wrong trace ({})", span, traceId); return; } + rootSpan.compareAndSet(null, new WeakReference(span)); synchronized (span) { if (null == span.ref) { span.ref = new WeakReference(span, referenceQueue); @@ -121,6 +125,11 @@ public class PendingTrace extends ConcurrentLinkedDeque { expireSpan(span); } + public DDSpan getRootSpan() { + WeakReference rootRef = rootSpan.get(); + return rootRef == null ? null : rootRef.get(); + } + /** * When using continuations, it's possible one may be used after all existing spans are otherwise * completed, so we need to wait till continuations are de-referenced before reporting.