Fix jax-rs/jakarta-rs span double-finish on synchronous AsyncResponse#resume()

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/ResourceMethodSpanTracker.java

@@ -0,0 +1,53 @@
+package datadog.trace.bootstrap;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import java.util.ArrayDeque;
+
+/**
+ * Tracks, per thread, which annotated resource-method invocation (identified by its span) is
+ * currently the innermost one still on the call stack -- i.e. whose exit advice has not yet run.
+ *
+ * <p>Used by the JAX-RS/Jakarta-RS instrumentation to tell a synchronous {@code
+ * AsyncResponse#resume()}/{@code cancel()} call -- one nested inside the still-running resource
+ * method that owns the response -- apart from a genuinely asynchronous one called later, from a
+ * different thread or after an intervening unrelated instrumented call. A generic "is any resource
+ * method open on this thread" counter is not enough for that: it can't tell one resource method's
+ * invocation apart from another's on a shared thread pool, and it can't see past a nested
+ * instrumented call (e.g. a {@code @Trace}-annotated helper) that becomes the current active span
+ * without popping this stack. Comparing the specific span object against the top of this stack
+ * answers the exact question that matters, without either failure mode.
+ *
+ * <p>Deliberately bootstrap-loaded (like {@link CallDepthThreadLocalMap}) rather than living on a
+ * per-instrumentation helper class: helper classes are injected once per target classloader, so a
+ * container that loads the resource-method advice and the AsyncResponse advice into different
+ * classloaders (e.g. a modular server where the JAX-RS runtime and the deployed application are in
+ * separate classloaders) would otherwise give each advice its own, disconnected copy of this state.
+ */
+public final class ResourceMethodSpanTracker {
+
+  private static final ThreadLocal<ArrayDeque<AgentSpan>> STACK = new ThreadLocal<>();
+
+  private ResourceMethodSpanTracker() {}
+
+  public static void enter(final AgentSpan span) {
+    ArrayDeque<AgentSpan> stack = STACK.get();
+    if (stack == null) {
+      stack = new ArrayDeque<>(4);
+      STACK.set(stack);
+    }
+    stack.push(span);
+  }
+
+  public static void exit() {
+    final ArrayDeque<AgentSpan> stack = STACK.get();
+    if (stack != null) {
+      stack.pop();
+    }
+  }
+
+  /** True if {@code span} is the innermost still-open resource-method invocation on this thread. */
+  public static boolean isInnermost(final AgentSpan span) {
+    final ArrayDeque<AgentSpan> stack = STACK.get();
+    return stack != null && stack.peek() == span;
+  }
+}

dd-java-agent/instrumentation/cxf-2.1/src/test/groovy/CxfContextPropagationTest.groovy

