Fix Netty 4.1 HTTP/1.1 pipelined response tracing

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

@@ -46,6 +46,7 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
   public String[] helperClassNames() {
     return new String[] {
       packageName + ".AttributeKeys",
+      packageName + ".ServerRequestContext",
       // client helpers
       packageName + ".client.NettyHttpClientDecorator",
       packageName + ".client.NettyResponseInjectAdapter",

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

@@ -14,7 +14,6 @@
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 
 import com.google.auto.service.AutoService;
-import datadog.context.Context;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
@@ -48,6 +47,7 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
   public String[] helperClassNames() {
     return new String[] {
       packageName + ".AttributeKeys",
+      packageName + ".ServerRequestContext",
       packageName + ".client.NettyHttpClientDecorator",
       packageName + ".server.ResponseExtractAdapter",
       packageName + ".server.NettyHttpServerDecorator",
@@ -71,8 +71,8 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class FireAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentScope scopeSpan(@Advice.This final ChannelHandlerContext ctx) {
-      final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get();
-      final AgentSpan channelSpan = spanFromContext(storedContext);
+      final AgentSpan channelSpan =
+          spanFromContext(ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get());
       if (channelSpan == null || channelSpan == activeSpan()) {
         // don't modify the scope
         return noopScope();

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

@@ -72,6 +72,7 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
   public String[] helperClassNames() {
     return new String[] {
       packageName + ".AttributeKeys",
+      packageName + ".ServerRequestContext",
       // client helpers
       packageName + ".client.NettyHttpClientDecorator",
       packageName + ".client.NettyResponseInjectAdapter",

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

@@ -1,16 +1,14 @@
 package datadog.trace.instrumentation.netty41.server;
 
-import static datadog.trace.instrumentation.netty41.AttributeKeys.ANALYZED_RESPONSE_KEY;
-import static datadog.trace.instrumentation.netty41.AttributeKeys.BLOCKED_RESPONSE_KEY;
 import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY;
 import static datadog.trace.instrumentation.netty41.AttributeKeys.PARENT_CONTEXT_ATTRIBUTE_KEY;
-import static datadog.trace.instrumentation.netty41.AttributeKeys.REQUEST_HEADERS_ATTRIBUTE_KEY;
 import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.DECORATE;
 
 import datadog.context.Context;
 import datadog.context.ContextScope;
 import datadog.trace.api.gateway.Flow;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.instrumentation.netty41.ServerRequestContext;
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelHandlerContext;
@@ -38,6 +36,12 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
     }
 
     final HttpRequest request = (HttpRequest) msg;
+    if (!ServerRequestContext.canTrackRequest(channel)) {
+      channel.attr(PARENT_CONTEXT_ATTRIBUTE_KEY).remove();
+      ctx.fireChannelRead(msg);
+      return;
+    }
+
     final HttpHeaders headers = request.headers();
     final Context storedParentContext = channel.attr(PARENT_CONTEXT_ATTRIBUTE_KEY).getAndRemove();
     final Context parentContext =
@@ -49,11 +53,8 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
       DECORATE.afterStart(span);
       DECORATE.onRequest(span, channel, request, parentContext);
 
-      channel.attr(ANALYZED_RESPONSE_KEY).set(null);
-      channel.attr(BLOCKED_RESPONSE_KEY).set(null);
-
-      channel.attr(CONTEXT_ATTRIBUTE_KEY).set(context);
-      channel.attr(REQUEST_HEADERS_ATTRIBUTE_KEY).set(request.headers());
+      final ServerRequestContext serverContext =
+          ServerRequestContext.add(channel, context, request.headers());
 
       Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction();
       if (rba != null) {
@@ -76,7 +77,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
         DECORATE.onError(span, throwable);
         DECORATE.beforeFinish(ignored.context());
         span.finish(); // Finish the span manually since finishSpanOnClose was false
-        ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).remove();
+        ServerRequestContext.remove(ctx.channel(), serverContext);
         throw throwable;
       }
     }
@@ -88,12 +89,7 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception {
       super.channelInactive(ctx);
     } finally {
       try {
-        final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).getAndRemove();
-        final AgentSpan span = AgentSpan.fromContext(storedContext);
-        if (span != null && span.phasedFinish()) {
-          // at this point we can just publish this span to avoid loosing the rest of the trace
-          span.publish();
-        }
+        ServerRequestContext.closeAll(ctx.channel());
       } catch (final Throwable ignored) {
       }
     }

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

@@ -1,13 +1,13 @@
 package datadog.trace.instrumentation.netty41.server;
 
-import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY;
 import static datadog.trace.instrumentation.netty41.AttributeKeys.WEBSOCKET_SENDER_HANDLER_CONTEXT;
 import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.DECORATE;
 
 import datadog.context.Context;
 import datadog.context.ContextScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext;
+import datadog.trace.instrumentation.netty41.ServerRequestContext;
 import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelOutboundHandlerAdapter;
@@ -22,10 +22,16 @@ public class HttpServerResponseTracingHandler extends ChannelOutboundHandlerAdap
 
   @Override
   public void write(final ChannelHandlerContext ctx, final Object msg, final ChannelPromise prm) {
-    final Context storedContext = ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).get();
+    if (!(msg instanceof HttpResponse)) {
+      ctx.write(msg, prm);
+      return;
+    }
+
+    final ServerRequestContext serverContext = ServerRequestContext.nextResponse(ctx.channel());
+    final Context storedContext = serverContext == null ? null : serverContext.tracingContext();
     final AgentSpan span = AgentSpan.fromContext(storedContext);
 
-    if (span == null || !(msg instanceof HttpResponse)) {
+    if (span == null) {
       ctx.write(msg, prm);
       return;
     }
@@ -39,7 +45,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann
         DECORATE.onError(span, throwable);
         span.setHttpStatusCode(500);
         span.finish(); // Finish the span manually since finishSpanOnClose was false
-        ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).remove();
+        ServerRequestContext.remove(ctx.channel(), serverContext);
         throw throwable;
       }
       final boolean isWebsocketUpgrade =
@@ -55,7 +61,7 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann
         DECORATE.onResponse(span, response);
         DECORATE.beforeFinish(scope.context());
         span.finish(); // Finish the span manually since finishSpanOnClose was false
-        ctx.channel().attr(CONTEXT_ATTRIBUTE_KEY).remove();
+        ServerRequestContext.remove(ctx.channel(), serverContext);
       }
     }
   }

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

