Keep Netty 4.1 server spans open until streaming responses complete

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

@@ -15,6 +15,7 @@
 import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.handler.codec.http.HttpHeaders;
 import io.netty.handler.codec.http.HttpRequest;
+import java.util.Deque;
 
 @ChannelHandler.Sharable
 public class HttpServerRequestTracingHandler extends ChannelInboundHandlerAdapter {
@@ -95,9 +96,43 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception {
       super.channelInactive(ctx);
     } finally {
       try {
-        ServerRequestContext.closeAll(ctx.channel());
+        final Deque<ServerRequestContext> storedContexts =
+            ServerRequestContext.removeAll(ctx.channel());
+        if (storedContexts != null) {
+          ServerRequestContext storedContext;
+          while ((storedContext = storedContexts.pollFirst()) != null) {
+            if (storedContext.isResponseStarted()) {
+              finishSpanOnChannelClose(storedContext);
+            } else {
+              publishSpanOnChannelClose(storedContext.tracingContext());
+            }
+          }
+        }
       } catch (final Throwable ignored) {
       }
     }
   }
+
+  private static void finishSpanOnChannelClose(final ServerRequestContext serverContext) {
+    final Context storedContext = serverContext.tracingContext();
+    final AgentSpan span = AgentSpan.fromContext(storedContext);
+    if (span == null) {
+      return;
+    }
+    try (final ContextScope ignored = storedContext.attach()) {
+      if (!serverContext.isBeforeFinishCalled()) {
+        serverContext.markBeforeFinishCalled();
+        DECORATE.beforeFinish(storedContext);
+      }
+      span.finish();
+    }
+  }
+
+  private static void publishSpanOnChannelClose(final Context storedContext) {
+    final AgentSpan span = AgentSpan.fromContext(storedContext);
+    if (span != null && span.phasedFinish()) {
+      // At this point we can just publish this span to avoid losing the rest of the trace.
+      span.publish();
+    }
+  }
 }

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

@@ -14,9 +14,11 @@
 import io.netty.channel.ChannelOutboundHandlerAdapter;
 import io.netty.channel.ChannelPromise;
 import io.netty.handler.codec.http.HttpHeaderNames;
+import io.netty.handler.codec.http.HttpHeaderValues;
 import io.netty.handler.codec.http.HttpResponse;
 import io.netty.handler.codec.http.HttpResponseStatus;
 import io.netty.handler.codec.http.LastHttpContent;
