Introduce built-in swift-distributed-tracing support (#857)

Co-authored-by: Moritz Lang <[email protected]>
Co-authored-by: George Barnett <[email protected]>

Author: 3 people

.licenseignore

@@ -16,6 +16,8 @@
 *.json
 Package.swift
 **/Package.swift
+Package@swift-*.swift
+**/Package@swift-*.swift
 Package@-*.swift
 **/Package@-*.swift
 Package.resolved

Package.swift

@@ -47,6 +47,7 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-log.git", from: "1.6.0"),
         .package(url: "https://github.com/apple/swift-atomics.git", from: "1.0.2"),
         .package(url: "https://github.com/apple/swift-algorithms.git", from: "1.0.0"),
+        .package(url: "https://github.com/apple/swift-distributed-tracing.git", from: "1.3.0"),
     ],
     targets: [
         .target(
@@ -70,9 +71,11 @@ let package = Package(
                 .product(name: "NIOHTTPCompression", package: "swift-nio-extras"),
                 .product(name: "NIOSOCKS", package: "swift-nio-extras"),
                 .product(name: "NIOTransportServices", package: "swift-nio-transport-services"),
-                .product(name: "Logging", package: "swift-log"),
                 .product(name: "Atomics", package: "swift-atomics"),
                 .product(name: "Algorithms", package: "swift-algorithms"),
+                // Observability support
+                .product(name: "Logging", package: "swift-log"),
+                .product(name: "Tracing", package: "swift-distributed-tracing"),
             ],
             swiftSettings: strictConcurrencySettings
         ),
@@ -89,9 +92,12 @@ let package = Package(
                 .product(name: "NIOSSL", package: "swift-nio-ssl"),
                 .product(name: "NIOHTTP2", package: "swift-nio-http2"),
                 .product(name: "NIOSOCKS", package: "swift-nio-extras"),
-                .product(name: "Logging", package: "swift-log"),
                 .product(name: "Atomics", package: "swift-atomics"),
                 .product(name: "Algorithms", package: "swift-algorithms"),
+                // Observability support
+                .product(name: "Logging", package: "swift-log"),
+                .product(name: "Tracing", package: "swift-distributed-tracing"),
+                .product(name: "InMemoryTracing", package: "swift-distributed-tracing"),
             ],
             resources: [
                 .copy("Resources/self_signed_cert.pem"),

Sources/AsyncHTTPClient/AsyncAwait/HTTPClient+execute.swift

@@ -15,6 +15,7 @@
 import Logging
 import NIOCore
 import NIOHTTP1
+import Tracing
 
 import struct Foundation.URL
 
@@ -36,12 +37,14 @@ extension HTTPClient {
         deadline: NIODeadline,
         logger: Logger? = nil
     ) async throws -> HTTPClientResponse {
-        try await self.executeAndFollowRedirectsIfNeeded(
-            request,
-            deadline: deadline,
-            logger: logger ?? Self.loggingDisabled,
-            redirectState: RedirectState(self.configuration.redirectConfiguration.mode, initialURL: request.url)
-        )
+        try await withRequestSpan(request) {
+            try await self.executeAndFollowRedirectsIfNeeded(
+                request,
+                deadline: deadline,
+                logger: logger ?? Self.loggingDisabled,
+                redirectState: RedirectState(self.configuration.redirectConfiguration.mode, initialURL: request.url)
+            )
+        }
     }
 }
 

Sources/AsyncHTTPClient/AsyncAwait/HTTPClient+tracing.swift

@@ -0,0 +1,41 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the AsyncHTTPClient open source project
+//
+// Copyright (c) 2021 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 NIOHTTP1
+import Tracing
+
+@available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
+extension HTTPClient {
+    @inlinable
+    func withRequestSpan(
+        _ request: HTTPClientRequest,
+        _ body: () async throws -> HTTPClientResponse
+    ) async rethrows -> HTTPClientResponse {
+        guard let tracer = self.tracer else {
+            return try await body()
+        }
+
+        return try await tracer.withSpan(request.method.rawValue, ofKind: .client) { span in
+            let keys = self.configuration.tracing.attributeKeys
+            span.attributes[keys.requestMethod] = request.method.rawValue
+            // TODO: set more attributes on the span
+            let response = try await body()
+
+            // set response span attributes
+            TracingSupport.handleResponseStatusCode(span, response.status, keys: tracing.attributeKeys)
+
+            return response
+        }
+    }
+}

Sources/AsyncHTTPClient/AsyncAwait/Transaction.swift

@@ -17,6 +17,7 @@ import NIOConcurrencyHelpers
 import NIOCore
 import NIOHTTP1
 import NIOSSL
+import Tracing
 
 @available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
 @usableFromInline

Sources/AsyncHTTPClient/HTTPClient.swift

@@ -23,6 +23,7 @@ import NIOPosix
 import NIOSSL
 import NIOTLS
 import NIOTransportServices
+import Tracing
 
 extension Logger {
     private func requestInfo(_ request: HTTPClient.Request) -> Logger.Metadata.Value {
@@ -62,15 +63,29 @@ public final class HTTPClient: Sendable {
     ///
     /// All HTTP transactions will occur on loops owned by this group.
     public let eventLoopGroup: EventLoopGroup
-    let configuration: Configuration
     let poolManager: HTTPConnectionPool.Manager
 
+    @usableFromInline
+    let configuration: Configuration
+
     /// Shared thread pool used for file IO. It is lazily created on first access of ``Task/fileIOThreadPool``.
     private let fileIOThreadPool: NIOLockedValueBox<NIOThreadPool?>
 
     private let state: NIOLockedValueBox<State>
     private let canBeShutDown: Bool
 
+    /// Tracer configured for this HTTPClient at configuration time.
+    @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
+    public var tracer: (any Tracer)? {
+        configuration.tracing.tracer
+    }
+
+    /// Access to tracing configuration in order to get configured attribute keys etc.
+    @usableFromInline
+    package var tracing: TracingConfiguration {
+        self.configuration.tracing
+    }
+
     static let loggingDisabled = Logger(label: "AHC-do-not-log", factory: { _ in SwiftLogNoOpLogHandler() })
 
     /// Create an ``HTTPClient`` with specified `EventLoopGroup` provider and configuration.
@@ -705,6 +720,7 @@ public final class HTTPClient: Sendable {
             request,
             requestID: globalRequestID.wrappingIncrementThenLoad(ordering: .relaxed)
         )
+
         let taskEL: EventLoop
         switch eventLoopPreference.preference {
         case .indifferent:
@@ -734,7 +750,7 @@ public final class HTTPClient: Sendable {
             ]
         )
 
-        let failedTask: Task<Delegate.Response>? = self.state.withLockedValue { state in
+        let failedTask: Task<Delegate.Response>? = self.state.withLockedValue { state -> (Task<Delegate.Response>?) in
             switch state {
             case .upAndRunning:
                 return nil
@@ -744,6 +760,7 @@ public final class HTTPClient: Sendable {
                     eventLoop: taskEL,
                     error: HTTPClientError.alreadyShutdown,
                     logger: logger,
+                    tracing: tracing,
                     makeOrGetFileIOThreadPool: self.makeOrGetFileIOThreadPool
                 )
             }
@@ -768,11 +785,14 @@ public final class HTTPClient: Sendable {
             }
         }()
 
-        let task = Task<Delegate.Response>(
-            eventLoop: taskEL,
-            logger: logger,
-            makeOrGetFileIOThreadPool: self.makeOrGetFileIOThreadPool
-        )
+        let task: HTTPClient.Task<Delegate.Response> =
+            Task<Delegate.Response>(
+                eventLoop: taskEL,
+                logger: logger,
+                tracing: self.tracing,
+                makeOrGetFileIOThreadPool: self.makeOrGetFileIOThreadPool
+            )
+
         do {
             let requestBag = try RequestBag(
                 request: request,
@@ -884,6 +904,9 @@ public final class HTTPClient: Sendable {
         /// A method with access to the HTTP/2 stream channel that is called when creating the stream.
         public var http2StreamChannelDebugInitializer: (@Sendable (Channel) -> EventLoopFuture<Void>)?
 
+        /// Configuration how distributed traces are created and handled.
+        public var tracing: TracingConfiguration = .init()
+
         public init(
             tlsConfiguration: TLSConfiguration? = nil,
             redirectConfiguration: RedirectConfiguration? = nil,
@@ -1012,6 +1035,84 @@ public final class HTTPClient: Sendable {
             self.http2ConnectionDebugInitializer = http2ConnectionDebugInitializer
             self.http2StreamChannelDebugInitializer = http2StreamChannelDebugInitializer
         }
+
+        public init(
+            tlsConfiguration: TLSConfiguration? = nil,
+            redirectConfiguration: RedirectConfiguration? = nil,
+            timeout: Timeout = Timeout(),
+            connectionPool: ConnectionPool = ConnectionPool(),
+            proxy: Proxy? = nil,
+            ignoreUncleanSSLShutdown: Bool = false,
+            decompression: Decompression = .disabled,
+            http1_1ConnectionDebugInitializer: (@Sendable (Channel) -> EventLoopFuture<Void>)? = nil,
+            http2ConnectionDebugInitializer: (@Sendable (Channel) -> EventLoopFuture<Void>)? = nil,
+            http2StreamChannelDebugInitializer: (@Sendable (Channel) -> EventLoopFuture<Void>)? = nil,
+            tracing: TracingConfiguration = .init()
+        ) {
+            self.init(
+                tlsConfiguration: tlsConfiguration,
+                redirectConfiguration: redirectConfiguration,
+                timeout: timeout,
+                connectionPool: connectionPool,
+                proxy: proxy,
+                ignoreUncleanSSLShutdown: ignoreUncleanSSLShutdown,
+                decompression: decompression
+            )
+            self.http1_1ConnectionDebugInitializer = http1_1ConnectionDebugInitializer
+            self.http2ConnectionDebugInitializer = http2ConnectionDebugInitializer
+            self.http2StreamChannelDebugInitializer = http2StreamChannelDebugInitializer
+            self.tracing = tracing
+        }
+    }
+
+    public struct TracingConfiguration: Sendable {
+
+        @usableFromInline
+        var _tracer: Optional<any Sendable>  // erasure trick so we don't have to make Configuration @available
+
+        /// Tracer that should be used by the HTTPClient.
+        ///
+        /// This is selected at configuration creation time, and if no tracer is passed explicitly,
+        /// (including `nil` in order to disable traces), the default global bootstrapped tracer will
+        /// be stored in this property, and used for all subsequent requests made by this client.
+        @inlinable
+        @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
+        public var tracer: (any Tracer)? {
+            get {
+                guard let _tracer else {
+                    return nil
+                }
+                return _tracer as! (any Tracer)?
+            }
+            set {
+                self._tracer = newValue
+            }
+        }
+
+        // TODO: Open up customization of keys we use?
+        /// Configuration for tracing attributes set by the HTTPClient.
+        @usableFromInline
+        package var attributeKeys: AttributeKeys
+
+        public init() {
+            self._tracer = nil
+            self.attributeKeys = .init()
+        }
+
+        /// Span attribute keys that the HTTPClient should set automatically.
+        /// This struct allows the configuration of the attribute names (keys) which will be used for the apropriate values.
+        @usableFromInline
+        package struct AttributeKeys: Sendable {
+            @usableFromInline package var requestMethod: String = "http.request.method"
+            @usableFromInline package var requestBodySize: String = "http.request.body.size"
+
+            @usableFromInline package var responseBodySize: String = "http.response.body.size"
+            @usableFromInline package var responseStatusCode: String = "http.status_code"
+
+            @usableFromInline package var httpFlavor: String = "http.flavor"
+
+            @usableFromInline package init() {}
+        }
     }
 
     /// Specifies how `EventLoopGroup` will be created and establishes lifecycle ownership.

Sources/AsyncHTTPClient/HTTPHandler.swift

@@ -19,6 +19,7 @@ import NIOCore
 import NIOHTTP1
 import NIOPosix
 import NIOSSL
+import Tracing
 
 #if compiler(>=6.0)
 import Foundation
@@ -924,6 +925,12 @@ extension HTTPClient {
         /// The `Logger` used by the `Task` for logging.
         public let logger: Logger  // We are okay to store the logger here because a Task is for only one request.
 
+        @available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
+        public var tracer: (any Tracer)? {
+            tracing.tracer
+        }
+        let tracing: TracingConfiguration
+
         let promise: EventLoopPromise<Response>
 
         struct State: Sendable {
@@ -953,10 +960,16 @@ extension HTTPClient {
             self.makeOrGetFileIOThreadPool()
         }
 
-        init(eventLoop: EventLoop, logger: Logger, makeOrGetFileIOThreadPool: @escaping @Sendable () -> NIOThreadPool) {
+        init(
+            eventLoop: EventLoop,
+            logger: Logger,
+            tracing: TracingConfiguration,
+            makeOrGetFileIOThreadPool: @escaping @Sendable () -> NIOThreadPool
+        ) {
             self.eventLoop = eventLoop
             self.promise = eventLoop.makePromise()
             self.logger = logger
+            self.tracing = tracing
             self.makeOrGetFileIOThreadPool = makeOrGetFileIOThreadPool
             self.state = NIOLockedValueBox(State(isCancelled: false, taskDelegate: nil))
         }
@@ -965,11 +978,13 @@ extension HTTPClient {
             eventLoop: EventLoop,
             error: Error,
             logger: Logger,
+            tracing: TracingConfiguration,
             makeOrGetFileIOThreadPool: @escaping @Sendable () -> NIOThreadPool
         ) -> Task<Response> {
             let task = self.init(
                 eventLoop: eventLoop,
                 logger: logger,
+                tracing: tracing,
                 makeOrGetFileIOThreadPool: makeOrGetFileIOThreadPool
             )
             task.promise.fail(error)