From 35405764394074a18b8efb58e83337ead33b62c2 Mon Sep 17 00:00:00 2001
From: Carter Kozak <ckozak@palantir.com>
Date: Wed, 2 Oct 2024 11:41:17 -0400
Subject: [PATCH] broaden the requestTrace attachment accessor (#1297)

broaden the requestTrace attachment accessor
---
 changelog/@unreleased/pr-1297.v2.yml                  |  5 +++++
 .../palantir/tracing/undertow/TracingAttachments.java | 11 ++++++++++-
 .../tracing/undertow/TracedOperationHandlerTest.java  |  2 ++
 3 files changed, 17 insertions(+), 1 deletion(-)
 create mode 100644 changelog/@unreleased/pr-1297.v2.yml

diff --git a/changelog/@unreleased/pr-1297.v2.yml b/changelog/@unreleased/pr-1297.v2.yml
new file mode 100644
index 000000000..927fc8044
--- /dev/null
+++ b/changelog/@unreleased/pr-1297.v2.yml
@@ -0,0 +1,5 @@
+type: improvement
+improvement:
+  description: broaden the requestTrace attachment accessor
+  links:
+  - https://github.com/palantir/tracing-java/pull/1297
diff --git a/tracing-undertow/src/main/java/com/palantir/tracing/undertow/TracingAttachments.java b/tracing-undertow/src/main/java/com/palantir/tracing/undertow/TracingAttachments.java
index 1968b0bca..348365261 100644
--- a/tracing-undertow/src/main/java/com/palantir/tracing/undertow/TracingAttachments.java
+++ b/tracing-undertow/src/main/java/com/palantir/tracing/undertow/TracingAttachments.java
@@ -38,6 +38,15 @@ public final class TracingAttachments {
      */
     static final AttachmentKey<DetachedSpan> REQUEST_SPAN = AttachmentKey.create(DetachedSpan.class);
 
+    /**
+     * 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.
+     */
+    @SuppressWarnings("unchecked")
+    public static final AttachmentKey<Detached> REQUEST_DETACHED_TRACE =
+            (AttachmentKey<Detached>) (AttachmentKey<?>) REQUEST_SPAN;
+
     /**
      * 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
@@ -45,7 +54,7 @@ public final class TracingAttachments {
      */
     @Nullable
     public static Detached requestTrace(HttpServerExchange exchange) {
-        return exchange.getAttachment(REQUEST_SPAN);
+        return exchange.getAttachment(REQUEST_DETACHED_TRACE);
     }
 
     private TracingAttachments() {}
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
index 3011399d1..663db1e5d 100644
--- a/tracing-undertow/src/test/java/com/palantir/tracing/undertow/TracedOperationHandlerTest.java
+++ b/tracing-undertow/src/test/java/com/palantir/tracing/undertow/TracedOperationHandlerTest.java
@@ -201,6 +201,8 @@ public void whenSamplingDecisionHasNotBeenMade_callsSampler() throws Exception {
     public void setsDetachedTrace() throws Exception {
         handler.handleRequest(exchange);
         assertThat(TracingAttachments.requestTrace(exchange)).isNotNull();
+        assertThat(exchange.getAttachment(TracingAttachments.REQUEST_DETACHED_TRACE))
+                .isNotNull();
     }
 
     @Test