Introduce requestId MDC property

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

@@ -35,6 +35,7 @@
 import javax.ws.rs.ext.Provider;
 import org.glassfish.jersey.server.ExtendedUriInfo;
 import org.glassfish.jersey.server.model.Resource;
+import org.slf4j.MDC;
 
 // Default is `Priorities.USER` == 5000. This filter needs to execute earlier to ensure traces are ready to use.
 @Priority(500)
@@ -46,6 +47,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
@@ -82,6 +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());
+        requestContext.setProperty(REQUEST_ID_PROPERTY_NAME, MDC.get(Tracers.REQUEST_ID_KEY));
         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("/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

@@ -29,6 +29,7 @@
 import io.undertow.util.AttachmentKey;
 import io.undertow.util.HeaderMap;
 import io.undertow.util.HttpString;
+import org.slf4j.MDC;
 
 /**
  * Extracts Zipkin-style trace information from the given HTTP request and sets up a corresponding
@@ -45,6 +46,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);
@@ -65,6 +67,7 @@ public void handleRequest(HttpServerExchange exchange) throws Exception {
         // 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, MDC.get(Tracers.REQUEST_ID_KEY));
         try {
             delegate.handleRequest(exchange);
         } finally {

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);
@@ -166,11 +181,19 @@ public void completesSpanEvenIfDelegateThrows() throws Exception {
     public void populatesSlf4jMdc() throws Exception {
         setRequestTraceId(traceId);
         AtomicReference<String> mdcTraceValue = new AtomicReference<>();
-        new TracedOperationHandler(exc -> mdcTraceValue.set(MDC.get(Tracers.TRACE_ID_KEY)), "GET /traced")
-                .handleRequest(exchange);
+        AtomicReference<String> mdcRequestIdValue = new AtomicReference<>();
+        new TracedOperationHandler(
+                exc -> {
+                    mdcTraceValue.set(MDC.get(Tracers.TRACE_ID_KEY));
+                    mdcRequestIdValue.set(MDC.get(Tracers.REQUEST_ID_KEY));
+                },
+                "GET /traced"
+        ).handleRequest(exchange);
         assertThat(mdcTraceValue).hasValue(traceId);
+        assertThat(mdcRequestIdValue.get()).isNotNull();
         // Value should be cleared when the handler returns
         assertThat(MDC.get(Tracers.TRACE_ID_KEY)).isNull();
+        assertThat(MDC.get(Tracers.REQUEST_ID_KEY)).isNull();
     }
 
     private void setRequestTraceId(String theTraceId) {

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. */

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);
     }
 }

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

@@ -35,6 +35,12 @@ public final class Tracers {
      * field can take the values of "1" or "0", where "1" indicates the trace was sampled.
      */
     public static final String TRACE_SAMPLED_KEY = "_sampled";
+    /**
+     * The key under which the outermost span ids are inserted into SLF4J {@link org.slf4j.MDC MDCs}.
+     * It is present only if the outermost span of the current trace is a SERVER_INCOMING span, and can be used
+     * to distinguish between incoming requests with the same traceid header.
+     */
+    public static final String REQUEST_ID_KEY = "_requestId";
     private static final String DEFAULT_ROOT_SPAN_OPERATION = "root";
     private static final char[] HEX_DIGITS =
             {'0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f'};