Add HttpServerResponseCustomizer support for Netty (#8094)

Add `HttpServerResponseCustomizer` support for Netty 3.8, 4.0 and 4.1
instrumentations and enable testing for it in their respective
`HttpServerTest` tests.

instrumentation/netty/netty-3.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/v3_8/server/HttpServerResponseTracingHandler.java

@@ -11,6 +11,7 @@
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.instrumentation.api.util.VirtualField;
 import io.opentelemetry.instrumentation.netty.common.internal.NettyErrorHolder;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizerHolder;
 import io.opentelemetry.javaagent.instrumentation.netty.v3_8.HttpRequestAndChannel;
 import org.jboss.netty.channel.Channel;
 import org.jboss.netty.channel.ChannelHandlerContext;
@@ -38,6 +39,7 @@ public void writeRequested(ChannelHandlerContext ctx, MessageEvent msg) throws E
 
     Throwable error = null;
     try (Scope ignored = context.makeCurrent()) {
+      customizeResponse(context, response);
       super.writeRequested(ctx, msg);
     } catch (Throwable t) {
       error = t;
@@ -47,4 +49,13 @@ public void writeRequested(ChannelHandlerContext ctx, MessageEvent msg) throws E
       instrumenter().end(context, request, response, error);
     }
   }
+
+  private static void customizeResponse(Context context, HttpResponse response) {
+    try {
+      HttpServerResponseCustomizerHolder.getCustomizer()
+          .customize(context, response, NettyHttpResponseMutator.INSTANCE);
+    } catch (Throwable ignore) {
+      // Ignore.
+    }
+  }
 }

instrumentation/netty/netty-3.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/v3_8/server/NettyHttpResponseMutator.java

@@ -0,0 +1,20 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.netty.v3_8.server;
+
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
+import org.jboss.netty.handler.codec.http.HttpResponse;
+
+public enum NettyHttpResponseMutator implements HttpServerResponseMutator<HttpResponse> {
+  INSTANCE;
+
+  NettyHttpResponseMutator() {}
+
+  @Override
+  public void appendHeader(HttpResponse response, String name, String value) {
+    response.headers().add(name, value);
+  }
+}

instrumentation/netty/netty-3.8/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/netty/v3_8/server/Netty38ServerTest.java

@@ -89,6 +89,7 @@ protected void configure(HttpServerTestOptions options) {
         });
 
     options.setExpectedException(new IllegalArgumentException(ServerEndpoint.EXCEPTION.getBody()));
+    options.setHasResponseCustomizer(serverEndpoint -> true);
   }
 
   private static ChannelPipeline channelPipeline() {

instrumentation/netty/netty-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/v4_0/server/HttpServerResponseTracingHandler.java

@@ -17,6 +17,7 @@
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.instrumentation.netty.common.internal.NettyErrorHolder;
 import io.opentelemetry.instrumentation.netty.v4.common.HttpRequestAndChannel;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizerHolder;
 import io.opentelemetry.javaagent.instrumentation.netty.v4_0.AttributeKeys;
 import javax.annotation.Nullable;
 
@@ -31,6 +32,7 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) {
     }
 
     try (Scope ignored = context.makeCurrent()) {
+      customizeResponse(context, (HttpResponse) msg);
       ctx.write(msg, prm);
       end(ctx.channel(), (HttpResponse) msg, null);
     } catch (Throwable throwable) {
@@ -46,4 +48,13 @@ private static void end(Channel channel, HttpResponse response, @Nullable Throwa
     error = NettyErrorHolder.getOrDefault(context, error);
     instrumenter().end(context, request, response, error);
   }
+
+  private static void customizeResponse(Context context, HttpResponse response) {
+    try {
+      HttpServerResponseCustomizerHolder.getCustomizer()
+          .customize(context, response, NettyHttpResponseMutator.INSTANCE);
+    } catch (Throwable ignore) {
+      // Ignore.
+    }
+  }
 }

instrumentation/netty/netty-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/v4_0/server/NettyHttpResponseMutator.java

@@ -0,0 +1,20 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.netty.v4_0.server;
+
+import io.netty.handler.codec.http.HttpResponse;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
+
+public enum NettyHttpResponseMutator implements HttpServerResponseMutator<HttpResponse> {
+  INSTANCE;
+
+  NettyHttpResponseMutator() {}
+
+  @Override
+  public void appendHeader(HttpResponse response, String name, String value) {
+    response.headers().add(name, value);
+  }
+}

instrumentation/netty/netty-4.0/javaagent/src/test/groovy/Netty40ServerTest.groovy

