fix connection keep alive on netty 4.0 (#386)

hypertrace · Jun 27, 2023 · c9cb1fe · c9cb1fe
1 parent b741811
commit c9cb1fe
Show file tree

Hide file tree

Showing 2 changed files with 100 additions and 0 deletions.
diff --git a/...aagent/instrumentation/hypertrace/netty/v4_0/server/HttpServerResponseTracingHandler.java b/...aagent/instrumentation/hypertrace/netty/v4_0/server/HttpServerResponseTracingHandler.java
@@ -104,6 +104,14 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) {
       span.setStatus(code >= 100 && code < 500 ? StatusCode.UNSET : StatusCode.ERROR);
     }
     if (msg instanceof LastHttpContent) {
+      // When we end the span, we should set the server context and request attr to null so that
+      // for the next request a new context and request is made and stored in channel.
+      // Else, when using a Connection: keep-alive header, the same server context and request attr
+      // is used in the subsequent requests.
+      ctx.channel()
+          .attr(io.opentelemetry.javaagent.instrumentation.netty.v4_0.AttributeKeys.SERVER_CONTEXT)
+          .set(null);
+      ctx.channel().attr(AttributeKeys.REQUEST).set(null);
       span.end();
     }
   }

diff --git a/...nstrumentation/hypertrace/netty/v4_0/server/AbstractNetty40ServerInstrumentationTest.java b/...nstrumentation/hypertrace/netty/v4_0/server/AbstractNetty40ServerInstrumentationTest.java
@@ -206,4 +206,96 @@ public void blocking() throws IOException, TimeoutException, InterruptedExceptio
             .getAttributes()
             .get(HypertraceSemanticAttributes.httpResponseHeader(RESPONSE_BODY)));
   }
+
+  @Test
+  public void connectionKeepAlive() throws IOException, TimeoutException, InterruptedException {
+    Request request =
+        new Request.Builder()
+            .url(String.format("http://localhost:%d/post", port))
+            .header(REQUEST_HEADER_NAME, REQUEST_HEADER_VALUE)
+            .header("first", "1st")
+            .header("connection", "keep-alive")
+            .get()
+            .build();
+
+    try (Response response = httpClient.newCall(request).execute()) {
+      Assertions.assertEquals(200, response.code());
+      Assertions.assertEquals(RESPONSE_BODY, response.body().string());
+    }
+
+    List<List<SpanData>> traces = TEST_WRITER.getTraces();
+    TEST_WRITER.waitForTraces(1);
+    Assertions.assertEquals(1, traces.size());
+    List<SpanData> trace = traces.get(0);
+    Assertions.assertEquals(1, trace.size());
+    SpanData spanData = trace.get(0);
+
+    Assertions.assertEquals(
+        REQUEST_HEADER_VALUE,
+        spanData
+            .getAttributes()
+            .get(HypertraceSemanticAttributes.httpRequestHeader(REQUEST_HEADER_NAME)));
+    Assertions.assertEquals(
+        "1st",
+        spanData.getAttributes().get(HypertraceSemanticAttributes.httpRequestHeader("first")));
+    Assertions.assertEquals(
+        "keep-alive",
+        spanData.getAttributes().get(HypertraceSemanticAttributes.httpRequestHeader("connection")));
+    Assertions.assertEquals(
+        RESPONSE_HEADER_VALUE,
+        spanData
+            .getAttributes()
+            .get(HypertraceSemanticAttributes.httpResponseHeader(RESPONSE_HEADER_NAME)));
+    Assertions.assertNull(
+        spanData.getAttributes().get(HypertraceSemanticAttributes.HTTP_REQUEST_BODY));
+    Assertions.assertEquals(
+        RESPONSE_BODY,
+        spanData.getAttributes().get(HypertraceSemanticAttributes.HTTP_RESPONSE_BODY));
+
+    RequestBody requestBody = blockedRequestBody(true, 3000, 75);
+    Request request2 =
+        new Request.Builder()
+            .url(String.format("http://localhost:%d/post", port))
+            .header(REQUEST_HEADER_NAME, "REQUEST_HEADER_VALUE")
+            .header("second", "2nd")
+            .header("connection", "keep-alive")
+            .post(requestBody)
+            .build();
+
+    try (Response response = httpClient.newCall(request2).execute()) {
+      Assertions.assertEquals(403, response.code());
+      Assertions.assertTrue(response.body().string().isEmpty());
+    }
+
+    List<List<SpanData>> traces2 = TEST_WRITER.getTraces();
+    TEST_WRITER.waitForTraces(2);
+    Assertions.assertEquals(2, traces2.size());
+    List<SpanData> trace2 = traces2.get(1);
+    Assertions.assertEquals(1, trace2.size());
+    SpanData spanData2 = trace2.get(0);
+
+    Assertions.assertEquals(
+        "REQUEST_HEADER_VALUE",
+        spanData2
+            .getAttributes()
+            .get(HypertraceSemanticAttributes.httpRequestHeader(REQUEST_HEADER_NAME)));
+    Assertions.assertEquals(
+        "2nd",
+        spanData2.getAttributes().get(HypertraceSemanticAttributes.httpRequestHeader("second")));
+    Assertions.assertEquals(
+        "keep-alive",
+        spanData2
+            .getAttributes()
+            .get(HypertraceSemanticAttributes.httpRequestHeader("connection")));
+    Assertions.assertNull(
+        spanData2.getAttributes().get(HypertraceSemanticAttributes.httpRequestHeader("first")));
+    Assertions.assertNull(
+        spanData2
+            .getAttributes()
+            .get(HypertraceSemanticAttributes.httpResponseHeader(RESPONSE_HEADER_NAME)));
+    Assertions.assertNull(
+        spanData2
+            .getAttributes()
+            .get(HypertraceSemanticAttributes.httpResponseHeader(RESPONSE_BODY)));
+  }
 }