Read requestId from MDC in handler/filter

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)
@@ -66,31 +67,25 @@ 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());
-            requestId = Tracer.startSpan(operation, SpanType.SERVER_INCOMING).getSpanId();
+            Tracer.fastStartSpan(operation, SpanType.SERVER_INCOMING);
         } else {
             Tracer.initTrace(getObservabilityFromHeader(requestContext), traceId);
             if (spanId == null) {
-                requestId = Tracer.startSpan(operation, SpanType.SERVER_INCOMING).getSpanId();
+                Tracer.fastStartSpan(operation, SpanType.SERVER_INCOMING);
             } else {
                 // caller's span is this span's parent.
-                requestId = Tracer.startSpan(operation, spanId, SpanType.SERVER_INCOMING).getSpanId();
+                Tracer.fastStartSpan(operation, spanId, SpanType.SERVER_INCOMING);
             }
         }
 
         // 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());
-        requestContext.setProperty(REQUEST_ID_PROPERTY_NAME, requestId);
     }
 
     // Handles outgoing response

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
@@ -61,32 +62,12 @@ public TracedOperationHandler(HttpHandler delegate, String operation) {
 
     @Override
     public void handleRequest(HttpServerExchange exchange) throws Exception {
-        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();
-            }
-        }
+        String traceId = initializeTrace(exchange);
 
         // 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);
+        exchange.putAttachment(REQUEST_ID_ATTACHMENT, MDC.get(Tracers.REQUEST_ID_KEY));
         try {
             delegate.handleRequest(exchange);
         } finally {
@@ -105,11 +86,39 @@ 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;
     }