@@ -47,6 +47,11 @@ class Netty40ServerTest extends HttpServerTest<EventLoopGroup> implements AgentT
 
   static final LoggingHandler LOGGING_HANDLER = new LoggingHandler(SERVER_LOGGER.name, LogLevel.DEBUG)
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    true
+  }
+
   @Override
   EventLoopGroup startServer(int port) {
     def eventLoopGroup = new NioEventLoopGroup()

instrumentation/netty/netty-4.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/v4_1/NettyChannelPipelineInstrumentation.java

@@ -105,11 +105,17 @@ public static void addHandler(
       ChannelHandler ourHandler = null;
       // Server pipeline handlers
       if (handler instanceof HttpServerCodec) {
-        ourHandler = new HttpServerTracingHandler(NettyServerSingletons.instrumenter());
+        ourHandler =
+            new HttpServerTracingHandler(
+                NettyServerSingletons.instrumenter(),
+                NettyHttpServerResponseBeforeCommitHandler.INSTANCE);
       } else if (handler instanceof HttpRequestDecoder) {
         ourHandler = new HttpServerRequestTracingHandler(NettyServerSingletons.instrumenter());
       } else if (handler instanceof HttpResponseEncoder) {
-        ourHandler = new HttpServerResponseTracingHandler(NettyServerSingletons.instrumenter());
+        ourHandler =
+            new HttpServerResponseTracingHandler(
+                NettyServerSingletons.instrumenter(),
+                NettyHttpServerResponseBeforeCommitHandler.INSTANCE);
         // Client pipeline handlers
       } else if (handler instanceof HttpClientCodec) {
         ourHandler = new HttpClientTracingHandler(NettyClientSingletons.instrumenter());

instrumentation/netty/netty-4.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/v4_1/NettyHttpServerResponseBeforeCommitHandler.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.netty.v4_1;
+
+import io.netty.handler.codec.http.HttpResponse;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.netty.v4_1.internal.server.HttpServerResponseBeforeCommitHandler;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizerHolder;
+
+public enum NettyHttpServerResponseBeforeCommitHandler
+    implements HttpServerResponseBeforeCommitHandler {
+  INSTANCE;
+
+  @Override
+  public void handle(Context context, HttpResponse response) {
+    HttpServerResponseCustomizerHolder.getCustomizer()
+        .customize(context, response, NettyHttpServerResponseMutator.INSTANCE);
+  }
+}

instrumentation/netty/netty-4.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/v4_1/NettyHttpServerResponseMutator.java

@@ -0,0 +1,18 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.netty.v4_1;
+
+import io.netty.handler.codec.http.HttpResponse;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
+
+public enum NettyHttpServerResponseMutator implements HttpServerResponseMutator<HttpResponse> {
+  INSTANCE;
+
+  @Override
+  public void appendHeader(HttpResponse response, String name, String value) {
+    response.headers().add(name, value);
+  }
+}

instrumentation/netty/netty-4.1/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/netty/v4_1/Netty41ServerTest.java

@@ -9,13 +9,20 @@
 import io.opentelemetry.instrumentation.netty.v4_1.AbstractNetty41ServerTest;
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
 import io.opentelemetry.instrumentation.testing.junit.http.HttpServerInstrumentationExtension;
+import io.opentelemetry.instrumentation.testing.junit.http.HttpServerTestOptions;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
 public class Netty41ServerTest extends AbstractNetty41ServerTest {
 
   @RegisterExtension
   static final InstrumentationExtension testing = HttpServerInstrumentationExtension.forAgent();
 
+  @Override
+  protected void configure(HttpServerTestOptions options) {
+    super.configure(options);
+    options.setHasResponseCustomizer(serverEndpoint -> true);
+  }
+
   @Override
   protected void configurePipeline(ChannelPipeline channelPipeline) {}
 }

instrumentation/netty/netty-4.1/library/src/main/java/io/opentelemetry/instrumentation/netty/v4_1/NettyServerTelemetry.java

@@ -13,6 +13,7 @@
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.netty.v4.common.HttpRequestAndChannel;
 import io.opentelemetry.instrumentation.netty.v4_1.internal.server.HttpServerRequestTracingHandler;
+import io.opentelemetry.instrumentation.netty.v4_1.internal.server.HttpServerResponseBeforeCommitHandler;
 import io.opentelemetry.instrumentation.netty.v4_1.internal.server.HttpServerResponseTracingHandler;
 import io.opentelemetry.instrumentation.netty.v4_1.internal.server.HttpServerTracingHandler;
 
@@ -51,7 +52,8 @@ public ChannelInboundHandlerAdapter createRequestHandler() {
    * responses. Must be paired with {@link #createRequestHandler()}.
    */
   public ChannelOutboundHandlerAdapter createResponseHandler() {
-    return new HttpServerResponseTracingHandler(instrumenter);
+    return new HttpServerResponseTracingHandler(
+        instrumenter, HttpServerResponseBeforeCommitHandler.Noop.INSTANCE);
   }
 
   /**
@@ -61,6 +63,7 @@ public ChannelOutboundHandlerAdapter createResponseHandler() {
   public CombinedChannelDuplexHandler<
           ? extends ChannelInboundHandlerAdapter, ? extends ChannelOutboundHandlerAdapter>
       createCombinedHandler() {
-    return new HttpServerTracingHandler(instrumenter);
+    return new HttpServerTracingHandler(
+        instrumenter, HttpServerResponseBeforeCommitHandler.Noop.INSTANCE);
   }
 }

instrumentation/netty/netty-4.1/library/src/main/java/io/opentelemetry/instrumentation/netty/v4_1/internal/server/HttpServerResponseBeforeCommitHandler.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.netty.v4_1.internal.server;
+
+import io.netty.handler.codec.http.HttpResponse;
+import io.opentelemetry.context.Context;
+
+/**
+ * This class is internal and is hence not for public use. Its APIs are unstable and can change at
+ * any time.
+ */
+public interface HttpServerResponseBeforeCommitHandler {
+  void handle(Context context, HttpResponse response);
+
+  /**
+   * This class is internal and is hence not for public use. Its APIs are unstable and can change at
+   * any time.
+   */
+  enum Noop implements HttpServerResponseBeforeCommitHandler {
+    INSTANCE;
+
+    @Override
+    public void handle(Context context, HttpResponse response) {}
+  }
+}

instrumentation/netty/netty-4.1/library/src/main/java/io/opentelemetry/instrumentation/netty/v4_1/internal/server/HttpServerResponseTracingHandler.java

@@ -35,10 +35,13 @@ public class HttpServerResponseTracingHandler extends ChannelOutboundHandlerAdap
       AttributeKey.valueOf(HttpServerResponseTracingHandler.class, "http-server-response");
 
   private final Instrumenter<HttpRequestAndChannel, HttpResponse> instrumenter;
+  private final HttpServerResponseBeforeCommitHandler beforeCommitHandler;
 
   public HttpServerResponseTracingHandler(
-      Instrumenter<HttpRequestAndChannel, HttpResponse> instrumenter) {
+      Instrumenter<HttpRequestAndChannel, HttpResponse> instrumenter,
+      HttpServerResponseBeforeCommitHandler beforeCommitHandler) {
     this.instrumenter = instrumenter;
+    this.beforeCommitHandler = beforeCommitHandler;
   }
 
   @Override
@@ -65,6 +68,7 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) thr
       // Going to finish the span after the write of the last content finishes.
       if (msg instanceof FullHttpResponse) {
         // Headers and body all sent together, we have the response information in the msg.
+        beforeCommitHandler.handle(context, (HttpResponse) msg);
         writePromise.addListener(
             future -> end(ctx.channel(), (FullHttpResponse) msg, writePromise));
       } else {
@@ -82,6 +86,7 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) thr
       writePromise = prm;
       if (msg instanceof HttpResponse) {
         // Headers before body has been sent, store them to use when finishing the span.
+        beforeCommitHandler.handle(context, (HttpResponse) msg);
         ctx.channel().attr(HTTP_SERVER_RESPONSE).set((HttpResponse) msg);
       }
     }

instrumentation/netty/netty-4.1/library/src/main/java/io/opentelemetry/instrumentation/netty/v4_1/internal/server/HttpServerTracingHandler.java

@@ -18,9 +18,11 @@ public class HttpServerTracingHandler
     extends CombinedChannelDuplexHandler<
         HttpServerRequestTracingHandler, HttpServerResponseTracingHandler> {
 
-  public HttpServerTracingHandler(Instrumenter<HttpRequestAndChannel, HttpResponse> instrumenter) {
+  public HttpServerTracingHandler(
+      Instrumenter<HttpRequestAndChannel, HttpResponse> instrumenter,
+      HttpServerResponseBeforeCommitHandler responseBeforeCommitHandler) {
     super(
         new HttpServerRequestTracingHandler(instrumenter),
-        new HttpServerResponseTracingHandler(instrumenter));
+        new HttpServerResponseTracingHandler(instrumenter, responseBeforeCommitHandler));
   }
 }