Skip to content

Trace thread state is removed when root spans are completed #22

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
iamdanfox merged 6 commits into from
Oct 29, 2018
Merged

Trace thread state is removed when root spans are completed #22

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
38b237a
Trace thread state is removed when root spans are completed
carterkozak Oct 25, 2018
10b89ba
Tracer.fastCompleteSpan fix
carterkozak Oct 26, 2018
de812a3
Prefer Trace.isObservable over static lookups
carterkozak Oct 26, 2018
1cfd152
lambda warning comments
carterkozak Oct 26, 2018
102b5df
Test that closing a root span completes the trace
carterkozak Oct 28, 2018
1bbfefa
Implement Tracer.hasTraceId
carterkozak Oct 29, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 13 additions & 4 deletions tracing/src/main/java/com/palantir/tracing/DeferredTracer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -36,7 +38,7 @@
* </pre>
*/
public final class DeferredTracer {
private final Trace trace;
private final Optional<Trace> trace;

public DeferredTracer() {
this.trace = Tracer.copyTrace();
Expand All @@ -47,12 +49,19 @@ public DeferredTracer() {
* the time of construction of this {@link DeferredTracer}.
*/
public <T, E extends Throwable> T withTrace(Tracers.ThrowingCallable<T, E> inner) throws E {
Trace originalTrace = Tracer.copyTrace();
Tracer.setTrace(trace);
if (!trace.isPresent()) {
return inner.call();
}
Optional<Trace> 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();
}
}
}
}
94 changes: 65 additions & 29 deletions tracing/src/main/java/com/palantir/tracing/Tracer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,11 +46,7 @@ public final class Tracer {
private Tracer() {}

// Thread-safe since thread-local
private static final ThreadLocal<Trace> currentTrace = ThreadLocal.withInitial(() -> {
Trace trace = createTrace(Optional.empty(), Tracers.randomId());
MDC.put(Tracers.TRACE_ID_KEY, trace.getTraceId());
return trace;
});
private static final ThreadLocal<Trace> currentTrace = new ThreadLocal<>();
Copy link
Contributor

@iamdanfox iamdanfox Oct 26, 2018 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK so I guess this is the key change - this guy now uses null to represent the idea that we're not in a trace?

I guess we just have to be super careful now everywhere that we use currentTrace.get() that it might actually be null!

Copy link
Contributor Author

@carterkozak carterkozak Oct 26, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That is correct, otherwise trace thread state leaks outside of root spans, which causes multiple root spans unless code directly creates a new traceId.


// Only access in a class-synchronized fashion
private static final Map<String, SpanObserver> observers = new HashMap<>();
Expand Down Expand Up @@ -85,7 +81,8 @@ public static void initTrace(Optional<Boolean> 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);
Expand All @@ -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;
}

Expand All @@ -119,13 +116,15 @@ private static OpenSpan startSpanInternal(String operation, SpanType type) {
.spanId(Tracers.randomId())
.type(type);

Optional<OpenSpan> prevState = currentTrace.get().top();
Trace trace = getOrCreateCurrentTrace();
Optional<OpenSpan> prevState = trace.top();
// Avoid lambda allocation in hot paths
if (prevState.isPresent()) {
spanBuilder.parentSpanId(prevState.get().getSpanId());
}
Copy link
Contributor

@schlosna schlosna Oct 26, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: could simplify

Suggested change
}
trace.top().ifPresent(openSpan -> spanBuilder.parentSpanId(openSpan.getSpanId()));

Copy link
Contributor Author

@carterkozak carterkozak Oct 26, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I opted to avoid the lambda allocation since this is a relatively hot path. We do a lot elsewhere, so I could go either way.

Copy link
Contributor

@schlosna schlosna Oct 26, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

makes sense


OpenSpan span = spanBuilder.build();
currentTrace.get().push(span);
trace.push(span);
return span;
}

