Use top span id as the local id

Author: JacekLach

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

@@ -73,7 +73,7 @@ public void filter(ContainerRequestContext requestContext) throws IOException {
             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,7 +84,7 @@ 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 ->
+        Tracer.getTopSpanId().ifPresent(localId ->
                 requestContext.setProperty(LOCAL_TRACE_ID_PROPERTY_NAME, localId));
         requestContext.setProperty(SAMPLED_PROPERTY_NAME, Tracer.isTraceObservable());
     }

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

@@ -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();
@@ -305,7 +294,7 @@ public void getFailingTraceOperation() {
 
         @Override
         public String getLocalTraceId() {
-            return Tracer.getLocalTraceId().orElse(null);
+            return Tracer.getTopSpanId().orElse(null);
         }
 
         @Override

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

@@ -45,13 +45,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,39 +140,31 @@ 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);
         }
 
 
@@ -216,6 +205,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 +223,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 +239,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 +291,7 @@ Optional<OpenSpan> pop() {
             }
             if (numberOfSpans == 0) {
                 originatingSpanId = Optional.empty();
+                topSpanId = Optional.empty();
             }
             return Optional.empty();
         }
@@ -303,14 +307,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. */