Add option to create span on new netty connection

instrumentation/netty/netty-4-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/common/FutureListenerWrappers.java

@@ -12,7 +12,6 @@
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.instrumentation.api.caching.Cache;
-import io.opentelemetry.javaagent.instrumentation.netty.common.client.AbstractNettyHttpClientTracer;
 
 public final class FutureListenerWrappers {
   // Instead of ContextStore use Cache with weak keys and weak values to store link between original
@@ -64,7 +63,6 @@ private WrappedFutureListener(Context context, GenericFutureListener<Future<?>>
 
     @Override
     public void operationComplete(Future<?> future) throws Exception {
-      AbstractNettyHttpClientTracer.operationComplete(future);
       try (Scope ignored = context.makeCurrent()) {
         delegate.operationComplete(future);
       }
@@ -93,7 +91,6 @@ public void operationProgressed(ProgressiveFuture<?> progressiveFuture, long l,
 
     @Override
     public void operationComplete(ProgressiveFuture<?> progressiveFuture) throws Exception {
-      AbstractNettyHttpClientTracer.operationComplete(progressiveFuture);
       try (Scope ignored = context.makeCurrent()) {
         delegate.operationComplete(progressiveFuture);
       }

instrumentation/netty/netty-4-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/common/client/AbstractNettyHttpClientTracer.java

@@ -6,32 +6,37 @@
 package io.opentelemetry.javaagent.instrumentation.netty.common.client;
 
 import static io.opentelemetry.api.trace.SpanKind.CLIENT;
+import static io.opentelemetry.api.trace.SpanKind.INTERNAL;
 import static io.opentelemetry.javaagent.instrumentation.netty.common.client.NettyResponseInjectAdapter.SETTER;
 import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.NetTransportValues.IP_TCP;
 import static io.opentelemetry.semconv.trace.attributes.SemanticAttributes.NetTransportValues.IP_UDP;
 
 import io.netty.channel.Channel;
-import io.netty.channel.ChannelFuture;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.socket.DatagramChannel;
 import io.netty.handler.codec.http.HttpHeaders;
 import io.netty.handler.codec.http.HttpResponse;
-import io.netty.util.concurrent.Future;
+import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanBuilder;
-import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.propagation.TextMapSetter;
+import io.opentelemetry.instrumentation.api.config.Config;
 import io.opentelemetry.instrumentation.api.tracer.HttpClientTracer;
 import io.opentelemetry.instrumentation.api.tracer.net.NetPeerAttributes;
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes;
 import java.net.InetSocketAddress;
+import java.net.SocketAddress;
 import java.net.URI;
 import java.net.URISyntaxException;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 public abstract class AbstractNettyHttpClientTracer<REQUEST extends AbstractNettyRequestWrapper>
     extends HttpClientTracer<REQUEST, HttpHeaders, HttpResponse> {
 
+  private static final boolean alwaysCreateConnectSpan =
+      Config.get()
+          .getBooleanProperty("otel.instrumentation.netty.always-create-connect-span", false);
+
   protected AbstractNettyHttpClientTracer() {
     super(NetPeerAttributes.INSTANCE);
   }
@@ -84,39 +89,60 @@ protected TextMapSetter<HttpHeaders> getSetter() {
     return SETTER;
   }
 
-  public static void operationComplete(Future<?> future) {
-    AbstractNettyHttpClientTracer tracer = NettyHttpClientTracerAccess.getTracer();
-    if (tracer == null) {
-      return;
+  public Context startConnectionSpan(Context parentContext, SocketAddress remoteAddress) {
+    if (!alwaysCreateConnectSpan) {
+      return null;
     }
 
-    if (!(future instanceof ChannelFuture)) {
-      return;
-    }
-    // If first call to GenericFutureListener#operationComplete has an exception then we
-    // treat it as the cause of connection failure and create a special span for it
-    ChannelFuture channelFuture = (ChannelFuture) future;
-    Context parentContext = tracer.getAndRemoveConnectContext(channelFuture);
-    if (parentContext == null) {
-      return;
-    }
-    Throwable cause = future.cause();
-    if (cause == null) {
-      return;
-    }
+    SpanBuilder spanBuilder = spanBuilder(parentContext, "CONNECT", INTERNAL);
+    NetPeerAttributes.INSTANCE.setNetPeer(spanBuilder, (InetSocketAddress) remoteAddress);
+
+    return parentContext.with(spanBuilder.startSpan());
+  }
 
-    if (tracer.shouldStartSpan(parentContext, SpanKind.CLIENT)) {
-      tracer.connectionFailure(parentContext, channelFuture.channel(), cause);
+  public void endConnectionSpan(
+      Context context,
+      Context parentContext,
+      SocketAddress remoteAddress,
+      Channel channel,
+      Throwable throwable) {
+    if (alwaysCreateConnectSpan) {
+      if (context != null) {
+        // if context is present we started span in startConnectionSpan
+        endConnectionSpan(context, channel, throwable);
+      }
+    } else if (throwable != null && shouldStartSpan(parentContext, CLIENT)) {
+      // if we didn't start span in startConnectionSpan create a span only when the request fails
+      // and when not inside a client span
+      connectionFailure(parentContext, remoteAddress, channel, throwable);
     }
   }
 
-  protected abstract Context getAndRemoveConnectContext(ChannelFuture channelFuture);
+  private void endConnectionSpan(Context context, Channel channel, Throwable throwable) {
+    if (channel != null) {
+      Span span = Span.fromContext(context);
+      span.setAttribute(
+          SemanticAttributes.NET_TRANSPORT, channel instanceof DatagramChannel ? IP_UDP : IP_TCP);
+      NetPeerAttributes.INSTANCE.setNetPeer(span, (InetSocketAddress) channel.remoteAddress());
+    }
+    if (throwable != null) {
+      endExceptionally(context, throwable);
+    } else {
+      end(context);
+    }
+  }
 
-  private void connectionFailure(Context parentContext, Channel channel, Throwable throwable) {
+  private void connectionFailure(
+      Context parentContext, SocketAddress remoteAddress, Channel channel, Throwable throwable) {
     SpanBuilder spanBuilder = spanBuilder(parentContext, "CONNECT", CLIENT);
-    spanBuilder.setAttribute(
-        SemanticAttributes.NET_TRANSPORT, channel instanceof DatagramChannel ? IP_UDP : IP_TCP);
-    NetPeerAttributes.INSTANCE.setNetPeer(spanBuilder, (InetSocketAddress) channel.remoteAddress());
+    if (channel != null) {
+      spanBuilder.setAttribute(
+          SemanticAttributes.NET_TRANSPORT, channel instanceof DatagramChannel ? IP_UDP : IP_TCP);
+      NetPeerAttributes.INSTANCE.setNetPeer(
+          spanBuilder, (InetSocketAddress) channel.remoteAddress());
+    } else if (remoteAddress != null) {
+      NetPeerAttributes.INSTANCE.setNetPeer(spanBuilder, (InetSocketAddress) remoteAddress);
+    }
 
     Context context = withClientSpan(parentContext, spanBuilder.startSpan());
     endExceptionally(context, throwable);

instrumentation/netty/netty-4-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/common/client/ConnectionCompleteListener.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.netty.common.client;
+
+import io.netty.channel.Channel;
+import io.netty.channel.ChannelFuture;
+import io.netty.util.concurrent.Future;
+import io.netty.util.concurrent.GenericFutureListener;
+import io.opentelemetry.context.Context;
+
+public class ConnectionCompleteListener implements GenericFutureListener<Future<Void>> {
+  private final Context context;
+  private final Context parentContext;
+
+  public ConnectionCompleteListener(Context context, Context parentContext) {
+    this.context = context;
+    this.parentContext = parentContext;
+  }
+
+  @Override
+  public void operationComplete(Future<Void> future) {
+    AbstractNettyHttpClientTracer tracer = NettyHttpClientTracerAccess.getTracer();
+    if (tracer == null) {
+      return;
+    }
+
+    Channel channel = null;
+    if (future instanceof ChannelFuture) {
+      channel = ((ChannelFuture) future).channel();
+    }
+    tracer.endConnectionSpan(context, parentContext, null, channel, future.cause());
+  }
+}

instrumentation/netty/netty-4.0/javaagent/build.gradle.kts

@@ -29,6 +29,25 @@ dependencies {
   latestDepTestLibrary("io.netty:netty-codec-http:4.0.56.Final")
 }
 
+tasks {
+  val testConnectionSpan by registering(Test::class) {
+    filter {
+      includeTestsMatching("Netty40ConnectionSpanTest")
+      isFailOnNoMatchingTests = false
+    }
+    include("**/Netty40ConnectionSpanTest.*")
+    jvmArgs("-Dotel.instrumentation.netty.always-create-connect-span=true")
+  }
+
+  named<Test>("test") {
+    dependsOn(testConnectionSpan)
+    filter {
+      excludeTestsMatching("Netty40ConnectionSpanTest")
+      isFailOnNoMatchingTests = false
+    }
+  }
+}
+
 // We need to force the dependency to the earliest supported version because other libraries declare newer versions.
 if (!(findProperty("testLatestDeps") as Boolean)) {
   configurations.configureEach {

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

@@ -8,9 +8,16 @@
 import static io.opentelemetry.javaagent.instrumentation.netty.v4_0.client.NettyHttpClientTracer.tracer;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
+import io.netty.channel.ChannelFuture;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
+import io.opentelemetry.javaagent.instrumentation.netty.common.client.ConnectionCompleteListener;
+import java.net.SocketAddress;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -25,6 +32,9 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   public void transform(TypeTransformer transformer) {
     transformer.applyAdviceToMethod(
         isConstructor(), BootstrapInstrumentation.class.getName() + "$InitAdvice");
+    transformer.applyAdviceToMethod(
+        named("doConnect").and(takesArgument(0, SocketAddress.class)),
+        BootstrapInstrumentation.class.getName() + "$ConnectAdvice");
   }
 
   @SuppressWarnings("unused")
@@ -37,4 +47,38 @@ public static void enter() {
       tracer();
     }
   }
+
+  public static class ConnectAdvice {
+    @Advice.OnMethodEnter
+    public static void startConnect(
+        @Advice.Argument(0) SocketAddress remoteAddress,
+        @Advice.Local("otelContext") Context context,
+        @Advice.Local("otelParentContext") Context parentContext,
+        @Advice.Local("otelScope") Scope scope) {
+      parentContext = Java8BytecodeBridge.currentContext();
+      context = tracer().startConnectionSpan(parentContext, remoteAddress);
+      if (context != null) {
+        scope = context.makeCurrent();
+      }
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void endConnect(
+        @Advice.Thrown Throwable throwable,
+        @Advice.Argument(0) SocketAddress remoteAddress,
+        @Advice.Return ChannelFuture channelFuture,
+        @Advice.Local("otelContext") Context context,
+        @Advice.Local("otelParentContext") Context parentContext,
+        @Advice.Local("otelScope") Scope scope) {
+      if (scope != null) {
+        scope.close();
+      }
+
+      if (throwable != null) {
+        tracer().endConnectionSpan(context, parentContext, remoteAddress, null, throwable);
+      } else {
+        channelFuture.addListener(new ConnectionCompleteListener(context, parentContext));
+      }
+    }
+  }
 }

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

@@ -8,7 +8,6 @@
 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.returns;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import io.netty.channel.ChannelHandler;
@@ -19,12 +18,9 @@
 import io.netty.handler.codec.http.HttpResponseDecoder;
 import io.netty.handler.codec.http.HttpResponseEncoder;
 import io.netty.handler.codec.http.HttpServerCodec;
-import io.netty.util.Attribute;
-import io.opentelemetry.context.Context;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import io.opentelemetry.javaagent.instrumentation.api.CallDepth;
 import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
-import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.instrumentation.netty.common.AbstractNettyChannelPipelineInstrumentation;
 import io.opentelemetry.javaagent.instrumentation.netty.v4_0.client.HttpClientRequestTracingHandler;
 import io.opentelemetry.javaagent.instrumentation.netty.v4_0.client.HttpClientResponseTracingHandler;
@@ -46,9 +42,6 @@ public void transform(TypeTransformer transformer) {
             .and(nameStartsWith("add"))
             .and(takesArgument(2, named("io.netty.channel.ChannelHandler"))),
         NettyChannelPipelineInstrumentation.class.getName() + "$ChannelPipelineAddAdvice");
-    transformer.applyAdviceToMethod(
-        isMethod().and(named("connect")).and(returns(named("io.netty.channel.ChannelFuture"))),
-        NettyChannelPipelineInstrumentation.class.getName() + "$ChannelPipelineConnectAdvice");
   }
 
   /**
@@ -103,15 +96,4 @@ public static void addHandler(
       }
     }
   }
-
-  @SuppressWarnings("unused")
-  public static class ChannelPipelineConnectAdvice {
-
-    @Advice.OnMethodEnter
-    public static void addParentSpan(@Advice.This ChannelPipeline pipeline) {
-      Context context = Java8BytecodeBridge.currentContext();
-      Attribute<Context> attribute = pipeline.channel().attr(AttributeKeys.CONNECT_CONTEXT);
-      attribute.compareAndSet(null, context);
-    }
-  }
 }

instrumentation/netty/netty-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/v4_0/client/NettyHttpClientTracer.java

@@ -5,12 +5,9 @@
 
 package io.opentelemetry.javaagent.instrumentation.netty.v4_0.client;
 
-import io.netty.channel.ChannelFuture;
 import io.netty.handler.codec.http.HttpResponse;
-import io.opentelemetry.context.Context;
 import io.opentelemetry.javaagent.instrumentation.netty.common.client.AbstractNettyHttpClientTracer;
 import io.opentelemetry.javaagent.instrumentation.netty.common.client.NettyHttpClientTracerAccess;
-import io.opentelemetry.javaagent.instrumentation.netty.v4_0.AttributeKeys;
 
 public class NettyHttpClientTracer extends AbstractNettyHttpClientTracer<NettyRequestWrapper> {
   private static final NettyHttpClientTracer TRACER = new NettyHttpClientTracer();
@@ -25,11 +22,6 @@ public static NettyHttpClientTracer tracer() {
     return TRACER;
   }
 
-  @Override
-  protected Context getAndRemoveConnectContext(ChannelFuture channelFuture) {
-    return channelFuture.channel().attr(AttributeKeys.CONNECT_CONTEXT).getAndRemove();
-  }
-
   @Override
   protected Integer status(HttpResponse httpResponse) {
     return httpResponse.getStatus().code();