Fix continuation leakage on netty http2 client pipeline

dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/Http2ConnectContinuationListener.java

@@ -0,0 +1,34 @@
+package datadog.trace.instrumentation.netty41;
+
+import static datadog.trace.instrumentation.netty41.AttributeKeys.CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import io.netty.channel.Channel;
+import io.netty.channel.ChannelFuture;
+import io.netty.channel.ChannelFutureListener;
+
+public final class Http2ConnectContinuationListener implements ChannelFutureListener {
+  public static final ChannelFutureListener INSTANCE = new Http2ConnectContinuationListener();
+
+  private Http2ConnectContinuationListener() {}
+
+  @Override
+  public void operationComplete(final ChannelFuture future) {
+    if (future.isSuccess() || future.isCancelled()) {
+      cancel(future.channel());
+    }
+    // Failed connects are left for ChannelFutureListenerInstrumentation, which creates the
+    // netty.connect error span under this continuation.
+  }
+
+  public static void cancel(final Channel channel) {
+    if (channel == null) {
+      return;
+    }
+    final AgentScope.Continuation continuation =
+        channel.attr(CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY).getAndRemove();
+    if (continuation != null) {
+      continuation.cancel();
+    }
+  }
+}

dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelPipelineInstrumentation.java

@@ -7,6 +7,7 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureActiveSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopContinuation;
 import static datadog.trace.instrumentation.netty41.AttributeKeys.CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY;
+import static datadog.trace.instrumentation.netty41.AttributeKeys.HTTP2_CONNECTION_CODEC_ATTRIBUTE_KEY;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.returns;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -29,6 +30,7 @@
 import datadog.trace.instrumentation.netty41.server.websocket.WebSocketServerInboundTracingHandler;
 import datadog.trace.instrumentation.netty41.server.websocket.WebSocketServerOutboundTracingHandler;
 import datadog.trace.instrumentation.netty41.server.websocket.WebSocketServerTracingHandler;
+import io.netty.channel.ChannelFuture;
 import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelPipeline;
 import io.netty.handler.codec.http.HttpClientCodec;
@@ -90,6 +92,7 @@ public String[] helperClassNames() {
       packageName + ".server.websocket.WebSocketServerTracingHandler",
       packageName + ".server.websocket.WebSocketServerOutboundTracingHandler",
       packageName + ".server.websocket.WebSocketServerInboundTracingHandler",
+      packageName + ".Http2ConnectContinuationListener",
       packageName + ".NettyHttp2Helper",
       packageName + ".NettyPipelineHelper",
     };
