Provide ability to add HTTP server response headers, with Tomcat 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) {
-    serverSpan(trace, index, traceID, parentID, method,
+                  ServerEndpoint endpoint = SUCCESS,
+                  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 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 static void onEnter(
     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 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 static void onEnter(
     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.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 @@ private static void installBytebuddyAgent(
     ResettableClassFileTransformer resettableClassFileTransformer = agentBuilder.installOn(inst);
     ClassFileTransformerHolder.setClassFileTransformer(resettableClassFileTransformer);
 
+    addHttpServerResponseCustomizers(extensionClassLoader);
+
     runAfterAgentListeners(agentListeners, autoConfiguredSdk);
   }
 
@@ -234,6 +239,23 @@ private static AgentBuilder configureIgnoredTypes(
             });
   }
 
+  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,16 +449,21 @@ 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)
     if (parentID == null) {
       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) {