[improvement] Add methods to set span operation when wrapping with a 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.

Author: pkoenig10

tracing/src/main/java/com/palantir/tracing/Tracers.java

@@ -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;
@@ -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();
@@ -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();
@@ -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();

tracing/src/test/java/com/palantir/tracing/TracersTest.java

@@ -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();
@@ -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();
@@ -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();
@@ -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 = () -> {
@@ -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";
@@ -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();
@@ -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);
             }
         };