Introduce localTraceId to disambiguate incoming requests that share traceid

Logs from concurrent incoming requests with the same traceid are currently difficult
to tell apart (one can look at thread names etc, but not across executors).

This introduces a new optional property on a Trace, localTraceId, that we set to
a new unique identifier if an incoming request already has a traceid assigned.

Author: JacekLach

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

@@ -46,6 +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 SAMPLED_PROPERTY_NAME = "com.palantir.tracing.sampled";
 
     @Context
@@ -71,7 +72,8 @@ public void filter(ContainerRequestContext requestContext) throws IOException {
             Tracer.initTrace(getObservabilityFromHeader(requestContext), Tracers.randomId());
             Tracer.fastStartSpan(operation, SpanType.SERVER_INCOMING);
         } else {
-            Tracer.initTrace(getObservabilityFromHeader(requestContext), traceId);
+            // propagate the traceid from request headers, but create a new local trace id to disambiguate
+            Tracer.initTrace(getObservabilityFromHeader(requestContext), traceId, Tracers.randomId());
             if (spanId == null) {
                 Tracer.fastStartSpan(operation, SpanType.SERVER_INCOMING);
             } else {
@@ -82,6 +84,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));
         requestContext.setProperty(SAMPLED_PROPERTY_NAME, Tracer.isTraceObservable());
     }
 

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

@@ -117,6 +117,36 @@ public void testTraceState_withHeaderUsesTraceId() {
         assertThat(spanCaptor.getValue().getOperation()).isEqualTo("Jersey: GET /trace");
     }
 
