Forward all read events in extracting NIO handler

Author: slashmo

Sources/NIOInstrumentation/HTTPHeadersExtractingHandler.swift renamed to Sources/NIOInstrumentation/HeaderExtractingHTTPServerHandler.swift

@@ -16,17 +16,24 @@ import Instrumentation
 import NIO
 import NIOHTTP1
 
-public final class HTTPHeadersExtractingHandler: ChannelInboundHandler {
+import Baggage
+import Instrumentation
+import NIO
+import NIOHTTP1
+
+public final class HeaderExtractingHTTPServerHandler: ChannelInboundHandler {
     public typealias InboundIn = HTTPServerRequestPart
     public typealias InboundOut = HTTPServerRequestPart
 
     public init() {}
 
     public func channelRead(context: ChannelHandlerContext, data: NIOAny) {
-        guard case .head(let head) = unwrapInboundIn(data) else { return }
-        var baggage = Baggage.topLevel
-        InstrumentationSystem.instrument.extract(head.headers, into: &baggage, using: HTTPHeadersExtractor())
-        context.baggage = baggage
+        if case .head(let head) = self.unwrapInboundIn(data) {
+            var baggage = Baggage.topLevel
+            InstrumentationSystem.instrument.extract(head.headers, into: &baggage, using: HTTPHeadersExtractor())
+            context.baggage = baggage
+        }
+
         context.fireChannelRead(data)
     }
 }

Tests/LinuxMain.swift

@@ -37,8 +37,8 @@ class LinuxMainRunnerImpl: LinuxMainRunner {
        XCTMain([
              testCase(HTTPHeadersCarrierTests.allTests),
              testCase(HTTPHeadersExtractInjectTests.allTests),
-             testCase(HTTPHeadersExtractingHandlerTests.allTests),
              testCase(HTTPHeadersInjectingHandlerTests.allTests),
+             testCase(HeaderExtractingHTTPServerHandlerTests.allTests),
              testCase(InstrumentTests.allTests),
              testCase(InstrumentationSystemTests.allTests),
              testCase(SpanAttributeSemanticsTests.allTests),

Tests/NIOInstrumentationTests/HTTPHeadersExtractInjectTests.swift

@@ -28,7 +28,7 @@ final class HTTPHeadersExtractInjectTests: XCTestCase {
 
         let outboundHandler = HTTPHeadersInjectingHandler()
         let requestHandler = MockRequestHandler()
-        let inboundHandler = HTTPHeadersExtractingHandler()
+        let inboundHandler = HeaderExtractingHTTPServerHandler()
 
         let channel = EmbeddedChannel(loop: EmbeddedEventLoop())
         XCTAssertNoThrow(try channel.pipeline.addHandlers([outboundHandler, requestHandler, inboundHandler]).wait())

Tests/NIOInstrumentationTests/HTTPHeadersExtractingHandlerTests+XCTest.swift renamed to Tests/NIOInstrumentationTests/HeaderExtractingHTTPServerHandlerTests+XCTest.swift

@@ -11,7 +11,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// HTTPHeadersExtractingHandlerTests+XCTest.swift
+// HeaderExtractingHTTPServerHandlerTests+XCTest.swift
 //
 import XCTest
 ///
@@ -20,12 +20,13 @@ import XCTest
 /// Do NOT edit this file directly as it will be regenerated automatically when needed.
 ///
 
-extension HTTPHeadersExtractingHandlerTests {
+extension HeaderExtractingHTTPServerHandlerTests {
 
    @available(*, deprecated, message: "not actually deprecated. Just deprecated to allow deprecated tests (which test deprecated functionality) without warnings")
-   static var allTests : [(String, (HTTPHeadersExtractingHandlerTests) -> () throws -> Void)] {
+   static var allTests : [(String, (HeaderExtractingHTTPServerHandlerTests) -> () throws -> Void)] {
       return [
                 ("test_extracts_http_request_headers_into_baggage", test_extracts_http_request_headers_into_baggage),
+                ("test_forwards_all_read_events", test_forwards_all_read_events),
            ]
    }
 }

Tests/NIOInstrumentationTests/HTTPHeadersExtractingHandlerTests.swift renamed to Tests/NIOInstrumentationTests/HeaderExtractingHTTPServerHandlerTests.swift

@@ -18,7 +18,7 @@ import NIOHTTP1
 import NIOInstrumentation
 import XCTest
 
-final class HTTPHeadersExtractingHandlerTests: XCTestCase {
+final class HeaderExtractingHTTPServerHandlerTests: XCTestCase {
     override class func tearDown() {
         super.tearDown()
         InstrumentationSystem.bootstrapInternal(nil)
@@ -28,7 +28,7 @@ final class HTTPHeadersExtractingHandlerTests: XCTestCase {
         InstrumentationSystem.bootstrapInternal(FakeTracer())
 
         let traceID = "abc"
-        let handler = HTTPHeadersExtractingHandler()
+        let handler = HeaderExtractingHTTPServerHandler()
         let loop = EmbeddedEventLoop()
         let channel = EmbeddedChannel(handler: handler, loop: loop)
 
@@ -41,4 +41,21 @@ final class HTTPHeadersExtractingHandlerTests: XCTestCase {
 
         XCTAssertEqual(channel._channelCore.baggage[FakeTracer.TraceIDKey.self], traceID)
     }
+
+    func test_forwards_all_read_events() throws {
+        InstrumentationSystem.bootstrapInternal(FakeTracer())
+
+        let channel = EmbeddedChannel(loop: EmbeddedEventLoop())
+        try channel.pipeline.addHandlers(HeaderExtractingHTTPServerHandler()).wait()
+
+        let requestHead = HTTPRequestHead(version: .init(major: 1, minor: 1), method: .GET, uri: "/")
+        try channel.writeInbound(HTTPServerRequestPart.head(requestHead))
+        XCTAssertNotNil(try channel.readInbound(as: HTTPServerRequestPart.self))
+
+        try channel.writeInbound(HTTPServerRequestPart.body(channel.allocator.buffer(string: "Test")))
+        XCTAssertNotNil(try channel.readInbound(as: HTTPServerRequestPart.self))
+
+        try channel.writeInbound(HTTPServerRequestPart.end(nil))
+        XCTAssertNotNil(try channel.readInbound(as: HTTPServerRequestPart.self))
+    }
 }