Provide ability to add HTTP server response headers, with Tomcat implementation (#7990)

This allows custom distributions of the agent to register
`HttpServerResponseCustomizer` implementations. When a supported HTTP
server instrumentation starts processing a response, the `onStart`
method of all registered implementations will be invoked with the
`Context` of the SERVER span, an instrumentation-specific response
object and `HttpServerResponseMutator` instance that allows appending
headers to that response.

The intent of this is to allow custom distributions to set a header
containing span context information, such as the trace and span IDs. As
such, the initial implementation only allows appending response headers
and nothing else.

The `HttpServerResponseCustomizer` and related classes are currently in
a subpackage of the `io.opentelemetry.javaagent.bootstrap` package in
`javaagent-extension-api`. This makes them get loaded in the bootstrap
classloader, thus directly accessible from instrumentations. I am not
aware if there is an elegant way to put it in the agent classloader
instead, yet have the same instance accessible from both
`AgentInstaller` and instrumentations.

This also includes Tomcat-specific implementation in order to be able to
demonstrate that it works, and add automated testing of this to
HttpServerTest including one implementation.

instrumentation/jaxrs/jaxrs-common/testing/src/main/groovy/AbstractJaxRsHttpServerTest.groovy

@@ -226,8 +226,9 @@ abstract class AbstractJaxRsHttpServerTest<S> extends HttpServerTest<S> implemen
                   String parentID = null,
                   String method = "GET",
                   Long responseContentLength = null,
-                  ServerEndpoint endpoint = SUCCESS) {
+                  ServerEndpoint endpoint = SUCCESS,
-    serverSpan(trace, index, traceID, parentID, method,
+                  String spanID = null) {
+    serverSpan(trace, index, traceID, parentID, spanID, method,
       endpoint == PATH_PARAM ? getContextPath() + "/path/{id}/param" : endpoint.resolvePath(address).path,
       endpoint.resolve(address),
       endpoint.status,
@@ -239,7 +240,7 @@ abstract class AbstractJaxRsHttpServerTest<S> extends HttpServerTest<S> implemen
                        String url,
                        int statusCode) {
     def rawUrl = URI.create(url).toURL()
-    serverSpan(trace, index, null, null, "GET",
+    serverSpan(trace, index, null, null, null, "GET",
       rawUrl.path,
       rawUrl.toURI(),
       statusCode,
@@ -250,6 +251,7 @@ abstract class AbstractJaxRsHttpServerTest<S> extends HttpServerTest<S> implemen
                   int index,
                   String traceID,
                   String parentID,
+                  String spanID,
                   String method,
                   String path,
                   URI fullUrl,
@@ -261,12 +263,17 @@ abstract class AbstractJaxRsHttpServerTest<S> extends HttpServerTest<S> implemen
       if (statusCode >= 500) {
         status ERROR
       }
-      if (parentID != null) {
+      if (traceID != null) {
         traceId traceID
+      }
+      if (parentID != null) {
         parentSpanId parentID
       } else {
         hasNoParent()
       }
+      if (spanID != null) {
+        spanId spanID
+      }
       attributes {
         "$SemanticAttributes.NET_HOST_NAME" fullUrl.host
         "$SemanticAttributes.NET_HOST_PORT" fullUrl.port

instrumentation/tomcat/tomcat-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/tomcat/v10_0/Tomcat10ResponseMutator.java

@@ -0,0 +1,20 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.tomcat.v10_0;
+
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
+import org.apache.coyote.Response;
+
+public enum Tomcat10ResponseMutator implements HttpServerResponseMutator<Response> {
+  INSTANCE;
+
+  Tomcat10ResponseMutator() {}
+
+  @Override
+  public void appendHeader(Response response, String name, String value) {
+    response.addHeader(name, value);
+  }
+}

instrumentation/tomcat/tomcat-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/tomcat/v10_0/Tomcat10ServerHandlerAdvice.java

@@ -10,6 +10,7 @@ import static io.opentelemetry.javaagent.instrumentation.tomcat.v10_0.Tomcat10Si
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizerHolder;
 import net.bytebuddy.asm.Advice;
 import org.apache.coyote.Request;
 import org.apache.coyote.Response;
@@ -32,6 +33,9 @@ public class Tomcat10ServerHandlerAdvice {
     context = helper().start(parentContext, request);
 
     scope = context.makeCurrent();
+
+    HttpServerResponseCustomizerHolder.getCustomizer()
+        .customize(context, response, Tomcat10ResponseMutator.INSTANCE);
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

instrumentation/tomcat/tomcat-10.0/javaagent/src/test/groovy/io/opentelemetry/javaagent/instrumentation/tomcat/v10_0/TomcatHandlerTest.groovy

@@ -33,6 +33,11 @@ class TomcatHandlerTest extends HttpServerTest<Tomcat> implements AgentTestTrait
     return "/app"
   }
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    true
+  }
+
   @Override
   boolean testCapturedRequestParameters() {
     true

instrumentation/tomcat/tomcat-7.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/tomcat/v7_0/Tomcat7ResponseMutator.java

@@ -0,0 +1,20 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.tomcat.v7_0;
+
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
+import org.apache.coyote.Response;
+
+public class Tomcat7ResponseMutator implements HttpServerResponseMutator<Response> {
+  public static final Tomcat7ResponseMutator INSTANCE = new Tomcat7ResponseMutator();
+
+  private Tomcat7ResponseMutator() {}
+
+  @Override
+  public void appendHeader(Response response, String name, String value) {
+    response.addHeader(name, value);
+  }
+}

instrumentation/tomcat/tomcat-7.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/tomcat/v7_0/Tomcat7ServerHandlerAdvice.java

@@ -10,6 +10,7 @@ import static io.opentelemetry.javaagent.instrumentation.tomcat.v7_0.Tomcat7Sing
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizerHolder;
 import net.bytebuddy.asm.Advice;
 import org.apache.coyote.Request;
 import org.apache.coyote.Response;
@@ -32,6 +33,9 @@ public class Tomcat7ServerHandlerAdvice {
     context = helper().start(parentContext, request);
 
     scope = context.makeCurrent();
+
+    HttpServerResponseCustomizerHolder.getCustomizer()
+        .customize(context, response, Tomcat7ResponseMutator.INSTANCE);
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

instrumentation/tomcat/tomcat-7.0/javaagent/src/test/groovy/io/opentelemetry/javaagent/instrumentation/tomcat/v7_0/TomcatHandlerTest.groovy

@@ -33,6 +33,11 @@ class TomcatHandlerTest extends HttpServerTest<Tomcat> implements AgentTestTrait
     return "/app"
   }
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    true
+  }
+
   @Override
   boolean testCapturedRequestParameters() {
     true

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/bootstrap/http/HttpServerResponseCustomizer.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.bootstrap.http;
+
+import io.opentelemetry.context.Context;
+
+/**
+ * {@link HttpServerResponseCustomizer} can be used to execute code after an HTTP server response is
+ * created for the purpose of mutating the response in some way, such as appending headers, that may
+ * depend on the context of the SERVER span.
+ *
+ * <p>This is a service provider interface that requires implementations to be registered in a
+ * provider-configuration file stored in the {@code META-INF/services} resource directory.
+ */
+public interface HttpServerResponseCustomizer {
+  /**
+   * Called for each HTTP server response with its SERVER span context provided. This is called at a
+   * time when response headers can already be set, but the response is not yet committed, which is
+   * typically at the start of request handling.
+   *
+   * @param serverContext Context of a SERVER span {@link io.opentelemetry.api.trace.SpanContext}
+   * @param response Response object specific to the library being instrumented
+   * @param responseMutator Mutator through which the provided response object can be mutated
+   */
+  <RESPONSE> void customize(
+      Context serverContext,
+      RESPONSE response,
+      HttpServerResponseMutator<RESPONSE> responseMutator);
+}

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/bootstrap/http/HttpServerResponseCustomizerHolder.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.bootstrap.http;
+
+import io.opentelemetry.context.Context;
+
+/**
+ * Holds the currently active response customizer. This is set during agent initialization to an
+ * instance that calls each {@link HttpServerResponseCustomizer} found in the agent classpath. It is
+ * intended to be used directly from HTTP server library instrumentations, which is why this package
+ * is inside the bootstrap package that gets loaded in the bootstrap classloader.
+ */
+public final class HttpServerResponseCustomizerHolder {
+  private static volatile HttpServerResponseCustomizer responseCustomizer = new NoOpCustomizer();
+
+  public static void setCustomizer(HttpServerResponseCustomizer customizer) {
+    HttpServerResponseCustomizerHolder.responseCustomizer = customizer;
+  }
+
+  public static HttpServerResponseCustomizer getCustomizer() {
+    return responseCustomizer;
+  }
+
+  private HttpServerResponseCustomizerHolder() {}
+
+  private static class NoOpCustomizer implements HttpServerResponseCustomizer {
+
+    @Override
+    public <T> void customize(
+        Context serverContext, T response, HttpServerResponseMutator<T> responseMutator) {}
+  }
+}

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/bootstrap/http/HttpServerResponseMutator.java

@@ -0,0 +1,11 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.bootstrap.http;
+
+/** Provides the ability to mutate an instrumentation library specific response. */
+public interface HttpServerResponseMutator<RESPONSE> {
+  void appendHeader(RESPONSE response, String name, String value);
+}

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentInstaller.java

@@ -22,6 +22,9 @@ import io.opentelemetry.javaagent.bootstrap.BootstrapPackagePrefixesHolder;
 import io.opentelemetry.javaagent.bootstrap.ClassFileTransformerHolder;
 import io.opentelemetry.javaagent.bootstrap.DefineClassHelper;
 import io.opentelemetry.javaagent.bootstrap.InstrumentedTaskClasses;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizer;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizerHolder;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
 import io.opentelemetry.javaagent.bootstrap.internal.InstrumentationConfig;
 import io.opentelemetry.javaagent.extension.AgentListener;
 import io.opentelemetry.javaagent.extension.ignore.IgnoredTypesConfigurer;
@@ -182,6 +185,8 @@ public class AgentInstaller {
     ResettableClassFileTransformer resettableClassFileTransformer = agentBuilder.installOn(inst);
     ClassFileTransformerHolder.setClassFileTransformer(resettableClassFileTransformer);
 
+    addHttpServerResponseCustomizers(extensionClassLoader);
+
     runAfterAgentListeners(agentListeners, autoConfiguredSdk);
   }
 
@@ -234,6 +239,23 @@ public class AgentInstaller {
             });
   }
 
+  private static void addHttpServerResponseCustomizers(ClassLoader extensionClassLoader) {
+    List<HttpServerResponseCustomizer> customizers =
+        load(HttpServerResponseCustomizer.class, extensionClassLoader);
+
+    HttpServerResponseCustomizerHolder.setCustomizer(
+        new HttpServerResponseCustomizer() {
+          @Override
+          public <T> void customize(
+              Context serverContext, T response, HttpServerResponseMutator<T> responseMutator) {
+
+            for (HttpServerResponseCustomizer modifier : customizers) {
+              modifier.customize(serverContext, response, responseMutator);
+            }
+          }
+        });
+  }
+
   private static void runAfterAgentListeners(
       Iterable<AgentListener> agentListeners, AutoConfiguredOpenTelemetrySdk autoConfiguredSdk) {
     // java.util.logging.LogManager maintains a final static LogManager, which is created during

testing-common/src/main/groovy/io/opentelemetry/instrumentation/test/asserts/SpanAssert.groovy

@@ -102,6 +102,11 @@ class SpanAssert {
     checked.parentId = true
   }
 
+  def spanId(String expected) {
+    assert span.spanId == expected
+    checked.spanId = true
+  }
+
   def traceId(String expected) {
     assert span.traceId == expected
     checked.traceId = true

testing-common/src/main/groovy/io/opentelemetry/instrumentation/test/base/HttpServerTest.groovy

@@ -100,6 +100,10 @@ abstract class HttpServerTest<SERVER> extends InstrumentationSpecification imple
     false
   }
 
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    false
+  }
+
   String sockPeerAddr(ServerEndpoint endpoint) {
     "127.0.0.1"
   }
@@ -199,6 +203,9 @@ abstract class HttpServerTest<SERVER> extends InstrumentationSpecification imple
       options.hasExceptionOnServerSpan = { endpoint ->
         HttpServerTest.this.hasExceptionOnServerSpan(endpoint)
       }
+      options.hasResponseCustomizer = { endpoint ->
+        HttpServerTest.this.hasResponseCustomizer(endpoint)
+      }
       options.sockPeerAddr = { endpoint ->
         HttpServerTest.this.sockPeerAddr(endpoint)
       }
@@ -221,10 +228,11 @@ abstract class HttpServerTest<SERVER> extends InstrumentationSpecification imple
       int size,
       String traceId,
       String parentId,
+      String spanId,
       String method,
       ServerEndpoint endpoint,
       AggregatedHttpResponse response) {
-      HttpServerTest.this.assertTheTraces(size, traceId, parentId, method, endpoint, response)
+      HttpServerTest.this.assertTheTraces(size, traceId, parentId, spanId, method, endpoint, response)
     }
 
     @Override
@@ -342,7 +350,7 @@ abstract class HttpServerTest<SERVER> extends InstrumentationSpecification imple
 
   //FIXME: add tests for POST with large/chunked data
 
-  void assertTheTraces(int size, String traceID = null, String parentID = null, String method = "GET", ServerEndpoint endpoint = SUCCESS, AggregatedHttpResponse response = null) {
+  void assertTheTraces(int size, String traceID = null, String parentID = null, String spanID = null, String method = "GET", ServerEndpoint endpoint = SUCCESS, AggregatedHttpResponse response = null) {
     def spanCount = 1 // server span
     if (hasResponseSpan(endpoint)) {
       spanCount++
@@ -368,7 +376,7 @@ abstract class HttpServerTest<SERVER> extends InstrumentationSpecification imple
               assert it.span(0).endEpochNanos - it.span(index).endEpochNanos >= 0
             }
           }
-          serverSpan(it, spanIndex++, traceID, parentID, method, response?.content()?.length(), endpoint)
+          serverSpan(it, spanIndex++, traceID, parentID, method, response?.content()?.length(), endpoint, spanID)
           if (hasHandlerSpan(endpoint)) {
             handlerSpan(it, spanIndex++, span(0), method, endpoint)
           }
@@ -441,7 +449,7 @@ abstract class HttpServerTest<SERVER> extends InstrumentationSpecification imple
     }
   }
 
-  void serverSpan(TraceAssert trace, int index, String traceID = null, String parentID = null, String method = "GET", Long responseContentLength = null, ServerEndpoint endpoint = SUCCESS) {
+  void serverSpan(TraceAssert trace, int index, String traceID = null, String parentID = null, String method = "GET", Long responseContentLength = null, ServerEndpoint endpoint = SUCCESS, String spanID = null) {
     trace.assertedIndexes.add(index)
     def spanData = trace.span(index)
     def assertion = junitTest.assertServerSpan(OpenTelemetryAssertions.assertThat(spanData), method, endpoint)
@@ -449,8 +457,13 @@ abstract class HttpServerTest<SERVER> extends InstrumentationSpecification imple
       assertion.hasParentSpanId(SpanId.invalid)
     } else {
       assertion.hasParentSpanId(parentID)
+    }
+    if (traceID != null) {
       assertion.hasTraceId(traceID)
     }
+    if (spanID != null) {
+      assertion.hasSpanId(spanID)
+    }
   }
 
   void indexedServerSpan(TraceAssert trace, int index, Object parent, int requestId) {

testing-common/src/main/java/io/opentelemetry/instrumentation/testing/junit/http/AbstractHttpServerTest.java

@@ -109,9 +109,10 @@ public abstract class AbstractHttpServerTest<SERVER> extends AbstractHttpServerU
     for (AggregatedHttpResponse response : responses) {
       assertThat(response.status().code()).isEqualTo(SUCCESS.getStatus());
       assertThat(response.contentUtf8()).isEqualTo(SUCCESS.getBody());
+      assertResponseHasCustomizedHeaders(response, SUCCESS, null);
     }
 
-    assertTheTraces(count, null, null, method, SUCCESS, responses.get(0));
+    assertTheTraces(count, null, null, null, method, SUCCESS, responses.get(0));
   }
 
   @Test
@@ -131,7 +132,8 @@ public abstract class AbstractHttpServerTest<SERVER> extends AbstractHttpServerU
     assertThat(response.status().code()).isEqualTo(SUCCESS.getStatus());
     assertThat(response.contentUtf8()).isEqualTo(SUCCESS.getBody());
 
-    assertTheTraces(1, traceId, parentId, "GET", SUCCESS, response);
+    String spanId = assertResponseHasCustomizedHeaders(response, SUCCESS, traceId);
+    assertTheTraces(1, traceId, parentId, spanId, "GET", SUCCESS, response);
   }
 
   @Test
@@ -149,7 +151,8 @@ public abstract class AbstractHttpServerTest<SERVER> extends AbstractHttpServerU
     assertThat(response.status().code()).isEqualTo(SUCCESS.getStatus());
     assertThat(response.contentUtf8()).isEqualTo(SUCCESS.getBody());
 
-    assertTheTraces(1, traceId, parentId, "GET", SUCCESS, response);
+    String spanId = assertResponseHasCustomizedHeaders(response, SUCCESS, traceId);
+    assertTheTraces(1, traceId, parentId, spanId, "GET", SUCCESS, response);
   }
 
   @ParameterizedTest
@@ -164,7 +167,8 @@ public abstract class AbstractHttpServerTest<SERVER> extends AbstractHttpServerU
     assertThat(response.status().code()).isEqualTo(endpoint.getStatus());
     assertThat(response.contentUtf8()).isEqualTo(endpoint.getBody());
 
-    assertTheTraces(1, null, null, method, endpoint, response);
+    String spanId = assertResponseHasCustomizedHeaders(response, endpoint, null);
+    assertTheTraces(1, null, null, spanId, method, endpoint, response);
   }
 
   @Test