@@ -175,6 +178,9 @@ public static void addHandler(
           handler2 instanceof ChannelHandler ? (ChannelHandler) handler2 : handler3;
 
       try {
+        if (NettyHttp2Helper.isHttp2ConnectionCodec(handler)) {
+          pipeline.channel().attr(HTTP2_CONNECTION_CODEC_ATTRIBUTE_KEY).set(Boolean.TRUE);
+        }
         // Server pipeline handlers
         if (handler instanceof HttpServerCodec) {
           NettyPipelineHelper.addHandlerAfter(
@@ -249,14 +255,33 @@ else if (handler instanceof HttpClientCodec) {
 
   public static class ConnectAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void addParentSpan(@Advice.This final ChannelPipeline pipeline) {
+    public static boolean addParentSpan(@Advice.This final ChannelPipeline pipeline) {
       AgentScope.Continuation continuation = captureActiveSpan();
       if (continuation != noopContinuation()) {
         final Attribute<AgentScope.Continuation> attribute =
             pipeline.channel().attr(CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY);
         if (!attribute.compareAndSet(null, continuation)) {
           continuation.cancel();
+          return false;
         }
+        return Boolean.TRUE.equals(
+            pipeline.channel().attr(HTTP2_CONNECTION_CODEC_ATTRIBUTE_KEY).get());
+      }
+      return false;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void cleanupHttp2ConnectParentContinuation(
+        @Advice.Enter final boolean cleanupHttp2Continuation,
+        @Advice.This final ChannelPipeline pipeline,
+        @Advice.Return final ChannelFuture future) {
+      if (!cleanupHttp2Continuation) {
+        return;
+      }
+      if (future == null) {
+        Http2ConnectContinuationListener.cancel(pipeline.channel());
+      } else {
+        future.addListener(Http2ConnectContinuationListener.INSTANCE);
       }
     }
   }

dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyHttp2Helper.java

@@ -10,11 +10,13 @@
 
 public class NettyHttp2Helper {
   private static final Class HTTP2_CODEC_CLS;
+  private static final Class HTTP2_FRAME_CODEC_CLS;
   private static final MethodHandle IS_SERVER_FIELD;
   private static final Logger LOGGER = LoggerFactory.getLogger(NettyHttp2Helper.class);
 
   static {
     Class codecClass;
+    Class frameCodecClass;
     MethodHandle isServerField;
     try {
       codecClass =
@@ -38,14 +40,33 @@ public class NettyHttp2Helper {
       isServerField = null;
       LOGGER.debug("Unable to setup netty http2 instrumentation", t);
     }
+    try {
+      frameCodecClass =
+          Class.forName(
+              "io.netty.handler.codec.http2.Http2FrameCodec",
+              false,
+              NettyHttp2Helper.class.getClassLoader());
+    } catch (final ClassNotFoundException cnfe) {
+      // can be expected
+      frameCodecClass = null;
+    } catch (Throwable t) {
+      // unexpected
+      frameCodecClass = null;
+      LOGGER.debug("Unable to setup netty http2 connection detection", t);
+    }
     HTTP2_CODEC_CLS = codecClass;
+    HTTP2_FRAME_CODEC_CLS = frameCodecClass;
     IS_SERVER_FIELD = isServerField;
   }
 
   public static boolean isHttp2FrameCodec(final ChannelHandler handler) {
     return HTTP2_CODEC_CLS != null && HTTP2_CODEC_CLS.isInstance(handler);
   }
 
+  public static boolean isHttp2ConnectionCodec(final ChannelHandler handler) {
+    return HTTP2_FRAME_CODEC_CLS != null && HTTP2_FRAME_CODEC_CLS.isInstance(handler);
+  }
+
   public static boolean isServer(final ChannelHandler handler) {
     try {
       return IS_SERVER_FIELD != null && (boolean) IS_SERVER_FIELD.invokeExact(handler);

dd-java-agent/instrumentation/netty/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java

@@ -18,6 +18,7 @@
 import datadog.trace.api.Config;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import io.netty.channel.Channel;
 import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelOutboundHandlerAdapter;
@@ -54,8 +55,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann
     }
 
     AgentScope parentScope = null;
-    final AgentScope.Continuation continuation =
-        ctx.channel().attr(CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY).getAndRemove();
+    final AgentScope.Continuation continuation = takeConnectParentContinuation(ctx);
     if (continuation != null) {
       parentScope = continuation.activate();
     }
@@ -111,4 +111,16 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann
       }
     }
   }
+
+  private static AgentScope.Continuation takeConnectParentContinuation(
+      final ChannelHandlerContext ctx) {
+    final Channel channel = ctx.channel();
+    AgentScope.Continuation continuation =
+        channel.attr(CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY).getAndRemove();
+    if (continuation == null && channel.parent() != null) {
+      continuation =
+          channel.parent().attr(CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY).getAndRemove();
+    }
+    return continuation;
+  }
 }

dd-java-agent/instrumentation/netty/netty-common/src/main/java/datadog/trace/instrumentation/netty41/AttributeKeys.java

@@ -27,6 +27,9 @@ public final class AttributeKeys {
       CONNECT_PARENT_CONTINUATION_ATTRIBUTE_KEY =
           attributeKey("datadog.connect.parent.continuation");
 
+  public static final AttributeKey<Boolean> HTTP2_CONNECTION_CODEC_ATTRIBUTE_KEY =
+      attributeKey("datadog.http2.connection.codec");
+
   public static final AttributeKey<Context> PARENT_CONTEXT_ATTRIBUTE_KEY =
       attributeKey("datadog.server.parent-context");
 

dd-java-agent/instrumentation/reactor-netty-1.0/src/test/groovy/ReactorNettyHttp2ClientTest.groovy

@@ -9,6 +9,7 @@ import reactor.netty.http.server.HttpServer
 import spock.lang.IgnoreIf
 import spock.lang.Shared
 
+import static datadog.trace.agent.test.utils.PortUtils.UNUSABLE_PORT
 import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
 import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
 
@@ -24,16 +25,50 @@ class ReactorNettyHttp2ClientTest extends InstrumentationSpecification {
   .handle { req, res -> res.status(200).send() }
   .bindNow()
 
-  @Override
-  boolean useStrictTraceWrites() {
-    false
-  }
-
   @Override
   def cleanupSpec() {
     server?.disposeNow()
   }
 
+  def "test http2 prior knowledge failed connect creates connect error span"() {
+    setup:
+    HttpClient httpClient = HttpClient.create()
+      .disableRetry(true)
+      .protocol(HttpProtocol.H2C)
+
+    when:
+    runUnderTrace("parent", {
+      httpClient.baseUrl("http://127.0.0.1:${UNUSABLE_PORT}")
+        .get()
+        .uri("/")
+        .response()
+        .block()
+    })
+
+    then:
+    def ex = thrown(Exception)
+    (ex instanceof ConnectException) || (ex.cause instanceof ConnectException)
+
+    and:
+    assertTraces(1) {
+      trace(2) {
+        basicSpan(it, "parent", null, ex)
+
+        span {
+          operationName "netty.connect"
+          resourceName "netty.connect"
+          childOf span(0)
+          errored true
+          tags {
+            "$Tags.COMPONENT" "netty"
+            errorTags Throwable, ~"Connection refused"
+            defaultTags()
+          }
+        }
+      }
+    }
+  }
+
   def "test http2 client/server propagation"() {
     setup:
     HttpClient httpClient = HttpClient.create()