Use top span id as the local id

tracing-jersey/src/main/java/com/palantir/tracing/jersey/TraceEnrichingFilter.java

@@ -46,7 +46,7 @@ public final class TraceEnrichingFilter implements ContainerRequestFilter, Conta
      * This is the name of the trace id property we set on {@link ContainerRequestContext}.
      */
     public static final String TRACE_ID_PROPERTY_NAME = "com.palantir.tracing.traceId";
-    public static final String LOCAL_TRACE_ID_PROPERTY_NAME = "com.palantir.tracing.localTraceId";
+    public static final String TOP_SPAN_ID_PROPERTY_NAME = "com.palantir.tracing.topSpanId";
     public static final String SAMPLED_PROPERTY_NAME = "com.palantir.tracing.sampled";
 
     @Context
@@ -72,8 +72,7 @@ public void filter(ContainerRequestContext requestContext) throws IOException {
             Tracer.initTrace(getObservabilityFromHeader(requestContext), Tracers.randomId());
             Tracer.fastStartSpan(operation, SpanType.SERVER_INCOMING);
         } else {
-            // propagate the traceid from request headers, but create a new local trace id to disambiguate
-            Tracer.initTrace(getObservabilityFromHeader(requestContext), traceId, Tracers.randomId());
+            Tracer.initTrace(getObservabilityFromHeader(requestContext), traceId);
             if (spanId == null) {
                 Tracer.fastStartSpan(operation, SpanType.SERVER_INCOMING);
             } else {
@@ -84,8 +83,8 @@ public void filter(ContainerRequestContext requestContext) throws IOException {
 
         // Give asynchronous downstream handlers access to the trace id
         requestContext.setProperty(TRACE_ID_PROPERTY_NAME, Tracer.getTraceId());
-        Tracer.getLocalTraceId().ifPresent(localId ->
-                requestContext.setProperty(LOCAL_TRACE_ID_PROPERTY_NAME, localId));
+        Tracer.getTopSpanId().ifPresent(localId ->
+                requestContext.setProperty(TOP_SPAN_ID_PROPERTY_NAME, localId));
         requestContext.setProperty(SAMPLED_PROPERTY_NAME, Tracer.isTraceObservable());
     }
 

tracing-jersey/src/test/java/com/palantir/tracing/jersey/TraceEnrichingFilterTest.java

@@ -119,7 +119,7 @@ public void testTraceState_withHeaderUsesTraceId() {
 
     @Test
     public void testTraceState_withHeaderUsesTraceIdWithDifferentLocalIds() {
-        Response response = target.path("/local-trace").request()
+        Response response = target.path("/top-span").request()
                 .header(TraceHttpHeaders.TRACE_ID, "traceId")
                 .header(TraceHttpHeaders.PARENT_SPAN_ID, "parentSpanId")
                 .header(TraceHttpHeaders.SPAN_ID, "spanId")
@@ -128,13 +128,13 @@ public void testTraceState_withHeaderUsesTraceIdWithDifferentLocalIds() {
         assertThat(response.getHeaderString(TraceHttpHeaders.PARENT_SPAN_ID)).isNull();
         assertThat(response.getHeaderString(TraceHttpHeaders.SPAN_ID)).isNull();
         verify(observer).consume(spanCaptor.capture());
-        assertThat(spanCaptor.getValue().getOperation()).isEqualTo("Jersey: GET /local-trace");
+        assertThat(spanCaptor.getValue().getOperation()).isEqualTo("Jersey: GET /top-span");
 
         String firstLocalTrace = response.readEntity(String.class);
         assertThat(firstLocalTrace).isNotEmpty();
 
         // make the query again
-        Response otherResponse = target.path("/local-trace").request()
+        Response otherResponse = target.path("/top-span").request()
                 .header(TraceHttpHeaders.TRACE_ID, "traceId")
                 .header(TraceHttpHeaders.PARENT_SPAN_ID, "parentSpanId")
                 .header(TraceHttpHeaders.SPAN_ID, "spanId")
@@ -185,17 +185,6 @@ public void testTraceState_withoutRequestHeadersGeneratesValidTraceResponseHeade
         assertThat(spanCaptor.getValue().getOperation()).isEqualTo("Jersey: GET /trace");
     }
 
-    @Test
-    public void testTraceState_withoutRequestHeadersDoesNotGenerateLocalTrace() {
-        Response response = target.path("/local-trace").request().get();
-        assertThat(response.getHeaderString(TraceHttpHeaders.TRACE_ID)).isNotNull();
-        assertThat(response.getHeaderString(TraceHttpHeaders.PARENT_SPAN_ID)).isNull();
-        assertThat(response.getHeaderString(TraceHttpHeaders.SPAN_ID)).isNull();
-        verify(observer).consume(spanCaptor.capture());
-        assertThat(spanCaptor.getValue().getOperation()).isEqualTo("Jersey: GET /local-trace");
-        assertThat(response.readEntity(String.class)).isNullOrEmpty();
-    }
-
     @Test
     public void testTraceState_withoutRequestHeadersGeneratesValidTraceResponseHeadersWhenFailing() {
         Response response = target.path("/failing-trace").request().get();
@@ -259,6 +248,8 @@ public void testFilter_setsMdcIfTraceIdHeaderIsPresent() throws Exception {
         TraceEnrichingFilter.INSTANCE.filter(request);
         assertThat(MDC.get(Tracers.TRACE_ID_KEY)).isEqualTo("traceId");
         verify(request).setProperty(TraceEnrichingFilter.TRACE_ID_PROPERTY_NAME, "traceId");
+        // the top span id is randomly generated
+        verify(request).setProperty(eq(TraceEnrichingFilter.TOP_SPAN_ID_PROPERTY_NAME), anyString());
         // Note: this will be set to a random value; we want to check whether the value is being set
         verify(request).setProperty(eq(TraceEnrichingFilter.SAMPLED_PROPERTY_NAME), anyBoolean());
     }
@@ -304,8 +295,8 @@ public void getFailingTraceOperation() {
         }
 
         @Override
-        public String getLocalTraceId() {
-            return Tracer.getLocalTraceId().orElse(null);
+        public String getTopSpanId() {
+            return Tracer.getTopSpanId().get();
         }
 
         @Override
@@ -343,8 +334,8 @@ public interface TracingTestService {
         void getFailingTraceOperation();
 
         @GET
-        @Path("/local-trace")
-        String getLocalTraceId();
+        @Path("/top-span")
+        String getTopSpanId();
 
         @GET
         @Path("/failing-streaming-trace")

tracing-undertow/src/main/java/com/palantir/tracing/undertow/TracedOperationHandler.java

@@ -100,8 +100,7 @@ private String initializeTrace(HttpServerExchange exchange) {
 
     /** Initializes trace state given a trace-id header from the client. */
     private void initializeTraceFromExisting(HeaderMap headers, String traceId) {
-        // propagate the traceid from request headers, but create a new local trace id to disambiguate
-        Tracer.initTrace(getObservabilityFromHeader(headers), traceId, Tracers.randomId());
+        Tracer.initTrace(getObservabilityFromHeader(headers), traceId);
         String spanId = headers.getFirst(SPAN_ID); // nullable
         if (spanId == null) {
             Tracer.fastStartSpan(operation, SpanType.SERVER_INCOMING);

tracing-undertow/src/test/java/com/palantir/tracing/undertow/TracedOperationHandlerTest.java

@@ -110,7 +110,7 @@ public void whenTraceIsInHeader_usesGivenTraceIdwithDifferentLocalIds() throws E
         TracedOperationHandler traceSettingHandler =
                 new TracedOperationHandler(exc -> {
                     traceIdValue.set(Tracer.getTraceId());
-                    localTraceIdValue.set(Tracer.getLocalTraceId().orElse(null));
+                    localTraceIdValue.set(Tracer.getTopSpanId().orElse(null));
 
                 }, "GET /traced");
         traceSettingHandler.handleRequest(exchange);

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

@@ -29,7 +29,6 @@
 import java.util.ArrayDeque;
 import java.util.Deque;
 import java.util.Optional;
-import javax.annotation.Nullable;
 
 /**
  * Represents a trace as an ordered list of non-completed spans. Supports adding and removing of spans. This class is
@@ -45,13 +44,10 @@
 public abstract class Trace {
 
     private final String traceId;
-    @Nullable
-    private final String localTraceId;
 
-    private Trace(String traceId, @Nullable  String localTraceId) {
+    private Trace(String traceId) {
         checkArgument(!Strings.isNullOrEmpty(traceId), "traceId must be non-empty");
         this.traceId = traceId;
-        this.localTraceId = localTraceId;
     }
 
     /**
@@ -143,42 +139,33 @@ final String getTraceId() {
     /**
      * The globally unique non-empty identifier for this call trace within a service (or another locality context).
      *
-     * While {@link #getTraceId()} is expected to be propagated across RPC calls, localTraceId distinguishes between
-     * two concurrent RPC calls made to a service with the same traceid.
+     * While {@link #getTraceId()} is expected to be propagated across RPC calls, the top span id distinguishes
+     * between two concurrent RPC calls made to a service with the same traceid.
      */
-    final Optional<String> getLocalTraceId() {
-        // XXX: should this be equal to traceId if localTraceId is not set?
-        return Optional.ofNullable(localTraceId);
-    }
+    abstract Optional<String> getTopSpanId();
 
     abstract Optional<String> getOriginatingSpanId();
 
     /** Returns a copy of this Trace which can be independently mutated. */
     abstract Trace deepCopy();
 
     static Trace of(boolean isObservable, String traceId) {
-        return isObservable ? new Sampled(traceId, null) : new Unsampled(traceId, null);
-    }
-
-    static Trace of(boolean isObservable, String traceId, String localTraceId) {
-        checkArgument(!Strings.isNullOrEmpty(localTraceId), "localTraceId must be non-empty");
-        return isObservable ? new Sampled(traceId, localTraceId) : new Unsampled(traceId, localTraceId);
+        return isObservable ? new Sampled(traceId) : new Unsampled(traceId);
     }
 
     private static final class Sampled extends Trace {
 
         private final Deque<OpenSpan> stack;
 
-        private Sampled(ArrayDeque<OpenSpan> stack, String traceId, @Nullable String localTraceId) {
-            super(traceId, localTraceId);
+        private Sampled(ArrayDeque<OpenSpan> stack, String traceId) {
+            super(traceId);
             this.stack = stack;
         }
 
-        private Sampled(String traceId, @Nullable String localTraceId) {
-            this(new ArrayDeque<>(), traceId, localTraceId);
+        private Sampled(String traceId) {
+            this(new ArrayDeque<>(), traceId);
         }
 
-
         @Override
         @SuppressWarnings("ResultOfMethodCallIgnored") // Sampled traces cannot optimize this path
         void fastStartSpan(String operation, String parentSpanId, SpanType type) {
@@ -216,6 +203,14 @@ boolean isObservable() {
             return true;
         }
 
+        @Override
+        Optional<String> getTopSpanId() {
+            if (stack.isEmpty()) {
+                return Optional.empty();
+            }
+            return Optional.of(stack.peekLast().getSpanId());
+        }
+
         @Override
         Optional<String> getOriginatingSpanId() {
             if (stack.isEmpty()) {
@@ -226,7 +221,7 @@ Optional<String> getOriginatingSpanId() {
 
         @Override
         Trace deepCopy() {
-            return new Sampled(new ArrayDeque<>(stack), getTraceId(), getLocalTraceId().orElse(null));
+            return new Sampled(new ArrayDeque<>(stack), getTraceId());
         }
 
         @Override
@@ -242,35 +237,41 @@ private static final class Unsampled extends Trace {
          */
         private int numberOfSpans;
         private Optional<String> originatingSpanId = Optional.empty();
+        private Optional<String> topSpanId = Optional.empty();
 
-        private Unsampled(int numberOfSpans, String traceId, @Nullable String localTraceId) {
-            super(traceId, localTraceId);
+        private Unsampled(int numberOfSpans, String traceId) {
+            super(traceId);
             this.numberOfSpans = numberOfSpans;
             validateNumberOfSpans();
         }
 
-        private Unsampled(String traceId, @Nullable String localTraceId) {
-            this(0, traceId, localTraceId);
+        private Unsampled(String traceId) {
+            this(0, traceId);
         }
 
         @Override
         void fastStartSpan(String _operation, String parentSpanId, SpanType _type) {
-            startSpan(Optional.of(parentSpanId));
+            if (numberOfSpans == 0) {
+                originatingSpanId = Optional.of(parentSpanId);
+                topSpanId = Optional.of(Tracers.randomId());
+            }
+            numberOfSpans++;
         }
 
         @Override
         void fastStartSpan(String _operation, SpanType _type) {
+            if (numberOfSpans == 0) {
+                topSpanId = Optional.of(Tracers.randomId());
+            }
             numberOfSpans++;
         }
 
         @Override
         protected void push(OpenSpan span) {
-            startSpan(span.getParentSpanId());
-        }
-
-        private void startSpan(Optional<String> parentSpanId) {
             if (numberOfSpans == 0) {
-                originatingSpanId = parentSpanId;
+                // TODO: shouldn't this be span.getOriginatingSpanId?
+                originatingSpanId = span.getParentSpanId();
+                topSpanId = Optional.of(span.getSpanId());
             }
             numberOfSpans++;
         }
@@ -288,6 +289,7 @@ Optional<OpenSpan> pop() {
             }
             if (numberOfSpans == 0) {
                 originatingSpanId = Optional.empty();
+                topSpanId = Optional.empty();
             }
             return Optional.empty();
         }
@@ -303,14 +305,20 @@ boolean isObservable() {
             return false;
         }
 
+        @Override
+        Optional<String> getTopSpanId() {
+            return topSpanId;
+        }
+
         @Override
         Optional<String> getOriginatingSpanId() {
             return originatingSpanId;
         }
 
         @Override
         Trace deepCopy() {
-            return new Unsampled(numberOfSpans, getTraceId(), getLocalTraceId().orElse(null));
+            // TODO: shouldn't this preserve originatingSpanId / topSpanId?
+            return new Unsampled(numberOfSpans, getTraceId());
         }
 
         /** Internal validation, this should never fail because {@link #pop()} only decrements positive values. */

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

@@ -78,15 +78,6 @@ private static Trace createTrace(Observability observability, String traceId) {
         return Trace.of(observable, traceId);
     }
 
-    /**
-     * Creates a new localized trace, but does not set it as the current trace.
-     */
-    private static Trace createTrace(Observability observability, String traceId, String localTraceId) {
-        checkArgument(!Strings.isNullOrEmpty(traceId), "traceId must be non-empty");
-        boolean observable = shouldObserve(observability);
-        return Trace.of(observable, traceId, localTraceId);
-    }
-
     private static boolean shouldObserve(Observability observability) {
         switch (observability) {
             case SAMPLE:
@@ -148,13 +139,6 @@ public static void initTrace(Observability observability, String traceId) {
         setTrace(createTrace(observability, traceId));
     }
 
-    /**
-     * Initializes the current thread's trace, erasing any previously accrued open spans.
-     */
-    public static void initTrace(Observability observability, String traceId, String localTraceId) {
-        setTrace(createTrace(observability, traceId, localTraceId));
-    }
-
     /**
      * Opens a new span for this thread's call trace, labeled with the provided operation and parent span. Only allowed
      * when the current trace is empty.
@@ -500,8 +484,8 @@ public static String getTraceId() {
     }
 
     /** Returns the globally unique identifier for this thread's trace specific to this call. */
-    public static Optional<String> getLocalTraceId() {
-        return checkNotNull(currentTrace.get(), "There is no trace").getLocalTraceId();
+    public static Optional<String> getTopSpanId() {
+        return checkNotNull(currentTrace.get(), "There is no trace").getTopSpanId();
     }
 
     /** Clears the current trace id and returns it if present. */