From 63b6ce55a3d90951e854776ba353ce3b97f38bb5 Mon Sep 17 00:00:00 2001 From: David Dunn Date: Thu, 10 Jan 2019 12:48:42 -0500 Subject: [PATCH] Restore trace state back to cleared when wrapWithNewTrace is called with no trace state. --- .../java/com/palantir/tracing/Tracer.java | 7 ++++ .../java/com/palantir/tracing/Tracers.java | 28 ++++++++----- .../java/com/palantir/tracing/TracerTest.java | 13 +++++++ .../com/palantir/tracing/TracersTest.java | 39 +++++++++++++++++++ 4 files changed, 78 insertions(+), 9 deletions(-) diff --git a/tracing/src/main/java/com/palantir/tracing/Tracer.java b/tracing/src/main/java/com/palantir/tracing/Tracer.java index 71b05810b..1aac77f8c 100644 --- a/tracing/src/main/java/com/palantir/tracing/Tracer.java +++ b/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 getAndClearTraceIfPresent() { + Optional 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(); diff --git a/tracing/src/main/java/com/palantir/tracing/Tracers.java b/tracing/src/main/java/com/palantir/tracing/Tracers.java index d64998e3f..c5088aac4 100644 --- a/tracing/src/main/java/com/palantir/tracing/Tracers.java +++ b/tracing/src/main/java/com/palantir/tracing/Tracers.java @@ -145,7 +145,7 @@ protected Callable wrapTask(Callable callable) { public static Callable wrapWithNewTrace(Callable delegate) { return () -> { // clear the existing trace and keep it around for restoration when we're done - Trace originalTrace = Tracer.getAndClearTrace(); + Optional originalTrace = Tracer.getAndClearTraceIfPresent(); try { Tracer.initTrace(Optional.empty(), Tracers.randomId()); @@ -153,8 +153,7 @@ public static Callable wrapWithNewTrace(Callable delegate) { return delegate.call(); } finally { Tracer.fastCompleteSpan(); - // restore the trace - Tracer.setTrace(originalTrace); + restoreTrace(originalTrace); } }; } @@ -165,7 +164,7 @@ public static Callable wrapWithNewTrace(Callable 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 originalTrace = Tracer.getAndClearTraceIfPresent(); try { Tracer.initTrace(Optional.empty(), Tracers.randomId()); @@ -173,8 +172,7 @@ public static Runnable wrapWithNewTrace(Runnable delegate) { delegate.run(); } finally { Tracer.fastCompleteSpan(); - // restore the trace - Tracer.setTrace(originalTrace); + restoreTrace(originalTrace); } }; } @@ -188,7 +186,7 @@ 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 originalTrace = Tracer.getAndClearTraceIfPresent(); try { Tracer.initTrace(Optional.empty(), traceId); @@ -196,12 +194,24 @@ public static Runnable wrapWithAlternateTraceId(String traceId, Runnable delegat 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) { + 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. diff --git a/tracing/src/test/java/com/palantir/tracing/TracerTest.java b/tracing/src/test/java/com/palantir/tracing/TracerTest.java index 335e20449..4e28f9da2 100644 --- a/tracing/src/test/java/com/palantir/tracing/TracerTest.java +++ b/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 nonEmptyTrace = Tracer.getAndClearTraceIfPresent(); + assertThat(nonEmptyTrace).hasValue(trace); + assertThat(Tracer.hasTraceId()).isFalse(); + + Optional emptyTrace = Tracer.getAndClearTraceIfPresent(); + assertThat(emptyTrace).isEmpty(); + } + @Test public void testClearAndGetTraceClearsMdc() { Tracer.startSpan("test"); diff --git a/tracing/src/test/java/com/palantir/tracing/TracersTest.java b/tracing/src/test/java/com/palantir/tracing/TracersTest.java index 0c9350c5d..9fe24b310 100644 --- a/tracing/src/test/java/com/palantir/tracing/TracersTest.java +++ b/tracing/src/test/java/com/palantir/tracing/TracersTest.java @@ -24,6 +24,7 @@ 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; @@ -31,6 +32,7 @@ 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