Provide an API to access the detached trace associated with an exchan…

…ge (#1287)

Provide an API to access the detached trace associated with an exchange

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

@@ -0,0 +1,5 @@
+type: improvement
+improvement:
+  description: Provide an API to access the detached trace associated with an exchange
+  links:
+  - https://github.com/palantir/tracing-java/pull/1287

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

@@ -16,7 +16,11 @@
 
 package com.palantir.tracing.undertow;
 
+import com.palantir.tracing.Detached;
+import com.palantir.tracing.DetachedSpan;
+import io.undertow.server.HttpServerExchange;
 import io.undertow.util.AttachmentKey;
+import javax.annotation.Nullable;
 
 /** Provides public tracing {@link AttachmentKey attachment keys}. */
 public final class TracingAttachments {
@@ -27,5 +31,22 @@ public final class TracingAttachments {
     /** Attachment providing the request identifier. */
     public static final AttachmentKey<String> REQUEST_ID = AttachmentKey.create(String.class);
 
+    /**
+     * Detached span object representing the entire request including asynchronous components.
+     * This is intentionally not public, we expose only the {@link Detached} component which critically does not
+     * include {@link DetachedSpan#complete()} APIs.
+     */
+    static final AttachmentKey<DetachedSpan> REQUEST_SPAN = AttachmentKey.create(DetachedSpan.class);
+
+    /**
+     * Gets the {@link Detached} trace state which represents the top-level request being processed. This may
+     * be used to apply thread state to code executing outside traced handlers, exchange completion
+     * listeners, for example.
+     */
+    @Nullable
+    public static Detached requestTrace(HttpServerExchange exchange) {
+        return exchange.getAttachment(REQUEST_SPAN);
+    }
+
     private TracingAttachments() {}
 }

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

@@ -56,13 +56,6 @@ final class UndertowTracing {
     @VisibleForTesting
     static final HttpString FETCH_USER_AGENT = HttpString.tryFromString("Fetch-User-Agent");
 
-    // Consider moving this to TracingAttachments and making it public. For now it's well encapsulated
-    // here because we expect the two handler implementations to be sufficient.
-    /**
-     * Detached span object representing the entire request including asynchronous components.
-     */
-    static final AttachmentKey<DetachedSpan> REQUEST_SPAN = AttachmentKey.create(DetachedSpan.class);
-
     private static final AttachmentKey<TagTranslator<? super HttpServerExchange>> TAG_TRANSLATOR_ATTACHMENT_KEY =
             AttachmentKey.create(TagTranslator.class);
 
@@ -71,7 +64,7 @@ final class UndertowTracing {
      */
     static DetachedSpan getOrInitializeRequestTrace(
             HttpServerExchange exchange, String operationName, TagTranslator<? super HttpServerExchange> translator) {
-        DetachedSpan detachedSpan = exchange.getAttachment(REQUEST_SPAN);
+        DetachedSpan detachedSpan = exchange.getAttachment(TracingAttachments.REQUEST_SPAN);
         if (detachedSpan == null) {
             return initializeRequestTrace(exchange, operationName, translator);
         }
@@ -104,7 +97,7 @@ private static void setExchangeState(
             throw new SafeIllegalStateException("No requestId is set", SafeArg.of("span", detachedSpan));
         }
         exchange.putAttachment(TracingAttachments.REQUEST_ID, requestId.get());
-        exchange.putAttachment(REQUEST_SPAN, detachedSpan);
+        exchange.putAttachment(TracingAttachments.REQUEST_SPAN, detachedSpan);
         exchange.putAttachment(TAG_TRANSLATOR_ATTACHMENT_KEY, translator);
         exchange.addExchangeCompleteListener(DetachedTraceCompletionListener.INSTANCE);
     }
@@ -130,7 +123,7 @@ private enum DetachedTraceCompletionListener implements ExchangeCompletionListen
         @Override
         public void exchangeEvent(HttpServerExchange exchange, NextListener nextListener) {
             try {
-                DetachedSpan detachedSpan = exchange.getAttachment(REQUEST_SPAN);
+                DetachedSpan detachedSpan = exchange.getAttachment(TracingAttachments.REQUEST_SPAN);
                 if (detachedSpan != null) {
                     detachedSpan.complete(tagTranslator(exchange), exchange);
                 }

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

@@ -118,7 +118,7 @@ public void whenParentSpanIsGiven_usesParentSpan() throws Exception {
 
         handler.handleRequest(exchange);
         // Since we're not running a full request, the completion handler cannot execute normally.
-        exchange.getAttachment(UndertowTracing.REQUEST_SPAN).complete();
+        exchange.getAttachment(TracingAttachments.REQUEST_SPAN).complete();
         verify(observer, times(2)).consume(spanCaptor.capture());
         List<Span> spans = spanCaptor.getAllValues();
         assertThat(spans).hasSize(2);
@@ -134,7 +134,7 @@ public void whenOnlyUserAgentIsProvided_setsItAsForUserAgent() throws Exception
         setUserAgent("userAgent");
         handler.handleRequest(exchange);
 
-        DetachedSpan detachedSpan = exchange.getAttachment(UndertowTracing.REQUEST_SPAN);
+        DetachedSpan detachedSpan = exchange.getAttachment(TracingAttachments.REQUEST_SPAN);
         assertThat(InternalTracers.getForUserAgent(detachedSpan)).contains("userAgent");
     }
 
@@ -145,7 +145,7 @@ public void whenFetchUserAgentIsProvided_setsItAsForUserAgent() throws Exception
         setFetchUserAgent("fetchUserAgent");
         handler.handleRequest(exchange);
 
-        DetachedSpan detachedSpan = exchange.getAttachment(UndertowTracing.REQUEST_SPAN);
+        DetachedSpan detachedSpan = exchange.getAttachment(TracingAttachments.REQUEST_SPAN);
         assertThat(InternalTracers.getForUserAgent(detachedSpan)).contains("fetchUserAgent");
     }
 
@@ -157,7 +157,7 @@ public void whenForUserAgentIsProvided_propagateItFurther() throws Exception {
         setForUserAgent("forUserAgent");
         handler.handleRequest(exchange);
 
-        DetachedSpan detachedSpan = exchange.getAttachment(UndertowTracing.REQUEST_SPAN);
+        DetachedSpan detachedSpan = exchange.getAttachment(TracingAttachments.REQUEST_SPAN);
         assertThat(InternalTracers.getForUserAgent(detachedSpan)).contains("forUserAgent");
     }
 
@@ -197,6 +197,12 @@ public void whenSamplingDecisionHasNotBeenMade_callsSampler() throws Exception {
         verify(traceSampler).sample();
     }
 
+    @Test
+    public void setsDetachedTrace() throws Exception {
+        handler.handleRequest(exchange);
+        assertThat(TracingAttachments.requestTrace(exchange)).isNotNull();
+    }
+
     @Test
     public void completesSpanEvenIfDelegateThrows() throws Exception {
         doThrow(new RuntimeException()).when(delegate).handleRequest(any());