Respect previous baggage values in extracting NIO handler

Author: slashmo

Sources/NIOInstrumentation/HeaderExtractingHTTPServerHandler.swift

@@ -29,9 +29,11 @@ public final class HeaderExtractingHTTPServerHandler: ChannelInboundHandler {
 
     public func channelRead(context: ChannelHandlerContext, data: NIOAny) {
         if case .head(let head) = self.unwrapInboundIn(data) {
-            var baggage = Baggage.topLevel
-            InstrumentationSystem.instrument.extract(head.headers, into: &baggage, using: HTTPHeadersExtractor())
-            context.baggage = baggage
+            InstrumentationSystem.instrument.extract(
+                head.headers,
+                into: &context.baggage,
+                using: HTTPHeadersExtractor()
+            )
         }
 
         context.fireChannelRead(data)

Tests/NIOInstrumentationTests/HeaderExtractingHTTPServerHandlerTests+XCTest.swift

@@ -26,6 +26,7 @@ extension HeaderExtractingHTTPServerHandlerTests {
    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/HeaderExtractingHTTPServerHandlerTests.swift

@@ -29,8 +29,7 @@ final class HeaderExtractingHTTPServerHandlerTests: XCTestCase {
 
         let traceID = "abc"
         let handler = HeaderExtractingHTTPServerHandler()
-        let loop = EmbeddedEventLoop()
-        let channel = EmbeddedChannel(handler: handler, loop: loop)
+        let channel = EmbeddedChannel(handler: handler, loop: EmbeddedEventLoop())
 
         var requestHead = HTTPRequestHead(version: .init(major: 1, minor: 1), method: .GET, uri: "/")
         requestHead.headers = [FakeTracer.headerName: traceID]
@@ -42,6 +41,25 @@ final class HeaderExtractingHTTPServerHandlerTests: XCTestCase {
         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 {
         InstrumentationSystem.bootstrapInternal(FakeTracer())
 
@@ -59,3 +77,7 @@ final class HeaderExtractingHTTPServerHandlerTests: XCTestCase {
         XCTAssertNotNil(try channel.readInbound(as: HTTPServerRequestPart.self))
     }
 }
+
+private enum TestKey: Baggage.Key {
+    typealias Value = String
+}