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-364.v2.yml

@@ -0,0 +1,6 @@
+type: feature
+feature:
+  description: Introduce a new MDC property, _requestId, that disambiguates span stacks
+    that belong to the same traceid but are different requests.
+  links:
+  - https://github.com/palantir/tracing-java/pull/364

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,32 @@ public TracedOperationHandler(HttpHandler delegate, String operation) {
 
     @Override
     public void handleRequest(HttpServerExchange exchange) throws Exception {
-        String traceId = initializeTrace(exchange);
+        HeaderMap headers = exchange.getRequestHeaders();
+        // TODO(rfink): Log/warn if we find multiple headers?
+        String traceIdHeader = headers.getFirst(TRACE_ID); // nullable
+        String traceId;
+        // XXX: again, nasty abstraction leak
+        String requestId;
+
+        // Set up thread-local span that inherits state from HTTP headers
+        if (Strings.isNullOrEmpty(traceIdHeader)) {
+            traceId = initializeNewTrace(headers);
+            requestId = Tracer.startSpan(operation, SpanType.SERVER_INCOMING).getSpanId();
+        } else {
+            traceId = traceIdHeader;
+            Tracer.initTrace(getObservabilityFromHeader(headers), traceIdHeader);
+            String spanId = headers.getFirst(SPAN_ID); // nullable
+            if (spanId == null) {
+                requestId = Tracer.startSpan(operation, SpanType.SERVER_INCOMING).getSpanId();
+            } else {
+                requestId = Tracer.startSpan(operation, spanId, 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 {
@@ -83,39 +105,11 @@ private static Observability getObservabilityFromHeader(HeaderMap headers) {
         }
     }
 
-    /** Initializes trace state and a root span for this request, returning the traceId. */
-    private String initializeTrace(HttpServerExchange exchange) {
-        HeaderMap headers = exchange.getRequestHeaders();
-        // TODO(rfink): Log/warn if we find multiple headers?
-        String traceId = headers.getFirst(TRACE_ID); // nullable
-
-        // Set up thread-local span that inherits state from HTTP headers
-        if (Strings.isNullOrEmpty(traceId)) {
-            return initializeNewTrace(headers);
-        } else {
-            initializeTraceFromExisting(headers, traceId);
-        }
-        return traceId;
-    }
-
-    /** Initializes trace state given a trace-id header from the client. */
-    private void initializeTraceFromExisting(HeaderMap headers, String traceId) {
-        Tracer.initTrace(getObservabilityFromHeader(headers), traceId);
-        String spanId = headers.getFirst(SPAN_ID); // nullable
-        if (spanId == null) {
-            Tracer.fastStartSpan(operation, SpanType.SERVER_INCOMING);
-        } else {
-            // caller's span is this span's parent.
-            Tracer.fastStartSpan(operation, spanId, SpanType.SERVER_INCOMING);
-        }
-    }
-
     /** Initializes trace state for a request without tracing headers. */
     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

@@ -136,6 +136,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 +197,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 +234,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 +297,7 @@ Optional<OpenSpan> pop() {
             }
             if (numberOfSpans == 0) {
                 originatingSpanId = Optional.empty();
+                requestId = Optional.empty();
             }
             return Optional.empty();
         }
@@ -282,14 +313,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. */