Add HttpServerResponseCustomizer support for Servlet and Jetty

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

@@ -219,6 +219,11 @@ abstract class AbstractJaxRsHttpServerTest<S> extends HttpServerTest<S> implemen
     Boolean.getBoolean("testLatestDeps")
   }
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    true
+  }
+
   @Override
   void serverSpan(TraceAssert trace,
                   int index,

instrumentation/jetty/jetty-11.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jetty/v11_0/Jetty11HandlerAdvice.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 io.opentelemetry.javaagent.instrumentation.servlet.ServletRequestContext;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
@@ -45,6 +46,9 @@ public static void onEnter(
 
     // Must be set here since Jetty handlers can use startAsync outside of servlet scope.
     helper().setAsyncListenerResponse(request, response);
+
+    HttpServerResponseCustomizerHolder.getCustomizer()
+        .customize(context, response, Jetty11ResponseMutator.INSTANCE);
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

instrumentation/jetty/jetty-11.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jetty/v11_0/Jetty11ResponseMutator.java

@@ -0,0 +1,18 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.jetty.v11_0;
+
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
+import jakarta.servlet.http.HttpServletResponse;
+
+public enum Jetty11ResponseMutator implements HttpServerResponseMutator<HttpServletResponse> {
+  INSTANCE;
+
+  @Override
+  public void appendHeader(HttpServletResponse response, String name, String value) {
+    response.addHeader(name, value);
+  }
+}

instrumentation/jetty/jetty-11.0/javaagent/src/test/Java/io/opentelemetry/javaagent/instrumentation/jetty/v11_0/JettyHandlerTest.java → instrumentation/jetty/jetty-11.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/jetty/v11_0/JettyHandlerTest.java

@@ -88,6 +88,7 @@ protected void configure(HttpServerTestOptions options) {
                 DEFAULT_HTTP_ATTRIBUTES, Collections.singleton(SemanticAttributes.HTTP_ROUTE)));
     options.setHasResponseSpan(endpoint -> endpoint == REDIRECT || endpoint == ERROR);
     options.setExpectedException(new IllegalStateException(EXCEPTION.getBody()));
+    options.setHasResponseCustomizer(endpoint -> true);
   }
 
   @Override

instrumentation/jetty/jetty-8.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jetty/v8_0/Jetty8HandlerAdvice.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 io.opentelemetry.javaagent.instrumentation.servlet.ServletRequestContext;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
@@ -45,6 +46,9 @@ public static void onEnter(
 
     // Must be set here since Jetty handlers can use startAsync outside of servlet scope.
     helper().setAsyncListenerResponse(request, response);
+
+    HttpServerResponseCustomizerHolder.getCustomizer()
+        .customize(context, response, Jetty8ResponseMutator.INSTANCE);
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

instrumentation/jetty/jetty-8.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jetty/v8_0/Jetty8ResponseMutator.java

@@ -0,0 +1,18 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.jetty.v8_0;
+
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
+import javax.servlet.http.HttpServletResponse;
+
+public enum Jetty8ResponseMutator implements HttpServerResponseMutator<HttpServletResponse> {
+  INSTANCE;
+
+  @Override
+  public void appendHeader(HttpServletResponse response, String name, String value) {
+    response.addHeader(name, value);
+  }
+}

instrumentation/jetty/jetty-8.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/jetty/v8_0/JettyHandlerTest.java

@@ -88,6 +88,7 @@ protected void configure(HttpServerTestOptions options) {
                 DEFAULT_HTTP_ATTRIBUTES, Collections.singleton(SemanticAttributes.HTTP_ROUTE)));
     options.setHasResponseSpan(endpoint -> endpoint == REDIRECT || endpoint == ERROR);
     options.setExpectedException(new IllegalStateException(EXCEPTION.getBody()));
+    options.setHasResponseCustomizer(endpoint -> endpoint != EXCEPTION);
   }
 
   @Override

