Implement HttpServerResponseCustomizer support for Restlet

instrumentation/restlet/restlet-1.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/restlet/v1_1/RestletResponseMutator.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.restlet.v1_1;
+
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
+import org.restlet.data.Form;
+import org.restlet.data.Response;
+
+public enum RestletResponseMutator implements HttpServerResponseMutator<Response> {
+  INSTANCE;
+
+  public static final String HEADERS_ATTRIBUTE = "org.restlet.http.headers";
+
+  @Override
+  public void appendHeader(Response response, String name, String value) {
+    Form headers =
+        (Form) response.getAttributes().computeIfAbsent(HEADERS_ATTRIBUTE, k -> new Form());
+    String existing = headers.getValues(name);
+    if (existing != null) {
+      value = existing + "," + value;
+    }
+    headers.set(name, value, /* ignoreCase= */ true);
+  }
+}

instrumentation/restlet/restlet-1.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/restlet/v1_1/ServerInstrumentation.java

@@ -16,6 +16,7 @@
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.instrumentation.api.instrumenter.http.HttpRouteHolder;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizerHolder;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import net.bytebuddy.asm.Advice;
@@ -79,6 +80,9 @@ public static void finishRequest(
         HttpRouteHolder.updateHttpRoute(context, CONTROLLER, serverSpanName(), "/*");
       }
 
+      HttpServerResponseCustomizerHolder.getCustomizer()
+          .customize(context, response, RestletResponseMutator.INSTANCE);
+
       if (exception != null) {
         instrumenter().end(context, request, response, exception);
         return;

instrumentation/restlet/restlet-1.1/javaagent/src/test/groovy/io/opentelemetry/javaagent/instrumentation/restlet/v1_1/RestletServerTest.groovy

@@ -23,4 +23,9 @@ class RestletServerTest extends AbstractRestletServerTest implements AgentTestTr
     }
   }
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    return true
+  }
+
 }

instrumentation/restlet/restlet-1.1/javaagent/src/test/groovy/io/opentelemetry/javaagent/instrumentation/restlet/v1_1/ServletServerTest.groovy

@@ -8,7 +8,13 @@ package io.opentelemetry.javaagent.instrumentation.restlet.v1_1
 
 import io.opentelemetry.instrumentation.restlet.v1_1.AbstractServletServerTest
 import io.opentelemetry.instrumentation.test.AgentTestTrait
+import io.opentelemetry.instrumentation.testing.junit.http.ServerEndpoint
 
 class ServletServerTest extends AbstractServletServerTest implements AgentTestTrait {
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    return true
+  }
+
 }

instrumentation/restlet/restlet-1.1/javaagent/src/test/groovy/io/opentelemetry/javaagent/instrumentation/restlet/v1_1/spring/SpringBeanRouterTest.groovy

@@ -7,11 +7,17 @@ package io.opentelemetry.javaagent.instrumentation.restlet.v1_1.spring
 
 import io.opentelemetry.instrumentation.restlet.v1_1.spring.AbstractSpringServerTest
 import io.opentelemetry.instrumentation.test.AgentTestTrait
+import io.opentelemetry.instrumentation.testing.junit.http.ServerEndpoint
 
 class SpringBeanRouterTest extends AbstractSpringServerTest implements AgentTestTrait {
   @Override
   String getConfigurationName() {
     return "springBeanRouterConf.xml"
   }
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    return true
+  }
+
 }

instrumentation/restlet/restlet-1.1/javaagent/src/test/groovy/io/opentelemetry/javaagent/instrumentation/restlet/v1_1/spring/SpringRouterTest.groovy

@@ -7,11 +7,17 @@ package io.opentelemetry.javaagent.instrumentation.restlet.v1_1.spring
 
 import io.opentelemetry.instrumentation.restlet.v1_1.spring.AbstractSpringServerTest
 import io.opentelemetry.instrumentation.test.AgentTestTrait
+import io.opentelemetry.instrumentation.testing.junit.http.ServerEndpoint
 
 class SpringRouterTest extends AbstractSpringServerTest implements AgentTestTrait {
   @Override
   String getConfigurationName() {
     return "springRouterConf.xml"
   }
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    return true
+  }
+
 }