@@ -183,7 +187,8 @@ public abstract class AbstractHttpServerTest<SERVER> extends AbstractHttpServerU
                 assertThat(new URI(location).normalize().toString())
                     .isEqualTo(address.resolve(REDIRECT.getBody()).toString()));
 
-    assertTheTraces(1, null, null, method, REDIRECT, response);
+    String spanId = assertResponseHasCustomizedHeaders(response, REDIRECT, null);
+    assertTheTraces(1, null, null, spanId, method, REDIRECT, response);
   }
 
   @Test
@@ -199,7 +204,8 @@ public abstract class AbstractHttpServerTest<SERVER> extends AbstractHttpServerU
       assertThat(response.contentUtf8()).isEqualTo(ERROR.getBody());
     }
 
-    assertTheTraces(1, null, null, method, ERROR, response);
+    String spanId = assertResponseHasCustomizedHeaders(response, ERROR, null);
+    assertTheTraces(1, null, null, spanId, method, ERROR, response);
   }
 
   @Test
@@ -216,7 +222,8 @@ public abstract class AbstractHttpServerTest<SERVER> extends AbstractHttpServerU
 
       assertThat(response.status().code()).isEqualTo(EXCEPTION.getStatus());
 
-      assertTheTraces(1, null, null, method, EXCEPTION, response);
+      String spanId = assertResponseHasCustomizedHeaders(response, EXCEPTION, null);
+      assertTheTraces(1, null, null, spanId, method, EXCEPTION, response);
     } finally {
       Awaitility.reset();
     }
