From 38b237a5723d03b5fdbb9c37378c9580bebd1c18 Mon Sep 17 00:00:00 2001 From: Carter Kozak Date: Thu, 25 Oct 2018 14:57:25 -0400 Subject: [PATCH 1/6] Trace thread state is removed when root spans are completed Multiple root spans are not allowed, traceId state should not leak between multiple root spans. --- .../com/palantir/tracing/DeferredTracer.java | 17 +++- .../java/com/palantir/tracing/Tracer.java | 83 ++++++++++++------- .../palantir/tracing/CloseableTracerTest.java | 4 +- .../java/com/palantir/tracing/TracerTest.java | 36 +++++--- .../com/palantir/tracing/TracersTest.java | 14 ++-- 5 files changed, 100 insertions(+), 54 deletions(-) diff --git a/tracing/src/main/java/com/palantir/tracing/DeferredTracer.java b/tracing/src/main/java/com/palantir/tracing/DeferredTracer.java index c395c52a8..fb75b6653 100644 --- a/tracing/src/main/java/com/palantir/tracing/DeferredTracer.java +++ b/tracing/src/main/java/com/palantir/tracing/DeferredTracer.java @@ -16,6 +16,8 @@ package com.palantir.tracing; +import java.util.Optional; + /** * Utility class for capturing the current trace at time of construction, and then * running callables at some later time with that captured trace. @@ -36,7 +38,7 @@ * */ public final class DeferredTracer { - private final Trace trace; + private final Optional trace; public DeferredTracer() { this.trace = Tracer.copyTrace(); @@ -47,12 +49,19 @@ public DeferredTracer() { * the time of construction of this {@link DeferredTracer}. */ public T withTrace(Tracers.ThrowingCallable inner) throws E { - Trace originalTrace = Tracer.copyTrace(); - Tracer.setTrace(trace); + if (!trace.isPresent()) { + return inner.call(); + } + Optional originalTrace = Tracer.copyTrace(); + Tracer.setTrace(trace.get()); try { return inner.call(); } finally { - Tracer.setTrace(originalTrace); + if (originalTrace.isPresent()) { + Tracer.setTrace(originalTrace.get()); + } else { + Tracer.getAndClearTrace(); + } } } } diff --git a/tracing/src/main/java/com/palantir/tracing/Tracer.java b/tracing/src/main/java/com/palantir/tracing/Tracer.java index 77ead25bb..5f0563122 100644 --- a/tracing/src/main/java/com/palantir/tracing/Tracer.java +++ b/tracing/src/main/java/com/palantir/tracing/Tracer.java @@ -46,11 +46,7 @@ public final class Tracer { private Tracer() {} // Thread-safe since thread-local - private static final ThreadLocal currentTrace = ThreadLocal.withInitial(() -> { - Trace trace = createTrace(Optional.empty(), Tracers.randomId()); - MDC.put(Tracers.TRACE_ID_KEY, trace.getTraceId()); - return trace; - }); + private static final ThreadLocal currentTrace = new ThreadLocal<>(); // Only access in a class-synchronized fashion private static final Map observers = new HashMap<>(); @@ -85,7 +81,8 @@ public static void initTrace(Optional isObservable, String traceId) { * when the current trace is empty. */ public static OpenSpan startSpan(String operation, String parentSpanId, SpanType type) { - Preconditions.checkState(currentTrace.get().isEmpty(), + Trace current = getOrCreateCurrentTrace(); + Preconditions.checkState(current.isEmpty(), "Cannot start a span with explicit parent if the current thread's trace is non-empty"); Preconditions.checkArgument(parentSpanId != null && !parentSpanId.isEmpty(), "parentTraceId must be non-empty: %s", parentSpanId); @@ -95,7 +92,7 @@ public static OpenSpan startSpan(String operation, String parentSpanId, SpanType .parentSpanId(parentSpanId) .type(type) .build(); - currentTrace.get().push(span); + current.push(span); return span; } @@ -119,13 +116,14 @@ private static OpenSpan startSpanInternal(String operation, SpanType type) { .spanId(Tracers.randomId()) .type(type); - Optional prevState = currentTrace.get().top(); + Trace trace = getOrCreateCurrentTrace(); + Optional prevState = trace.top(); if (prevState.isPresent()) { spanBuilder.parentSpanId(prevState.get().getSpanId()); } OpenSpan span = spanBuilder.build(); - currentTrace.get().push(span); + trace.push(span); return span; } @@ -143,10 +141,12 @@ public static void fastCompleteSpan() { * Like {@link #fastCompleteSpan()}, but adds {@code metadata} to the current span being completed. */ public static void fastCompleteSpan(Map metadata) { - popCurrentSpan() - .filter(openSpan -> currentTrace.get().isObservable()) - .map(openSpan -> toSpan(openSpan, metadata)) - .ifPresent(Tracer::notifyObservers); + if (isTraceObservable()) { + String traceId = getTraceId(); + popCurrentSpan() + .map(openSpan -> toSpan(openSpan, metadata, traceId)) + .ifPresent(Tracer::notifyObservers); + } } /** @@ -161,15 +161,15 @@ public static Optional completeSpan() { * Like {@link #completeSpan()}, but adds {@code metadata} to the current span being completed. */ public static Optional completeSpan(Map metadata) { + boolean observable = isTraceObservable(); + Trace trace = currentTrace.get(); Optional maybeSpan = popCurrentSpan() - .map(openSpan -> toSpan(openSpan, metadata)); + .map(openSpan -> toSpan(openSpan, metadata, trace.getTraceId())); // Notify subscribers iff trace is observable - maybeSpan.ifPresent(span -> { - if (currentTrace.get().isObservable()) { - notifyObservers(span); - } - }); + if (maybeSpan.isPresent() && observable) { + notifyObservers(maybeSpan.get()); + } return maybeSpan; } @@ -181,12 +181,20 @@ private static void notifyObservers(Span span) { } private static Optional popCurrentSpan() { - return currentTrace.get().pop(); + Trace trace = currentTrace.get(); + if (trace != null) { + Optional span = trace.pop(); + if (trace.isEmpty()) { + clearCurrentTrace(); + } + return span; + } + return Optional.empty(); } - private static Span toSpan(OpenSpan openSpan, Map metadata) { + private static Span toSpan(OpenSpan openSpan, Map metadata, String traceId) { return Span.builder() - .traceId(getTraceId()) + .traceId(traceId) .spanId(openSpan.getSpanId()) .type(openSpan.type()) .parentSpanId(openSpan.getParentSpanId()) @@ -239,14 +247,13 @@ public static void setSampler(TraceSampler sampler) { /** Returns the globally unique identifier for this thread's trace. */ public static String getTraceId() { - return currentTrace.get().getTraceId(); + return Preconditions.checkNotNull(currentTrace.get(), "There is no root span").getTraceId(); } /** Clears the current trace id and returns (a copy of) it. */ public static Trace getAndClearTrace() { - Trace trace = currentTrace.get(); - currentTrace.remove(); - MDC.remove(Tracers.TRACE_ID_KEY); + Trace trace = getOrCreateCurrentTrace(); + clearCurrentTrace(); return trace; } @@ -255,12 +262,17 @@ public static Trace getAndClearTrace() { * Tracer#completeSpan span completion}. */ public static boolean isTraceObservable() { - return currentTrace.get().isObservable(); + Trace trace = currentTrace.get(); + return trace != null && trace.isObservable(); } /** Returns an independent copy of this thread's {@link Trace}. */ - static Trace copyTrace() { - return currentTrace.get().deepCopy(); + static Optional copyTrace() { + Trace trace = currentTrace.get(); + if (trace != null) { + return Optional.of(trace.deepCopy()); + } + return Optional.empty(); } /** @@ -274,4 +286,17 @@ static void setTrace(Trace trace) { MDC.put(Tracers.TRACE_ID_KEY, trace.getTraceId()); } + private static Trace getOrCreateCurrentTrace() { + Trace trace = currentTrace.get(); + if (trace == null) { + trace = createTrace(Optional.empty(), Tracers.randomId()); + setTrace(trace); + } + return trace; + } + + private static void clearCurrentTrace() { + currentTrace.remove(); + MDC.remove(Tracers.TRACE_ID_KEY); + } } diff --git a/tracing/src/test/java/com/palantir/tracing/CloseableTracerTest.java b/tracing/src/test/java/com/palantir/tracing/CloseableTracerTest.java index 3ef302750..3692afd2e 100644 --- a/tracing/src/test/java/com/palantir/tracing/CloseableTracerTest.java +++ b/tracing/src/test/java/com/palantir/tracing/CloseableTracerTest.java @@ -33,8 +33,8 @@ public void before() { @Test public void startsAndClosesSpan() { try (CloseableTracer tracer = CloseableTracer.startSpan("foo")) { - assertThat(Tracer.copyTrace().top()).isNotEmpty(); + assertThat(Tracer.copyTrace().get().top()).isNotEmpty(); } - assertThat(Tracer.copyTrace().top()).isEmpty(); + assertThat(Tracer.getAndClearTrace().top()).isEmpty(); } } diff --git a/tracing/src/test/java/com/palantir/tracing/TracerTest.java b/tracing/src/test/java/com/palantir/tracing/TracerTest.java index 71ff22d76..e68365bdf 100644 --- a/tracing/src/test/java/com/palantir/tracing/TracerTest.java +++ b/tracing/src/test/java/com/palantir/tracing/TracerTest.java @@ -63,6 +63,7 @@ public void after() { Tracer.setSampler(AlwaysSampler.INSTANCE); Tracer.unsubscribe("1"); Tracer.unsubscribe("2"); + Tracer.getAndClearTrace(); } @Test @@ -164,23 +165,26 @@ public void testInitTraceCallsSampler() throws Exception { when(sampler.sample()).thenReturn(true, false); Tracer.subscribe("1", observer1); - Tracer.initTrace(Optional.empty(), Tracers.randomId()); - verify(sampler).sample(); Span span = startAndCompleteSpan(); + verify(sampler).sample(); verify(observer1).consume(span); verifyNoMoreInteractions(observer1, sampler); Mockito.reset(observer1, sampler); - Tracer.initTrace(Optional.empty(), Tracers.randomId()); - verify(sampler).sample(); startAndCompleteSpan(); // not sampled, see above + verify(sampler).sample(); verifyNoMoreInteractions(observer1, sampler); } @Test public void testTraceCopyIsIndependent() throws Exception { - Trace trace = Tracer.copyTrace(); - trace.push(mock(OpenSpan.class)); + Tracer.startSpan("span"); + try { + Trace trace = Tracer.copyTrace().get(); + trace.push(mock(OpenSpan.class)); + } finally { + Tracer.fastCompleteSpan(); + } assertThat(Tracer.completeSpan().isPresent()).isFalse(); } @@ -189,8 +193,9 @@ public void testSetTraceSetsCurrentTraceAndMdcTraceIdKey() throws Exception { Tracer.startSpan("operation"); Tracer.setTrace(new Trace(true, "newTraceId")); assertThat(Tracer.getTraceId()).isEqualTo("newTraceId"); - assertThat(Tracer.completeSpan()).isEmpty(); assertThat(MDC.get(Tracers.TRACE_ID_KEY)).isEqualTo("newTraceId"); + assertThat(Tracer.completeSpan()).isEmpty(); + assertThat(MDC.get(Tracers.TRACE_ID_KEY)).isNull(); } @Test @@ -245,12 +250,17 @@ public void testFastCompleteSpanWithMetadata() { @Test public void testClearAndGetTraceClearsMdc() { - String startTrace = Tracer.getTraceId(); - assertThat(MDC.get(Tracers.TRACE_ID_KEY)).isEqualTo(startTrace); - - Trace oldTrace = Tracer.getAndClearTrace(); - assertThat(oldTrace.getTraceId()).isEqualTo(startTrace); - assertThat(MDC.get(Tracers.TRACE_ID_KEY)).isNull(); // after clearing, it's empty + Tracer.startSpan("test"); + try { + String startTrace = Tracer.getTraceId(); + assertThat(MDC.get(Tracers.TRACE_ID_KEY)).isEqualTo(startTrace); + + Trace oldTrace = Tracer.getAndClearTrace(); + assertThat(oldTrace.getTraceId()).isEqualTo(startTrace); + assertThat(MDC.get(Tracers.TRACE_ID_KEY)).isNull(); // after clearing, it's empty + } finally { + Tracer.fastCompleteSpan(); + } } private static Span startAndCompleteSpan() { diff --git a/tracing/src/test/java/com/palantir/tracing/TracersTest.java b/tracing/src/test/java/com/palantir/tracing/TracersTest.java index 79c71bb6a..826f6f47c 100644 --- a/tracing/src/test/java/com/palantir/tracing/TracersTest.java +++ b/tracing/src/test/java/com/palantir/tracing/TracersTest.java @@ -21,6 +21,7 @@ import com.google.common.collect.Lists; import com.palantir.tracing.api.OpenSpan; +import java.util.Collections; import java.util.HashSet; import java.util.List; import java.util.Set; @@ -374,11 +375,12 @@ public void run() { } private static List getCurrentFullTrace() { - Trace trace = Tracer.copyTrace(); - List spans = Lists.newArrayList(); - while (!trace.isEmpty()) { - spans.add(trace.pop().get()); - } - return spans; + return Tracer.copyTrace().map(trace -> { + List spans = Lists.newArrayList(); + while (!trace.isEmpty()) { + spans.add(trace.pop().get()); + } + return spans; + }).orElse(Collections.emptyList()); } } From 10b89ba65329373d2941c227d5b190bd21be14b4 Mon Sep 17 00:00:00 2001 From: Carter Kozak Date: Fri, 26 Oct 2018 13:08:29 -0400 Subject: [PATCH 2/6] Tracer.fastCompleteSpan fix --- .../src/main/java/com/palantir/tracing/Tracer.java | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/tracing/src/main/java/com/palantir/tracing/Tracer.java b/tracing/src/main/java/com/palantir/tracing/Tracer.java index 5f0563122..c70c9ed28 100644 --- a/tracing/src/main/java/com/palantir/tracing/Tracer.java +++ b/tracing/src/main/java/com/palantir/tracing/Tracer.java @@ -141,11 +141,14 @@ public static void fastCompleteSpan() { * Like {@link #fastCompleteSpan()}, but adds {@code metadata} to the current span being completed. */ public static void fastCompleteSpan(Map metadata) { - if (isTraceObservable()) { - String traceId = getTraceId(); - popCurrentSpan() - .map(openSpan -> toSpan(openSpan, metadata, traceId)) - .ifPresent(Tracer::notifyObservers); + Trace trace = currentTrace.get(); + if (trace != null) { + boolean observable = isTraceObservable(); + Optional span = popCurrentSpan(); + if (observable) { + span.map(openSpan -> toSpan(openSpan, metadata, trace.getTraceId())) + .ifPresent(Tracer::notifyObservers); + } } } From de812a376006e8a72d67d91fb674231515fd11b8 Mon Sep 17 00:00:00 2001 From: Carter Kozak Date: Fri, 26 Oct 2018 13:54:34 -0400 Subject: [PATCH 3/6] Prefer Trace.isObservable over static lookups --- tracing/src/main/java/com/palantir/tracing/Tracer.java | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tracing/src/main/java/com/palantir/tracing/Tracer.java b/tracing/src/main/java/com/palantir/tracing/Tracer.java index c70c9ed28..83d9f3aec 100644 --- a/tracing/src/main/java/com/palantir/tracing/Tracer.java +++ b/tracing/src/main/java/com/palantir/tracing/Tracer.java @@ -143,9 +143,8 @@ public static void fastCompleteSpan() { public static void fastCompleteSpan(Map metadata) { Trace trace = currentTrace.get(); if (trace != null) { - boolean observable = isTraceObservable(); Optional span = popCurrentSpan(); - if (observable) { + if (trace.isObservable()) { span.map(openSpan -> toSpan(openSpan, metadata, trace.getTraceId())) .ifPresent(Tracer::notifyObservers); } @@ -164,13 +163,15 @@ public static Optional completeSpan() { * Like {@link #completeSpan()}, but adds {@code metadata} to the current span being completed. */ public static Optional completeSpan(Map metadata) { - boolean observable = isTraceObservable(); Trace trace = currentTrace.get(); + if (trace == null) { + return Optional.empty(); + } Optional maybeSpan = popCurrentSpan() .map(openSpan -> toSpan(openSpan, metadata, trace.getTraceId())); // Notify subscribers iff trace is observable - if (maybeSpan.isPresent() && observable) { + if (maybeSpan.isPresent() && trace.isObservable()) { notifyObservers(maybeSpan.get()); } From 1cfd1528c8d62414b3dd66432734ef1644b05cb7 Mon Sep 17 00:00:00 2001 From: Carter Kozak Date: Fri, 26 Oct 2018 14:05:28 -0400 Subject: [PATCH 4/6] lambda warning comments --- tracing/src/main/java/com/palantir/tracing/Tracer.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tracing/src/main/java/com/palantir/tracing/Tracer.java b/tracing/src/main/java/com/palantir/tracing/Tracer.java index 83d9f3aec..7589c4528 100644 --- a/tracing/src/main/java/com/palantir/tracing/Tracer.java +++ b/tracing/src/main/java/com/palantir/tracing/Tracer.java @@ -118,6 +118,7 @@ private static OpenSpan startSpanInternal(String operation, SpanType type) { Trace trace = getOrCreateCurrentTrace(); Optional prevState = trace.top(); + // Avoid lambda allocation in hot paths if (prevState.isPresent()) { spanBuilder.parentSpanId(prevState.get().getSpanId()); } @@ -172,6 +173,7 @@ public static Optional completeSpan(Map metadata) { // Notify subscribers iff trace is observable if (maybeSpan.isPresent() && trace.isObservable()) { + // Avoid lambda allocation in hot paths notifyObservers(maybeSpan.get()); } From 102b5dfd1bc45851f666ab20211ee6f5af25721b Mon Sep 17 00:00:00 2001 From: Carter Kozak Date: Sun, 28 Oct 2018 15:16:19 -0400 Subject: [PATCH 5/6] Test that closing a root span completes the trace --- .../java/com/palantir/tracing/TracerTest.java | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tracing/src/test/java/com/palantir/tracing/TracerTest.java b/tracing/src/test/java/com/palantir/tracing/TracerTest.java index e68365bdf..fed3db952 100644 --- a/tracing/src/test/java/com/palantir/tracing/TracerTest.java +++ b/tracing/src/test/java/com/palantir/tracing/TracerTest.java @@ -31,6 +31,8 @@ import com.palantir.tracing.api.SpanType; import java.util.Map; import java.util.Optional; +import java.util.Set; +import org.assertj.core.util.Sets; import org.junit.After; import org.junit.Before; import org.junit.Test; @@ -263,6 +265,22 @@ public void testClearAndGetTraceClearsMdc() { } } + @Test + public void testCompleteRootSpanCompletesTrace() { + Set traceIds = Sets.newHashSet(); + SpanObserver traceIdCaptor = span -> traceIds.add(span.getTraceId()); + Tracer.subscribe("traceIds", traceIdCaptor); + try { + // Only one root span is allowed, the second span + // is expected to generate a second traceId + startAndCompleteSpan(); + startAndCompleteSpan(); + } finally { + Tracer.unsubscribe("traceIds"); + } + assertThat(traceIds.size()).isEqualTo(2); + } + private static Span startAndCompleteSpan() { Tracer.startSpan("operation"); return Tracer.completeSpan().get(); From 1bbfefa10b744d1b9c2e93ea9e85041ba9157f7f Mon Sep 17 00:00:00 2001 From: Carter Kozak Date: Mon, 29 Oct 2018 08:57:55 -0400 Subject: [PATCH 6/6] Implement Tracer.hasTraceId --- .../src/main/java/com/palantir/tracing/Tracer.java | 5 +++++ .../test/java/com/palantir/tracing/TracerTest.java | 12 ++++++++++++ 2 files changed, 17 insertions(+) diff --git a/tracing/src/main/java/com/palantir/tracing/Tracer.java b/tracing/src/main/java/com/palantir/tracing/Tracer.java index 7589c4528..82d475d26 100644 --- a/tracing/src/main/java/com/palantir/tracing/Tracer.java +++ b/tracing/src/main/java/com/palantir/tracing/Tracer.java @@ -251,6 +251,11 @@ public static void setSampler(TraceSampler sampler) { Tracer.sampler = sampler; } + /** Returns true if there is an active trace on this thread. */ + public static boolean hasTraceId() { + return currentTrace.get() != null; + } + /** Returns the globally unique identifier for this thread's trace. */ public static String getTraceId() { return Preconditions.checkNotNull(currentTrace.get(), "There is no root span").getTraceId(); diff --git a/tracing/src/test/java/com/palantir/tracing/TracerTest.java b/tracing/src/test/java/com/palantir/tracing/TracerTest.java index fed3db952..a69b1b2b0 100644 --- a/tracing/src/test/java/com/palantir/tracing/TracerTest.java +++ b/tracing/src/test/java/com/palantir/tracing/TracerTest.java @@ -281,6 +281,18 @@ public void testCompleteRootSpanCompletesTrace() { assertThat(traceIds.size()).isEqualTo(2); } + @Test + public void testHasTraceId() { + assertThat(Tracer.hasTraceId()).isEqualTo(false); + Tracer.startSpan("testSpan"); + try { + assertThat(Tracer.hasTraceId()).isEqualTo(true); + } finally { + Tracer.completeSpan(); + } + assertThat(Tracer.hasTraceId()).isEqualTo(false); + } + private static Span startAndCompleteSpan() { Tracer.startSpan("operation"); return Tracer.completeSpan().get();