Instrument Netty 4.0 to add Server-Timing header

instrumentation/netty-3.8/build.gradle

@@ -7,6 +7,9 @@ dependencies {
     implementation project(':instrumentation:common')
 
     testInstrumentation group: 'io.opentelemetry.javaagent.instrumentation', name: 'opentelemetry-javaagent-netty-3.8', version: versions.opentelemetryJavaagent
+    testInstrumentation group: 'io.opentelemetry.javaagent.instrumentation', name: 'opentelemetry-javaagent-netty-4.0', version: versions.opentelemetryJavaagent
+    testInstrumentation group: 'io.opentelemetry.javaagent.instrumentation', name: 'opentelemetry-javaagent-netty-4.1', version: versions.opentelemetryJavaagent
+    testInstrumentation project(':instrumentation:netty-4.0')
 
     testImplementation group: 'io.netty', name: 'netty', version: '3.8.0.Final'
 }

instrumentation/netty-3.8/src/main/java/com/splunk/opentelemetry/netty/v3_8/NettyChannelPipelineInstrumentation.java → instrumentation/netty-3.8/src/main/java/com/splunk/opentelemetry/netty/v3_8/ChannelPipelineInstrumentation.java

@@ -40,7 +40,7 @@
 import org.jboss.netty.handler.codec.http.HttpResponseEncoder;
 import org.jboss.netty.handler.codec.http.HttpServerCodec;
 
