Restore trace state back to cleared when wrapWithNewTrace is called with no trace state

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

@@ -270,6 +270,13 @@ public static String getTraceId() {
         return Preconditions.checkNotNull(currentTrace.get(), "There is no root span").getTraceId();
     }
 
+    /** Clears the current trace id and returns it if present. */
+    static Optional<Trace> getAndClearTraceIfPresent() {
+        Optional<Trace> trace = Optional.ofNullable(currentTrace.get());
+        clearCurrentTrace();
+        return trace;
+    }
+
     /** Clears the current trace id and returns (a copy of) it. */
     public static Trace getAndClearTrace() {
         Trace trace = getOrCreateCurrentTrace();

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

@@ -145,16 +145,15 @@ protected <T> Callable<T> wrapTask(Callable<T> callable) {
     public static <V> Callable<V> wrapWithNewTrace(Callable<V> delegate) {
         return () -> {
             // clear the existing trace and keep it around for restoration when we're done
-            Trace originalTrace = Tracer.getAndClearTrace();
+            Optional<Trace> originalTrace = Tracer.getAndClearTraceIfPresent();
 
             try {
                 Tracer.initTrace(Optional.empty(), Tracers.randomId());
                 Tracer.startSpan(ROOT_SPAN_OPERATION);
                 return delegate.call();
             } finally {
                 Tracer.fastCompleteSpan();
-                // restore the trace
-                Tracer.setTrace(originalTrace);
+                restoreTrace(originalTrace);
             }
         };
     }
@@ -165,16 +164,15 @@ public static <V> Callable<V> wrapWithNewTrace(Callable<V> delegate) {
     public static Runnable wrapWithNewTrace(Runnable delegate) {
         return () -> {
             // clear the existing trace and keep it around for restoration when we're done
-            Trace originalTrace = Tracer.getAndClearTrace();
+            Optional<Trace> originalTrace = Tracer.getAndClearTraceIfPresent();
 
             try {
                 Tracer.initTrace(Optional.empty(), Tracers.randomId());
                 Tracer.startSpan(ROOT_SPAN_OPERATION);
                 delegate.run();
             } finally {
                 Tracer.fastCompleteSpan();
-                // restore the trace
-                Tracer.setTrace(originalTrace);
+                restoreTrace(originalTrace);
             }
         };
     }
@@ -188,20 +186,32 @@ public static Runnable wrapWithNewTrace(Runnable delegate) {
     public static Runnable wrapWithAlternateTraceId(String traceId, Runnable delegate) {
         return () -> {
             // clear the existing trace and keep it around for restoration when we're done
-            Trace originalTrace = Tracer.getAndClearTrace();
+            Optional<Trace> originalTrace = Tracer.getAndClearTraceIfPresent();
 
             try {
                 Tracer.initTrace(Optional.empty(), traceId);
                 Tracer.startSpan(ROOT_SPAN_OPERATION);
                 delegate.run();
             } finally {
                 Tracer.fastCompleteSpan();
-                // restore the trace
-                Tracer.setTrace(originalTrace);
+                restoreTrace(originalTrace);
             }
         };
     }
 
+    /**
+     * Restores or clears trace state based on provided {@link Trace}. Used to cleanup trace state for
+     * {@link #wrapWithNewTrace} calls.
+     */
+    private static void restoreTrace(Optional<Trace> trace) {
+        if (trace.isPresent()) {
+            Tracer.setTrace(trace.get());
+        } else {
+            // Ignoring returned value, used to clear trace only
+            Tracer.getAndClearTraceIfPresent();
+        }
+    }
+
     /**
      * Wraps a given callable such that its execution operates with the {@link Trace thread-local Trace} of the thread
      * that constructs the {@link TracingAwareCallable} instance rather than the thread that executes the callable.

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

@@ -251,6 +251,19 @@ public void testFastCompleteSpanWithMetadata() {
         assertThat(spanCaptor.getValue().getMetadata()).isEqualTo(metadata);
     }
 
+    @Test
+    public void testGetAndClearTraceIfPresent() {
+        Trace trace = new Trace(true, "newTraceId");
+        Tracer.setTrace(trace);
+
+        Optional<Trace> nonEmptyTrace = Tracer.getAndClearTraceIfPresent();
+        assertThat(nonEmptyTrace).hasValue(trace);
+        assertThat(Tracer.hasTraceId()).isFalse();
+
+        Optional<Trace> emptyTrace = Tracer.getAndClearTraceIfPresent();
+        assertThat(emptyTrace).isEmpty();
+    }
+
     @Test
     public void testClearAndGetTraceClearsMdc() {
         Tracer.startSpan("test");

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

@@ -24,13 +24,15 @@
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.MockitoAnnotations;
@@ -42,6 +44,15 @@ public final class TracersTest {
     public void before() {
         MockitoAnnotations.initMocks(this);
         MDC.clear();
+
+        // Initialize a new trace for each test
+        Tracer.initTrace(Optional.empty(), "defaultTraceId");
+    }
+
+    @After
+    public void after() {
+        // Clear out the old trace from each test
+        Tracer.getAndClearTraceIfPresent();
     }
 
     @Test
@@ -243,6 +254,14 @@ public void testWrapCallableWithNewTrace_traceStateRestoredWhenThrows() throws E
         assertThat(Tracer.getTraceId()).isEqualTo(traceIdBeforeConstruction);
     }
 
+    @Test
+    public void testWrapCallableWithNewTrace_traceStateRestoredToCleared() throws Exception {
+        // Clear out the default initialized trace
+        Tracer.getAndClearTraceIfPresent();
+        Tracers.wrapWithNewTrace(() -> null).call();
+        assertThat(Tracer.hasTraceId()).isFalse();
+    }
+
     @Test
     public void testWrapRunnableWithNewTrace_traceStateInsideRunnableIsIsolated() throws Exception {
         String traceIdBeforeConstruction = Tracer.getTraceId();
@@ -306,6 +325,16 @@ public void testWrapRunnableWithNewTrace_traceStateRestoredWhenThrows() throws E
         assertThat(Tracer.getTraceId()).isEqualTo(traceIdBeforeConstruction);
     }
 
+    @Test
+    public void testWrapRunnableWithNewTrace_traceStateRestoredToCleared() {
+        // Clear out the default initialized trace
+        Tracer.getAndClearTraceIfPresent();
+        Tracers.wrapWithNewTrace(() -> {
+            // no-op
+        }).run();
+        assertThat(Tracer.hasTraceId()).isFalse();
+    }
+
     @Test
     public void testWrapRunnableWithAlternateTraceId_traceStateInsideRunnableUsesGivenTraceId() {
         String traceIdBeforeConstruction = Tracer.getTraceId();
@@ -358,6 +387,16 @@ public void testWrapRunnableWithAlternateTraceId_traceStateRestoredWhenThrows()
         assertThat(Tracer.getTraceId()).isEqualTo(traceIdBeforeConstruction);
     }
 
+    @Test
+    public void testWrapRunnableWithAlternateTraceId_traceStateRestoredToCleared() {
+        // Clear out the default initialized trace
+        Tracer.getAndClearTraceIfPresent();
+        Tracers.wrapWithAlternateTraceId("someTraceId", () -> {
+            // no-op
+        }).run();
+        assertThat(Tracer.hasTraceId()).isFalse();
+    }
+
     @Test
     public void testTraceIdGeneration() throws Exception {
         assertThat(Tracers.randomId()).hasSize(16); // fails with p=1/16 if generated string is not padded