Tracing uniquely identifies requests (#574)

Tracing uniquely identifies requests

Author: carterkozak

.palantir/revapi.yml

@@ -0,0 +1,6 @@
+acceptedBreaks:
+  "4.5.2":
+    com.palantir.tracing:tracing:
+    - code: "java.method.addedToInterface"
+      new: "method java.util.Optional<java.lang.String> com.palantir.tracing.TraceMetadata::getRequestId()"
+      justification: "immutables are not for extension"

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

@@ -0,0 +1,6 @@
+type: improvement
+improvement:
+  description: Tracing uniquely identifies requests with request metadata and a `_requestId`
+    MDC property.
+  links:
+  - https://github.com/palantir/tracing-java/pull/574

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

@@ -18,6 +18,7 @@
 
 import com.google.common.base.Strings;
 import com.palantir.tracing.Observability;
+import com.palantir.tracing.TraceMetadata;
 import com.palantir.tracing.Tracer;
 import com.palantir.tracing.Tracers;
 import com.palantir.tracing.api.Span;
@@ -45,6 +46,8 @@ 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
@@ -67,21 +70,26 @@ public void filter(ContainerRequestContext requestContext) throws IOException {
         // 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);
+            Tracer.initTraceWithSpan(
+                    getObservabilityFromHeader(requestContext),
+                    Tracers.randomId(),
+                    operation,
+                    SpanType.SERVER_INCOMING);
+        } else if (spanId == null) {
+            Tracer.initTraceWithSpan(
+                    getObservabilityFromHeader(requestContext), traceId, operation, SpanType.SERVER_INCOMING);
         } else {
-            Tracer.initTrace(getObservabilityFromHeader(requestContext), traceId);
-            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);
-            }
+            // caller's span is this span's parent.
+            Tracer.initTraceWithSpan(
+                    getObservabilityFromHeader(requestContext), traceId, operation, spanId, SpanType.SERVER_INCOMING);
         }
 
         // Give asynchronous downstream handlers access to the trace id
         requestContext.setProperty(TRACE_ID_PROPERTY_NAME, Tracer.getTraceId());
         requestContext.setProperty(SAMPLED_PROPERTY_NAME, Tracer.isTraceObservable());
+        Tracer.maybeGetTraceMetadata()
+                .flatMap(TraceMetadata::getRequestId)
+                .ifPresent(requestId -> requestContext.setProperty(REQUEST_ID_PROPERTY_NAME, requestId));
     }
 
     // Handles outgoing response

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

@@ -225,6 +225,7 @@ public void testFilter_setsMdcIfTraceIdHeaderIsPresent() throws Exception {
         verify(request).setProperty(TraceEnrichingFilter.TRACE_ID_PROPERTY_NAME, "traceId");
         // 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());
+        verify(request).setProperty(eq(TraceEnrichingFilter.REQUEST_ID_PROPERTY_NAME), anyString());
     }
 
     @Test
