Merge pull request #152 from slashmo/fix/nio-handlers

Correctly forward channel events in NIO instrumentation handlers

Author: slashmo

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

@@ -16,17 +16,26 @@ 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) {
+            InstrumentationSystem.instrument.extract(
+                head.headers,
+                into: &context.baggage,
+                using: HTTPHeadersExtractor()
+            )
+        }
+
         context.fireChannelRead(data)
     }
 }

Sources/NIOInstrumentation/HTTPHeadersInjectingHandler.swift renamed to Sources/NIOInstrumentation/HeaderInjectingHTTPClientHandler.swift

@@ -16,15 +16,18 @@ import Instrumentation
 import NIO
 import NIOHTTP1
 
-public final class HTTPHeadersInjectingHandler: ChannelOutboundHandler {
+public final class HeaderInjectingHTTPClientHandler: ChannelOutboundHandler {
     public typealias OutboundIn = HTTPClientRequestPart
     public typealias OutboundOut = HTTPClientRequestPart
 
     public init() {}
 
     public func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
-        let requestPart = unwrapOutboundIn(data)
-        guard case .head(var head) = requestPart else { return }
+        guard case .head(var head) = self.unwrapOutboundIn(data) else {
+            context.write(data, promise: promise)
+            return
+        }
+
         InstrumentationSystem.instrument.inject(context.baggage, into: &head.headers, using: HTTPHeadersInjector())
         context.write(self.wrapOutboundOut(.head(head)), promise: promise)
     }

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(HeaderInjectingHTTPClientHandlerTests.allTests),
              testCase(InstrumentTests.allTests),
              testCase(InstrumentationSystemTests.allTests),
              testCase(SpanAttributeSemanticsTests.allTests),

Tests/NIOInstrumentationTests/HTTPHeadersExtractInjectTests.swift

@@ -26,9 +26,9 @@ final class HTTPHeadersExtractInjectTests: XCTestCase {
     func test_extracted_baggage_into_subsequent_request_headers() throws {
         InstrumentationSystem.bootstrapInternal(FakeTracer())
 
-        let outboundHandler = HTTPHeadersInjectingHandler()
+        let outboundHandler = HeaderInjectingHTTPClientHandler()
         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.swift

@@ -11,7 +11,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// HTTPHeadersExtractingHandlerTests+XCTest.swift
+// HeaderExtractingHTTPServerHandlerTests+XCTest.swift
 //
 import XCTest
 ///
@@ -20,12 +20,14 @@ 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_respects_previous_baggage_values", test_respects_previous_baggage_values),
+                ("test_forwards_all_read_events", test_forwards_all_read_events),
            ]
    }
 }

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

@@ -0,0 +1,86 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift Distributed Tracing open source project
+//
+// Copyright (c) 2020 Moritz Lang and the Swift Tracing project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import Baggage
+@testable import Instrumentation
+import NIO
+import NIOHTTP1
+import NIOInstrumentation
+import XCTest
+
+final class HeaderExtractingHTTPServerHandlerTests: XCTestCase {
+    override class func tearDown() {
+        super.tearDown()
+        InstrumentationSystem.bootstrapInternal(nil)
+    }
+
+    func test_extracts_http_request_headers_into_baggage() throws {
+        InstrumentationSystem.bootstrapInternal(FakeTracer())
+
+        let traceID = "abc"
+        let handler = HeaderExtractingHTTPServerHandler()
+        let channel = EmbeddedChannel(handler: handler, loop: EmbeddedEventLoop())
+
+        var requestHead = HTTPRequestHead(version: .init(major: 1, minor: 1), method: .GET, uri: "/")
+        requestHead.headers = [FakeTracer.headerName: traceID]
+
+        XCTAssertNil(channel._channelCore.baggage[FakeTracer.TraceIDKey.self])
+
+        try channel.writeInbound(HTTPServerRequestPart.head(requestHead))
+
+        XCTAssertEqual(channel._channelCore.baggage[FakeTracer.TraceIDKey.self], traceID)
+    }
+
+    func test_respects_previous_baggage_values() throws {
+        InstrumentationSystem.bootstrapInternal(FakeTracer())
+
+        let traceID = "abc"
+        let handler = HeaderExtractingHTTPServerHandler()
+        let channel = EmbeddedChannel(handler: handler, loop: EmbeddedEventLoop())
+        channel._channelCore.baggage[TestKey.self] = "test"
+
+        var requestHead = HTTPRequestHead(version: .init(major: 1, minor: 1), method: .GET, uri: "/")
+        requestHead.headers = [FakeTracer.headerName: traceID]
+
+        XCTAssertNil(channel._channelCore.baggage[FakeTracer.TraceIDKey.self])
+
+        try channel.writeInbound(HTTPServerRequestPart.head(requestHead))
+
+        XCTAssertEqual(channel._channelCore.baggage[FakeTracer.TraceIDKey.self], traceID)
+        XCTAssertEqual(channel._channelCore.baggage[TestKey.self], "test")
+    }
+
+    func test_forwards_all_read_events() throws {
+        let channel = EmbeddedChannel(
+            handler: HeaderExtractingHTTPServerHandler(),
+            loop: EmbeddedEventLoop()
+        )
+
+        let requestHead = HTTPRequestHead(version: .init(major: 1, minor: 1), method: .GET, uri: "/")
+        let head = HTTPServerRequestPart.head(requestHead)
+        try channel.writeInbound(head)
+        XCTAssertEqual(try channel.readInbound(), head)
+
+        let body = HTTPServerRequestPart.body(channel.allocator.buffer(string: "Test"))
+        try channel.writeInbound(body)
+        XCTAssertEqual(try channel.readInbound(), body)
+
+        let end = HTTPServerRequestPart.end(nil)
+        try channel.writeInbound(end)
+        XCTAssertEqual(try channel.readInbound(as: HTTPServerRequestPart.self), end)
+    }
+}
+
+private enum TestKey: Baggage.Key {
+    typealias Value = String
+}