@@ -1,9 +1,5 @@
 package datadog.trace.instrumentation.netty41.server;
 
-import static datadog.trace.instrumentation.netty41.AttributeKeys.ANALYZED_RESPONSE_KEY;
-import static datadog.trace.instrumentation.netty41.AttributeKeys.BLOCKED_RESPONSE_KEY;
-import static datadog.trace.instrumentation.netty41.AttributeKeys.CONTEXT_ATTRIBUTE_KEY;
-import static datadog.trace.instrumentation.netty41.AttributeKeys.REQUEST_HEADERS_ATTRIBUTE_KEY;
 import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.DECORATE;
 import static io.netty.handler.codec.http.HttpHeaders.setContentLength;
 
@@ -13,6 +9,7 @@
 import datadog.trace.api.gateway.RequestContext;
 import datadog.trace.bootstrap.blocking.BlockingActionHelper;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.instrumentation.netty41.ServerRequestContext;
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelHandlerContext;
@@ -36,36 +33,21 @@ public class MaybeBlockResponseHandler extends ChannelOutboundHandlerAdapter {
 
   private MaybeBlockResponseHandler() {}
 
-  private static boolean isAnalyzedResponse(Channel ch) {
-    return ch.attr(ANALYZED_RESPONSE_KEY).get() != null;
-  }
-
-  private static void markAnalyzedResponse(Channel ch) {
-    ch.attr(ANALYZED_RESPONSE_KEY).set(Boolean.TRUE);
-  }
-
-  private static boolean isBlockedResponse(Channel ch) {
-    return ch.attr(BLOCKED_RESPONSE_KEY).get() != null;
-  }
-
-  private static void markBlockedResponse(Channel ch) {
-    ch.attr(BLOCKED_RESPONSE_KEY).set(Boolean.TRUE);
-  }
-
   @Override
   public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) throws Exception {
     Channel channel = ctx.channel();
 
-    Context storedContext = channel.attr(CONTEXT_ATTRIBUTE_KEY).get();
+    ServerRequestContext serverContext = ServerRequestContext.nextResponse(channel);
+    Context storedContext = serverContext == null ? null : serverContext.tracingContext();
     AgentSpan span = AgentSpan.fromContext(storedContext);
     RequestContext requestContext;
     if (span == null || (requestContext = span.getRequestContext()) == null) {
       super.write(ctx, msg, prm);
       return;
     }
 
-    if (isAnalyzedResponse(channel)) {
-      if (isBlockedResponse(channel)) {
+    if (serverContext.isResponseAnalyzed()) {
+      if (serverContext.isResponseBlocked()) {
         // block further writes
         log.debug("Write suppressed, msg {} dropped", msg);
         ReferenceCountUtil.release(msg);
@@ -88,14 +70,14 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) thr
     Flow<Void> flow =
         DECORATE.callIGCallbackResponseAndHeaders(
             span, origResponse, origResponse.getStatus().code(), ResponseExtractAdapter.GETTER);
-    markAnalyzedResponse(channel);
+    serverContext.markResponseAnalyzed();
     Flow.Action action = flow.getAction();
     if (!(action instanceof Flow.Action.RequestBlockingAction)) {
       super.write(ctx, msg, prm);
       return;
     }
 
-    markBlockedResponse(channel);
+    serverContext.markResponseBlocked();
     Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) action;
     int httpCode = BlockingActionHelper.getHttpCode(rba.getStatusCode());
     HttpResponseStatus httpResponseStatus = HttpResponseStatus.valueOf(httpCode);
@@ -112,7 +94,7 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) thr
 
     BlockingContentType bct = rba.getBlockingContentType();
     if (bct != BlockingContentType.NONE) {
-      HttpHeaders reqHeaders = ctx.attr(REQUEST_HEADERS_ATTRIBUTE_KEY).get();
+      HttpHeaders reqHeaders = serverContext.requestHeaders();
       String acceptHeader = reqHeaders != null ? reqHeaders.get("accept") : null;
       BlockingActionHelper.TemplateType type =
           BlockingActionHelper.determineTemplateType(bct, acceptHeader);

dd-java-agent/instrumentation/netty/netty-4.1/src/test/java/datadog/trace/instrumentation/netty41/ServerRequestContextTest.java

@@ -0,0 +1,28 @@
+package datadog.trace.instrumentation.netty41;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+import datadog.context.Context;
+import io.netty.handler.codec.http.DefaultHttpHeaders;
+import io.netty.util.DefaultAttributeMap;
+import org.junit.jupiter.api.Test;
+
+class ServerRequestContextTest {
+  private static final int PIPELINING_LIMIT = 1000;
+
+  @Test
+  void disablesTrackingWhenPipeliningLimitIsExceeded() {
+    DefaultAttributeMap attributes = new DefaultAttributeMap();
+
+    for (int i = 0; i < PIPELINING_LIMIT; i++) {
+      assertNotNull(ServerRequestContext.add(attributes, Context.root(), new DefaultHttpHeaders()));
+    }
+
+    assertFalse(ServerRequestContext.canTrackRequest(attributes));
+    assertNull(ServerRequestContext.nextResponse(attributes));
+    assertNull(attributes.attr(AttributeKeys.CONTEXT_ATTRIBUTE_KEY).get());
+    assertNull(ServerRequestContext.add(attributes, Context.root(), new DefaultHttpHeaders()));
+  }
+}