@@ -240,6 +241,7 @@ public void testFilter_setsMdcIfTraceIdHeaderIsNotePresent() throws Exception {
         TraceEnrichingFilter.INSTANCE.filter(request);
         assertThat(MDC.get(Tracers.TRACE_ID_KEY)).hasSize(16);
         verify(request).setProperty(eq(TraceEnrichingFilter.TRACE_ID_PROPERTY_NAME), anyString());
+        verify(request).setProperty(eq(TraceEnrichingFilter.REQUEST_ID_PROPERTY_NAME), anyString());
     }
 
     public static class TracingTestServer extends Application<Configuration> {

tracing-okhttp3/src/test/java/com/palantir/tracing/okhttp3/OkhttpTraceInterceptorTest.java

@@ -67,12 +67,13 @@ public void before() {
 
     @After
     public void after() {
-        Tracer.initTrace(Observability.SAMPLE, Tracers.randomId());
+        Tracer.initTraceWithSpan(Observability.SAMPLE, Tracers.randomId(), "op", SpanType.LOCAL);
         Tracer.unsubscribe("");
     }
 
     @Test
     public void testPopulatesNewTrace_whenNoTraceIsPresentInGlobalState() throws IOException {
+        Tracer.getAndClearTrace();
         OkhttpTraceInterceptor.INSTANCE.intercept(chain);
         verify(chain).request();
         verify(chain).proceed(requestCaptor.capture());
@@ -88,7 +89,7 @@ public void testPopulatesNewTrace_whenNoTraceIsPresentInGlobalState() throws IOE
     @Test
     public void testPopulatesNewTrace_whenParentTraceIsPresent() throws IOException {
         String originatingSpanId = "originating Span";
-        OpenSpan parentState = Tracer.startSpan("operation", originatingSpanId, SpanType.SERVER_INCOMING);
+        Tracer.initTraceWithSpan(Observability.SAMPLE, "id", "operation", originatingSpanId, SpanType.SERVER_INCOMING);
         String traceId = Tracer.getTraceId();
         try {
             OkhttpTraceInterceptor.INSTANCE.intercept(chain);
@@ -102,15 +103,13 @@ public void testPopulatesNewTrace_whenParentTraceIsPresent() throws IOException
 
         Request intercepted = requestCaptor.getValue();
         assertThat(intercepted.headers(TraceHttpHeaders.SPAN_ID)).isNotNull();
-        assertThat(intercepted.headers(TraceHttpHeaders.SPAN_ID)).doesNotContain(parentState.getSpanId());
         assertThat(intercepted.headers(TraceHttpHeaders.TRACE_ID)).containsOnly(traceId);
-        assertThat(intercepted.headers(TraceHttpHeaders.PARENT_SPAN_ID)).containsOnly(parentState.getSpanId());
         assertThat(intercepted.headers(TraceHttpHeaders.ORIGINATING_SPAN_ID)).containsOnly(originatingSpanId);
     }
 
     @Test
     public void testAddsIsSampledHeader_whenTraceIsObservable() throws IOException {
-        Tracer.initTrace(Observability.SAMPLE, Tracers.randomId());
+        Tracer.initTraceWithSpan(Observability.SAMPLE, Tracers.randomId(), "op", SpanType.LOCAL);
         OkhttpTraceInterceptor.INSTANCE.intercept(chain);
         verify(chain).proceed(requestCaptor.capture());
         assertThat(requestCaptor.getValue().headers(TraceHttpHeaders.IS_SAMPLED))
@@ -119,7 +118,7 @@ public void testAddsIsSampledHeader_whenTraceIsObservable() throws IOException {
 
     @Test
     public void testHeaders_whenTraceIsNotObservable() throws IOException {
-        Tracer.initTrace(Observability.DO_NOT_SAMPLE, Tracers.randomId());
+        Tracer.initTraceWithSpan(Observability.DO_NOT_SAMPLE, Tracers.randomId(), "op", SpanType.LOCAL);
         String traceId = Tracer.getTraceId();
         OkhttpTraceInterceptor.INSTANCE.intercept(chain);
         verify(chain).proceed(requestCaptor.capture());
@@ -175,7 +174,7 @@ public void testHeaders_noMultiValue() throws IOException {
                 .build();
         when(chain.request()).thenReturn(request);
 
-        Tracer.initTrace(Observability.SAMPLE, Tracers.randomId());
+        Tracer.initTraceWithSpan(Observability.SAMPLE, Tracers.randomId(), "op", SpanType.LOCAL);
         OkhttpTraceInterceptor.INSTANCE.intercept(chain);
         verify(chain).proceed(requestCaptor.capture());
 

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

@@ -24,5 +24,8 @@ public final class TracingAttachments {
     /** Attachment to check whether the current request is being traced. */
     public static final AttachmentKey<Boolean> IS_SAMPLED = AttachmentKey.create(Boolean.class);
 
+    /** Attachment providing the request identifier. */
+    public static final AttachmentKey<String> REQUEST_ID = AttachmentKey.create(String.class);
+
     private TracingAttachments() {}
 }

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

@@ -18,6 +18,8 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Strings;
+import com.palantir.logsafe.SafeArg;
+import com.palantir.logsafe.exceptions.SafeIllegalStateException;
 import com.palantir.tracing.DetachedSpan;
 import com.palantir.tracing.InternalTracers;
 import com.palantir.tracing.Observability;
@@ -73,6 +75,11 @@ private static void setExchangeState(HttpServerExchange exchange, DetachedSpan d
         // Populate response before proceeding since later operations might commit the response.
         exchange.getResponseHeaders().put(TRACE_ID, traceId);
         exchange.putAttachment(TracingAttachments.IS_SAMPLED, InternalTracers.isSampled(detachedSpan));
+        Optional<String> requestId = InternalTracers.getRequestId(detachedSpan);
+        if (!requestId.isPresent()) {
+            throw new SafeIllegalStateException("No requestId is set", SafeArg.of("span", detachedSpan));
+        }
+        exchange.putAttachment(TracingAttachments.REQUEST_ID, requestId.get());
         exchange.putAttachment(REQUEST_SPAN, detachedSpan);
         exchange.addExchangeCompleteListener(DetachedTraceCompletionListener.INSTANCE);
     }

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

@@ -28,7 +28,7 @@
 import com.palantir.tracing.api.TraceHttpHeaders;
 import io.undertow.Undertow;
 import io.undertow.server.HttpHandler;
-import io.undertow.server.handlers.ResponseCodeHandler;
+import io.undertow.util.HttpString;
 import java.io.IOException;
 import java.net.HttpURLConnection;
 import java.net.URL;
@@ -70,7 +70,8 @@ public void before() {
 
         traceReportedLatch = new CountDownLatch(1);
         port = portSelector.incrementAndGet();
-        HttpHandler nextHandler = new TracedRequestHandler(ResponseCodeHandler.HANDLE_200);
+        HttpHandler nextHandler = new TracedRequestHandler(exchange -> exchange.getResponseHeaders()
+                .put(HttpString.tryFromString("requestId"), exchange.getAttachment(TracingAttachments.REQUEST_ID)));
         server = Undertow.builder()
                 .addHttpListener(port, null)
                 .setHandler(exchange -> {
@@ -101,6 +102,7 @@ public void testRequestTracing_sampled() throws Exception {
         con.setRequestProperty(TraceHttpHeaders.TRACE_ID, "1234");
         assertThat(con.getResponseCode()).isEqualTo(200);
         assertThat(con.getHeaderField(TraceHttpHeaders.TRACE_ID)).isEqualTo("1234");
+        assertThat(con.getHeaderField("requestId")).isNotEmpty();
         assertThat(traceReportedLatch.await(5, TimeUnit.SECONDS)).isTrue();
         verifyNoMoreInteractions(traceSampler);
         verify(observer).consume(spanCaptor.capture());

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

@@ -65,6 +65,9 @@ public final class DeferredTracer implements Serializable {
     @Nullable
     private final String parentSpanId;
 
+    @Nullable
+    private final transient String requestId;
+
     /**
      * Deprecated.
      *
@@ -90,11 +93,13 @@ public DeferredTracer(String operation) {
         if (maybeTrace.isPresent()) {
             Trace trace = maybeTrace.get();
             this.traceId = trace.getTraceId();
+            this.requestId = trace.getRequestId().orElse(null);
             this.isObservable = trace.isObservable();
             this.parentSpanId = trace.top().map(OpenSpan::getSpanId).orElse(null);
             this.operation = operation;
         } else {
             this.traceId = null;
+            this.requestId = null;
             this.isObservable = false;
             this.parentSpanId = null;
             this.operation = null;
@@ -117,7 +122,7 @@ CloseableTrace withTrace() {
 
         Optional<Trace> originalTrace = Tracer.copyTrace();
 
-        Tracer.setTrace(Trace.of(isObservable, traceId));
+        Tracer.setTrace(Trace.of(isObservable, traceId, Optional.ofNullable(requestId)));
         if (parentSpanId != null) {
             Tracer.fastStartSpan(operation, parentSpanId, SpanType.LOCAL);
         } else {

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

@@ -16,6 +16,8 @@
 
 package com.palantir.tracing;
 
+import java.util.Optional;
+
 /**
  * Internal utilities meant for consumption only inside of the tracing codebase.
  *
@@ -29,5 +31,10 @@ public static boolean isSampled(DetachedSpan detachedSpan) {
         return Tracer.isSampled(detachedSpan);
     }
 
+    /** Returns true if the provided detachedSpan is sampled. */
+    public static Optional<String> getRequestId(DetachedSpan detachedSpan) {
+        return Tracer.getRequestId(detachedSpan);
+    }
+
     private InternalTracers() {}
 }