@@ -232,7 +239,8 @@ public abstract class AbstractHttpServerTest<SERVER> extends AbstractHttpServerU
 
     assertThat(response.status().code()).isEqualTo(NOT_FOUND.getStatus());
 
-    assertTheTraces(1, null, null, method, NOT_FOUND, response);
+    String spanId = assertResponseHasCustomizedHeaders(response, NOT_FOUND, null);
+    assertTheTraces(1, null, null, spanId, method, NOT_FOUND, response);
   }
 
   @Test
@@ -246,7 +254,8 @@ public abstract class AbstractHttpServerTest<SERVER> extends AbstractHttpServerU
     assertThat(response.status().code()).isEqualTo(PATH_PARAM.getStatus());
     assertThat(response.contentUtf8()).isEqualTo(PATH_PARAM.getBody());
 
-    assertTheTraces(1, null, null, method, PATH_PARAM, response);
+    String spanId = assertResponseHasCustomizedHeaders(response, PATH_PARAM, null);
+    assertTheTraces(1, null, null, spanId, method, PATH_PARAM, response);
   }
 
   @Test
@@ -264,7 +273,8 @@ public abstract class AbstractHttpServerTest<SERVER> extends AbstractHttpServerU
     assertThat(response.contentUtf8()).isEqualTo(CAPTURE_HEADERS.getBody());
     assertThat(response.headers().get("X-Test-Response")).isEqualTo("test");
 
