Implement HttpServerResponseCustomizer support for Grizzly (#8263)

open-telemetry · Apr 12, 2023 · 271c72b · 271c72b
1 parent 66f4c80
commit 271c72b
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 0 deletions.
diff --git a/...n/java/io/opentelemetry/javaagent/instrumentation/grizzly/GrizzlyHttpResponseMutator.java b/...n/java/io/opentelemetry/javaagent/instrumentation/grizzly/GrizzlyHttpResponseMutator.java
@@ -0,0 +1,29 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.grizzly;
+
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
+import org.glassfish.grizzly.http.HttpResponsePacket;
+import org.glassfish.grizzly.http.util.DataChunk;
+import org.glassfish.grizzly.http.util.MimeHeaders;
+
+public enum GrizzlyHttpResponseMutator implements HttpServerResponseMutator<HttpResponsePacket> {
+  INSTANCE;
+
+  @Override
+  public void appendHeader(HttpResponsePacket response, String name, String value) {
+    MimeHeaders headers = response.getHeaders();
+    DataChunk data = headers.getValue(name);
+    if (data == null) {
+      data = headers.addValue(name);
+    }
+    if (data.getLength() > 0) {
+      data.setString(data.toString() + "," + value);
+    } else {
+      data.setString(value);
+    }
+  }
+}
diff --git a/...a/io/opentelemetry/javaagent/instrumentation/grizzly/HttpServerFilterInstrumentation.java b/...a/io/opentelemetry/javaagent/instrumentation/grizzly/HttpServerFilterInstrumentation.java
@@ -11,6 +11,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import io.opentelemetry.context.Context;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizerHolder;
 import io.opentelemetry.javaagent.bootstrap.servlet.AppServerBridge;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
@@ -43,6 +44,15 @@ public void transform(TypeTransformer transformer) {
   @SuppressWarnings("unused")
   public static class PrepareResponseAdvice {
 
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.Argument(0) FilterChainContext ctx,
+        @Advice.Argument(2) HttpResponsePacket response) {
+      Context context = GrizzlyStateStorage.getContext(ctx);
+      HttpServerResponseCustomizerHolder.getCustomizer()
+          .customize(context, response, GrizzlyHttpResponseMutator.INSTANCE);
+    }
+
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void onExit(
         @Advice.Argument(0) FilterChainContext ctx,

diff --git a/instrumentation/grizzly-2.0/javaagent/src/test/groovy/GrizzlyTest.groovy b/instrumentation/grizzly-2.0/javaagent/src/test/groovy/GrizzlyTest.groovy
@@ -58,6 +58,11 @@ class GrizzlyTest extends HttpServerTest<HttpServer> implements AgentTestTrait {
     server.stop()
   }
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    true
+  }
+
   @Override
   boolean testCapturedHttpHeaders() {
     false