Tests/NIOInstrumentationTests/HeaderExtractingHTTPServerHandlerTests.swift

@@ -11,7 +11,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// HTTPHeadersInjectingHandlerTests+XCTest.swift
+// HeaderInjectingHTTPClientHandlerTests+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 HTTPHeadersInjectingHandlerTests {
+extension HeaderInjectingHTTPClientHandlerTests {
 
    @available(*, deprecated, message: "not actually deprecated. Just deprecated to allow deprecated tests (which test deprecated functionality) without warnings")
-   static var allTests : [(String, (HTTPHeadersInjectingHandlerTests) -> () throws -> Void)] {
+   static var allTests : [(String, (HeaderInjectingHTTPClientHandlerTests) -> () throws -> Void)] {
       return [
                 ("test_injects_baggage_into_http_request_headers", test_injects_baggage_into_http_request_headers),
+                ("test_forwards_all_write_events", test_forwards_all_write_events),
            ]
    }
 }

Tests/NIOInstrumentationTests/HTTPHeadersInjectingHandlerTests+XCTest.swift renamed to Tests/NIOInstrumentationTests/HeaderInjectingHTTPClientHandlerTests+XCTest.swift

@@ -18,7 +18,9 @@ import NIOHTTP1
 import NIOInstrumentation
 import XCTest
 
-final class HTTPHeadersInjectingHandlerTests: XCTestCase {
+final class HeaderInjectingHTTPClientHandlerTests: XCTestCase {
+    private let httpVersion = HTTPVersion(major: 1, minor: 1)
+
     override class func tearDown() {
         super.tearDown()
         InstrumentationSystem.bootstrapInternal(nil)
@@ -32,19 +34,37 @@ final class HTTPHeadersInjectingHandlerTests: XCTestCase {
         var baggage = Baggage.topLevel
         baggage[FakeTracer.TraceIDKey.self] = traceID
 
-        let httpVersion = HTTPVersion(major: 1, minor: 1)
-        let handler = HTTPHeadersInjectingHandler()
+        let handler = HeaderInjectingHTTPClientHandler()
         let loop = EmbeddedEventLoop()
         let channel = EmbeddedChannel(handler: handler, loop: loop)
         channel._channelCore.baggage = baggage
         let requestHead = HTTPRequestHead(version: httpVersion, method: .GET, uri: "/")
 
         try channel.writeOutbound(HTTPClientRequestPart.head(requestHead))
-        let modifiedRequestPart = try channel.readOutbound(as: HTTPClientRequestPart.self)
 
         XCTAssertEqual(
-            modifiedRequestPart,
-            .head(.init(version: httpVersion, method: .GET, uri: "/", headers: [FakeTracer.headerName: traceID]))
+            try channel.readOutbound(as: HTTPClientRequestPart.self),
+            .head(.init(version: self.httpVersion, method: .GET, uri: "/", headers: [FakeTracer.headerName: traceID]))
+        )
+    }
+
+    func test_forwards_all_write_events() throws {
+        let channel = EmbeddedChannel(
+            handler: HeaderInjectingHTTPClientHandler(),
+            loop: EmbeddedEventLoop()
         )
+
+        let requestHead = HTTPRequestHead(version: httpVersion, method: .GET, uri: "/")
+        let head = HTTPClientRequestPart.head(requestHead)
+        try channel.writeOutbound(head)
+        XCTAssertEqual(try channel.readOutbound(), head)
+
+        let body = HTTPClientRequestPart.body(.byteBuffer(channel.allocator.buffer(string: "test")))
+        try channel.writeOutbound(body)
+        XCTAssertEqual(try channel.readOutbound(), body)
+
+        let end = HTTPClientRequestPart.end(nil)
+        try channel.writeOutbound(end)
+        XCTAssertEqual(try channel.readOutbound(), end)
     }
 }