@@ -22,12 +22,20 @@ class CxfContextPropagationTest extends InstrumentationSpecification {
   @Override
   void setupSpec() {
     JAXRSServerFactoryBean sf = new JAXRSServerFactoryBean()
-    sf.setResourceClasses(TestResource)
+    sf.setResourceClasses(TestResource, AsyncResumeResource, TrueAsyncResumeResource, AsyncCancelResource, NestedResumeResource)
     List<Object> providers = [new TestExceptionMapper()]
     sf.setProviders(providers)
 
     sf.setResourceProvider(TestResource,
       new SingletonResourceProvider(new TestResource(), true))
+    sf.setResourceProvider(AsyncResumeResource,
+      new SingletonResourceProvider(new AsyncResumeResource(), true))
+    sf.setResourceProvider(TrueAsyncResumeResource,
+      new SingletonResourceProvider(new TrueAsyncResumeResource(), true))
+    sf.setResourceProvider(AsyncCancelResource,
+      new SingletonResourceProvider(new AsyncCancelResource(), true))
+    sf.setResourceProvider(NestedResumeResource,
+      new SingletonResourceProvider(new NestedResumeResource(), true))
     sf.setAddress("http://localhost:0")
 
     server = sf.create()
@@ -40,6 +48,15 @@ class CxfContextPropagationTest extends InstrumentationSpecification {
     server?.stop()
   }
 
+  @Override
+  protected boolean enabledFinishTimingChecks() {
+    // Regression guard for https://github.com/DataDog/dd-trace-java/issues/12597:
+    // fails the test with the exact "finished more than once" stack traces if the
+    // jax-rs.request span is ever finished twice (e.g. once from AsyncResponse#resume()
+    // and again from the resource method's own exit advice).
+    return true
+  }
+
   def "should propagate context on async request resume"() {
     setup:
     def client = OkHttpUtils.client()
@@ -97,4 +114,251 @@ class CxfContextPropagationTest extends InstrumentationSpecification {
       }
     }
   }
+
+  def "resume() called synchronously from within the resource method finishes the span only once"() {
+    // Regression test for https://github.com/DataDog/dd-trace-java/issues/12597: when
+    // AsyncResponse#resume() is called synchronously (not truly suspended, the resource
+    // method keeps running), the resource method's own jax-rs.request scope is still on
+    // top of the scope stack. Before the fix, JakartaRsAsyncResponseInstrumentation /
+    // JaxRsAsyncResponseInstrumentation would eagerly finish the span there, so any work
+    // done afterwards (here: doWorkAfterResume()) would be attributed as a child of an
+    // already-finished span, and the resource method's own exit advice would finish the
+    // same span a second time (caught by enabledFinishTimingChecks()).
+    setup:
+    def client = OkHttpUtils.client()
+    when:
+    def response = client.newCall(new Request.Builder()
+      .url("http://localhost:$port/asyncresume")
+      .get().build()).execute()
+    then:
+    assert response.code() == 200
+    assert response.body().string() == "OK"
+
+    assertTraces(1) {
+      trace(3) {
+        sortSpansByStart()
+        span {
+          operationName "servlet.request"
+          resourceName "GET /asyncresume"
+          spanType DDSpanTypes.HTTP_SERVER
+          errored false
+          parent()
+          tags {
+            "$Tags.COMPONENT" "jax-rs"
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_SERVER
+            "$Tags.PEER_HOST_IPV4" "127.0.0.1"
+            "$Tags.PEER_PORT" Integer
+            "$Tags.HTTP_URL" "http://localhost:$port/asyncresume"
+            "$Tags.HTTP_HOSTNAME" "localhost"
+            "$Tags.HTTP_METHOD" "GET"
+            "$Tags.HTTP_STATUS" 200
+            "$Tags.HTTP_ROUTE" String
+            "servlet.path" { it == null || it == "/asyncresume" }
+            "$Tags.HTTP_USER_AGENT" String
+            "$Tags.HTTP_CLIENT_IP" "127.0.0.1"
+            "$Tags.NETWORK_CLIENT_IP" "127.0.0.1"
+            withCustomIntegrationName("jetty-server")
+            defaultTags()
+          }
+        }
+        span {
+          operationName "jax-rs.request"
+          resourceName "AsyncResumeResource.resumeThenWork"
+          spanType DDSpanTypes.HTTP_SERVER
+          errored false
+          childOfPrevious()
+          tags {
+            "$Tags.COMPONENT" "jax-rs-controller"
+            defaultTags()
+          }
+        }
+        // Still parented under jax-rs.request: proves that span wasn't finished (and its
+        // scope wasn't popped) by resume() itself, before the resource method returned.
+        TraceUtils.basicSpan(it, "trace.annotation", "AsyncResumeResource.doWorkAfterResume", span(1), null, ["component": "trace"])
+      }
+    }
+  }
+
+  def "cancel() called synchronously from within the resource method finishes the span only once"() {
+    // Same regression as above, but for AsyncResponseCancelAdvice: cancel() is called
+    // synchronously and the resource method keeps running afterwards.
+    setup:
+    def client = OkHttpUtils.client()
+    when:
+    def response = client.newCall(new Request.Builder()
+      .url("http://localhost:$port/asynccancel")
+      .get().build()).execute()
+    then:
+    assert response.code() == 503
+
+    assertTraces(1) {
+      trace(3) {
+        sortSpansByStart()
+        span {
+          operationName "servlet.request"
+          resourceName "GET /asynccancel"
+          spanType DDSpanTypes.HTTP_SERVER
+          errored true
+          parent()
+          tags {
+            "$Tags.COMPONENT" "jax-rs"
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_SERVER
+            "$Tags.PEER_HOST_IPV4" "127.0.0.1"
+            "$Tags.PEER_PORT" Integer
+            "$Tags.HTTP_URL" "http://localhost:$port/asynccancel"
+            "$Tags.HTTP_HOSTNAME" "localhost"
+            "$Tags.HTTP_METHOD" "GET"
+            "$Tags.HTTP_STATUS" 503
+            "$Tags.HTTP_ROUTE" String
+            "servlet.path" { it == null || it == "/asynccancel" }
+            "$Tags.HTTP_USER_AGENT" String
+            "$Tags.HTTP_CLIENT_IP" "127.0.0.1"
+            "$Tags.NETWORK_CLIENT_IP" "127.0.0.1"
+            withCustomIntegrationName("jetty-server")
+            defaultTags()
+          }
+        }
+        span {
+          operationName "jax-rs.request"
+          resourceName "AsyncCancelResource.cancelThenWork"
+          spanType DDSpanTypes.HTTP_SERVER
+          errored false
+          childOfPrevious()
+          tags {
+            "$Tags.COMPONENT" "jax-rs-controller"
+            "canceled" true
+            defaultTags()
+          }
+        }
+        // Still parented under jax-rs.request: proves the span wasn't finished (and its
+        // scope wasn't popped) by cancel() itself, before the resource method returned.
+        TraceUtils.basicSpan(it, "trace.annotation", "AsyncCancelResource.doWorkAfterCancel", span(1), null, ["component": "trace"])
+      }
+    }
+  }
+
+  def "resume() called from a genuinely different thread (textbook async pattern) is unaffected"() {
+    // Regression guard the other way: the fix must not change the standard cross-thread
+    // async pattern, where the resource method returns without resolving anything and a
+    // completely different thread calls resume() later. Here the span IS finished by the
+    // resume() advice (activeSpan() on that other thread is not this span), exactly as
+    // before the fix.
+    setup:
+    def client = OkHttpUtils.client()
+    when:
+    def response = client.newCall(new Request.Builder()
+      .url("http://localhost:$port/trueasyncresume")
+      .get().build()).execute()
+    then:
+    assert response.code() == 200
+    assert response.body().string() == "OK"
+
+    assertTraces(1) {
+      trace(3) {
+        sortSpansByStart()
+        span {
+          operationName "servlet.request"
+          resourceName "GET /trueasyncresume"
+          spanType DDSpanTypes.HTTP_SERVER
+          errored false
+          parent()
+          tags {
+            "$Tags.COMPONENT" "jax-rs"
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_SERVER
+            "$Tags.PEER_HOST_IPV4" "127.0.0.1"
+            "$Tags.PEER_PORT" Integer
+            "$Tags.HTTP_URL" "http://localhost:$port/trueasyncresume"
+            "$Tags.HTTP_HOSTNAME" "localhost"
+            "$Tags.HTTP_METHOD" "GET"
+            "$Tags.HTTP_STATUS" 200
+            "$Tags.HTTP_ROUTE" String
+            "servlet.path" { it == null || it == "/trueasyncresume" }
+            "$Tags.HTTP_USER_AGENT" String
+            "$Tags.HTTP_CLIENT_IP" "127.0.0.1"
+            "$Tags.NETWORK_CLIENT_IP" "127.0.0.1"
+            withCustomIntegrationName("jetty-server")
+            defaultTags()
+          }
+        }
+        span {
+          operationName "jax-rs.request"
+          resourceName "TrueAsyncResumeResource.suspendThenResumeFromAnotherThread"
+          spanType DDSpanTypes.HTTP_SERVER
+          errored false
+          childOfPrevious()
+          tags {
+            "$Tags.COMPONENT" "jax-rs-controller"
+            defaultTags()
+          }
+        }
+        // Runs on the background thread, before resume() -- still correctly parented
+        // under the (still-open, cross-thread-propagated) jax-rs.request span.
+        TraceUtils.basicSpan(it, "trace.annotation", "TrueAsyncResumeResource.doWorkOnBackgroundThread", span(1), null, ["component": "trace"])
+      }
+    }
+  }
+
+  def "resume() called synchronously from a nested @Trace helper finishes the span only once"() {
+    // Regression test for a gap found reviewing the fix for GH-12597: resume() is called
+    // synchronously, but from a @Trace-annotated helper method rather than directly from
+    // the resource method's own body. At that moment, the *helper's* span is the active
+    // one on this thread, not the resource method's -- checking activeSpan() against the
+    // resource method's span directly (an earlier version of this fix) would wrongly treat
+    // this as a genuinely-async resume and finish the span right there, then finish it
+    // again when the resource method itself returns (caught by enabledFinishTimingChecks()).
+    setup:
+    def client = OkHttpUtils.client()
+    when:
+    def response = client.newCall(new Request.Builder()
+      .url("http://localhost:$port/nestedresume")
+      .get().build()).execute()
+    then:
+    assert response.code() == 200
+    assert response.body().string() == "OK"
+
+    assertTraces(1) {
+      trace(3) {
+        sortSpansByStart()
+        span {
+          operationName "servlet.request"
+          resourceName "GET /nestedresume"
+          spanType DDSpanTypes.HTTP_SERVER
+          errored false
+          parent()
+          tags {
+            "$Tags.COMPONENT" "jax-rs"
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_SERVER
+            "$Tags.PEER_HOST_IPV4" "127.0.0.1"
+            "$Tags.PEER_PORT" Integer
+            "$Tags.HTTP_URL" "http://localhost:$port/nestedresume"
+            "$Tags.HTTP_HOSTNAME" "localhost"
+            "$Tags.HTTP_METHOD" "GET"
+            "$Tags.HTTP_STATUS" 200
+            "$Tags.HTTP_ROUTE" String
+            "servlet.path" { it == null || it == "/nestedresume" }
+            "$Tags.HTTP_USER_AGENT" String
+            "$Tags.HTTP_CLIENT_IP" "127.0.0.1"
+            "$Tags.NETWORK_CLIENT_IP" "127.0.0.1"
+            withCustomIntegrationName("jetty-server")
+            defaultTags()
+          }
+        }
+        span {
+          operationName "jax-rs.request"
+          resourceName "NestedResumeResource.resumeViaHelper"
+          spanType DDSpanTypes.HTTP_SERVER
+          errored false
+          childOfPrevious()
+          tags {
+            "$Tags.COMPONENT" "jax-rs-controller"
+            defaultTags()
+          }
+        }
+        // The helper that actually calls resume() -- still parented under jax-rs.request,
+        // proving the resource-method span wasn't finished/popped while the helper (and
+        // resume() inside it) was still running.
+        TraceUtils.basicSpan(it, "trace.annotation", "NestedResumeResource.resumeFromHelper", span(1), null, ["component": "trace"])
+      }
+    }
+  }
 }

