Introduce requestId MDC property

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 MDC property managed by Tracers, which stores a unique identifier
when the top span of the current trace is a SERVER_INCOMING span.

changelog/@unreleased/pr-352.v2.yml

@@ -0,0 +1,7 @@
+type: feature
+feature:
+  description: Introduce a new property on a Trace, outermostSpanId, that can be read to
+    disambiguate span stacks that belong to the same traceid but had different entry
+    points.
+  links:
+  - https://github.com/palantir/tracing-java/pull/352

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 REQUEST_ID_PROPERTY_NAME = "com.palantir.tracing.requestId";
     public static final String SAMPLED_PROPERTY_NAME = "com.palantir.tracing.sampled";
 
     @Context
@@ -65,24 +66,31 @@ public void filter(ContainerRequestContext requestContext) throws IOException {
         String traceId = requestContext.getHeaderString(TraceHttpHeaders.TRACE_ID);
         String spanId = requestContext.getHeaderString(TraceHttpHeaders.SPAN_ID);
 
+        // XXX: leaks abstraction and ties us to requestId === spanId of initial span
+        // however the leak is contained without this repository, so maybe OK?
+        // alternative is to expose requestId as a property on the current trace to all clients.
+        // or read the property from the MDC at this point
+        String requestId;
+
         // Set up thread-local span that inherits state from HTTP headers
         if (Strings.isNullOrEmpty(traceId)) {
             // HTTP request did not indicate a trace; initialize trace state and create a span.
             Tracer.initTrace(getObservabilityFromHeader(requestContext), Tracers.randomId());
-            Tracer.fastStartSpan(operation, SpanType.SERVER_INCOMING);
+            requestId = Tracer.startSpan(operation, SpanType.SERVER_INCOMING).getSpanId();
         } else {
             Tracer.initTrace(getObservabilityFromHeader(requestContext), traceId);
             if (spanId == null) {
-                Tracer.fastStartSpan(operation, SpanType.SERVER_INCOMING);
+                requestId = Tracer.startSpan(operation, SpanType.SERVER_INCOMING).getSpanId();
             } else {
                 // caller's span is this span's parent.
-                Tracer.fastStartSpan(operation, spanId, SpanType.SERVER_INCOMING);
+                requestId = Tracer.startSpan(operation, spanId, SpanType.SERVER_INCOMING).getSpanId();
             }
         }
 
         // Give asynchronous downstream handlers access to the trace id
         requestContext.setProperty(TRACE_ID_PROPERTY_NAME, Tracer.getTraceId());
         requestContext.setProperty(SAMPLED_PROPERTY_NAME, Tracer.isTraceObservable());
+        requestContext.setProperty(REQUEST_ID_PROPERTY_NAME, requestId);
     }
 
     // Handles outgoing response

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("/top-span").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 /top-span");
+
+        String firstLocalTrace = response.readEntity(String.class);
+        assertThat(firstLocalTrace).isNotEmpty();
+
+        // make the query again
+        Response otherResponse = target.path("/top-span").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()
@@ -218,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.REQUEST_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());
     }
@@ -262,6 +294,11 @@ public void getFailingTraceOperation() {
             throw new RuntimeException();
         }
 
