Fix #177: Add APIs for Supplier<String> operation names

palantir · Sep 10, 2019 · ab5a181 · ab5a181
1 parent a942e27
commit ab5a181
Show file tree

Hide file tree

Showing 7 changed files with 185 additions and 3 deletions.
diff --git a/tracing-jersey/src/main/java/com/palantir/tracing/jersey/TraceEnrichingFilter.java b/tracing-jersey/src/main/java/com/palantir/tracing/jersey/TraceEnrichingFilter.java
@@ -25,6 +25,7 @@
 import com.palantir.tracing.api.TraceHttpHeaders;
 import java.io.IOException;
 import java.util.Optional;
+import java.util.function.Supplier;
 import javax.annotation.Priority;
 import javax.ws.rs.container.ContainerRequestContext;
 import javax.ws.rs.container.ContainerRequestFilter;
@@ -55,12 +56,10 @@ public final class TraceEnrichingFilter implements ContainerRequestFilter, Conta
     // Handles incoming request
     @Override
     public void filter(ContainerRequestContext requestContext) throws IOException {
-        String path = Optional.ofNullable(uriInfo)
+        Supplier<String> operation = () -> "Jersey: " + requestContext.getMethod() + " " + Optional.ofNullable(uriInfo)
                 .map(ExtendedUriInfo::getMatchedModelResource)
                 .map(Resource::getPath)
                 .orElse("(unknown)");
-
-        String operation = "Jersey: " + requestContext.getMethod() + " " + path;
         // The following strings are all nullable
         String traceId = requestContext.getHeaderString(TraceHttpHeaders.TRACE_ID);
         String spanId = requestContext.getHeaderString(TraceHttpHeaders.SPAN_ID);

diff --git a/tracing/src/main/java/com/palantir/tracing/CloseableTracer.java b/tracing/src/main/java/com/palantir/tracing/CloseableTracer.java
@@ -17,6 +17,7 @@
 package com.palantir.tracing;
 
 import com.palantir.tracing.api.SpanType;
+import java.util.function.Supplier;
 
 /**
  * Wraps the {@link Tracer} methods in a closeable resource to enable the usage of the try-with-resources pattern.
@@ -39,6 +40,13 @@ public static CloseableTracer startSpan(String operation) {
         return startSpan(operation, SpanType.LOCAL);
     }
 
+    /**
+     * Opens a new {@link SpanType#LOCAL LOCAL} span for this thread's call trace, labeled with the provided operation.
+     */
+    public static CloseableTracer startSpan(Supplier<String> operation) {
+        return startSpan(operation, SpanType.LOCAL);
+    }
+
     /**
      * Opens a new span for this thread's call trace with the provided {@link SpanType},
      * labeled with the provided operation.
@@ -50,6 +58,17 @@ public static CloseableTracer startSpan(String operation, SpanType spanType) {
         return INSTANCE;
     }
 
+    /**
+     * Opens a new span for this thread's call trace with the provided {@link SpanType},
+     * labeled with the provided operation.
+     *
+     * If you need to a span that may complete on another thread, use {@link DetachedSpan#start} instead.
+     */
+    public static CloseableTracer startSpan(Supplier<String> operation, SpanType spanType) {
+        Tracer.fastStartSpan(operation, spanType);
+        return INSTANCE;
+    }
+
     @Override
     public void close() {
         Tracer.fastCompleteSpan();

diff --git a/tracing/src/main/java/com/palantir/tracing/DetachedSpan.java b/tracing/src/main/java/com/palantir/tracing/DetachedSpan.java
@@ -18,6 +18,7 @@
 
 import com.google.errorprone.annotations.MustBeClosed;
 import com.palantir.tracing.api.SpanType;
+import java.util.function.Supplier;
 import javax.annotation.CheckReturnValue;
 
 /** Span which is not bound to thread state, and can be completed on any other thread. */
@@ -38,6 +39,12 @@ static DetachedSpan start(String operation) {
         return start(operation, SpanType.LOCAL);
     }
 
+    /** Equivalent to {@link #start(String)} with a lazy supplier for the operation name. */
+    @CheckReturnValue
+    static DetachedSpan start(Supplier<String> operation) {
+        return start(operation, SpanType.LOCAL);
+    }
+
     /**
      * Marks the beginning of a span, which you can {@link #complete} on any other thread.
      *
@@ -48,13 +55,26 @@ static DetachedSpan start(String operation, SpanType type) {
         return Tracer.detachInternal(operation, type);
     }
 
+    /** Equivalent to {@link #start(String, SpanType)} with a lazy supplier for the operation name. */
+    @CheckReturnValue
+    static DetachedSpan start(Supplier<String> operation, SpanType type) {
+        return Tracer.detachInternal(operation, type);
+    }
+
     /**
      * Equivalent to {@link Tracer#startSpan(String, SpanType)}, but using this {@link DetachedSpan}
      * as the parent instead of thread state.
      */
     @MustBeClosed
     CloseableSpan childSpan(String operationName, SpanType type);
 
+    /**
+     * Equivalent to {@link Tracer#fastStartSpan(Supplier, SpanType)}, but using this {@link DetachedSpan}
+     * as the parent instead of thread state.
+     */
+    @MustBeClosed
+    CloseableSpan childSpan(Supplier<String> operationName, SpanType type);
+
     /**
      * Equivalent to {@link Tracer#startSpan(String)}, but using this {@link DetachedSpan} as the parent instead
      * of thread state.
@@ -72,15 +92,32 @@ default CloseableSpan completeAndStartChild(String operationName, SpanType type)
         return child;
     }
 
+    @MustBeClosed
+    @SuppressWarnings("MustBeClosedChecker")
+    default CloseableSpan completeAndStartChild(Supplier<String> operationName, SpanType type) {
+        CloseableSpan child = childSpan(operationName, type);
+        complete();
+        return child;
+    }
+
     @MustBeClosed
     default CloseableSpan completeAndStartChild(String operationName) {
         return completeAndStartChild(operationName, SpanType.LOCAL);
     }
 
+    @MustBeClosed
+    default CloseableSpan completeAndStartChild(Supplier<String> operationName) {
+        return completeAndStartChild(operationName, SpanType.LOCAL);
+    }
+
     /** Starts a child {@link DetachedSpan} using this instance as the parent. */
     @CheckReturnValue
     DetachedSpan childDetachedSpan(String operation, SpanType type);
 
+    /** Starts a child {@link DetachedSpan} using this instance as the parent. */
+    @CheckReturnValue
+    DetachedSpan childDetachedSpan(Supplier<String> operation, SpanType type);
+
     /**
      * Starts a child {@link DetachedSpan} using this instance as the parent.
      * Equivalent to {@link #childDetachedSpan(String, SpanType)} using {@link SpanType#LOCAL}.
@@ -90,6 +127,15 @@ default DetachedSpan childDetachedSpan(String operation) {
         return childDetachedSpan(operation, SpanType.LOCAL);
     }
 
+    /**
+     * Starts a child {@link DetachedSpan} using this instance as the parent.
+     * Equivalent to {@link #childDetachedSpan(Supplier, SpanType)} using {@link SpanType#LOCAL}.
+     */
+    @CheckReturnValue
+    default DetachedSpan childDetachedSpan(Supplier<String> operation) {
+        return childDetachedSpan(operation, SpanType.LOCAL);
+    }
+
     /**
      * Completes this span. After complete is invoked, other methods are not expected to produce spans, but
      * they must not throw either in order to avoid confusing failures.

diff --git a/tracing/src/main/java/com/palantir/tracing/Trace.java b/tracing/src/main/java/com/palantir/tracing/Trace.java
@@ -29,6 +29,7 @@
 import java.util.ArrayDeque;
 import java.util.Deque;
 import java.util.Optional;
+import java.util.function.Supplier;
 
 /**
  * Represents a trace as an ordered list of non-completed spans. Supports adding and removing of spans. This class is
@@ -115,6 +116,16 @@ private static <T> Optional<T> orElse(Optional<T> left, Optional<T> right) {
      */
     abstract void fastStartSpan(String operation, SpanType type);
 
+    /**
+     * Like {@link #startSpan(String, String, SpanType)}, but does not return an {@link OpenSpan}.
+     */
+    abstract void fastStartSpan(Supplier<String> operation, String parentSpanId, SpanType type);
+
+    /**
+     * Like {@link #startSpan(String, SpanType)}, but does not return an {@link OpenSpan}.
+     */
+    abstract void fastStartSpan(Supplier<String> operation, SpanType type);
+
     protected abstract void push(OpenSpan span);
 
     abstract Optional<OpenSpan> top();
@@ -170,6 +181,16 @@ void fastStartSpan(String operation, SpanType type) {
             startSpan(operation, type);
         }
 
+        @Override
+        void fastStartSpan(Supplier<String> operation, String parentSpanId, SpanType type) {
+            fastStartSpan(operation.get(), parentSpanId, type);
+        }
+
+        @Override
+        void fastStartSpan(Supplier<String> operation, SpanType type) {
+            fastStartSpan(operation.get(), type);
+        }
+
         @Override
         protected void push(OpenSpan span) {
             stack.push(span);
@@ -242,6 +263,16 @@ void fastStartSpan(String operation, SpanType type) {
             numberOfSpans++;
         }
 
+        @Override
+        void fastStartSpan(Supplier<String> operation, String parentSpanId, SpanType type) {
+            numberOfSpans++;
+        }
+
+        @Override
+        void fastStartSpan(Supplier<String> operation, SpanType type) {
+            numberOfSpans++;
+        }
+
         @Override
         protected void push(OpenSpan span) {
             startSpan(span.getParentSpanId());

diff --git a/tracing/src/main/java/com/palantir/tracing/Tracer.java b/tracing/src/main/java/com/palantir/tracing/Tracer.java
@@ -40,6 +40,7 @@
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -188,6 +189,27 @@ public static void fastStartSpan(String operation) {
         fastStartSpan(operation, SpanType.LOCAL);
     }
 
+    /**
+     * Like {@link #startSpan(String, String, SpanType)}, but does not return an {@link OpenSpan}.
+     */
+    public static void fastStartSpan(Supplier<String> operation, String parentSpanId, SpanType type) {
+        getOrCreateCurrentTrace().fastStartSpan(operation, parentSpanId, type);
+    }
+
+    /**
+     * Like {@link #startSpan(String, SpanType)}, but does not return an {@link OpenSpan}.
+     */
+    public static void fastStartSpan(Supplier<String> operation, SpanType type) {
+        getOrCreateCurrentTrace().fastStartSpan(operation, type);
+    }
+
+    /**
+     * Like {@link #startSpan(String)}, but does not return an {@link OpenSpan}.
+     */
+    public static void fastStartSpan(Supplier<String> operation) {
+        fastStartSpan(operation, SpanType.LOCAL);
+    }
+
     /**
      * Like {@link #startSpan(String, SpanType)}, but does not set or modify tracing thread state.
      * This is an internal utility that should not be called directly outside of {@link DetachedSpan}.
@@ -203,6 +225,17 @@ static DetachedSpan detachInternal(String operation, SpanType type) {
                 : new UnsampledDetachedSpan(traceId);
     }
 
+    static DetachedSpan detachInternal(Supplier<String> operation, SpanType type) {
+        Trace maybeCurrentTrace = currentTrace.get();
+        String traceId = maybeCurrentTrace != null
+                ? maybeCurrentTrace.getTraceId() : Tracers.randomId();
+        boolean sampled = maybeCurrentTrace != null
+                ? maybeCurrentTrace.isObservable() : sampler.sample();
+        return sampled
+                ? new SampledDetachedSpan(operation.get(), type, traceId, getParentSpanId(maybeCurrentTrace))
+                : new UnsampledDetachedSpan(traceId);
+    }
+
     private static Optional<String> getParentSpanId(@Nullable Trace trace) {
         if (trace != null) {
             Optional<OpenSpan> maybeOpenSpan = trace.top();
@@ -240,12 +273,23 @@ public CloseableSpan childSpan(String operationName, SpanType type) {
             return TraceRestoringCloseableSpan.of(maybeCurrentTrace);
         }
 
+        @Override
+        @MustBeClosed
+        public CloseableSpan childSpan(Supplier<String> operationName, SpanType type) {
+            return childSpan(operationName.get(), type);
+        }
+
         @Override
         public DetachedSpan childDetachedSpan(String operation, SpanType type) {
             warnIfCompleted("startDetachedSpan");
             return new SampledDetachedSpan(operation, type, traceId, Optional.of(openSpan.getSpanId()));
         }
 
+        @Override
+        public DetachedSpan childDetachedSpan(Supplier<String> operation, SpanType type) {
+            return childDetachedSpan(operation.get(), type);
+        }
+
         @Override
         public void complete() {
             if (completed.compareAndSet(false, true)) {
@@ -284,11 +328,24 @@ public CloseableSpan childSpan(String operationName, SpanType type) {
             return TraceRestoringCloseableSpan.of(maybeCurrentTrace);
         }
 
+        @Override
+        public CloseableSpan childSpan(Supplier<String> operationName, SpanType type) {
+            Trace maybeCurrentTrace = currentTrace.get();
+            setTrace(Trace.of(false, traceId));
+            Tracer.fastStartSpan(operationName, type);
+            return TraceRestoringCloseableSpan.of(maybeCurrentTrace);
+        }
+
         @Override
         public DetachedSpan childDetachedSpan(String operation, SpanType type) {
             return this;
         }
 
+        @Override
+        public DetachedSpan childDetachedSpan(Supplier<String> operation, SpanType type) {
+            return this;
+        }
+
         @Override
         public void complete() {
             // nop

diff --git a/tracing/src/test/java/com/palantir/tracing/CloseableTracerTest.java b/tracing/src/test/java/com/palantir/tracing/CloseableTracerTest.java
@@ -44,6 +44,16 @@ public void startsAndClosesSpan() {
         assertThat(Tracer.getAndClearTrace().top()).isEmpty();
     }
 
+    @Test
+    public void startsAndClosesSpan_supplier() {
+        try (CloseableTracer tracer = CloseableTracer.startSpan(() -> "foo")) {
+            OpenSpan openSpan = Tracer.copyTrace().get().top().get();
+            assertThat(openSpan.getOperation()).isEqualTo("foo");
+            assertThat(openSpan.type()).isEqualTo(SpanType.LOCAL);
+        }
+        assertThat(Tracer.getAndClearTrace().top()).isEmpty();
+    }
+
     @Test
     public void startsAndClosesSpanWithType() {
         try (CloseableTracer tracer = CloseableTracer.startSpan("foo", SpanType.CLIENT_OUTGOING)) {
@@ -53,4 +63,14 @@ public void startsAndClosesSpanWithType() {
         }
         assertThat(Tracer.getAndClearTrace().top()).isEmpty();
     }
+
+    @Test
+    public void startsAndClosesSpanWithType_supplier() {
+        try (CloseableTracer tracer = CloseableTracer.startSpan(() -> "foo", SpanType.CLIENT_OUTGOING)) {
+            OpenSpan openSpan = Tracer.copyTrace().get().top().get();
+            assertThat(openSpan.getOperation()).isEqualTo("foo");
+            assertThat(openSpan.type()).isEqualTo(SpanType.CLIENT_OUTGOING);
+        }
+        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
@@ -270,6 +270,16 @@ public void testFastCompleteSpan() {
         assertThat(spanCaptor.getValue().getOperation()).isEqualTo(operation);
     }
 
+    @Test
+    public void testFastCompleteSpan_supplier() {
+        Tracer.subscribe("1", observer1);
+        String operation = "operation";
+        Tracer.fastStartSpan(() -> operation);
+        Tracer.fastCompleteSpan();
+        verify(observer1).consume(spanCaptor.capture());
+        assertThat(spanCaptor.getValue().getOperation()).isEqualTo(operation);
+    }
+
     @Test
     public void testFastCompleteSpanWithMetadata() {
         Tracer.subscribe("1", observer1);