Introduce requestId MDC property

JacekLach · JacekLach · commit 5ed48513d74f · 2019-11-25T15:09:48.000Z
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.
diff --git a/changelog/@unreleased/pr-364.v2.yml b/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
diff --git a/tracing-jersey/src/main/java/com/palantir/tracing/jersey/TraceEnrichingFilter.java b/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
diff --git a/tracing-jersey/src/test/java/com/palantir/tracing/jersey/TraceEnrichingFilterTest.java b/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)
diff --git a/tracing-undertow/src/main/java/com/palantir/tracing/undertow/TracedOperationHandler.java b/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;
     }
 
diff --git a/tracing-undertow/src/test/java/com/palantir/tracing/undertow/TracedOperationHandlerTest.java b/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) {
diff --git a/tracing/src/main/java/com/palantir/tracing/Trace.java b/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. */
diff --git a/tracing/src/main/java/com/palantir/tracing/Tracer.java b/tracing/src/main/java/com/palantir/tracing/Tracer.java
diff --git a/tracing/src/main/java/com/palantir/tracing/Tracers.java b/tracing/src/main/java/com/palantir/tracing/Tracers.java