+    @Test
+    public void testTraceState_withHeaderUsesTraceIdWithDifferentLocalIds() {
+        Response response = target.path("/local-trace").request()
+                .header(TraceHttpHeaders.TRACE_ID, "traceId")
+                .header(TraceHttpHeaders.PARENT_SPAN_ID, "parentSpanId")
+                .header(TraceHttpHeaders.SPAN_ID, "spanId")
+                .get();
+        assertThat(response.getHeaderString(TraceHttpHeaders.TRACE_ID)).isEqualTo("traceId");
+        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");
+
+        String firstLocalTrace = response.readEntity(String.class);
+        assertThat(firstLocalTrace).isNotEmpty();
+
+        // make the query again
+        Response otherResponse = target.path("/local-trace").request()
+                .header(TraceHttpHeaders.TRACE_ID, "traceId")
+                .header(TraceHttpHeaders.PARENT_SPAN_ID, "parentSpanId")
+                .header(TraceHttpHeaders.SPAN_ID, "spanId")
+                .get();
+        assertThat(otherResponse.getHeaderString(TraceHttpHeaders.TRACE_ID)).isEqualTo("traceId");
+        assertThat(otherResponse.getHeaderString(TraceHttpHeaders.PARENT_SPAN_ID)).isNull();
+        assertThat(otherResponse.getHeaderString(TraceHttpHeaders.SPAN_ID)).isNull();
+        String otherLocalTrace = otherResponse.readEntity(String.class);
+        assertThat(otherLocalTrace).isNotEmpty();
+        assertThat(otherLocalTrace).isNotEqualTo(firstLocalTrace);
+    }
+
     @Test
     public void testTraceState_respectsMethod() {
         Response response = target.path("/trace").request()
@@ -155,6 +185,17 @@ 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();
@@ -262,6 +303,11 @@ public void getFailingTraceOperation() {
             throw new RuntimeException();
         }
 
+        @Override
+        public String getLocalTraceId() {
+            return Tracer.getLocalTraceId().orElse(null);
+        }
+
         @Override
         public StreamingOutput getFailingStreamingTraceOperation() {
             return os -> {
@@ -296,6 +342,10 @@ public interface TracingTestService {
         @Path("/failing-trace")
         void getFailingTraceOperation();
 
+        @GET
+        @Path("/local-trace")
+        String getLocalTraceId();
+
         @GET
         @Path("/failing-streaming-trace")
         @Produces(MediaType.APPLICATION_OCTET_STREAM)

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

@@ -100,7 +100,8 @@ private String initializeTrace(HttpServerExchange exchange) {
 
     /** Initializes trace state given a trace-id header from the client. */
     private void initializeTraceFromExisting(HeaderMap headers, String traceId) {
-        Tracer.initTrace(getObservabilityFromHeader(headers), traceId);
+        // propagate the traceid from request headers, but create a new local trace id to disambiguate
+        Tracer.initTrace(getObservabilityFromHeader(headers), traceId, Tracers.randomId());
         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

@@ -102,6 +102,29 @@ public void whenTraceIsInHeader_usesGivenTraceId() throws Exception {
         assertThat(exchange.getResponseHeaders().getFirst(TraceHttpHeaders.TRACE_ID)).isEqualTo(traceId);
     }
 
+    @Test
+    public void whenTraceIsInHeader_usesGivenTraceIdwithDifferentLocalIds() throws Exception {
+        setRequestTraceId(traceId);
+        AtomicReference<String> traceIdValue = new AtomicReference<>();
+        AtomicReference<String> localTraceIdValue = new AtomicReference<>();
+        TracedOperationHandler traceSettingHandler =
+                new TracedOperationHandler(exc -> {
+                    traceIdValue.set(Tracer.getTraceId());
+                    localTraceIdValue.set(Tracer.getLocalTraceId().orElse(null));
+
+                }, "GET /traced");
+        traceSettingHandler.handleRequest(exchange);
+
+        assertThat(traceIdValue.get()).isEqualTo(traceId);
+        assertThat(localTraceIdValue).isNotNull();
+        String firstLocalId = localTraceIdValue.get();
+
+        traceSettingHandler.handleRequest(exchange);
+        assertThat(traceIdValue.get()).isEqualTo(traceId);
+        assertThat(localTraceIdValue.get()).isNotNull();
+        assertThat(localTraceIdValue.get()).isNotEqualTo(firstLocalId);
+    }
+
     @Test
     public void whenParentSpanIsGiven_usesParentSpan() throws Exception {
         setRequestTraceId(traceId);

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 javax.annotation.Nullable;
 
 /**
  * Represents a trace as an ordered list of non-completed spans. Supports adding and removing of spans. This class is
@@ -44,10 +45,13 @@
 public abstract class Trace {
 
     private final String traceId;
+    @Nullable
+    private final String localTraceId;
 
-    private Trace(String traceId) {
+    private Trace(String traceId, @Nullable  String localTraceId) {
         checkArgument(!Strings.isNullOrEmpty(traceId), "traceId must be non-empty");
         this.traceId = traceId;
+        this.localTraceId = localTraceId;
     }
 
     /**
@@ -136,28 +140,45 @@ final String getTraceId() {
         return traceId;
     }
 
+    /**
+     * 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.
+     */
+    final Optional<String> getLocalTraceId() {
+        // XXX: should this be equal to traceId if localTraceId is not set?
+        return Optional.ofNullable(localTraceId);
+    }
+
     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) : new Unsampled(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);
     }
 
     private static final class Sampled extends Trace {
 
         private final Deque<OpenSpan> stack;
 
-        private Sampled(ArrayDeque<OpenSpan> stack, String traceId) {
-            super(traceId);
+        private Sampled(ArrayDeque<OpenSpan> stack, String traceId, @Nullable String localTraceId) {
+            super(traceId, localTraceId);
             this.stack = stack;
         }
 
-        private Sampled(String traceId) {
-            this(new ArrayDeque<>(), traceId);
+        private Sampled(String traceId, @Nullable String localTraceId) {
+            this(new ArrayDeque<>(), traceId, localTraceId);
         }
 
+
         @Override
         @SuppressWarnings("ResultOfMethodCallIgnored") // Sampled traces cannot optimize this path
         void fastStartSpan(String operation, String parentSpanId, SpanType type) {
@@ -205,7 +226,7 @@ Optional<String> getOriginatingSpanId() {
 
         @Override
         Trace deepCopy() {
-            return new Sampled(new ArrayDeque<>(stack), getTraceId());
+            return new Sampled(new ArrayDeque<>(stack), getTraceId(), getLocalTraceId().orElse(null));
         }
 
         @Override
@@ -222,14 +243,14 @@ private static final class Unsampled extends Trace {
         private int numberOfSpans;
         private Optional<String> originatingSpanId = Optional.empty();
 
-        private Unsampled(int numberOfSpans, String traceId) {
-            super(traceId);
+        private Unsampled(int numberOfSpans, String traceId, @Nullable String localTraceId) {
+            super(traceId, localTraceId);
             this.numberOfSpans = numberOfSpans;
             validateNumberOfSpans();
         }
 
-        private Unsampled(String traceId) {
-            this(0, traceId);
+        private Unsampled(String traceId, @Nullable String localTraceId) {
+            this(0, traceId, localTraceId);
         }
 
         @Override
@@ -289,7 +310,7 @@ Optional<String> getOriginatingSpanId() {
 
         @Override
         Trace deepCopy() {
-            return new Unsampled(numberOfSpans, getTraceId());
+            return new Unsampled(numberOfSpans, getTraceId(), getLocalTraceId().orElse(null));
         }
 
         /** Internal validation, this should never fail because {@link #pop()} only decrements positive values. */

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

@@ -78,6 +78,15 @@ 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:
@@ -139,6 +148,13 @@ 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.
@@ -483,6 +499,11 @@ public static String getTraceId() {
         return checkNotNull(currentTrace.get(), "There is no trace").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();
+    }
+
     /** Clears the current trace id and returns it if present. */
     static Optional<Trace> getAndClearTraceIfPresent() {
         Optional<Trace> trace = Optional.ofNullable(currentTrace.get());