Skip to content
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.

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

[improvement] Add methods to set span operation when wrapping with a new trace #59

Merged
merged 4 commits into from
Feb 6, 2019
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
76 changes: 58 additions & 18 deletions tracing/src/main/java/com/palantir/tracing/Tracers.java
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
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