+        @Override
+        public String getRequestId() {
+            return MDC.get(Tracers.REQUEST_ID_KEY);
+        }
+
         @Override
         public StreamingOutput getFailingStreamingTraceOperation() {
             return os -> {
@@ -296,6 +333,10 @@ public interface TracingTestService {
         @Path("/failing-trace")
         void getFailingTraceOperation();
 
+        @GET
+        @Path("/top-span")
+        String getRequestId();
+
         @GET
         @Path("/failing-streaming-trace")
         @Produces(MediaType.APPLICATION_OCTET_STREAM)

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

@@ -45,6 +45,7 @@ public final class TracedOperationHandler implements HttpHandler {
      * Attachment to check whether the current request is being traced.
      */
     public static final AttachmentKey<Boolean> IS_SAMPLED_ATTACHMENT = AttachmentKey.create(Boolean.class);
+    public static final AttachmentKey<String> REQUEST_ID_ATTACHMENT = AttachmentKey.create(String.class);
 
     private static final HttpString TRACE_ID = HttpString.tryFromString(TraceHttpHeaders.TRACE_ID);
     private static final HttpString SPAN_ID = HttpString.tryFromString(TraceHttpHeaders.SPAN_ID);
@@ -60,11 +61,14 @@ public TracedOperationHandler(HttpHandler delegate, String operation) {
 
     @Override
     public void handleRequest(HttpServerExchange exchange) throws Exception {
-        String traceId = initializeTrace(exchange);
+        String traceId = initializeTraceWithNoSpans(exchange);
+        // XXX: again, nasty abstraction leak
+        String requestId = Tracer.startSpan(operation, SpanType.SERVER_INCOMING).getSpanId();
 
         // Populate response before calling delegate since delegate might commit the response.
         exchange.getResponseHeaders().put(TRACE_ID, traceId);
         exchange.putAttachment(IS_SAMPLED_ATTACHMENT, Tracer.isTraceObservable());
+        exchange.putAttachment(REQUEST_ID_ATTACHMENT, requestId);
         try {
             delegate.handleRequest(exchange);
         } finally {
@@ -84,7 +88,7 @@ private static Observability getObservabilityFromHeader(HeaderMap headers) {
     }
 
     /** Initializes trace state and a root span for this request, returning the traceId. */
-    private String initializeTrace(HttpServerExchange exchange) {
+    private String initializeTraceWithNoSpans(HttpServerExchange exchange) {
         HeaderMap headers = exchange.getRequestHeaders();
         // TODO(rfink): Log/warn if we find multiple headers?
         String traceId = headers.getFirst(TRACE_ID); // nullable
@@ -115,7 +119,6 @@ private String initializeNewTrace(HeaderMap headers) {
         // HTTP request did not indicate a trace; initialize trace state and create a span.
         String newTraceId = Tracers.randomId();
         Tracer.initTrace(getObservabilityFromHeader(headers), newTraceId);
-        Tracer.fastStartSpan(operation, SpanType.SERVER_INCOMING);
         return newTraceId;
     }
 

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

@@ -102,6 +102,21 @@ public void whenTraceIsInHeader_usesGivenTraceId() throws Exception {
         assertThat(exchange.getResponseHeaders().getFirst(TraceHttpHeaders.TRACE_ID)).isEqualTo(traceId);
     }
 
+    @Test
+    public void whenTraceIsInHeader_usesGivenTraceIdWithDifferentRequestIds() throws Exception {
+        setRequestTraceId(traceId);
+        exchange.getRequestHeaders().put(HttpString.tryFromString(TraceHttpHeaders.IS_SAMPLED), "1");
+
+        handler.handleRequest(exchange);
+        String firstRequestId = exchange.getAttachment(TracedOperationHandler.REQUEST_ID_ATTACHMENT);
+        assertThat(firstRequestId).isNotEmpty();
+
+        handler.handleRequest(exchange);
+        String secondRequestId = exchange.getAttachment(TracedOperationHandler.REQUEST_ID_ATTACHMENT);
+        assertThat(secondRequestId).isNotEmpty();
+        assertThat(firstRequestId).isNotEqualTo(secondRequestId);
+    }
+
     @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 org.slf4j.MDC;
 
 /**
  * Represents a trace as an ordered list of non-completed spans. Supports adding and removing of spans. This class is
@@ -136,6 +137,8 @@ final String getTraceId() {
         return traceId;
     }
 
+    abstract Optional<String> getRequestId();
+
     abstract Optional<String> getOriginatingSpanId();
 
     /** Returns a copy of this Trace which can be independently mutated. */
@@ -195,6 +198,18 @@ boolean isObservable() {
             return true;
         }
 
+        @Override
+        Optional<String> getRequestId() {
+            if (stack.isEmpty()) {
+                return Optional.empty();
+            }
+            OpenSpan top = stack.peekFirst();
+            if (top.type() != SpanType.SERVER_INCOMING) {
+                return Optional.empty();
+            }
+            return Optional.of(top.getSpanId());
+        }
+
         @Override
         Optional<String> getOriginatingSpanId() {
             if (stack.isEmpty()) {
@@ -220,36 +235,52 @@ private static final class Unsampled extends Trace {
          * This allows thread trace state to be cleared when all "started" spans have been "removed".
          */
         private int numberOfSpans;
-        private Optional<String> originatingSpanId = Optional.empty();
-
-        private Unsampled(int numberOfSpans, String traceId) {
+        private Optional<String> originatingSpanId;
+        private Optional<String> requestId;
+
+        private Unsampled(
+                int numberOfSpans,
+                String traceId,
+                Optional<String> originatingSpanId,
+                Optional<String> requestId) {
             super(traceId);
             this.numberOfSpans = numberOfSpans;
+            this.originatingSpanId = originatingSpanId;
+            this.requestId = requestId;
             validateNumberOfSpans();
         }
 
         private Unsampled(String traceId) {
-            this(0, traceId);
+            this(0, traceId, Optional.empty(), Optional.empty());
         }
 
         @Override
-        void fastStartSpan(String _operation, String parentSpanId, SpanType _type) {
-            startSpan(Optional.of(parentSpanId));
+        void fastStartSpan(String _operation, String parentSpanId, SpanType type) {
+            if (numberOfSpans == 0) {
+                originatingSpanId = Optional.of(parentSpanId);
+                if (type == SpanType.SERVER_INCOMING) {
+                    requestId = Optional.of(Tracers.randomId());
+                }
+            }
+            numberOfSpans++;
         }
 
         @Override
-        void fastStartSpan(String _operation, SpanType _type) {
+        void fastStartSpan(String _operation, SpanType type) {
+            if (numberOfSpans == 0 && type == SpanType.SERVER_INCOMING) {
+                requestId = 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;
+                originatingSpanId = span.getParentSpanId();
+
+                if (span.type() == SpanType.SERVER_INCOMING) {
+                    requestId = Optional.of(span.getSpanId());
+                }
             }
             numberOfSpans++;
         }
@@ -267,6 +298,7 @@ Optional<OpenSpan> pop() {
             }
             if (numberOfSpans == 0) {
                 originatingSpanId = Optional.empty();
+                requestId = Optional.empty();
             }
             return Optional.empty();
         }
@@ -282,14 +314,19 @@ boolean isObservable() {
             return false;
         }
 
+        @Override
+        Optional<String> getRequestId() {
+            return requestId;
+        }
+
         @Override
         Optional<String> getOriginatingSpanId() {
             return originatingSpanId;
         }
 
         @Override
         Trace deepCopy() {
-            return new Unsampled(numberOfSpans, getTraceId());
+            return new Unsampled(numberOfSpans, getTraceId(), getOriginatingSpanId(), getRequestId());
         }
 
         /** Internal validation, this should never fail because {@link #pop()} only decrements positive values. */

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

@@ -146,7 +146,13 @@ public static void initTrace(Observability observability, String traceId) {
      */
     @CheckReturnValue
     public static OpenSpan startSpan(String operation, String parentSpanId, SpanType type) {
-        return getOrCreateCurrentTrace().startSpan(operation, parentSpanId, type);
+        Trace trace = getOrCreateCurrentTrace();
+        boolean wasEmpty = trace.isEmpty();
+        OpenSpan result = trace.startSpan(operation, parentSpanId, type);
+        if (wasEmpty) {
+            setRequestIdIfPresent(trace.getRequestId());
+        }
+        return result;
     }
 
     /**
@@ -155,7 +161,13 @@ public static OpenSpan startSpan(String operation, String parentSpanId, SpanType
      */
     @CheckReturnValue
     public static OpenSpan startSpan(String operation, SpanType type) {
-        return getOrCreateCurrentTrace().startSpan(operation, type);
+        Trace trace = getOrCreateCurrentTrace();
+        boolean wasEmpty = trace.isEmpty();
+        OpenSpan result = trace.startSpan(operation, type);
+        if (wasEmpty) {
+            setRequestIdIfPresent(trace.getRequestId());
+        }
+        return result;
     }
 
     /**
@@ -171,14 +183,24 @@ public static OpenSpan startSpan(String operation) {
      * Like {@link #startSpan(String, String, SpanType)}, but does not return an {@link OpenSpan}.
      */
     public static void fastStartSpan(String operation, String parentSpanId, SpanType type) {
-        getOrCreateCurrentTrace().fastStartSpan(operation, parentSpanId, type);
+        Trace trace = getOrCreateCurrentTrace();
+        boolean wasEmpty = trace.isEmpty();
+        trace.fastStartSpan(operation, parentSpanId, type);
+        if (wasEmpty) {
+            setRequestIdIfPresent(trace.getRequestId());
+        }
     }
 
     /**
      * Like {@link #startSpan(String, SpanType)}, but does not return an {@link OpenSpan}.
      */
     public static void fastStartSpan(String operation, SpanType type) {
-        getOrCreateCurrentTrace().fastStartSpan(operation, type);
+        Trace trace = getOrCreateCurrentTrace();
+        boolean wasEmpty = trace.isEmpty();
+        trace.fastStartSpan(operation, type);
+        if (wasEmpty) {
+            setRequestIdIfPresent(trace.getRequestId());
+        }
     }
 
     /**
@@ -525,6 +547,7 @@ static void setTrace(Trace trace) {
         // Give log appenders access to the trace id and whether the trace is being sampled
         MDC.put(Tracers.TRACE_ID_KEY, trace.getTraceId());
         setTraceSampledMdcIfObservable(trace.isObservable());
+        setRequestIdIfPresent(trace.getRequestId());
     }
 
     private static void setTraceSampledMdcIfObservable(boolean observable) {
@@ -537,6 +560,14 @@ private static void setTraceSampledMdcIfObservable(boolean observable) {
         }
     }
 
+    private static void setRequestIdIfPresent(Optional<String> requestId) {
+        if (requestId.isPresent()) {
+            MDC.put(Tracers.REQUEST_ID_KEY, requestId.get());
+        } else {
+            MDC.remove(Tracers.REQUEST_ID_KEY);
+        }
+    }
+
     private static Trace getOrCreateCurrentTrace() {
         Trace trace = currentTrace.get();
         if (trace == null) {
@@ -551,5 +582,6 @@ static void clearCurrentTrace() {
         currentTrace.remove();
         MDC.remove(Tracers.TRACE_ID_KEY);
         MDC.remove(Tracers.TRACE_SAMPLED_KEY);
+        MDC.remove(Tracers.REQUEST_ID_KEY);
     }
 }