[WIP] Implement Distributed Tracing

Author: slashmo

Package.swift

@@ -27,6 +27,8 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-nio-extras.git", from: "1.13.0"),
         .package(url: "https://github.com/apple/swift-nio-transport-services.git", from: "1.19.0"),
         .package(url: "https://github.com/apple/swift-log.git", from: "1.4.4"),
+        .package(url: "https://github.com/apple/swift-distributed-tracing.git", from: "1.0.0"),
+        .package(url: "https://github.com/apple/swift-service-context.git", from: "1.0.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"),
     ],
@@ -53,6 +55,8 @@ let package = Package(
                 .product(name: "NIOSOCKS", package: "swift-nio-extras"),
                 .product(name: "NIOTransportServices", package: "swift-nio-transport-services"),
                 .product(name: "Logging", package: "swift-log"),
+                .product(name: "Tracing", package: "swift-distributed-tracing"),
+                .product(name: "ServiceContextModule", package: "swift-service-context"),
                 .product(name: "Atomics", package: "swift-atomics"),
                 .product(name: "Algorithms", package: "swift-algorithms"),
             ]

Sources/AsyncHTTPClient/AsyncAwait/HTTPClient+execute.swift

@@ -15,6 +15,8 @@
 import Logging
 import NIOCore
 import NIOHTTP1
+import ServiceContextModule
+import Tracing
 
 import struct Foundation.URL
 
@@ -34,18 +36,20 @@ extension HTTPClient {
     public func execute(
         _ request: HTTPClientRequest,
         deadline: NIODeadline,
-        logger: Logger? = nil
+        logger: Logger? = nil,
+        context: ServiceContext? = nil
     ) async throws -> HTTPClientResponse {
         try await self.executeAndFollowRedirectsIfNeeded(
             request,
             deadline: deadline,
             logger: logger ?? Self.loggingDisabled,
+            context: context ?? .current ?? .topLevel,
             redirectState: RedirectState(self.configuration.redirectConfiguration.mode, initialURL: request.url)
         )
     }
 }
 
