From e5418641b0256bbff5e6678cdc0202624f8149b6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= <raphael.vandon@datadoghq.com>
Date: Mon, 15 Sep 2025 11:29:19 +0200
Subject: [PATCH 1/2] Revert "add waiting hook in HttpServerTest to try to
 solve Pekko's flakiness (#9493)"

This reverts commit 46441980ca83b9491dff93735ef814e87d53aaad.
---
 .../trace/agent/test/base/HttpServerTest.groovy   | 11 -----------
 .../PekkoHttpServerInstrumentationTest.groovy     | 15 ---------------
 2 files changed, 26 deletions(-)

diff --git a/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy b/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy
index b207cbd5f48..9bafc945800 100644
--- a/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy
+++ b/dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/base/HttpServerTest.groovy
@@ -174,15 +174,6 @@ abstract class HttpServerTest<SERVER> extends WithHttpServer<SERVER> {
     }
   }
 
-  /**
-   * Hook method to allow subclasses to wait for request processing to complete.
-   * This is particularly useful for asynchronous HTTP server implementations
-   * where spans may still not be closed after the HTTP response is sent.
-   */
-  protected void waitForRequestToComplete() {
-    // Default implementation does nothing - override in subclasses as needed
-  }
-
   // used in blocking tests to check if the handler was skipped
   volatile boolean handlerRan
 
@@ -1204,8 +1195,6 @@ abstract class HttpServerTest<SERVER> extends WithHttpServer<SERVER> {
       TEST_DATA_STREAMS_WRITER.waitForGroups(1)
     }
 
-    waitForRequestToComplete()
-
     expect:
     response.code() == EXCEPTION.status
     if (testExceptionBody()) {
diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/baseTest/groovy/PekkoHttpServerInstrumentationTest.groovy b/dd-java-agent/instrumentation/pekko-http-1.0/src/baseTest/groovy/PekkoHttpServerInstrumentationTest.groovy
index 52fe53cb39d..664ee4f0e7f 100644
--- a/dd-java-agent/instrumentation/pekko-http-1.0/src/baseTest/groovy/PekkoHttpServerInstrumentationTest.groovy
+++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/baseTest/groovy/PekkoHttpServerInstrumentationTest.groovy
@@ -40,21 +40,6 @@ abstract class PekkoHttpServerInstrumentationTest extends HttpServerTest<PekkoHt
     true
   }
 
-  @Override
-  protected void waitForRequestToComplete() {
-    // Pekko HTTP is asynchronous, and spans may not be completed
-    // immediately after the HTTP response is sent. Wait for the trace
-    // to be written to avoid race conditions in trace assertions.
-    try {
-      // Wait up to 1 sec for the trace to be written
-      TEST_WRITER.waitForTracesMax(1, 1)
-    } catch (InterruptedException e) {
-      Thread.currentThread().interrupt()
-      // If interrupted, proceed anyway - test may still pass
-    }
-  }
-
-
   //@Ignore("https://github.com/DataDog/dd-trace-java/pull/5213")
   @Override
   boolean testBadUrl() {

From fc8d09bc7aa199d00454eae31c3c1756a7c8b56c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= <raphael.vandon@datadoghq.com>
Date: Mon, 15 Sep 2025 11:30:18 +0200
Subject: [PATCH 2/2] close span on exception handling in pekko test server

---
 .../src/baseTest/scala/PekkoHttpTestWebServer.scala         | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/baseTest/scala/PekkoHttpTestWebServer.scala b/dd-java-agent/instrumentation/pekko-http-1.0/src/baseTest/scala/PekkoHttpTestWebServer.scala
index 583aa86e4d7..b7b0c30460b 100644
--- a/dd-java-agent/instrumentation/pekko-http-1.0/src/baseTest/scala/PekkoHttpTestWebServer.scala
+++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/baseTest/scala/PekkoHttpTestWebServer.scala
@@ -136,7 +136,11 @@ object PekkoHttpTestWebServer {
   private val exceptionHandler = ExceptionHandler {
     case e: Exception =>
       val span = activeSpan()
-      TraceUtils.handleException(span, e)
+      if (span != null) {
+        // The exception handler is bypassing the normal instrumentation flow, so we need to handle things here
+        TraceUtils.handleException(span, e)
+        span.finish()
+      }
       complete(
         HttpResponse(status = EXCEPTION.getStatus, entity = e.getMessage)
       )