-public class NettyChannelPipelineInstrumentation implements TypeInstrumentation {
+public class ChannelPipelineInstrumentation implements TypeInstrumentation {
 
   @Override
   public ElementMatcher<ClassLoader> classLoaderOptimization() {

instrumentation/netty-3.8/src/main/java/com/splunk/opentelemetry/netty/v3_8/NettyInstrumentationModule.java

@@ -16,6 +16,8 @@
 
 package com.splunk.opentelemetry.netty.v3_8;
 
+import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.ClassLoaderMatcher.hasClassesNamed;
+
 import com.google.auto.service.AutoService;
 import com.splunk.opentelemetry.servertiming.ServerTimingHeader;
 import io.opentelemetry.javaagent.instrumentation.netty.v3_8.ChannelTraceContext;
@@ -24,21 +26,26 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(InstrumentationModule.class)
 public class NettyInstrumentationModule extends InstrumentationModule {
   public NettyInstrumentationModule() {
     super("netty", "netty-3.8");
   }
 
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    return hasClassesNamed("org.jboss.netty.channel.ChannelPipeline");
+  }
+
   @Override
   protected String[] additionalHelperClassNames() {
     return new String[] {
       ServerTimingHeader.class.getName(),
       getClass().getPackage().getName() + ".ServerTimingHandler",
       getClass().getPackage().getName() + ".ServerTimingHandler$HeadersSetter",
-      getClass().getPackage().getName()
-          + ".NettyChannelPipelineInstrumentation$ChannelPipelineUtil",
+      getClass().getPackage().getName() + ".ChannelPipelineInstrumentation$ChannelPipelineUtil",
     };
   }
 
@@ -62,6 +69,6 @@ public Map<String, String> contextStore() {
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return Collections.singletonList(new NettyChannelPipelineInstrumentation());
+    return Collections.singletonList(new ChannelPipelineInstrumentation());
   }
 }

instrumentation/netty-4.0/build.gradle

@@ -0,0 +1,19 @@
+apply from: "$rootDir/gradle/instrumentation.gradle"
+
+dependencies {
+    compileOnly group: 'io.netty', name: 'netty-codec-http', version: '4.0.0.Final'
+    compileOnly group: 'io.opentelemetry.javaagent.instrumentation', name: 'opentelemetry-javaagent-netty-4.0', version: versions.opentelemetryJavaagent
+
+    implementation project(':instrumentation:common')
+
+    testInstrumentation group: 'io.opentelemetry.javaagent.instrumentation', name: 'opentelemetry-javaagent-netty-3.8', version: versions.opentelemetryJavaagent
+    testInstrumentation group: 'io.opentelemetry.javaagent.instrumentation', name: 'opentelemetry-javaagent-netty-4.0', version: versions.opentelemetryJavaagent
+    testInstrumentation group: 'io.opentelemetry.javaagent.instrumentation', name: 'opentelemetry-javaagent-netty-4.1', version: versions.opentelemetryJavaagent
+    testInstrumentation project(':instrumentation:netty-3.8')
+
+    testImplementation group: 'io.netty', name: 'netty-codec-http', version: '4.0.0.Final'
+}
+
+tasks.withType(Test) {
+    jvmArgs '-Dsplunk.context.server-timing.enabled=true'
+}

instrumentation/netty-4.0/src/main/java/com/splunk/opentelemetry/netty/v4_0/ChannelPipelineInstrumentation.java

@@ -0,0 +1,84 @@
+/*
+ * Copyright Splunk Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.splunk.opentelemetry.netty.v4_0;
+
+import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.AgentElementMatchers.implementsInterface;
+import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.ClassLoaderMatcher.hasClassesNamed;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import io.netty.channel.ChannelHandler;
+import io.netty.channel.ChannelPipeline;
+import io.netty.handler.codec.http.HttpResponseEncoder;
+import io.netty.handler.codec.http.HttpServerCodec;
+import io.opentelemetry.javaagent.instrumentation.api.CallDepthThreadLocalMap;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.util.Collections;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class ChannelPipelineInstrumentation implements TypeInstrumentation {
+  @Override
+  public ElementMatcher<ClassLoader> classLoaderOptimization() {
+    return hasClassesNamed("io.netty.channel.ChannelPipeline");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return implementsInterface(named("io.netty.channel.ChannelPipeline"));
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return Collections.singletonMap(
+        isMethod()
+            .and(nameStartsWith("add"))
+            .and(takesArgument(2, named("io.netty.channel.ChannelHandler"))),
+        this.getClass().getName() + "$ChannelPipelineAddAdvice");
+  }
+
+  public static class ChannelPipelineAddAdvice {
+    @Advice.OnMethodEnter
+    public static int trackCallDepth() {
+      return CallDepthThreadLocalMap.incrementCallDepth(ServerTimingHandler.class);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void addHandler(
+        @Advice.Enter int callDepth,
+        @Advice.This ChannelPipeline pipeline,
+        @Advice.Argument(2) ChannelHandler handler) {
+      if (callDepth > 0) {
+        return;
+      }
+      CallDepthThreadLocalMap.reset(ServerTimingHandler.class);
+
+      try {
+        if (handler instanceof HttpServerCodec || handler instanceof HttpResponseEncoder) {
+          pipeline.addLast(ServerTimingHandler.class.getName(), new ServerTimingHandler());
+        }
+      } catch (IllegalArgumentException e) {
+        // Prevented adding duplicate handlers.
+      }
+    }
+  }
+}

instrumentation/netty-4.0/src/main/java/com/splunk/opentelemetry/netty/v4_0/NettyInstrumentationModule.java

@@ -0,0 +1,68 @@
+/*
+ * Copyright Splunk Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.splunk.opentelemetry.netty.v4_0;
+
+import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.ClassLoaderMatcher.hasClassesNamed;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+
+import com.google.auto.service.AutoService;
+import com.splunk.opentelemetry.servertiming.ServerTimingHeader;
+import io.opentelemetry.javaagent.tooling.InstrumentationModule;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.util.Collections;
+import java.util.List;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumentationModule.class)
+public class NettyInstrumentationModule extends InstrumentationModule {
+  public NettyInstrumentationModule() {
+    super("netty", "netty-4.0");
+  }
+
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    // Class added in 4.1.0 and not in 4.0.56 to avoid resolving this instrumentation completely
+    // when using 4.1.
+    return not(hasClassesNamed("io.netty.handler.codec.http.CombinedHttpHeaders"));
+  }
+
+  @Override
+  protected String[] additionalHelperClassNames() {
+    return new String[] {
+      ServerTimingHeader.class.getName(),
+      this.getClass().getPackage().getName() + ".ServerTimingHandler",
+      this.getClass().getPackage().getName() + ".ServerTimingHandler$HeadersSetter"
+    };
+  }
+
+  // run after the upstream netty instrumentation
+  @Override
+  public int getOrder() {
+    return 1;
+  }
+
+  // enable the instrumentation only if the server-timing header flag is on
+  @Override
+  protected boolean defaultEnabled() {
+    return super.defaultEnabled() && ServerTimingHeader.shouldEmitServerTimingHeader();
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return Collections.singletonList(new ChannelPipelineInstrumentation());
+  }
+}

instrumentation/netty-4.0/src/main/java/com/splunk/opentelemetry/netty/v4_0/ServerTimingHandler.java

@@ -0,0 +1,52 @@
+/*
+ * Copyright Splunk Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.splunk.opentelemetry.netty.v4_0;
+
+import static io.opentelemetry.javaagent.instrumentation.netty.v4_0.server.NettyHttpServerTracer.tracer;
+
+import com.splunk.opentelemetry.servertiming.ServerTimingHeader;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelOutboundHandlerAdapter;
+import io.netty.channel.ChannelPromise;
+import io.netty.handler.codec.http.HttpHeaders;
+import io.netty.handler.codec.http.HttpResponse;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapPropagator;
+
+public class ServerTimingHandler extends ChannelOutboundHandlerAdapter {
+  @Override
+  public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) {
+    Context context = tracer().getServerContext(ctx.channel());
+    if (context == null || !(msg instanceof HttpResponse)) {
+      ctx.write(msg, prm);
+      return;
+    }
+
+    HttpResponse response = (HttpResponse) msg;
+    ServerTimingHeader.setHeaders(context, response.headers(), HeadersSetter.INSTANCE);
+    ctx.write(msg, prm);
+  }
+
+  public static final class HeadersSetter implements TextMapPropagator.Setter<HttpHeaders> {
+    private static final HeadersSetter INSTANCE = new HeadersSetter();
+
+    @Override
+    public void set(HttpHeaders carrier, String key, String value) {
+      carrier.add(key, value);
+    }
+  }
+}