instrumentation/restlet/restlet-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/restlet/v2_0/RestletResponseMutator.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.restlet.v2_0;
+
+import io.opentelemetry.instrumentation.restlet.v2_0.internal.MessageAttributesAccessor;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseMutator;
+import java.util.Map;
+import org.restlet.Response;
+import org.restlet.util.Series;
+
+public enum RestletResponseMutator implements HttpServerResponseMutator<Response> {
+  INSTANCE;
+
+  public static final String HEADERS_ATTRIBUTE = "org.restlet.http.headers";
+
+  @Override
+  public void appendHeader(Response response, String name, String value) {
+    Map<String, Object> attributes = MessageAttributesAccessor.getAttributes(response);
+    if (attributes == null) {
+      // should never happen in practice
+      return;
+    }
+    Series<?> headers = (Series<?>) attributes.get(HEADERS_ATTRIBUTE);
+    if (headers == null) {
+      headers = MessageAttributesAccessor.createHeaderSeries();
+      if (headers == null) {
+        // should never happen in practice; abort
+        return;
+      }
+      attributes.put(HEADERS_ATTRIBUTE, headers);
+    }
+    String existing = headers.getValues(name);
+    if (existing != null) {
+      value = existing + "," + value;
+    }
+    MessageAttributesAccessor.setSeriesValue(headers, name, value);
+  }
+}

instrumentation/restlet/restlet-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/restlet/v2_0/ServerInstrumentation.java

@@ -16,6 +16,7 @@
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.instrumentation.api.instrumenter.http.HttpRouteHolder;
+import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizerHolder;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import net.bytebuddy.asm.Advice;
@@ -79,6 +80,9 @@ public static void finishRequest(
         HttpRouteHolder.updateHttpRoute(context, CONTROLLER, serverSpanName(), "/*");
       }
 
+      HttpServerResponseCustomizerHolder.getCustomizer()
+          .customize(context, response, RestletResponseMutator.INSTANCE);
+
       if (exception != null) {
         instrumenter().end(context, request, response, exception);
         return;

instrumentation/restlet/restlet-2.0/javaagent/src/test/groovy/io/opentelemetry/javaagent/instrumentation/restlet/v2_0/RestletServerTest.groovy

@@ -23,4 +23,9 @@ class RestletServerTest extends AbstractRestletServerTest implements AgentTestTr
     }
   }
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    return true
+  }
+
 }

instrumentation/restlet/restlet-2.0/javaagent/src/test/groovy/io/opentelemetry/javaagent/instrumentation/restlet/v2_0/spring/SpringBeanRouterTest.groovy

@@ -7,11 +7,17 @@ package io.opentelemetry.javaagent.instrumentation.restlet.v2_0.spring
 
 import io.opentelemetry.instrumentation.restlet.v2_0.spring.AbstractSpringServerTest
 import io.opentelemetry.instrumentation.test.AgentTestTrait
+import io.opentelemetry.instrumentation.testing.junit.http.ServerEndpoint
 
 class SpringBeanRouterTest extends AbstractSpringServerTest implements AgentTestTrait {
   @Override
   String getConfigurationName() {
     return "springBeanRouterConf.xml"
   }
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    return true
+  }
+
 }

instrumentation/restlet/restlet-2.0/javaagent/src/test/groovy/io/opentelemetry/javaagent/instrumentation/restlet/v2_0/spring/SpringRouterTest.groovy

@@ -7,11 +7,17 @@ package io.opentelemetry.javaagent.instrumentation.restlet.v2_0.spring
 
 import io.opentelemetry.instrumentation.restlet.v2_0.spring.AbstractSpringServerTest
 import io.opentelemetry.instrumentation.test.AgentTestTrait
+import io.opentelemetry.instrumentation.testing.junit.http.ServerEndpoint
 
 class SpringRouterTest extends AbstractSpringServerTest implements AgentTestTrait {
   @Override
   String getConfigurationName() {
     return "springRouterConf.xml"
   }
 
+  @Override
+  boolean hasResponseCustomizer(ServerEndpoint endpoint) {
+    return true
+  }
+
 }

instrumentation/restlet/restlet-2.0/library/src/main/java/io/opentelemetry/instrumentation/restlet/v2_0/internal/MessageAttributesAccessor.java