-    assertTheTraces(1, null, null, "GET", CAPTURE_HEADERS, response);
+    String spanId = assertResponseHasCustomizedHeaders(response, CAPTURE_HEADERS, null);
+    assertTheTraces(1, null, null, spanId, "GET", CAPTURE_HEADERS, response);
   }
 
   @Test
@@ -283,7 +293,8 @@ public abstract class AbstractHttpServerTest<SERVER> extends AbstractHttpServerU
     assertThat(response.status().code()).isEqualTo(CAPTURE_PARAMETERS.getStatus());
     assertThat(response.contentUtf8()).isEqualTo(CAPTURE_PARAMETERS.getBody());
 
-    assertTheTraces(1, null, null, "POST", CAPTURE_PARAMETERS, response);
+    String spanId = assertResponseHasCustomizedHeaders(response, CAPTURE_PARAMETERS, null);
+    assertTheTraces(1, null, null, spanId, "POST", CAPTURE_PARAMETERS, response);
   }
 
   /**
@@ -373,12 +384,32 @@ public abstract class AbstractHttpServerTest<SERVER> extends AbstractHttpServerU
     testing.waitAndAssertTraces(assertions);
   }
 
+  protected String assertResponseHasCustomizedHeaders(
+      AggregatedHttpResponse response, ServerEndpoint endpoint, String expectedTraceId) {
+    if (!options.hasResponseCustomizer.test(endpoint)) {
+      return null;
+    }
+
+    String responseHeaderTraceId = response.headers().get("x-test-traceid");
+    String responseHeaderSpanId = response.headers().get("x-test-spanid");
+
+    if (expectedTraceId != null) {
+      assertThat(responseHeaderTraceId).matches(expectedTraceId);
+    } else {
+      assertThat(responseHeaderTraceId).isNotNull();
+    }
+
+    assertThat(responseHeaderSpanId).isNotNull();
+    return responseHeaderSpanId;
+  }
+
   // NOTE: this method does not currently implement asserting all the span types that groovy
   // HttpServerTest does
   protected void assertTheTraces(
       int size,
       String traceId,
       String parentId,
+      String spanId,
       String method,
       ServerEndpoint endpoint,
       AggregatedHttpResponse response) {
@@ -390,8 +421,14 @@ public abstract class AbstractHttpServerTest<SERVER> extends AbstractHttpServerU
             spanAssertions.add(
                 span -> {
                   assertServerSpan(span, method, endpoint);
+                  if (traceId != null) {
+                    span.hasTraceId(traceId);
+                  }
+                  if (spanId != null) {
+                    span.hasSpanId(spanId);
+                  }
                   if (parentId != null) {
-                    span.hasTraceId(traceId).hasParentSpanId(parentId);
+                    span.hasParentSpanId(parentId);
                   } else {
                     span.hasNoParent();
                   }

testing-common/src/main/java/io/opentelemetry/instrumentation/testing/junit/http/HttpServerTestOptions.java

@@ -39,6 +39,7 @@ public final class HttpServerTestOptions {
   Predicate<ServerEndpoint> hasHandlerSpan = unused -> false;
   Predicate<ServerEndpoint> hasResponseSpan = unused -> false;
   Predicate<ServerEndpoint> hasErrorPageSpans = unused -> false;
+  Predicate<ServerEndpoint> hasResponseCustomizer = unused -> false;
 
   Predicate<ServerEndpoint> hasExceptionOnServerSpan = endpoint -> !hasHandlerSpan.test(endpoint);
 
@@ -117,6 +118,13 @@ public final class HttpServerTestOptions {
     return this;
   }
 
+  @CanIgnoreReturnValue
+  public HttpServerTestOptions setHasResponseCustomizer(
+      Predicate<ServerEndpoint> hasResponseCustomizer) {
+    this.hasResponseCustomizer = hasResponseCustomizer;
+    return this;
+  }
+
   @CanIgnoreReturnValue
   public HttpServerTestOptions setTestRedirect(boolean testRedirect) {
     this.testRedirect = testRedirect;

testing/agent-exporter/src/main/java/io/opentelemetry/javaagent/testing/http/TestAgentHttpResponseCustomizer.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.testing.http;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizer;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
+
+@AutoService(HttpServerResponseCustomizer.class)
+public class TestAgentHttpResponseCustomizer implements HttpServerResponseCustomizer {
+
+  @Override
+  public <T> void customize(
+      Context serverContext, T response, HttpServerResponseMutator<T> responseMutator) {
+
+    SpanContext spanContext = Span.fromContext(serverContext).getSpanContext();
+    String traceId = spanContext.getTraceId();
+    String spanId = spanContext.getSpanId();
+
+    responseMutator.appendHeader(response, "X-Test-TraceId", traceId);
+    responseMutator.appendHeader(response, "X-Test-SpanId", spanId);
+  }
+}