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 extends WithHttpServer { } } - /** - * 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 extends WithHttpServer { 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 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) )