+import io.netty.util.concurrent.Future;
 
 @ChannelHandler.Sharable
 public class HttpServerResponseTracingHandler extends ChannelOutboundHandlerAdapter {
@@ -25,7 +27,8 @@ public class HttpServerResponseTracingHandler extends ChannelOutboundHandlerAdap
   @Override
   public void write(final ChannelHandlerContext ctx, final Object msg, final ChannelPromise prm) {
     final boolean isResponse = msg instanceof HttpResponse;
-    if (!isResponse && !(msg instanceof LastHttpContent)) {
+    final boolean isLastContent = msg instanceof LastHttpContent;
+    if (!isResponse && !isLastContent) {
       ctx.write(msg, prm);
       return;
     }
@@ -50,43 +53,99 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann
       return;
     }
 
-    try (final ContextScope scope = storedContext.attach()) {
+    try (final ContextScope ignored = storedContext.attach()) {
       final HttpResponse response = isResponse ? (HttpResponse) msg : null;
-      final boolean responseComplete = msg instanceof LastHttpContent;
-
+      final boolean websocketUpgrade = response != null && isWebsocketUpgrade(response);
+      final boolean informationalResponse =
+          response != null && isInformationalResponse(response) && !websocketUpgrade;
+      final boolean finishResponseOnWrite = isLastContent && !informationalResponse;
+      final ChannelPromise writePromise =
+          finishResponseOnWrite && prm.isVoid() ? ctx.newPromise() : prm;
       try {
-        ctx.write(msg, prm);
-      } catch (final Throwable throwable) {
-        DECORATE.onError(span, throwable);
-        span.setHttpStatusCode(500);
-        span.finish(); // Finish the span manually since finishSpanOnClose was false
-        removeServerContext(ctx, serverContext);
-        throw throwable;
-      }
-      if (response != null) {
-        final boolean isWebsocketUpgrade =
-            response.status() == HttpResponseStatus.SWITCHING_PROTOCOLS
-                && "websocket".equals(response.headers().get(HttpHeaderNames.UPGRADE));
-        if (isWebsocketUpgrade) {
-          ctx.channel()
-              .attr(WEBSOCKET_SENDER_HANDLER_CONTEXT)
-              .set(new HandlerContext.Sender(span, ctx.channel().id().asShortText()));
+        if (response != null && !informationalResponse) {
+          onResponse(ctx, span, serverContext, response, websocketUpgrade);
         }
-        if (isInformational(response) && !isWebsocketUpgrade) {
-          return;
+        if (finishResponseOnWrite) {
+          removeServerContext(ctx, serverContext);
+          writePromise.addListener(
+              future -> finishSpan(serverContext, storedContext, span, future));
         }
-        if (serverContext != null) {
-          serverContext.markResponseStarted();
+        ctx.write(msg, writePromise);
+        if (finishResponseOnWrite && (!writePromise.isDone() || writePromise.isSuccess())) {
+          final ServerRequestContext nextResponse =
+              ServerRequestContext.nextResponse(ctx.channel());
+          BlockingResponseHandler.maybeWriteDeferredBlockResponse(ctx, nextResponse);
         }
-        DECORATE.onResponse(span, response);
+      } catch (final Throwable throwable) {
+        if (!finishResponseOnWrite || !writePromise.isDone()) {
+          DECORATE.onError(span, throwable);
+          span.setHttpStatusCode(500);
+          if (!finishResponseOnWrite) {
+            removeServerContext(ctx, serverContext);
+          }
+          finishSpan(serverContext, storedContext, span);
+        }
+        throw throwable;
       }
-      if (responseComplete) {
-        DECORATE.beforeFinish(scope.context());
-        span.finish(); // Finish the span manually since finishSpanOnClose was false
-        removeServerContext(ctx, serverContext);
-        final ServerRequestContext nextResponse = ServerRequestContext.nextResponse(ctx.channel());
-        BlockingResponseHandler.maybeWriteDeferredBlockResponse(ctx, nextResponse);
+    }
+  }
+
+  private static void onResponse(
+      final ChannelHandlerContext ctx,
+      final AgentSpan span,
+      final ServerRequestContext serverContext,
+      final HttpResponse response,
+      final boolean websocketUpgrade) {
+    if (websocketUpgrade) {
+      ctx.channel()
+          .attr(WEBSOCKET_SENDER_HANDLER_CONTEXT)
+          .set(new HandlerContext.Sender(span, ctx.channel().id().asShortText()));
+    }
+    DECORATE.onResponse(span, response);
+    if (serverContext != null) {
+      serverContext.markResponseStarted();
+    }
+  }
+
+  private static boolean isInformationalResponse(final HttpResponse response) {
+    final int statusCode = response.status().code();
+    return statusCode >= 100 && statusCode < 200;
+  }
+
+  private static boolean isWebsocketUpgrade(final HttpResponse response) {
+    return response.status().code() == HttpResponseStatus.SWITCHING_PROTOCOLS.code()
+        && response
+            .headers()
+            .containsValue(HttpHeaderNames.UPGRADE, HttpHeaderValues.WEBSOCKET, true);
+  }
+
+  private static void finishSpan(
+      final ServerRequestContext serverContext,
+      final Context storedContext,
+      final AgentSpan span,
+      final Future<?> future) {
+    if (!future.isSuccess()) {
+      DECORATE.onError(span, future.cause());
+      span.setHttpStatusCode(500);
+    }
+    finishSpan(serverContext, storedContext, span);
+  }
+
+  private static void finishSpan(
+      final ServerRequestContext serverContext, final Context storedContext, final AgentSpan span) {
+    try (final ContextScope ignored = storedContext.attach()) {
+      beforeFinish(serverContext, storedContext);
+      span.finish(); // Finish the span manually since finishSpanOnClose was false
+    }
+  }
+
+  private static void beforeFinish(
+      final ServerRequestContext serverContext, final Context storedContext) {
+    if (serverContext == null || !serverContext.isBeforeFinishCalled()) {
+      if (serverContext != null) {
+        serverContext.markBeforeFinishCalled();
       }
+      DECORATE.beforeFinish(storedContext);
     }
   }
 
@@ -98,9 +157,4 @@ private static void removeServerContext(
       ServerRequestContext.remove(ctx.channel(), serverContext);
     }
   }
-
-  private static boolean isInformational(final HttpResponse response) {
-    final int statusCode = response.status().code();
-    return statusCode >= 100 && statusCode < 200;
-  }
 }

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

@@ -70,6 +70,10 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) thr
     }
     HttpResponse origResponse = (HttpResponse) msg;
     int statusCode = origResponse.status().code();
+    // Interim 1xx responses (e.g. 100 Continue, 103 Early Hints) precede the final response.
+    // Analyzing one here would consume the one-shot response analysis before the final response is
+    // written, so its status and headers would never be inspected. Switching Protocols (101) is
+    // terminal, so it is still analyzed.
     if (statusCode >= 100
         && statusCode < 200
         && statusCode != HttpResponseStatus.SWITCHING_PROTOCOLS.code()) {

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

@@ -13,6 +13,7 @@
 import static io.netty.handler.codec.http.HttpResponseStatus.OK;
 import static io.netty.handler.codec.http.HttpResponseStatus.SWITCHING_PROTOCOLS;
 import static io.netty.handler.codec.http.HttpVersion.HTTP_1_1;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertSame;
@@ -39,16 +40,22 @@
 class HttpServerResponseTracingHandlerTest extends AbstractInstrumentationTest {
 
   @Test
-  void finishesMirroredContextWhenRequestQueueIsAbsent() {
+  void finishesMirroredContextOnLastContentWhenRequestQueueIsAbsent() {
     EmbeddedChannel channel = new EmbeddedChannel(HttpServerResponseTracingHandler.INSTANCE);
     AgentSpan span = startSpan("netty", "mirrored-http2-server");
     channel.attr(CONTEXT_ATTRIBUTE_KEY).set(span);
 
-    assertTrue(channel.writeOutbound(new DefaultFullHttpResponse(HTTP_1_1, OK)));
+    HttpResponse response = new DefaultHttpResponse(HTTP_1_1, OK);
+    assertTrue(channel.writeOutbound(response));
 
-    FullHttpResponse response = channel.readOutbound();
-    assertNotNull(response);
-    response.release();
+    assertSame(span, channel.attr(CONTEXT_ATTRIBUTE_KEY).get());
+    HttpResponse forwarded = channel.readOutbound();
+    assertSame(response, forwarded);
+    ReferenceCountUtil.release(forwarded);
+
+    assertTrue(channel.writeOutbound(LastHttpContent.EMPTY_LAST_CONTENT));
+
+    ReferenceCountUtil.release(channel.readOutbound());
     assertNull(channel.attr(CONTEXT_ATTRIBUTE_KEY).get());
     channel.finishAndReleaseAll();
     assertTraces(trace(span().root().operationName("mirrored-http2-server")));
@@ -70,7 +77,7 @@ void forwardsLastContentBeforeFinalResponseWithoutCompletingContext() {
   }
 
   @Test
-  void headerOnlyResponseDoesNotCompleteContextBeforeLastContent() {
+  void headerOnlyResponseWaitsForLastContent() {
     EmbeddedChannel channel = new EmbeddedChannel(HttpServerResponseTracingHandler.INSTANCE);
     AgentSpan span = startSpan("netty", "header-only-server");
     ServerRequestContext serverContext = ServerRequestContext.add(channel, span, null);
@@ -84,6 +91,7 @@ void headerOnlyResponseDoesNotCompleteContextBeforeLastContent() {
     HttpResponse forwarded = channel.readOutbound();
     assertSame(response, forwarded);
     ReferenceCountUtil.release(forwarded);
+    assertNull(channel.readOutbound());
 
     assertTrue(channel.writeOutbound(LastHttpContent.EMPTY_LAST_CONTENT));
 
@@ -95,10 +103,12 @@ void headerOnlyResponseDoesNotCompleteContextBeforeLastContent() {
   }
 
   @Test
-  void rawFixedLengthBodyDoesNotCompleteContextBeforeLastContent() {
+  void rawFixedLengthBodyWaitsForLastContent() {
     EmbeddedChannel channel = new EmbeddedChannel(HttpServerResponseTracingHandler.INSTANCE);
     AgentSpan span = startSpan("netty", "raw-body-server");
     ServerRequestContext serverContext = ServerRequestContext.add(channel, span, null);
+    ServerRequestContext nextServerContext =
+        ServerRequestContext.add(channel, Context.root(), null);
     HttpResponse response = new DefaultHttpResponse(HTTP_1_1, OK);
     response.headers().set(CONTENT_LENGTH, 4);
 
@@ -108,15 +118,34 @@ void rawFixedLengthBodyDoesNotCompleteContextBeforeLastContent() {
 
     assertSame(serverContext, ServerRequestContext.nextResponse(channel));
     ReferenceCountUtil.release(channel.readOutbound());
+    assertNull(channel.readOutbound());
 
     assertTrue(channel.writeOutbound(LastHttpContent.EMPTY_LAST_CONTENT));
 
-    assertNull(ServerRequestContext.nextResponse(channel));
+    assertSame(nextServerContext, ServerRequestContext.nextResponse(channel));
     ReferenceCountUtil.release(channel.readOutbound());
+    ServerRequestContext.remove(channel, nextServerContext);
     channel.finishAndReleaseAll();
     assertTraces(trace(span().root().operationName("raw-body-server")));
   }
 
+  @Test
+  void doesNotThrowOnMalformedContentLength() {
+    EmbeddedChannel channel = new EmbeddedChannel(HttpServerResponseTracingHandler.INSTANCE);
+    AgentSpan span = startSpan("netty", "malformed-content-length-server");
+    ServerRequestContext.add(channel, span, null);
+    FullHttpResponse response = new DefaultFullHttpResponse(HTTP_1_1, OK);
+    response.headers().set(CONTENT_LENGTH, "malformed");
+
+    assertDoesNotThrow(() -> assertTrue(channel.writeOutbound(response)));
+
+    FullHttpResponse forwarded = channel.readOutbound();
+    assertSame(response, forwarded);
+    forwarded.release();
+    channel.finishAndReleaseAll();
+    assertTraces(trace(span().root().operationName("malformed-content-length-server")));
+  }
+
   @Test
   void nettyEncoderRequiresLastContentBeforeNextKeepAliveResponse() {
     EmbeddedChannel channel = new EmbeddedChannel(new HttpResponseEncoder());
@@ -159,4 +188,27 @@ void fullWebsocketUpgradeCompletesContextAndPreservesHandshakeSpan() {
     channel.finishAndReleaseAll();
     assertTraces(trace(span().root().operationName("websocket-handshake-server")));
   }
+
+  @Test
+  void fullNonWebsocketUpgradeWaitsForFinalResponse() {
+    EmbeddedChannel channel = new EmbeddedChannel(HttpServerResponseTracingHandler.INSTANCE);
+    AgentSpan span = startSpan("netty", "h2c-upgrade-server");
+    ServerRequestContext serverContext = ServerRequestContext.add(channel, span, null);
+    FullHttpResponse upgradeResponse = new DefaultFullHttpResponse(HTTP_1_1, SWITCHING_PROTOCOLS);
+    upgradeResponse.headers().set(UPGRADE, "h2c");
+
+    assertTrue(channel.writeOutbound(upgradeResponse));
+
+    assertSame(serverContext, ServerRequestContext.nextResponse(channel));
+    assertNull(channel.attr(WEBSOCKET_SENDER_HANDLER_CONTEXT).get());
+    ReferenceCountUtil.release(channel.readOutbound());
+
+    assertTrue(channel.writeOutbound(new DefaultFullHttpResponse(HTTP_1_1, OK)));
+
+    assertNull(ServerRequestContext.nextResponse(channel));
+    ReferenceCountUtil.release(channel.readOutbound());
+    channel.finishAndReleaseAll();
+
+    assertTraces(trace(span().root().operationName("h2c-upgrade-server")));
+  }
 }