Inject instrumentation metadata into HTTP headers

Motivation:

In order to instrument distributed systems, metadata such as trace ids
must be propagated across network boundaries.
As HTTPClient operates at one such boundary, it should take care of
injecting metadata into HTTP headers automatically using the configured
instrument.

Modifications:

HTTPClient gains new method overloads accepting LoggingContext.

Result:

- New HTTPClient method overloads accepting LoggingContext
- Existing overloads accepting Logger construct a DefaultLoggingContext
- Existing methods that neither take Logger nor LoggingContext construct
  a DefaultLoggingContext

Author: slashmo

Package.swift

@@ -26,12 +26,13 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-nio-extras.git", from: "1.3.0"),
         .package(url: "https://github.com/apple/swift-nio-transport-services.git", from: "1.5.1"),
         .package(url: "https://github.com/apple/swift-log.git", from: "1.4.0"),
+        .package(url: "https://github.com/apple/swift-distributed-tracing.git", from: "0.1.1"),
     ],
     targets: [
         .target(
             name: "AsyncHTTPClient",
             dependencies: ["NIO", "NIOHTTP1", "NIOSSL", "NIOConcurrencyHelpers", "NIOHTTPCompression",
-                           "NIOFoundationCompat", "NIOTransportServices", "Logging"]
+                           "NIOFoundationCompat", "NIOTransportServices", "Logging", "Instrumentation"]
         ),
         .testTarget(
             name: "AsyncHTTPClientTests",

Sources/AsyncHTTPClient/HTTPClient.swift

@@ -12,7 +12,9 @@
 //
 //===----------------------------------------------------------------------===//
 
+import Baggage
 import Foundation
+import Instrumentation
 import Logging
 import NIO
 import NIOConcurrencyHelpers
@@ -71,6 +73,7 @@ public class HTTPClient {
     private let stateLock = Lock()
 
     internal static let loggingDisabled = Logger(label: "AHC-do-not-log", factory: { _ in SwiftLogNoOpLogHandler() })
+    internal static let topLevelContextLoggingDisabled = DefaultLoggingContext(logger: loggingDisabled, baggage: .topLevel)
 
     /// Create an `HTTPClient` with specified `EventLoopGroup` provider and configuration.
     ///
@@ -227,7 +230,17 @@ public class HTTPClient {
     ///     - url: Remote URL.
     ///     - deadline: Point in time by which the request must complete.
     public func get(url: String, deadline: NIODeadline? = nil) -> EventLoopFuture<Response> {
-        return self.get(url: url, deadline: deadline, logger: HTTPClient.loggingDisabled)
+        return self.get(url: url, context: HTTPClient.topLevelContextLoggingDisabled, deadline: deadline)
+    }
+
+    /// Execute `GET` request using specified URL.
+    ///
+    /// - parameters:
+    ///     - url: Remote URL.
+    ///     - context: The logging context carrying a logger and instrumentation metadata.
+    ///     - deadline: Point in time by which the request must complete.
+    public func get(url: String, context: LoggingContext, deadline: NIODeadline? = nil) -> EventLoopFuture<Response> {
+        return self.execute(.GET, url: url, context: context, deadline: deadline)
     }
 
     /// Execute `GET` request using specified URL.
@@ -237,7 +250,7 @@ public class HTTPClient {
     ///     - deadline: Point in time by which the request must complete.
     ///     - logger: The logger to use for this request.
     public func get(url: String, deadline: NIODeadline? = nil, logger: Logger) -> EventLoopFuture<Response> {
-        return self.execute(.GET, url: url, deadline: deadline, logger: logger)
+        return self.get(url: url, context: DefaultLoggingContext(logger: logger, baggage: .topLevel), deadline: deadline)
     }
 
     /// Execute `POST` request using specified URL.
@@ -247,7 +260,18 @@ public class HTTPClient {
     ///     - body: Request body.
     ///     - deadline: Point in time by which the request must complete.
     public func post(url: String, body: Body? = nil, deadline: NIODeadline? = nil) -> EventLoopFuture<Response> {
-        return self.post(url: url, body: body, deadline: deadline, logger: HTTPClient.loggingDisabled)
+        return self.post(url: url, context: HTTPClient.topLevelContextLoggingDisabled, body: body, deadline: deadline)
+    }
+
+    /// Execute `POST` request using specified URL.
+    ///
+    /// - parameters:
+    ///     - url: Remote URL.
+    ///     - context: The logging context carrying a logger and instrumentation metadata.
+    ///     - body: Request body.
+    ///     - deadline: Point in time by which the request must complete.
+    public func post(url: String, context: LoggingContext, body: Body? = nil, deadline: NIODeadline? = nil) -> EventLoopFuture<Response> {
+        return self.execute(.POST, url: url, context: context, body: body, deadline: deadline)
     }
 
     /// Execute `POST` request using specified URL.
@@ -258,7 +282,12 @@ public class HTTPClient {
     ///     - deadline: Point in time by which the request must complete.
     ///     - logger: The logger to use for this request.
     public func post(url: String, body: Body? = nil, deadline: NIODeadline? = nil, logger: Logger) -> EventLoopFuture<Response> {
-        return self.execute(.POST, url: url, body: body, deadline: deadline, logger: logger)
+        return self.post(
+            url: url,
+            context: DefaultLoggingContext(logger: logger, baggage: .topLevel),
+            body: body,
+            deadline: deadline
+        )
     }
 
     /// Execute `PATCH` request using specified URL.
@@ -268,7 +297,18 @@ public class HTTPClient {
     ///     - body: Request body.
     ///     - deadline: Point in time by which the request must complete.
     public func patch(url: String, body: Body? = nil, deadline: NIODeadline? = nil) -> EventLoopFuture<Response> {
-        return self.patch(url: url, body: body, deadline: deadline, logger: HTTPClient.loggingDisabled)
+        return self.patch(url: url, context: HTTPClient.topLevelContextLoggingDisabled, body: body, deadline: deadline)
+    }
+
+    /// Execute `PATCH` request using specified URL.
+    ///
+    /// - parameters:
+    ///     - url: Remote URL.
+    ///     - context: The logging context carrying a logger and instrumentation metadata.
+    ///     - body: Request body.
+    ///     - deadline: Point in time by which the request must complete.
+    public func patch(url: String, context: LoggingContext, body: Body? = nil, deadline: NIODeadline? = nil) -> EventLoopFuture<Response> {
+        return self.execute(.PATCH, url: url, context: context, body: body, deadline: deadline)
     }
 
     /// Execute `PATCH` request using specified URL.
@@ -279,7 +319,12 @@ public class HTTPClient {
     ///     - deadline: Point in time by which the request must complete.
     ///     - logger: The logger to use for this request.
     public func patch(url: String, body: Body? = nil, deadline: NIODeadline? = nil, logger: Logger) -> EventLoopFuture<Response> {
-        return self.execute(.PATCH, url: url, body: body, deadline: deadline, logger: logger)
+        return self.patch(
+            url: url,
+            context: DefaultLoggingContext(logger: logger, baggage: .topLevel),
+            body: body,
+            deadline: deadline
+        )
     }
 
     /// Execute `PUT` request using specified URL.
@@ -289,7 +334,18 @@ public class HTTPClient {
     ///     - body: Request body.
     ///     - deadline: Point in time by which the request must complete.
     public func put(url: String, body: Body? = nil, deadline: NIODeadline? = nil) -> EventLoopFuture<Response> {
-        return self.put(url: url, body: body, deadline: deadline, logger: HTTPClient.loggingDisabled)
+        return self.put(url: url, context: HTTPClient.topLevelContextLoggingDisabled, body: body, deadline: deadline)
+    }
+
+    /// Execute `PUT` request using specified URL.
+    ///
+    /// - parameters:
+    ///     - url: Remote URL.
+    ///     - context: The logging context carrying a logger and instrumentation metadata.
+    ///     - body: Request body.
+    ///     - deadline: Point in time by which the request must complete.
+    public func put(url: String, context: LoggingContext, body: Body? = nil, deadline: NIODeadline? = nil) -> EventLoopFuture<Response> {
+        return self.execute(.PUT, url: url, context: context, body: body, deadline: deadline)
     }
 
     /// Execute `PUT` request using specified URL.
@@ -300,7 +356,12 @@ public class HTTPClient {
     ///     - deadline: Point in time by which the request must complete.
     ///     - logger: The logger to use for this request.
     public func put(url: String, body: Body? = nil, deadline: NIODeadline? = nil, logger: Logger) -> EventLoopFuture<Response> {
-        return self.execute(.PUT, url: url, body: body, deadline: deadline, logger: logger)
+        return self.put(
+            url: url,
+            context: DefaultLoggingContext(logger: logger, baggage: .topLevel),
+            body: body,
+            deadline: deadline
+        )
     }
 
     /// Execute `DELETE` request using specified URL.
@@ -309,7 +370,17 @@ public class HTTPClient {
     ///     - url: Remote URL.
     ///     - deadline: The time when the request must have been completed by.
     public func delete(url: String, deadline: NIODeadline? = nil) -> EventLoopFuture<Response> {
-        return self.delete(url: url, deadline: deadline, logger: HTTPClient.loggingDisabled)
+        return self.delete(url: url, context: HTTPClient.topLevelContextLoggingDisabled, deadline: deadline)
+    }
+
+    /// Execute `DELETE` request using specified URL.
+    ///
+    /// - parameters:
+    ///     - url: Remote URL.
+    ///     - context: The logging context carrying a logger and instrumentation metadata.
+    ///     - deadline: Point in time by which the request must complete.
+    public func delete(url: String, context: LoggingContext, deadline: NIODeadline? = nil) -> EventLoopFuture<Response> {
+        return self.execute(.DELETE, url: url, context: context, deadline: deadline)
     }
 
     /// Execute `DELETE` request using specified URL.
@@ -319,7 +390,24 @@ public class HTTPClient {
     ///     - deadline: The time when the request must have been completed by.
     ///     - logger: The logger to use for this request.
     public func delete(url: String, deadline: NIODeadline? = nil, logger: Logger) -> EventLoopFuture<Response> {
-        return self.execute(.DELETE, url: url, deadline: deadline, logger: logger)
+        return self.delete(url: url, context: DefaultLoggingContext(logger: logger, baggage: .topLevel), deadline: deadline)
+    }
+
+    /// Execute arbitrary HTTP request using specified URL.
+    ///
+    /// - parameters:
+    ///     - method: Request method.
+    ///     - url: Request url.
+    ///     - body: Request body.
+    ///     - deadline: Point in time by which the request must complete.
+    ///     - context: The logging context carrying a logger and instrumentation metadata.
+    public func execute(_ method: HTTPMethod = .GET, url: String, context: LoggingContext?, body: Body? = nil, deadline: NIODeadline? = nil) -> EventLoopFuture<Response> {
+        do {
+            let request = try Request(url: url, method: method, body: body)
+            return self.execute(request: request, context: context, deadline: deadline)
+        } catch {
+            return self.eventLoopGroup.next().makeFailedFuture(error)
+        }
     }
 
     /// Execute arbitrary HTTP request using specified URL.
@@ -390,6 +478,17 @@ public class HTTPClient {
         return self.execute(request: request, deadline: deadline, logger: HTTPClient.loggingDisabled)
     }
 
+    /// Execute arbitrary HTTP request using specified URL.
+    ///
+    /// - parameters:
+    ///     - request: HTTP request to execute.
+    ///     - context: The logging context carrying a logger and instrumentation metadata.
+    ///     - deadline: Point in time by which the request must complete.
+    public func execute(request: Request, context: LoggingContext?, deadline: NIODeadline? = nil) -> EventLoopFuture<Response> {
+        let accumulator = ResponseAccumulator(request: request)
+        return self.execute(request: request, delegate: accumulator, context: context, deadline: deadline).futureResult
+    }
+
     /// Execute arbitrary HTTP request using specified URL.
     ///
     /// - parameters:
@@ -441,6 +540,20 @@ public class HTTPClient {
         return self.execute(request: request, delegate: delegate, deadline: deadline, logger: HTTPClient.loggingDisabled)
     }
 
+    /// Execute arbitrary HTTP request and handle response processing using provided delegate.
+    ///
+    /// - parameters:
+    ///     - request: HTTP request to execute.
+    ///     - delegate: Delegate to process response parts.
+    ///     - deadline: Point in time by which the request must complete.
+    ///     - context: The logging context carrying a logger and instrumentation metadata.
+    public func execute<Delegate: HTTPClientResponseDelegate>(request: Request,
+                                                              delegate: Delegate,
+                                                              context: LoggingContext?,
+                                                              deadline: NIODeadline? = nil) -> Task<Delegate.Response> {
+        return self.execute(request: request, delegate: delegate, eventLoop: .indifferent, context: context, deadline: deadline)
+    }
+
     /// Execute arbitrary HTTP request and handle response processing using provided delegate.
     ///
     /// - parameters:
@@ -474,6 +587,30 @@ public class HTTPClient {
                             logger: HTTPClient.loggingDisabled)
     }
 
+    /// Execute arbitrary HTTP request and handle response processing using provided delegate.
+    ///
+    /// - parameters:
+    ///     - request: HTTP request to execute.
+    ///     - delegate: Delegate to process response parts.
+    ///     - eventLoop: NIO Event Loop preference.
+    ///     - deadline: Point in time by which the request must complete.
+    ///     - context: The logging context carrying a logger and instrumentation metadata.
+    public func execute<Delegate: HTTPClientResponseDelegate>(request: Request,
+                                                              delegate: Delegate,
+                                                              eventLoop eventLoopPreference: EventLoopPreference,
+                                                              context: LoggingContext?,
+                                                              deadline: NIODeadline? = nil) -> Task<Delegate.Response> {
+        var request = request
+        if let baggage = context?.baggage {
+            InstrumentationSystem.instrument.inject(baggage, into: &request.headers, using: HTTPHeadersInjector())
+        }
+        return self.execute(request: request,
+                            delegate: delegate,
+                            eventLoop: eventLoopPreference,
+                            deadline: deadline,
+                            logger: context?.logger)
+    }
+
     /// Execute arbitrary HTTP request and handle response processing using provided delegate.
     ///
     /// - parameters:

Sources/AsyncHTTPClient/HTTPHeadersInjector.swift

@@ -0,0 +1,25 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the AsyncHTTPClient open source project
+//
+// Copyright (c) 2020 Apple Inc. and the AsyncHTTPClient project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of AsyncHTTPClient project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import Instrumentation
+import NIOHTTP1
+
+/// Injects values into `NIOHTTP1.HTTPHeaders`.
+struct HTTPHeadersInjector: Injector {
+    init() {}
+
+    func inject(_ value: String, forKey key: String, into headers: inout HTTPHeaders) {
+        headers.replaceOrAdd(name: key, value: value)
+    }
+}

Tests/AsyncHTTPClientTests/HTTPClientTestUtils.swift

@@ -13,7 +13,9 @@
 //===----------------------------------------------------------------------===//
 
 import AsyncHTTPClient
+import CoreBaggage
 import Foundation
+import Instrumentation
 import Logging
 import NIO
 import NIOConcurrencyHelpers
@@ -849,6 +851,52 @@ struct CollectEverythingLogHandler: LogHandler {
     }
 }
 
+internal enum TestInstrumentIDKey: Baggage.Key {
+    typealias Value = String
+    static let nameOverride: String? = "instrumentation-test-id"
+    static let headerName = "x-instrumentation-test-id"
+}
+
+internal extension Baggage {
+    var testInstrumentID: String? {
+        get {
+            return self[TestInstrumentIDKey.self]
+        }
+        set {
+            self[TestInstrumentIDKey.self] = newValue
+        }
+    }
+}
+
+internal final class TestInstrument: Instrument {
+    private(set) var carrierAfterInjection: Any?
+
+    func inject<Carrier, Inject>(
+        _ baggage: Baggage,
+        into carrier: inout Carrier,
+        using injector: Inject
+    ) where Carrier == Inject.Carrier, Inject: Injector {
+        if let testID = baggage.testInstrumentID {
+            injector.inject(testID, forKey: TestInstrumentIDKey.headerName, into: &carrier)
+            self.carrierAfterInjection = carrier
+        }
+    }
+
+    func extract<Carrier, Extract>(
+        _ carrier: Carrier,
+        into baggage: inout Baggage,
+        using extractor: Extract
+    ) where Carrier == Extract.Carrier, Extract: Extractor {
+        // no-op
+    }
+}
+
+internal extension InstrumentationSystem {
+    static var testInstrument: TestInstrument? {
+        InstrumentationSystem.instrument as? TestInstrument
+    }
+}
+
 private let cert = """
 -----BEGIN CERTIFICATE-----
 MIICmDCCAYACCQCPC8JDqMh1zzANBgkqhkiG9w0BAQsFADANMQswCQYDVQQGEwJ1

Tests/AsyncHTTPClientTests/HTTPClientTests+XCTest.swift

@@ -116,6 +116,7 @@ extension HTTPClientTests {
             ("testLoggingCorrectlyAttachesRequestInformation", testLoggingCorrectlyAttachesRequestInformation),
             ("testNothingIsLoggedAtInfoOrHigher", testNothingIsLoggedAtInfoOrHigher),
             ("testAllMethodsLog", testAllMethodsLog),
+            ("testRequestWithBaggage", testRequestWithBaggage),
             ("testClosingIdleConnectionsInPoolLogsInTheBackground", testClosingIdleConnectionsInPoolLogsInTheBackground),
             ("testUploadStreamingNoLength", testUploadStreamingNoLength),
             ("testConnectErrorPropagatedToDelegate", testConnectErrorPropagatedToDelegate),