Support attaching DetachedSpan state to a thread without creating children

.palantir/revapi.yml

@@ -33,3 +33,8 @@ acceptedBreaks:
     - code: "java.class.removed"
       old: "class com.palantir.tracing.Observers"
       justification: "Sourcegraph search shows no more prod usages of this class"
+  "6.0.0":
+    com.palantir.tracing:tracing:
+    - code: "java.method.addedToInterface"
+      new: "method com.palantir.tracing.CloseableSpan com.palantir.tracing.DetachedSpan::attach()"
+      justification: "DetachedSpan is not meant for external implementations"

changelog/@unreleased/pr-769.v2.yml

@@ -0,0 +1,5 @@
+type: improvement
+improvement:
+  description: Support attaching DetachedSpan state to a thread without creating children
+  links:
+  - https://github.com/palantir/tracing-java/pull/769

tracing/src/main/java/com/palantir/tracing/DetachedSpan.java

@@ -138,6 +138,17 @@ default DetachedSpan childDetachedSpan(@Safe String operation) {
         return childDetachedSpan(operation, SpanType.LOCAL);
     }
 
+    /**
+     * Attaches the current {@link DetachedSpan} state to the current thread without creating additional spans.
+     * This is useful when a long-lived {@link DetachedSpan} measures many smaller operations (like async-io)
+     * in which we don't want to produce spans for each task, but do need tracing state associated for logging
+     * and potential child traces.
+     * @apiNote This must be executed within a try-with-resources block, and the parent detached span must still be
+     * completed separately.
+     */
+    @MustBeClosed
+    CloseableSpan attach();
+
     /**
      * Completes this span. After complete is invoked, other methods are not expected to produce spans, but they must
      * not throw either in order to avoid confusing failures.

tracing/src/main/java/com/palantir/tracing/Tracer.java

@@ -362,6 +362,22 @@ public DetachedSpan childDetachedSpan(String operation, SpanType type) {
             return new SampledDetachedSpan(operation, type, traceId, requestId, Optional.of(openSpan.getSpanId()));
         }
 
+        @Override
+        @MustBeClosed
+        public CloseableSpan attach() {
+            warnIfCompleted("startSpanOnCurrentThread");
+            Trace maybeCurrentTrace = currentTrace.get();
+            Trace newTrace = Trace.of(true, traceId, requestId);
+            // Push the DetachedSpan OpenSpan to provide the correct parent information
+            // to child spans created within the context of this attach.
+            // It is VITAL that this span is never completed, it exists only for attribution.
+            newTrace.push(openSpan);
+            setTrace(newTrace);
+            // Do not complete the synthetic root span, it simply prevents nested spans from removing trace state, and
+            // allows
+            return maybeCurrentTrace == null ? REMOVE_TRACE : () -> Tracer.setTrace(maybeCurrentTrace);
+        }
+
         @Override
         public void complete() {
             complete(NoTagTranslator.INSTANCE, NoTagTranslator.INSTANCE);
@@ -463,6 +479,14 @@ public DetachedSpan childDetachedSpan(String _operation, SpanType _type) {
             return this;
         }
 
+        @Override
+        @MustBeClosed
+        public CloseableSpan attach() {
+            // In the unsampled case this method is equivalent to 'childSpan' because spans are neither
+            // measured nor emitted.
+            return childSpan("SYNTHETIC_ATTACH");
+        }
+
         @Override
         public void complete() {
             // nop
@@ -481,6 +505,7 @@ public String toString() {
 
     // Complete the current span.
     private static final CloseableSpan DEFAULT_CLOSEABLE_SPAN = Tracer::fastCompleteSpan;
+    private static final CloseableSpan REMOVE_TRACE = Tracer::clearCurrentTrace;
 
     private static final class TraceRestoringCloseableSpan implements CloseableSpan {
 

tracing/src/main/java/com/palantir/tracing/Tracers.java

@@ -246,9 +246,7 @@ public static <T, U extends ListenableFuture<T>> U wrapListenableFuture(
         U result = null;
         // n.b. This span is required to apply tracing thread state to an initial request. Otherwise if there is
         // no active trace, the detached span would not be associated with work initiated by delegateFactory.
-        try (CloseableSpan ignored =
-                // This could be more efficient using https://github.com/palantir/tracing-java/issues/177
-                span.childSpan(operation + " initial", metadata)) {
+        try (CloseableSpan ignored = span.attach()) {
             result = Preconditions.checkNotNull(delegateFactory.get(), "Expected a ListenableFuture");
         } finally {
             if (result != null) {

tracing/src/test/java/com/palantir/tracing/DetachedSpanTest.java

@@ -0,0 +1,151 @@
+/*
+ * (c) Copyright 2021 Palantir Technologies Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.palantir.tracing;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import com.palantir.tracing.api.Span;
+import java.util.List;
+import java.util.UUID;
+import java.util.concurrent.CopyOnWriteArrayList;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.slf4j.MDC;
+
+public final class DetachedSpanTest {
+
+    @Before
+    public void before() {
+        MDC.clear();
+        Tracer.clearCurrentTrace();
+    }
+
+    @After
+    public void after() {
+        // Clear out the old trace from each test
+        Tracer.clearCurrentTrace();
+    }
+
+    @Test
+    public void testAttach_sampled() {
+        Tracer.setSampler(AlwaysSampler.INSTANCE);
+        String subscription = UUID.randomUUID().toString();
+        List<Span> spans = new CopyOnWriteArrayList<>();
+        Tracer.subscribe(subscription, spans::add);
+        try {
+            DetachedSpan span = DetachedSpan.start("root");
+            assertThat(Tracer.hasTraceId()).isFalse();
+            try (CloseableSpan ignored = span.attach()) {
+                assertThat(Tracer.hasTraceId()).as("Expected a traceId").isTrue();
+                String traceId = Tracer.getTraceId();
+                Tracer.fastStartSpan("one");
+                assertThat(Tracer.getTraceId()).isEqualTo(traceId);
+                Tracer.fastCompleteSpan();
+
+                // Completing the first nested span should leave tracing state in-tact
+                assertThat(Tracer.getTraceId()).isEqualTo(traceId);
+
+                Tracer.fastStartSpan("two");
+                assertThat(Tracer.getTraceId()).isEqualTo(traceId);
+                Tracer.fastCompleteSpan();
+            }
+            assertThat(Tracer.hasTraceId())
+                    .as("Completing the detached span clears thread state")
+                    .isFalse();
+            span.complete();
+            assertThat(spans).hasSize(3);
+            Span oneSpan = spans.get(0);
+            Span twoSpan = spans.get(1);
+            Span rootSpan = spans.get(2);
+            assertThat(rootSpan.getParentSpanId()).isEmpty();
+            assertThat(rootSpan.getOperation()).isEqualTo("root");
+            assertThat(rootSpan.getTraceId()).isEqualTo(oneSpan.getTraceId()).isEqualTo(twoSpan.getTraceId());
+            assertThat(oneSpan.getOperation()).isEqualTo("one");
+            assertThat(oneSpan.getParentSpanId()).hasValue(rootSpan.getSpanId());
+            assertThat(twoSpan.getOperation()).isEqualTo("two");
+            assertThat(twoSpan.getParentSpanId()).hasValue(rootSpan.getSpanId());
+        } finally {
+            Tracer.unsubscribe(subscription);
+        }
+    }
+
+    @Test
+    public void testAttachThreadState_sampled() {
+        Tracer.setSampler(AlwaysSampler.INSTANCE);
+        testAttachThreadState();
+    }
+
+    @Test
+    public void testAttachThreadState_unsampled() {
+        Tracer.setSampler(NeverSampler.INSTANCE);
+        testAttachThreadState();
+    }
+
+    private void testAttachThreadState() {
+        DetachedSpan span = DetachedSpan.start("root");
+        assertThat(Tracer.hasTraceId()).isFalse();
+        try (CloseableSpan ignored = span.attach()) {
+            assertThat(Tracer.hasTraceId()).as("Expected a traceId").isTrue();
+            String traceId = Tracer.getTraceId();
+            Tracer.fastStartSpan("one");
+            assertThat(Tracer.getTraceId()).isEqualTo(traceId);
+            Tracer.fastCompleteSpan();
+
+            // Completing the first nested span should leave tracing state in-tact
+            assertThat(Tracer.getTraceId()).isEqualTo(traceId);
+
+            Tracer.fastStartSpan("two");
+            assertThat(Tracer.getTraceId()).isEqualTo(traceId);
+            Tracer.fastCompleteSpan();
+        }
+        assertThat(Tracer.hasTraceId())
+                .as("Completing the detached span clears thread state")
+                .isFalse();
+    }
+
+    @Test
+    public void testAttachRestoresPreviousTrace_sampled() {
+        Tracer.setSampler(AlwaysSampler.INSTANCE);
+        testAttachRestoresPreviousTrace();
+    }
+
+    @Test
+    public void testAttachRestoresPreviousTrace_unsampled() {
+        Tracer.setSampler(NeverSampler.INSTANCE);
+        testAttachRestoresPreviousTrace();
+    }
+
+    private void testAttachRestoresPreviousTrace() {
+        // This detached span has a unique identifier which differs from the trace
+        // initialized by 'otherTrace'
+        DetachedSpan span = DetachedSpan.start("root");
+
+        Tracer.fastStartSpan("otherTrace");
+        String otherTraceId = Tracer.getTraceId();
+        // begin
+        try (CloseableSpan ignored = span.attach()) {
+            String detachedSpanTraceId = Tracer.getTraceId();
+            assertThat(detachedSpanTraceId).isNotEqualTo(otherTraceId);
+        }
+        assertThat(Tracer.getTraceId())
+                .as("The previous thread state should be restored")
+                .isEqualTo(otherTraceId);
+        Tracer.fastCompleteSpan();
+        assertThat(Tracer.hasTraceId()).isFalse();
+    }
+}

tracing/src/test/java/com/palantir/tracing/TracersTest.java

@@ -48,15 +48,13 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
-import org.mockito.MockitoAnnotations;
 import org.slf4j.MDC;
 
 @SuppressWarnings("deprecation")
 public final class TracersTest {
 
     @Before
     public void before() {
-        MockitoAnnotations.initMocks(this);
         MDC.clear();
 
         Tracer.setSampler(AlwaysSampler.INSTANCE);
@@ -383,10 +381,8 @@ public void testWrapFuture_immediate() {
             ListenableFuture<String> traced =
                     Tracers.wrapListenableFuture(operationName, () -> Futures.immediateFuture("result"));
             assertThat(traced).isDone();
-            assertThat(observed).hasSize(2);
-            // Inner operation must complete first to avoid confusion
-            assertThat(observed.get(0)).extracting(Span::getOperation).isEqualTo(operationName + " initial");
-            assertThat(observed.get(1)).extracting(Span::getOperation).isEqualTo(operationName);
+            assertThat(observed).hasSize(1);
+            assertThat(observed.get(0)).extracting(Span::getOperation).isEqualTo(operationName);
             assertThat(observed)
                     .allSatisfy(span ->
                             assertThat(span).extracting(Span::getTraceId).isEqualTo("defaultTraceId"));
@@ -404,10 +400,8 @@ public void testWrapFuture_failure() {
             ListenableFuture<String> traced = Tracers.wrapListenableFuture(
                     operationName, () -> Futures.immediateFailedFuture(new SafeRuntimeException("result")));
             assertThat(traced).isDone();
-            assertThat(observed).hasSize(2);
-            // Inner operation must complete first to avoid confusion
-            assertThat(observed.get(0)).extracting(Span::getOperation).isEqualTo(operationName + " initial");
-            assertThat(observed.get(1)).extracting(Span::getOperation).isEqualTo(operationName);
+            assertThat(observed).hasSize(1);
+            assertThat(observed.get(0)).extracting(Span::getOperation).isEqualTo(operationName);
             assertThat(observed)
                     .allSatisfy(span ->
                             assertThat(span).extracting(Span::getTraceId).isEqualTo("defaultTraceId"));
@@ -424,14 +418,13 @@ public void testWrapFuture_delayed() {
         SettableFuture<String> rawFuture = SettableFuture.create();
         SettableFuture<String> traced = Tracers.wrapListenableFuture(operationName, () -> rawFuture);
         assertThat(traced).isNotDone();
-        // Inner operation has completed
-        assertThat(observed).hasSize(1);
-        assertThat(observed.get(0)).extracting(Span::getOperation).isEqualTo(operationName + " initial");
+        // attached component has completed but does not emit any spans
+        assertThat(observed).isEmpty();
         // Complete the future
         rawFuture.set("complete");
         assertThat(traced).isDone();
-        assertThat(observed).hasSize(2);
-        assertThat(observed.get(1)).extracting(Span::getOperation).isEqualTo(operationName);
+        assertThat(observed).hasSize(1);
+        assertThat(observed.get(0)).extracting(Span::getOperation).isEqualTo(operationName);
         assertThat(observed)
                 .allSatisfy(
                         span -> assertThat(span).extracting(Span::getTraceId).isEqualTo("defaultTraceId"));
@@ -450,9 +443,8 @@ public void testWrapFuture_throws() {
                     }))
                     .isInstanceOf(SafeRuntimeException.class)
                     .hasMessage("initial operation failure");
-            assertThat(observed).hasSize(2);
-            assertThat(observed.get(0)).extracting(Span::getOperation).isEqualTo(operationName + " initial");
-            assertThat(observed.get(1)).extracting(Span::getOperation).isEqualTo(operationName);
+            assertThat(observed).hasSize(1);
+            assertThat(observed.get(0)).extracting(Span::getOperation).isEqualTo(operationName);
             assertThat(observed)
                     .allSatisfy(span ->
                             assertThat(span).extracting(Span::getTraceId).isEqualTo("defaultTraceId"));