dd-java-agent/instrumentation/cxf-2.1/src/test/java/AsyncCancelResource.java

@@ -0,0 +1,22 @@
+import datadog.trace.api.Trace;
+import javax.ws.rs.GET;
+import javax.ws.rs.Path;
+import javax.ws.rs.container.AsyncResponse;
+import javax.ws.rs.container.Suspended;
+
+/**
+ * Same GH-12597 pattern as {@link AsyncResumeResource}, but exercising {@code
+ * AsyncResponse#cancel()} instead of {@code resume()} -- the third advice touched by the fix
+ * ({@code AsyncResponseCancelAdvice}).
+ */
+@Path("/asynccancel")
+public class AsyncCancelResource {
+  @GET
+  public void cancelThenWork(@Suspended final AsyncResponse response) {
+    response.cancel();
+    doWorkAfterCancel();
+  }
+
+  @Trace
+  private void doWorkAfterCancel() {}
+}

dd-java-agent/instrumentation/cxf-2.1/src/test/java/AsyncResumeResource.java

@@ -0,0 +1,17 @@
+import datadog.trace.api.Trace;
+import javax.ws.rs.GET;
+import javax.ws.rs.Path;
+import javax.ws.rs.container.AsyncResponse;
+import javax.ws.rs.container.Suspended;
+
+@Path("/asyncresume")
+public class AsyncResumeResource {
+  @GET
+  public void resumeThenWork(@Suspended final AsyncResponse response) {
+    response.resume("OK");
+    doWorkAfterResume();
+  }
+
+  @Trace
+  private void doWorkAfterResume() {}
+}