instrumentation/servlet/servlet-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v2_2/Servlet2Accessor.java

@@ -5,14 +5,16 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.v2_2;
 
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletAsyncListener;
 import io.opentelemetry.javaagent.instrumentation.servlet.javax.JavaxServletAccessor;
 import java.util.Collections;
 import java.util.List;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 
-public class Servlet2Accessor extends JavaxServletAccessor<HttpServletResponse> {
+public class Servlet2Accessor extends JavaxServletAccessor<HttpServletResponse>
+    implements HttpServerResponseMutator<HttpServletResponse> {
   public static final Servlet2Accessor INSTANCE = new Servlet2Accessor();
 
   private Servlet2Accessor() {}
@@ -55,4 +57,9 @@ public List<String> getResponseHeaderValues(
   public boolean isResponseCommitted(HttpServletResponse httpServletResponse) {
     return httpServletResponse.isCommitted();
   }
+
+  @Override
+  public void appendHeader(HttpServletResponse response, String name, String value) {
+    response.addHeader(name, value);
+  }
 }

instrumentation/servlet/servlet-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v2_2/Servlet2Advice.java

@@ -12,6 +12,7 @@
 import io.opentelemetry.instrumentation.api.util.VirtualField;
 import io.opentelemetry.javaagent.bootstrap.CallDepth;
 import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizerHolder;
 import io.opentelemetry.javaagent.bootstrap.servlet.AppServerBridge;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletRequestContext;
 import javax.servlet.ServletRequest;
@@ -64,6 +65,9 @@ public static void onEnter(
     // reset response status from previous request
     // (some servlet containers reuse response objects to reduce memory allocations)
     VirtualField.find(ServletResponse.class, Integer.class).set(response, null);
+
+    HttpServerResponseCustomizerHolder.getCustomizer()
+        .customize(context, (HttpServletResponse) response, Servlet2Accessor.INSTANCE);
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

instrumentation/servlet/servlet-2.2/javaagent/src/test/groovy/JettyServlet2Test.groovy

@@ -111,6 +111,11 @@ class JettyServlet2Test extends HttpServerTest<Server> implements AgentTestTrait
     endpoint == REDIRECT || endpoint == ERROR
   }
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    true
+  }
+
   @Override
   void responseSpan(TraceAssert trace, int index, Object parent, String method = "GET", ServerEndpoint endpoint = SUCCESS) {
     def responseMethod = endpoint == REDIRECT ? "sendRedirect" : "sendError"

instrumentation/servlet/servlet-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/Servlet3Accessor.java

@@ -5,6 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.v3_0;
 
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletAsyncListener;
 import io.opentelemetry.javaagent.instrumentation.servlet.javax.JavaxServletAccessor;
 import java.util.ArrayList;
@@ -16,7 +17,8 @@
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 
-public class Servlet3Accessor extends JavaxServletAccessor<HttpServletResponse> {
+public class Servlet3Accessor extends JavaxServletAccessor<HttpServletResponse>
+    implements HttpServerResponseMutator<HttpServletResponse> {
   public static final Servlet3Accessor INSTANCE = new Servlet3Accessor();
 
   private Servlet3Accessor() {}
@@ -70,6 +72,11 @@ public boolean isResponseCommitted(HttpServletResponse response) {
     return response.isCommitted();
   }
 
+  @Override
+  public void appendHeader(HttpServletResponse response, String name, String value) {
+    response.addHeader(name, value);
+  }
+
   private static class Listener implements AsyncListener {
     private final ServletAsyncListener<HttpServletResponse> listener;
 

instrumentation/servlet/servlet-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/Servlet3Advice.java

@@ -11,6 +11,7 @@
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.bootstrap.CallDepth;
 import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizerHolder;
 import io.opentelemetry.javaagent.bootstrap.servlet.AppServerBridge;
 import io.opentelemetry.javaagent.bootstrap.servlet.MappingResolver;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletRequestContext;
@@ -73,6 +74,12 @@ && helper().needsRescoping(currentContext, attachedContext)) {
     contextToUpdate =
         helper().updateContext(contextToUpdate, httpServletRequest, mappingResolver, servlet);
     scope = contextToUpdate.makeCurrent();
+
+    if (context != null) {
+      // Only trigger response customizer once, so only if server span was created here
+      HttpServerResponseCustomizerHolder.getCustomizer()
+          .customize(contextToUpdate, (HttpServletResponse) response, Servlet3Accessor.INSTANCE);
+    }
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

instrumentation/servlet/servlet-3.0/javaagent/src/test/groovy/AbstractServlet3Test.groovy

@@ -79,6 +79,11 @@ abstract class AbstractServlet3Test<SERVER, CONTEXT> extends HttpServerTest<SERV
     true
   }
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    true
+  }
+
   @Override
   boolean hasResponseSpan(ServerEndpoint endpoint) {
     endpoint == REDIRECT || (endpoint == ERROR && errorEndpointUsesSendError())

instrumentation/servlet/servlet-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/Servlet5Accessor.java

@@ -5,6 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.v5_0;
 
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletAccessor;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletAsyncListener;
 import jakarta.servlet.AsyncEvent;
@@ -20,7 +21,9 @@
 import java.util.Enumeration;
 import java.util.List;
 
-public class Servlet5Accessor implements ServletAccessor<HttpServletRequest, HttpServletResponse> {
+public class Servlet5Accessor
+    implements ServletAccessor<HttpServletRequest, HttpServletResponse>,
+        HttpServerResponseMutator<HttpServletResponse> {
   public static final Servlet5Accessor INSTANCE = new Servlet5Accessor();
 
   private Servlet5Accessor() {}
@@ -172,6 +175,11 @@ public boolean isServletException(Throwable throwable) {
     return throwable instanceof ServletException;
   }
 
+  @Override
+  public void appendHeader(HttpServletResponse response, String name, String value) {
+    response.addHeader(name, value);
+  }
+
   private static class Listener implements AsyncListener {
     private final ServletAsyncListener<HttpServletResponse> listener;
 

instrumentation/servlet/servlet-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/service/JakartaServletServiceAdvice.java

@@ -11,9 +11,11 @@
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.bootstrap.CallDepth;
 import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizerHolder;
 import io.opentelemetry.javaagent.bootstrap.servlet.AppServerBridge;
 import io.opentelemetry.javaagent.bootstrap.servlet.MappingResolver;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletRequestContext;
+import io.opentelemetry.javaagent.instrumentation.servlet.v5_0.Servlet5Accessor;
 import io.opentelemetry.javaagent.instrumentation.servlet.v5_0.Servlet5Singletons;
 import jakarta.servlet.Servlet;
 import jakarta.servlet.ServletRequest;
@@ -74,6 +76,12 @@ && helper().needsRescoping(currentContext, attachedContext)) {
     contextToUpdate =
         helper().updateContext(contextToUpdate, httpServletRequest, mappingResolver, servlet);
     scope = contextToUpdate.makeCurrent();
+
+    if (context != null) {
+      // Only trigger response customizer once, so only if server span was created here
+      HttpServerResponseCustomizerHolder.getCustomizer()
+          .customize(contextToUpdate, (HttpServletResponse) response, Servlet5Accessor.INSTANCE);
+    }
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

instrumentation/servlet/servlet-5.0/javaagent/src/test/groovy/AbstractServlet5Test.groovy

@@ -78,6 +78,11 @@ abstract class AbstractServlet5Test<SERVER, CONTEXT> extends HttpServerTest<SERV
     }
   }
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    true
+  }
+
   @Override
   boolean hasResponseSpan(ServerEndpoint endpoint) {
     endpoint == REDIRECT || (endpoint == ERROR && errorEndpointUsesSendError())