Skip to content

Commit

Permalink
[improvement] Add methods to set span operation when wrapping with a …
Browse files Browse the repository at this point in the history 
…new trace (#59)

## Before this PR
When calling any of the `wrapWithNewTrace` methods, the operation of the initial span would be set to "root". This makes it difficult to determine the origin of traces if `wrapWithNewTrace` is used in more than one place.

## After this PR
`Tracers` now provides variants of the `wrapWithNewTrace` methods that allow callers to explicitly set the initial span operation.

Follow up to #49.
  • Loading branch information
@pkoenig10 @bulldozer-bot
pkoenig10 authored and bulldozer-bot[bot] committed Feb 6, 2019
1 parent 2597bae commit 74f46d4
Show file tree
Hide file tree
Showing 2 changed files with 129 additions and 30 deletions.
76 changes: 58 additions & 18 deletions tracing/src/main/java/com/palantir/tracing/Tracers.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

package com.palantir.tracing;

import com.google.errorprone.annotations.CompileTimeConstant;
import java.util.Optional;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
Expand Down Expand Up @@ -104,52 +105,75 @@ public static Runnable wrap(Runnable delegate) {
return new TracingAwareRunnable(delegate);
}

/**
* Like {@link #wrapWithNewTrace(String, ExecutorService)}, but with a default initial span operation.
*/
public static ExecutorService wrapWithNewTrace(ExecutorService executorService) {
return wrapWithNewTrace(ROOT_SPAN_OPERATION, executorService);
}

/**
* Wraps the provided executor service to make submitted tasks traceable with a fresh {@link Trace trace}
* for each execution, see {@link #wrapWithNewTrace(ExecutorService)}. This method should not be used to
* for each execution, see {@link #wrapWithNewTrace(String, ExecutorService)}. This method should not be used to
* wrap a ScheduledExecutorService that has already been {@link #wrap(ExecutorService) wrapped}. If this is
* done, a new trace will be generated for each execution, effectively bypassing the intent of the previous
* wrapping.
* wrapping. The given {@link String operation} is used to create the initial span.
*/
public static ExecutorService wrapWithNewTrace(ExecutorService executorService) {
public static ExecutorService wrapWithNewTrace(@CompileTimeConstant String operation,
ExecutorService executorService) {
return new WrappingExecutorService(executorService) {
@Override
protected <T> Callable<T> wrapTask(Callable<T> callable) {
return wrapWithNewTrace(callable);
return wrapWithNewTrace(operation, callable);
}
};
}

/**
* Wraps the provided scheduled executor service to make submitted tasks traceable with a fresh {@link Trace trace}
* for each execution, see {@link #wrapWithNewTrace(ScheduledExecutorService)}. This method should not be used to
* wrap a ScheduledExecutorService that has already been {@link #wrap(ScheduledExecutorService) wrapped}. If this is
* done, a new trace will be generated for each execution, effectively bypassing the intent of the previous
* wrapping.
* Like {@link #wrapWithNewTrace(String, ScheduledExecutorService)}, but with a default initial span operation.
*/
public static ScheduledExecutorService wrapWithNewTrace(ScheduledExecutorService executorService) {
return wrapWithNewTrace(ROOT_SPAN_OPERATION, executorService);
}

/**
* Wraps the provided scheduled executor service to make submitted tasks traceable with a fresh {@link Trace trace}
* for each execution, see {@link #wrapWithNewTrace(String, ScheduledExecutorService)}. This method should not be
* used to wrap a ScheduledExecutorService that has already been {@link #wrap(ScheduledExecutorService) wrapped}.
* If this is done, a new trace will be generated for each execution, effectively bypassing the intent of the
* previous wrapping. The given {@link String operation} is used to create the initial span.
*/
public static ScheduledExecutorService wrapWithNewTrace(@CompileTimeConstant String operation,
ScheduledExecutorService executorService) {
return new WrappingScheduledExecutorService(executorService) {
@Override
protected <T> Callable<T> wrapTask(Callable<T> callable) {
return wrapWithNewTrace(callable);
return wrapWithNewTrace(operation, callable);
}
};
}

/**
* Like {@link #wrapWithNewTrace(String, Callable)}, but with a default initial span operation.
*/
public static <V> Callable<V> wrapWithNewTrace(Callable<V> delegate) {
return wrapWithNewTrace(ROOT_SPAN_OPERATION, delegate);
}

/**
* Wraps the given {@link Callable} such that it creates a fresh {@link Trace tracing state} for its execution.
* That is, the trace during its {@link Callable#call() execution} is entirely separate from the trace at
* construction or any trace already set on the thread used to execute the callable. Each execution of the callable
* will have a fresh trace.
* will have a fresh trace. The given {@link String operation} is used to create the initial span.
*/
public static <V> Callable<V> wrapWithNewTrace(Callable<V> delegate) {
public static <V> Callable<V> wrapWithNewTrace(@CompileTimeConstant String operation, Callable<V> delegate) {
return () -> {
// clear the existing trace and keep it around for restoration when we're done
Optional<Trace> originalTrace = Tracer.getAndClearTraceIfPresent();

try {
Tracer.initTrace(Optional.empty(), Tracers.randomId());
Tracer.startSpan(ROOT_SPAN_OPERATION);
Tracer.startSpan(operation);
return delegate.call();
} finally {
Tracer.fastCompleteSpan();
Expand All @@ -159,16 +183,23 @@ public static <V> Callable<V> wrapWithNewTrace(Callable<V> delegate) {
}

/**
* Like {@link #wrapWithNewTrace(Callable)}, but for Runnables.
* Like {@link #wrapWithAlternateTraceId(String, Runnable)}, but with a default initial span operation.
*/
public static Runnable wrapWithNewTrace(Runnable delegate) {
return wrapWithNewTrace(ROOT_SPAN_OPERATION, delegate);
}

/**
* Like {@link #wrapWithNewTrace(String, Callable)}, but for Runnables.
*/
public static Runnable wrapWithNewTrace(@CompileTimeConstant String operation, Runnable delegate) {
return () -> {
// clear the existing trace and keep it around for restoration when we're done
Optional<Trace> originalTrace = Tracer.getAndClearTraceIfPresent();

try {
Tracer.initTrace(Optional.empty(), Tracers.randomId());
Tracer.startSpan(ROOT_SPAN_OPERATION);
Tracer.startSpan(operation);
delegate.run();
} finally {
Tracer.fastCompleteSpan();
Expand All @@ -177,20 +208,29 @@ public static Runnable wrapWithNewTrace(Runnable delegate) {
};
}

/**
* Like {@link #wrapWithAlternateTraceId(String, String, Runnable)}, but with a default initial span operation.
*/
public static Runnable wrapWithAlternateTraceId(String traceId, Runnable delegate) {
return wrapWithAlternateTraceId(traceId, ROOT_SPAN_OPERATION, delegate);
}

/**
* Wraps the given {@link Runnable} such that it creates a fresh {@link Trace tracing state with the given traceId}
* for its execution. That is, the trace during its {@link Runnable#run() execution} will use the traceId provided
* instead of any trace already set on the thread used to execute the runnable. Each execution of the runnable
* will use a new {@link Trace tracing state} with the same given traceId.
* will use a new {@link Trace tracing state} with the same given traceId. The given {@link String operation} is
* used to create the initial span.
*/
public static Runnable wrapWithAlternateTraceId(String traceId, Runnable delegate) {
public static Runnable wrapWithAlternateTraceId(String traceId, @CompileTimeConstant String operation, Runnable
delegate) {
return () -> {
// clear the existing trace and keep it around for restoration when we're done
Optional<Trace> originalTrace = Tracer.getAndClearTraceIfPresent();

try {
Tracer.initTrace(Optional.empty(), traceId);
Tracer.startSpan(ROOT_SPAN_OPERATION);
Tracer.startSpan(operation);
delegate.run();
} finally {
Tracer.fastCompleteSpan();
Expand Down
83 changes: 71 additions & 12 deletions tracing/src/test/java/com/palantir/tracing/TracersTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,10 +96,10 @@ public void testScheduledExecutorServiceWrapsCallables() throws Exception {
@Test
public void testScheduledExecutorServiceWrapsCallablesWithNewTraces() throws Exception {
ScheduledExecutorService wrappedService =
Tracers.wrapWithNewTrace(Executors.newSingleThreadScheduledExecutor());
Tracers.wrapWithNewTrace("operationToUse", Executors.newSingleThreadScheduledExecutor());

Callable<Void> callable = newTraceExpectingCallable();
Runnable runnable = newTraceExpectingRunnable();
Callable<Void> callable = newTraceExpectingCallable("operationToUse");
Runnable runnable = newTraceExpectingRunnable("operationToUse");

// Empty trace
wrappedService.schedule(callable, 0, TimeUnit.SECONDS).get();
Expand All @@ -122,10 +122,10 @@ public void testScheduledExecutorServiceWrapsCallablesWithNewTraces() throws Exc
@Test
public void testExecutorServiceWrapsCallablesWithNewTraces() throws Exception {
ExecutorService wrappedService =
Tracers.wrapWithNewTrace(Executors.newSingleThreadExecutor());
Tracers.wrapWithNewTrace("operationToUse", Executors.newSingleThreadExecutor());

Callable<Void> callable = newTraceExpectingCallable();
Runnable runnable = newTraceExpectingRunnable();
Callable<Void> callable = newTraceExpectingCallable("operationToUse");
Runnable runnable = newTraceExpectingRunnable("operationToUse");

// Empty trace
wrappedService.submit(callable).get();
Expand Down Expand Up @@ -241,6 +241,22 @@ public void testWrapCallableWithNewTrace_traceStateInsideCallableHasSpan() throw
assertThat(span.getParentSpanId()).isEmpty();
}

@Test
public void testWrapCallableWithNewTrace_traceStateInsideCallableHasGivenSpan() throws Exception {
Callable<List<OpenSpan>> wrappedCallable = Tracers.wrapWithNewTrace("someOperation", () -> {
return getCurrentFullTrace();
});

List<OpenSpan> spans = wrappedCallable.call();

assertThat(spans).hasSize(1);

OpenSpan span = spans.get(0);

assertThat(span.getOperation()).isEqualTo("someOperation");
assertThat(span.getParentSpanId()).isEmpty();
}

@Test
public void testWrapCallableWithNewTrace_traceStateRestoredWhenThrows() throws Exception {
String traceIdBeforeConstruction = Tracer.getTraceId();
Expand Down Expand Up @@ -312,7 +328,25 @@ public void testWrapRunnableWithNewTrace_traceStateInsideRunnableHasSpan() throw
}

@Test
public void testWrapRunnableWithNewTrace_traceStateRestoredWhenThrows() throws Exception {
public void testWrapRunnableWithNewTrace_traceStateInsideRunnableHasGivenSpan() throws Exception {
List<List<OpenSpan>> spans = Lists.newArrayList();

Runnable wrappedRunnable = Tracers.wrapWithNewTrace("someOperation", () -> {
spans.add(getCurrentFullTrace());
});

wrappedRunnable.run();

assertThat(spans.get(0)).hasSize(1);

OpenSpan span = spans.get(0).get(0);

assertThat(span.getOperation()).isEqualTo("someOperation");
assertThat(span.getParentSpanId()).isEmpty();
}

@Test
public void testWrapRunnableWithNewTrace_traceStateRestoredWhenThrows() {
String traceIdBeforeConstruction = Tracer.getTraceId();

Runnable rawRunnable = () -> {
Expand Down Expand Up @@ -357,7 +391,7 @@ public void testWrapRunnableWithAlternateTraceId_traceStateInsideRunnableUsesGiv
}

@Test
public void testWrapRunnableWithAlternateTraceId_traceStateInsideRunnableHasSpan() throws Exception {
public void testWrapRunnableWithAlternateTraceId_traceStateInsideRunnableHasSpan() {
List<List<OpenSpan>> spans = Lists.newArrayList();

String traceIdToUse = "someTraceId";
Expand All @@ -375,6 +409,25 @@ public void testWrapRunnableWithAlternateTraceId_traceStateInsideRunnableHasSpan
assertThat(span.getParentSpanId()).isEmpty();
}

@Test
public void testWrapRunnableWithAlternateTraceId_traceStateInsideRunnableHasGivenSpan() {
List<List<OpenSpan>> spans = Lists.newArrayList();

String traceIdToUse = "someTraceId";
Runnable wrappedRunnable = Tracers.wrapWithAlternateTraceId(traceIdToUse, "someOperation", () -> {
spans.add(getCurrentFullTrace());
});

wrappedRunnable.run();

assertThat(spans.get(0)).hasSize(1);

OpenSpan span = spans.get(0).get(0);

assertThat(span.getOperation()).isEqualTo("someOperation");
assertThat(span.getParentSpanId()).isEmpty();
}

@Test
public void testWrapRunnableWithAlternateTraceId_traceStateRestoredWhenThrows() {
String traceIdBeforeConstruction = Tracer.getTraceId();
Expand Down Expand Up @@ -406,36 +459,42 @@ public void testTraceIdGeneration() throws Exception {
assertThat(Tracers.longToPaddedHex(123456789L)).isEqualTo("00000000075bcd15");
}

private static Callable<Void> newTraceExpectingCallable() {
private static Callable<Void> newTraceExpectingCallable(String expectedOperation) {
final Set<String> seenTraceIds = new HashSet<>();
seenTraceIds.add(Tracer.getTraceId());

return new Callable<Void>() {
@Override
public Void call() throws Exception {
String newTraceId = Tracer.getTraceId();
List<OpenSpan> spans = getCurrentFullTrace();

assertThat(MDC.get(Tracers.TRACE_ID_KEY)).isEqualTo(newTraceId);
assertThat(seenTraceIds).doesNotContain(newTraceId);
assertThat(getCurrentFullTrace()).hasSize(1);
assertThat(spans).hasSize(1);
assertThat(spans.get(0).getOperation()).isEqualTo(expectedOperation);
assertThat(spans.get(0).getParentSpanId()).isEmpty();
seenTraceIds.add(newTraceId);
return null;
}
};
}

private static Runnable newTraceExpectingRunnable() {
private static Runnable newTraceExpectingRunnable(String expectedOperation) {
final Set<String> seenTraceIds = new HashSet<>();
seenTraceIds.add(Tracer.getTraceId());

return new Runnable() {
@Override
public void run() {
String newTraceId = Tracer.getTraceId();
List<OpenSpan> spans = getCurrentFullTrace();

assertThat(MDC.get(Tracers.TRACE_ID_KEY)).isEqualTo(newTraceId);
assertThat(seenTraceIds).doesNotContain(newTraceId);
assertThat(getCurrentFullTrace()).hasSize(1);
assertThat(spans).hasSize(1);
assertThat(spans.get(0).getOperation()).isEqualTo(expectedOperation);
assertThat(spans.get(0).getParentSpanId()).isEmpty();
seenTraceIds.add(newTraceId);
}
};
Expand Down

0 comments on commit 74f46d4

Please sign in to comment.