Expand All @@ -143,10 +142,14 @@ public static void fastCompleteSpan() {
* Like {@link #fastCompleteSpan()}, but adds {@code metadata} to the current span being completed.
*/
public static void fastCompleteSpan(Map<String, String> metadata) {
popCurrentSpan()
.filter(openSpan -> currentTrace.get().isObservable())
.map(openSpan -> toSpan(openSpan, metadata))
.ifPresent(Tracer::notifyObservers);
Trace trace = currentTrace.get();
if (trace != null) {
Optional<OpenSpan> span = popCurrentSpan();
if (trace.isObservable()) {
span.map(openSpan -> toSpan(openSpan, metadata, trace.getTraceId()))
.ifPresent(Tracer::notifyObservers);
}
}
}

/**
Expand All @@ -161,15 +164,18 @@ public static Optional<Span> completeSpan() {
* Like {@link #completeSpan()}, but adds {@code metadata} to the current span being completed.
*/
public static Optional<Span> completeSpan(Map<String, String> metadata) {
Trace trace = currentTrace.get();
if (trace == null) {
return Optional.empty();
}
Optional<Span> 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() && trace.isObservable()) {
// Avoid lambda allocation in hot paths
notifyObservers(maybeSpan.get());
}

return maybeSpan;
}
Expand All @@ -181,12 +187,20 @@ private static void notifyObservers(Span span) {
}

private static Optional<OpenSpan> popCurrentSpan() {
return currentTrace.get().pop();
Trace trace = currentTrace.get();
if (trace != null) {
Optional<OpenSpan> span = trace.pop();
if (trace.isEmpty()) {
clearCurrentTrace();
}
return span;
}
return Optional.empty();
}

private static Span toSpan(OpenSpan openSpan, Map<String, String> metadata) {
private static Span toSpan(OpenSpan openSpan, Map<String, String> metadata, String traceId) {
return Span.builder()
.traceId(getTraceId())
.traceId(traceId)
.spanId(openSpan.getSpanId())
.type(openSpan.type())
.parentSpanId(openSpan.getParentSpanId())
Expand Down Expand Up @@ -237,16 +251,20 @@ 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 currentTrace.get().getTraceId();
return Preconditions.checkNotNull(currentTrace.get(), "There is no root span").getTraceId();
Copy link
Contributor Author

@carterkozak carterkozak Oct 25, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I considered returning a newly generated ID when no trace is available, but that could potentially result in bad data if we initialize a new span assuming that traceId will be maintained.

Copy link
Contributor

@iamdanfox iamdanfox Oct 26, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Possibly want to use the safelogging Preconditions here so these messages will definitely not get redacted?

https://github.com/palantir/safe-logging#preconditions

Copy link
Contributor Author

@carterkozak carterkozak Oct 26, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree that would be ideal, but we don't currently depend on safe-logging. I'd prefer not to add deps in this change.

Copy link
Contributor

@schlosna schlosna Oct 26, 2018 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should we expose a public static boolean hasTraceId() to allow checking if tracing is currently enabled?

Once this PR goes through we can obviously clean up https://github.com/palantir/tritium/blob/develop/tritium-tracing/src/main/java/com/palantir/tritium/tracing/TracingInvocationEventHandler.java#L71 as getAndClearTrace will now do the right thing and clear state, but there may be cases where we want to know if there is a root trace ID already set.

Copy link
Contributor

@iamdanfox iamdanfox Oct 29, 2018 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1 to a hasTraceId or a new method to return Optional<String>... otherwise you can't eyeball I method like this and be confident it's safe:

    @Override
    public HttpConnection create(URL url) throws IOException {
        HttpConnection connection = delegate.create(url);
        String traceId = Tracer.getTraceId();
        connection.setRequestProperty(TraceHttpHeaders.TRACE_ID, traceId);
        return connection;
    }

Copy link
Contributor Author

@carterkozak carterkozak Oct 29, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added.

I searched for usages of getTraceId, and everything that I found called it after creating a new span, so it should be safe, though there may be cases that I'm not aware of. For what it's worth, anything that does throw here is also responsible for leaking traceIds.

}

/** 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;
}

Expand All @@ -255,12 +273,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();
Copy link
Contributor

@schlosna schlosna Oct 26, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

related to above comments, could add isObservable(Trace) to avoid additional thread local lookups

Suggested change
return trace != null && trace.isObservable();
return isObservable(currentTrace.get());
}
private static boolean isObservable(Trace trace) {
return trace != null && trace.isObservable();
}

}

/** Returns an independent copy of this thread's {@link Trace}. */
static Trace copyTrace() {
return currentTrace.get().deepCopy();
static Optional<Trace> copyTrace() {
Trace trace = currentTrace.get();
if (trace != null) {
return Optional.of(trace.deepCopy());
}
return Optional.empty();
}

/**
Expand All @@ -274,4 +297,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() {
Copy link
Contributor

@iamdanfox iamdanfox Oct 26, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like the only two places this can be called are:

  • public Trace getAndClearTrace()
  • when a call to fastCompleteSpan() -> popCurrentSpan() (pops the last remaining span)

Copy link
Contributor Author

@carterkozak carterkozak Oct 26, 2018

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That is correct

currentTrace.remove();
MDC.remove(Tracers.TRACE_ID_KEY);
}
}
4 changes: 2 additions & 2 deletions tracing/src/test/java/com/palantir/tracing/CloseableTracerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
64 changes: 52 additions & 12 deletions tracing/src/test/java/com/palantir/tracing/TracerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -63,6 +65,7 @@ public void after() {
Tracer.setSampler(AlwaysSampler.INSTANCE);
Tracer.unsubscribe("1");
Tracer.unsubscribe("2");
Tracer.getAndClearTrace();
}

@Test
Expand Down Expand Up @@ -164,23 +167,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();
}

Expand All @@ -189,8 +195,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
Expand Down Expand Up @@ -245,12 +252,45 @@ public void testFastCompleteSpanWithMetadata() {

@Test
public void testClearAndGetTraceClearsMdc() {
String startTrace = Tracer.getTraceId();
assertThat(MDC.get(Tracers.TRACE_ID_KEY)).isEqualTo(startTrace);
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();
}
}

Trace oldTrace = Tracer.getAndClearTrace();
assertThat(oldTrace.getTraceId()).isEqualTo(startTrace);
assertThat(MDC.get(Tracers.TRACE_ID_KEY)).isNull(); // after clearing, it's empty
@Test
public void testCompleteRootSpanCompletesTrace() {
Set<String> 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);
}

@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() {
Expand Down
14 changes: 8 additions & 6 deletions tracing/src/test/java/com/palantir/tracing/TracersTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -374,11 +375,12 @@ public void run() {
}

private static List<OpenSpan> getCurrentFullTrace() {
Trace trace = Tracer.copyTrace();
List<OpenSpan> spans = Lists.newArrayList();
while (!trace.isEmpty()) {
spans.add(trace.pop().get());
}
return spans;
return Tracer.copyTrace().map(trace -> {
List<OpenSpan> spans = Lists.newArrayList();
while (!trace.isEmpty()) {
spans.add(trace.pop().get());
}
return spans;
}).orElse(Collections.emptyList());
}
}