-// MARK: Connivence methods
+// MARK: Convenience methods
 
 @available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
 extension HTTPClient {
@@ -63,12 +67,14 @@ extension HTTPClient {
     public func execute(
         _ request: HTTPClientRequest,
         timeout: TimeAmount,
-        logger: Logger? = nil
+        logger: Logger? = nil,
+        context: ServiceContext? = nil
     ) async throws -> HTTPClientResponse {
         try await self.execute(
             request,
             deadline: .now() + timeout,
-            logger: logger
+            logger: logger,
+            context: context
         )
     }
 }
@@ -81,6 +87,7 @@ extension HTTPClient {
         _ request: HTTPClientRequest,
         deadline: NIODeadline,
         logger: Logger,
+        context: ServiceContext,
         redirectState: RedirectState?
     ) async throws -> HTTPClientResponse {
         var currentRequest = request
@@ -91,7 +98,13 @@ extension HTTPClient {
         while true {
             let preparedRequest = try HTTPClientRequest.Prepared(currentRequest, dnsOverride: configuration.dnsOverride)
             let response = try await {
-                var response = try await self.executeCancellable(preparedRequest, deadline: deadline, logger: logger)
+                var response = try await self.executeCancellable(
+                    preparedRequest,
+                    deadline: deadline,
+                    logger: logger,
+                    context: context,
+                    resendCount: history.isEmpty ? nil : history.count
+                )
 
                 history.append(
                     .init(
@@ -149,39 +162,68 @@ extension HTTPClient {
     private func executeCancellable(
         _ request: HTTPClientRequest.Prepared,
         deadline: NIODeadline,
-        logger: Logger
+        logger: Logger,
+        context: ServiceContext,
+        resendCount: Int?
     ) async throws -> HTTPClientResponse {
         let cancelHandler = TransactionCancelHandler()
 
-        return try await withTaskCancellationHandler(
-            operation: { () async throws -> HTTPClientResponse in
-                let eventLoop = self.eventLoopGroup.any()
-                let deadlineTask = eventLoop.scheduleTask(deadline: deadline) {
-                    cancelHandler.cancel(reason: .deadlineExceeded)
-                }
-                defer {
-                    deadlineTask.cancel()
-                }
-                return try await withCheckedThrowingContinuation {
-                    (continuation: CheckedContinuation<HTTPClientResponse, Swift.Error>) -> Void in
-                    let transaction = Transaction(
-                        request: request,
-                        requestOptions: .fromClientConfiguration(self.configuration),
-                        logger: logger,
-                        connectionDeadline: .now() + (self.configuration.timeout.connectionCreationTimeout),
-                        preferredEventLoop: eventLoop,
-                        responseContinuation: continuation
-                    )
+        return try await withSpan(request.head.method.rawValue, context: context, ofKind: .client) { span in
+            var request = request
+            request.head.headers.propagate(span.context)
+            span.updateAttributes { attributes in
+                attributes["http.request.method"] = request.head.method.rawValue
+                attributes["server.address"] = request.poolKey.connectionTarget.host
+                attributes["server.port"] = request.poolKey.connectionTarget.port
+                attributes["url.full"] = request.url.absoluteString
+                attributes["http.request.resend_count"] = resendCount
+            }
+
+            do {
+                let response = try await withTaskCancellationHandler(
+                    operation: { () async throws -> HTTPClientResponse in
+                        let eventLoop = self.eventLoopGroup.any()
+                        let deadlineTask = eventLoop.scheduleTask(deadline: deadline) {
+                            cancelHandler.cancel(reason: .deadlineExceeded)
+                        }
+                        defer {
+                            deadlineTask.cancel()
+                        }
+                        let response = try await withCheckedThrowingContinuation {
+                            (continuation: CheckedContinuation<HTTPClientResponse, Swift.Error>) -> Void in
+                            let transaction = Transaction(
+                                request: request,
+                                requestOptions: .fromClientConfiguration(self.configuration),
+                                logger: logger,
+                                connectionDeadline: .now() + (self.configuration.timeout.connectionCreationTimeout),
+                                preferredEventLoop: eventLoop,
+                                responseContinuation: continuation
+                            )
 
-                    cancelHandler.registerTransaction(transaction)
+                            cancelHandler.registerTransaction(transaction)
 
-                    self.poolManager.executeRequest(transaction)
-                }
-            },
-            onCancel: {
-                cancelHandler.cancel(reason: .taskCanceled)
+                            self.poolManager.executeRequest(transaction)
+                        }
+                        if response.status.code >= 400 {
+                            span.setStatus(.init(code: .error))
+                            span.attributes["error.type"] = "\(response.status.code)"
+                        }
+                        span.attributes["http.response.status_code"] = "\(response.status.code)"
+                        return response
+                    },
+                    onCancel: {
+                        span.setStatus(.init(code: .error))
+                        span.attributes["error.type"] = "\(CancellationError.self)"
+                        cancelHandler.cancel(reason: .taskCanceled)
+                    }
+                )
+                return response
+            } catch {
+                span.setStatus(.init(code: .error))
+                span.attributes["error.type"] = "\(type(of: error))"
+                throw error
             }
-        )
+        }
     }
 }
 

Sources/AsyncHTTPClient/ServiceContextPropagation.swift

@@ -0,0 +1,32 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the AsyncHTTPClient open source project
+//
+// Copyright (c) 2024 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
+import ServiceContextModule
+
+@available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
+extension HTTPHeaders {
+    mutating func propagate(_ context: ServiceContext) {
+        InstrumentationSystem.instrument.inject(context, into: &self, using: Self.injector)
+    }
+
+    private static let injector = HTTPHeadersInjector()
+}
+
+private struct HTTPHeadersInjector: Injector {
+    func inject(_ value: String, forKey name: String, into headers: inout HTTPHeaders) {
+        headers.add(name: name, value: value)
+    }
+}