@@ -0,0 +1,137 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.restlet.v2_0.internal;
+
+import static java.lang.invoke.MethodType.methodType;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.util.Map;
+import java.util.concurrent.ConcurrentMap;
+import javax.annotation.Nullable;
+import org.restlet.Message;
+import org.restlet.data.Form;
+import org.restlet.util.Series;
+
+/**
+ * This class is internal and is hence not for public use. Its APIs are unstable and can change at
+ * any time.
+ */
+public final class MessageAttributesAccessor {
+
+  private static final MethodHandle GET_ATTRIBUTES;
+  private static final MethodHandle SET_VALUE;
+  private static final Class<?> HEADER_CLASS;
+  private static final MethodHandle NEW_SERIES;
+
+  static {
+    MethodHandle getAttributes = null;
+    MethodHandle setValue = null;
+    Class<?> headerClass = null;
+    MethodHandle newSeries = null;
+
+    MethodHandles.Lookup lookup = MethodHandles.lookup();
+    try {
+      getAttributes = lookup.findVirtual(Message.class, "getAttributes", methodType(Map.class));
+    } catch (NoSuchMethodException | IllegalAccessException e) {
+      // changed the return type to ConcurrentMap in version 2.1
+      try {
+        getAttributes =
+            lookup.findVirtual(Message.class, "getAttributes", methodType(ConcurrentMap.class));
+      } catch (NoSuchMethodException | IllegalAccessException ex) {
+        // ignored
+      }
+    }
+
+    Class<?> setValueReturnType = null;
+    try {
+      // changed the generic bound to NamedValue in version 2.1; earlier than that it's Parameter
+      setValueReturnType = Class.forName("org.restlet.util.NamedValue");
+    } catch (ClassNotFoundException e) {
+      try {
+        setValueReturnType = Class.forName("org.restlet.data.Parameter");
+      } catch (ClassNotFoundException ex) {
+        // ignored
+      }
+    }
+    if (setValueReturnType != null) {
+      try {
+        setValue =
+            lookup.findVirtual(
+                Series.class, "set", methodType(setValueReturnType, String.class, String.class));
+      } catch (NoSuchMethodException | IllegalAccessException e) {
+        // ignored
+      }
+    }
+
+    try {
+      // restlet 2.3+
+      headerClass = Class.forName("org.restlet.data.Header");
+    } catch (ClassNotFoundException e) {
+      try {
+        // restlet 2.1-2.2
+        headerClass = Class.forName("org.restlet.engine.header.Header");
+      } catch (ClassNotFoundException ex) {
+        // restlet 2.0 does not have Header
+      }
+    }
+    if (headerClass != null) {
+      // restlet 2.1+ Series has different constructor
+      try {
+        newSeries = lookup.findConstructor(Series.class, methodType(void.class, Class.class));
+      } catch (NoSuchMethodException | IllegalAccessException e) {
+        // ignored
+      }
+    }
+
+    GET_ATTRIBUTES = getAttributes;
+    SET_VALUE = setValue;
+    HEADER_CLASS = headerClass;
+    NEW_SERIES = newSeries;
+  }
+
+  @SuppressWarnings("unchecked")
+  @Nullable
+  public static Map<String, Object> getAttributes(Message message) {
+    if (GET_ATTRIBUTES == null) {
+      return null;
+    }
+    try {
+      return (Map<String, Object>) GET_ATTRIBUTES.invoke(message);
+    } catch (Throwable e) {
+      return null;
+    }
+  }
+
+  @Nullable
+  public static Series<?> createHeaderSeries() {
+    if (HEADER_CLASS == null) {
+      return new Form();
+    }
+    if (NEW_SERIES == null) {
+      // should never really happen
+      return null;
+    }
+    try {
+      return (Series<?>) NEW_SERIES.invoke(HEADER_CLASS);
+    } catch (Throwable e) {
+      return null;
+    }
+  }
+
+  public static void setSeriesValue(Series<?> series, String name, String value) {
+    if (SET_VALUE == null) {
+      return;
+    }
+    try {
+      SET_VALUE.invoke(series, name, value);
+    } catch (Throwable ignored) {
+      // just do nothing
+    }
+  }
+
+  private MessageAttributesAccessor() {}
+}