Instrument AWSClient WIP

Package.swift

@@ -26,7 +26,11 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-nio.git", .upToNextMajor(from: "2.16.1")),
         .package(url: "https://github.com/apple/swift-nio-ssl.git", .upToNextMajor(from: "2.7.2")),
         .package(url: "https://github.com/apple/swift-nio-transport-services.git", .upToNextMajor(from: "1.0.0")),
-        .package(url: "https://github.com/swift-server/async-http-client.git", .upToNextMajor(from: "1.2.0")),
+        // .package(url: "https://github.com/swift-server/async-http-client.git", .upToNextMajor(from: "1.2.0")),
+        .package(url: "https://github.com/pokryfka/async-http-client.git", .branch("feature/instrumentation")),
+        .package(url: "https://github.com/slashmo/gsoc-swift-baggage-context.git", .upToNextMinor(from: "0.3.0")),
+        // TODO: use version when released
+        .package(url: "https://github.com/slashmo/gsoc-swift-tracing.git", .revision("fe80d764ad225b1dfd06dcb57d08b5e3485662f9")),
     ],
     targets: [
         .target(name: "AWSSDKSwiftCore", dependencies: [
@@ -41,6 +45,9 @@ let package = Package(
             .product(name: "NIOSSL", package: "swift-nio-ssl"),
             .product(name: "NIOTransportServices", package: "swift-nio-transport-services"),
             .product(name: "NIOFoundationCompat", package: "swift-nio"),
+            .product(name: "Baggage", package: "swift-baggage-context"),
+            .product(name: "BaggageLogging", package: "swift-baggage-context"),
+            .product(name: "TracingInstrumentation", package: "gsoc-swift-tracing"),
         ]),
         .target(name: "AWSCrypto", dependencies: []),
         .target(name: "AWSSignerV4", dependencies: [
@@ -53,6 +60,9 @@ let package = Package(
             .product(name: "NIOHTTP1", package: "swift-nio"),
             .product(name: "NIOFoundationCompat", package: "swift-nio"),
             .product(name: "NIOTestUtils", package: "swift-nio"),
+            .product(name: "Baggage", package: "swift-baggage-context"),
+            .product(name: "BaggageLogging", package: "swift-baggage-context"),
+            .product(name: "TracingInstrumentation", package: "gsoc-swift-tracing"),
         ]),
         .target(name: "AWSXML", dependencies: [
             .byName(name: "CAWSExpat"),

Sources/AWSSDKSwiftCore/AWSClient+Paginate.swift

@@ -34,17 +34,17 @@ extension AWSClient {
     ///   - onPage: closure called with each block of entries
     public func paginate<Input: AWSPaginateToken, Output: AWSShape>(
         input: Input,
-        command: @escaping (Input, EventLoop?, Logger) -> EventLoopFuture<Output>,
+        command: @escaping (Input, AWSClient.Context, EventLoop?) -> EventLoopFuture<Output>,
         tokenKey: KeyPath<Output, Input.Token?>,
+        context: AWSClient.Context,
         on eventLoop: EventLoop? = nil,
-        logger: Logger = AWSClient.loggingDisabled,
         onPage: @escaping (Output, EventLoop) -> EventLoopFuture<Bool>
     ) -> EventLoopFuture<Void> {
         let eventLoop = eventLoop ?? eventLoopGroup.next()
         let promise = eventLoop.makePromise(of: Void.self)
 
         func paginatePart(input: Input) {
-            let responseFuture = command(input, eventLoop, logger)
+            let responseFuture = command(input, context, eventLoop)
                 .flatMap { response in
                     return onPage(response, eventLoop)
                         .map { (rt) -> Void in

Sources/AWSSDKSwiftCore/AWSService.swift

@@ -36,7 +36,7 @@ extension AWSService {
     ///     - expires: How long before the signed URL expires
     /// - returns:
     ///     A signed URL
-    public func signURL(url: URL, httpMethod: String, expires: Int = 86400, logger: Logger = AWSClient.loggingDisabled) -> EventLoopFuture<URL> {
-        return self.client.signURL(url: url, httpMethod: httpMethod, expires: expires, serviceConfig: self.config, logger: logger)
+    public func signURL(url: URL, httpMethod: String, expires: Int = 86400, context: AWSClient.Context) -> EventLoopFuture<URL> {
+        return self.client.signURL(url: url, httpMethod: httpMethod, expires: expires, serviceConfig: self.config, context: context)
     }
 }

Sources/AWSSDKSwiftCore/Credential/ConfigFileCredentialProvider.swift

@@ -43,7 +43,7 @@ struct AWSConfigFileCredentialProvider: CredentialProvider {
         self.profile = profile ?? Environment["AWS_PROFILE"] ?? "default"
     }
 
-    func getCredential(on eventLoop: EventLoop, logger: Logger) -> EventLoopFuture<Credential> {
+    func getCredential(on eventLoop: EventLoop, context: CredentialProvider.Context) -> EventLoopFuture<Credential> {
         return AWSConfigFileCredentialProvider.fromSharedCredentials(credentialsFilePath: self.credentialsFilePath, profile: self.profile, on: eventLoop)
             .map { $0 }
     }

Sources/AWSSDKSwiftCore/Credential/CredentialProvider.swift

@@ -13,13 +13,16 @@
 //===----------------------------------------------------------------------===//
 
 import AWSSignerV4
-import Logging
+import Baggage
+import BaggageLogging
 import NIO
 import NIOConcurrencyHelpers
 
 /// Protocol providing future holding a credential
 public protocol CredentialProvider: CustomStringConvertible {
-    func getCredential(on eventLoop: EventLoop, logger: Logger) -> EventLoopFuture<Credential>
+    typealias Context = BaggageLogging.LoggingBaggageContextCarrier
+
+    func getCredential(on eventLoop: EventLoop, context: CredentialProvider.Context) -> EventLoopFuture<Credential>
     func shutdown(on eventLoop: EventLoop) -> EventLoopFuture<Void>
 }
 
@@ -34,13 +37,22 @@ extension CredentialProvider {
 /// A helper struct to defer the creation of a `CredentialProvider` until after the AWSClient has been created.
 public struct CredentialProviderFactory {
     /// The initialization context for a `ContextProvider`
-    public struct Context {
+    public struct Context: BaggageLogging.LoggingBaggageContextCarrier {
         /// The `AWSClient`s internal `HTTPClient`
         public let httpClient: AWSHTTPClient
         /// The `EventLoop` that the `CredentialProvider` should use for credential refreshs
         public let eventLoop: EventLoop
         /// The `Logger` attached to the AWSClient
-        public let logger: Logger
+        public var logger: Logger // TODO: should not need to be mutable
+        /// The context baggage.
+        public var baggage: BaggageContext
+
+        public init(httpClient: AWSHTTPClient, eventLoop: EventLoop, context: BaggageLogging.LoggingBaggageContextCarrier) {
+            self.httpClient = httpClient
+            self.eventLoop = eventLoop
+            self.logger = context.logger
+            self.baggage = context.baggage
+        }
     }
 
     private let cb: (Context) -> CredentialProvider

Sources/AWSSDKSwiftCore/Credential/DeferredCredentialProvider.swift

@@ -38,7 +38,7 @@ public class DeferredCredentialProvider: CredentialProvider {
     public init(context: CredentialProviderFactory.Context, provider: CredentialProvider) {
         self.startupPromise = context.eventLoop.makePromise(of: Credential.self)
         self.provider = provider
-        provider.getCredential(on: context.eventLoop, logger: context.logger)
+        provider.getCredential(on: context.eventLoop, context: context)
             .flatMapErrorThrowing { _ in throw CredentialProviderError.noProvider }
             .map { credential in
                 self.internalCredential = credential
@@ -59,7 +59,7 @@ public class DeferredCredentialProvider: CredentialProvider {
     /// otherwise return credentials store in class
     /// - Parameter eventLoop: EventLoop to run off
     /// - Returns: EventLoopFuture that will hold credentials
-    public func getCredential(on eventLoop: EventLoop, logger: Logger) -> EventLoopFuture<Credential> {
+    public func getCredential(on eventLoop: EventLoop, context: CredentialProvider.Context) -> EventLoopFuture<Credential> {
         if let credential = self.credential {
             return eventLoop.makeSucceededFuture(credential)
         }

Sources/AWSSDKSwiftCore/Credential/MetaDataCredentialProvider.swift

@@ -21,7 +21,7 @@ import struct Foundation.TimeZone
 import struct Foundation.URL
 
 import AWSSignerV4
-import Logging
+// import Logging
 import NIO
 import NIOConcurrencyHelpers
 import NIOHTTP1
@@ -30,12 +30,12 @@ import NIOHTTP1
 protocol MetaDataClient: CredentialProvider {
     associatedtype MetaData: ExpiringCredential & Decodable
 
-    func getMetaData(on eventLoop: EventLoop, logger: Logger) -> EventLoopFuture<MetaData>
+    func getMetaData(on eventLoop: EventLoop, context: CredentialProvider.Context) -> EventLoopFuture<MetaData>
 }
 
 extension MetaDataClient {
-    func getCredential(on eventLoop: EventLoop, logger: Logger) -> EventLoopFuture<Credential> {
-        self.getMetaData(on: eventLoop, logger: logger).map { metaData in
+    func getCredential(on eventLoop: EventLoop, context: CredentialProvider.Context) -> EventLoopFuture<Credential> {
+        self.getMetaData(on: eventLoop, context: context).map { metaData in
             metaData
         }
     }
@@ -107,8 +107,8 @@ struct ECSMetaDataClient: MetaDataClient {
         self.endpointURL = "\(host)\(relativeURL)"
     }
 
-    func getMetaData(on eventLoop: EventLoop, logger: Logger) -> EventLoopFuture<ECSMetaData> {
-        return request(url: endpointURL, timeout: 2, on: eventLoop, logger: logger)
+    func getMetaData(on eventLoop: EventLoop, context: CredentialProvider.Context) -> EventLoopFuture<ECSMetaData> {
+        return request(url: endpointURL, timeout: 2, on: eventLoop, context: context)
             .flatMapThrowing { response in
                 guard let body = response.body else {
                     throw MetaDataClientError.missingMetaData
@@ -117,9 +117,14 @@ struct ECSMetaDataClient: MetaDataClient {
             }
     }
 
-    private func request(url: String, timeout: TimeInterval, on eventLoop: EventLoop, logger: Logger) -> EventLoopFuture<AWSHTTPResponse> {
+    private func request(
+        url: String,
+        timeout: TimeInterval,
+        on eventLoop: EventLoop,
+        context: CredentialProvider.Context
+    ) -> EventLoopFuture<AWSHTTPResponse> {
         let request = AWSHTTPRequest(url: URL(string: url)!, method: .GET, headers: [:], body: .empty)
-        return httpClient.execute(request: request, timeout: TimeAmount.seconds(2), on: eventLoop, logger: logger)
+        return httpClient.execute(request: request, timeout: TimeAmount.seconds(2), on: eventLoop, context: context)
     }
 }
 
@@ -180,17 +185,17 @@ struct InstanceMetaDataClient: MetaDataClient {
         self.host = host
     }
 
-    func getMetaData(on eventLoop: EventLoop, logger: Logger) -> EventLoopFuture<InstanceMetaData> {
-        return getToken(on: eventLoop, logger: logger)
+    func getMetaData(on eventLoop: EventLoop, context: CredentialProvider.Context) -> EventLoopFuture<InstanceMetaData> {
+        return getToken(on: eventLoop, context: context)
             .map { token in
-                logger.info("Found IMDSv2 token")
+                context.logger.info("Found IMDSv2 token")
                 return HTTPHeaders([(Self.TokenHeaderName, token)])
             }
             .flatMapErrorThrowing { _ in
                 // If we didn't find a session key then assume we are running IMDSv1.
                 // (we could be running from a Docker container and the hop count for the PUT
                 // request is still set to 1)
-                logger.info("Did not find IMDSv2 token, use IMDSv1")
+                context.logger.info("Did not find IMDSv2 token, use IMDSv1")
                 return HTTPHeaders()
             }
             .flatMap { (headers) -> EventLoopFuture<(AWSHTTPResponse, HTTPHeaders)> in
@@ -200,7 +205,7 @@ struct InstanceMetaDataClient: MetaDataClient {
                     method: .GET,
                     headers: headers,
                     on: eventLoop,
-                    logger: logger
+                    context: context
                 ).map { ($0, headers) }
             }
             .flatMapThrowing { (response, headers) -> (String, HTTPHeaders) in
@@ -218,7 +223,7 @@ struct InstanceMetaDataClient: MetaDataClient {
             .flatMap { (roleName, headers) -> EventLoopFuture<AWSHTTPResponse> in
                 // request credentials with the rolename
                 let url = self.credentialURL.appendingPathComponent(roleName)
-                return self.request(url: url, headers: headers, on: eventLoop, logger: logger)
+                return self.request(url: url, headers: headers, on: eventLoop, context: context)
             }
             .flatMapThrowing { response in
                 // decode the repsonse payload into the metadata object
@@ -230,13 +235,13 @@ struct InstanceMetaDataClient: MetaDataClient {
             }
     }
 
-    func getToken(on eventLoop: EventLoop, logger: Logger) -> EventLoopFuture<String> {
+    func getToken(on eventLoop: EventLoop, context: CredentialProvider.Context) -> EventLoopFuture<String> {
         return request(
             url: self.tokenURL,
             method: .PUT,
             headers: HTTPHeaders([Self.TokenTimeToLiveHeader]), timeout: .seconds(2),
             on: eventLoop,
-            logger: logger
+            context: context
         ).flatMapThrowing { response in
             guard response.status == .ok else {
                 throw MetaDataClientError.unexpectedTokenResponseStatus(status: response.status)
@@ -255,9 +260,9 @@ struct InstanceMetaDataClient: MetaDataClient {
         headers: HTTPHeaders = .init(),
         timeout: TimeAmount = .seconds(2),
         on eventLoop: EventLoop,
-        logger: Logger
+        context: CredentialProvider.Context
     ) -> EventLoopFuture<AWSHTTPResponse> {
         let request = AWSHTTPRequest(url: url, method: method, headers: headers, body: .empty)
-        return httpClient.execute(request: request, timeout: timeout, on: eventLoop, logger: logger)
+        return httpClient.execute(request: request, timeout: timeout, on: eventLoop, context: context)
     }
 }

Sources/AWSSDKSwiftCore/Credential/NullCredentialProvider.swift

@@ -18,7 +18,7 @@ import NIO
 
 /// Credential provider that always fails
 public struct NullCredentialProvider: CredentialProvider {
-    public func getCredential(on eventLoop: EventLoop, logger: Logger) -> EventLoopFuture<Credential> {
+    public func getCredential(on eventLoop: EventLoop, context: CredentialProvider.Context) -> EventLoopFuture<Credential> {
         return eventLoop.makeFailedFuture(CredentialProviderError.noProvider)
     }
 }

Sources/AWSSDKSwiftCore/Credential/RotatingCredentialProvider.swift

@@ -33,7 +33,7 @@ public final class RotatingCredentialProvider: CredentialProvider {
     public init(context: CredentialProviderFactory.Context, provider: CredentialProvider, remainingTokenLifetimeForUse: TimeInterval? = nil) {
         self.provider = provider
         self.remainingTokenLifetimeForUse = remainingTokenLifetimeForUse ?? 3 * 60
-        _ = refreshCredentials(on: context.eventLoop, logger: context.logger)
+        _ = refreshCredentials(on: context.eventLoop, context: context)
     }
 
     public func shutdown(on eventLoop: EventLoop) -> EventLoopFuture<Void> {
@@ -45,18 +45,18 @@ public final class RotatingCredentialProvider: CredentialProvider {
         }
     }
 
-    public func getCredential(on eventLoop: EventLoop, logger: Logger) -> EventLoopFuture<Credential> {
+    public func getCredential(on eventLoop: EventLoop, context: CredentialProvider.Context) -> EventLoopFuture<Credential> {
         self.lock.lock()
         let cred = credential
         self.lock.unlock()
 
         switch cred {
         case .none:
-            return self.refreshCredentials(on: eventLoop, logger: logger)
+            return self.refreshCredentials(on: eventLoop, context: context)
         case .some(let cred as ExpiringCredential):
             if cred.isExpiring(within: remainingTokenLifetimeForUse) {
                 // the credentials are expiring... let's refresh
-                return self.refreshCredentials(on: eventLoop, logger: logger)
+                return self.refreshCredentials(on: eventLoop, context: context)
             }
 
             return eventLoop.makeSucceededFuture(cred)
@@ -66,7 +66,7 @@ public final class RotatingCredentialProvider: CredentialProvider {
         }
     }
 
-    private func refreshCredentials(on eventLoop: EventLoop, logger: Logger) -> EventLoopFuture<Credential> {
+    private func refreshCredentials(on eventLoop: EventLoop, context: CredentialProvider.Context) -> EventLoopFuture<Credential> {
         self.lock.lock()
         defer { self.lock.unlock() }
 
@@ -80,15 +80,15 @@ public final class RotatingCredentialProvider: CredentialProvider {
             return future
         }
 
-        logger.info("Refeshing AWS credentials", metadata: ["aws-credential-provider": .string("\(self)")])
+        context.logger.info("Refeshing AWS credentials", metadata: ["aws-credential-provider": .string("\(self)")])
 
-        credentialFuture = self.provider.getCredential(on: eventLoop, logger: logger)
+        credentialFuture = self.provider.getCredential(on: eventLoop, context: context)
             .map { (credential) -> (Credential) in
                 // update the internal credential locked
                 self.lock.withLock {
                     self.credentialFuture = nil
                     self.credential = credential
-                    logger.info("AWS credentials ready", metadata: ["aws-credential-provider": .string("\(self)")])
+                    context.logger.info("AWS credentials ready", metadata: ["aws-credential-provider": .string("\(self)")])
                 }
                 return credential
             }

Sources/AWSSDKSwiftCore/Credential/RuntimeSelectorCredentialProvider.swift

@@ -43,13 +43,13 @@ class RuntimeSelectorCredentialProvider: CredentialProvider {
         return self.startupPromise.futureResult.map { _ in }.hop(to: eventLoop)
     }
 
-    func getCredential(on eventLoop: EventLoop, logger: Logger) -> EventLoopFuture<Credential> {
+    func getCredential(on eventLoop: EventLoop, context: CredentialProvider.Context) -> EventLoopFuture<Credential> {
         if let provider = internalProvider {
-            return provider.getCredential(on: eventLoop, logger: logger)
+            return provider.getCredential(on: eventLoop, context: context)
         }
 
         return self.startupPromise.futureResult.hop(to: eventLoop).flatMap { provider in
-            return provider.getCredential(on: eventLoop, logger: logger)
+            return provider.getCredential(on: eventLoop, context: context)
         }
     }
 
@@ -63,7 +63,7 @@ class RuntimeSelectorCredentialProvider: CredentialProvider {
             }
             let providerFactory = providers[index]
             let provider = providerFactory.createProvider(context: context)
-            provider.getCredential(on: context.eventLoop, logger: context.logger).whenComplete { result in
+            provider.getCredential(on: context.eventLoop, context: context).whenComplete { result in
                 switch result {
                 case .success:
                     context.logger.info("Select credential provider", metadata: ["aws-credential-provider": .string("\(provider)")])

Sources/AWSSDKSwiftCore/Credential/StaticCredential+CredentialProvider.swift

@@ -16,7 +16,7 @@ import AWSSignerV4
 import Logging
 
 extension StaticCredential: CredentialProvider {
-    public func getCredential(on eventLoop: EventLoop, logger: Logger) -> EventLoopFuture<Credential> {
+    public func getCredential(on eventLoop: EventLoop, context: CredentialProvider.Context) -> EventLoopFuture<Credential> {
         eventLoop.makeSucceededFuture(self)
     }
 }