feat: use BaggageContext to get and propagate context

Closes #23
Affects #19

Examples/Package.swift

@@ -29,7 +29,6 @@ let package = Package(
             name: "AWSXRayRecorderExampleSDK",
             dependencies: [
                 .product(name: "AWSXRayRecorder", package: "aws-xray-sdk-swift"),
-                .product(name: "AWSXRayRecorderSDK", package: "aws-xray-sdk-swift"),
                 .product(name: "AsyncHTTPClient", package: "async-http-client"),
                 .product(name: "AWSS3", package: "aws-sdk-swift"),
             ]

Examples/Sources/AWSXRayRecorderExample/main.swift

@@ -22,8 +22,10 @@ enum ExampleError: Error {
 
 let recorder = XRayRecorder()
 
+let context = XRayRecorder.TraceContext()
+
 // begin and end (sub)segments explicitly
-let segment = recorder.beginSegment(name: "Segment 1")
+let segment = recorder.beginSegment(name: "Segment 1", context: context)
 segment.setAnnotation(98101, forKey: "zip_code")
 segment.setMetadata(["debug": ["test": "Metadata string"]])
 _ = segment.beginSubsegment(name: "Subsegment 1.1 in progress")
@@ -40,7 +42,7 @@ usleep(100_000)
 subsegment.end()
 
 // use closures for convenience
-recorder.segment(name: "Segment 2") { segment in
+recorder.segment(name: "Segment 2", context: context) { segment in
     try? segment.subsegment(name: "Subsegment 2.1") { segment in
         _ = segment.subsegment(name: "Subsegment 2.1.1 with Result") { _ -> String in
             usleep(100_000)

Examples/Sources/AWSXRayRecorderExampleSDK/main.swift

@@ -14,7 +14,6 @@
 import AsyncHTTPClient
 import AWSS3
 import AWSXRayRecorder
-import AWSXRayRecorderSDK
 import NIO
 
 func env(_ name: String) -> String? {
@@ -46,10 +45,10 @@ let recorder = XRayRecorder(
     eventLoopGroup: group
 )
 
-// TODO: WIP
+// TODO: WIP https://github.com/pokryfka/aws-xray-sdk-swift/issues/19
 
 let awsClient = AWSClient(
-    middlewares: [XRayMiddleware(recorder: recorder, name: "S3")],
+    //    middlewares: [XRayMiddleware(recorder: recorder, name: "S3")],
     httpClientProvider: .shared(httpClient)
 )
 let s3 = S3(client: awsClient)

Package.swift

@@ -10,13 +10,12 @@ let package = Package(
     products: [
         .library(name: "AWSXRayRecorder", targets: ["AWSXRayRecorder"]),
         .library(name: "AWSXRayRecorderLambda", targets: ["AWSXRayRecorderLambda"]),
-        .library(name: "AWSXRayRecorderSDK", targets: ["AWSXRayRecorderSDK"]),
     ],
     dependencies: [
         .package(url: "https://github.com/apple/swift-nio.git", .upToNextMajor(from: "2.17.0")),
         .package(url: "https://github.com/apple/swift-log.git", .upToNextMajor(from: "1.0.0")),
         .package(url: "https://github.com/Flight-School/AnyCodable.git", .upToNextMajor(from: "0.3.0")),
-        .package(url: "https://github.com/swift-aws/aws-sdk-swift-core.git", .upToNextMinor(from: "5.0.0-alpha.5")),
+        .package(name: "swift-baggage-context", url: "https://github.com/slashmo/gsoc-swift-baggage-context.git", .upToNextMinor(from: "0.1.0")),
         .package(url: "https://github.com/swift-server/swift-aws-lambda-runtime.git", .upToNextMajor(from: "0.2.0")),
     ],
     targets: [
@@ -26,6 +25,7 @@ let package = Package(
                 .product(name: "NIO", package: "swift-nio"),
                 .product(name: "Logging", package: "swift-log"),
                 .product(name: "AnyCodable", package: "AnyCodable"),
+                .product(name: "Baggage", package: "swift-baggage-context"),
             ]
         ),
         .target(
@@ -35,13 +35,6 @@ let package = Package(
                 .product(name: "AWSLambdaRuntime", package: "swift-aws-lambda-runtime"),
             ]
         ),
-        .target(
-            name: "AWSXRayRecorderSDK",
-            dependencies: [
-                .byName(name: "AWSXRayRecorder"),
-                .product(name: "AWSSDKSwiftCore", package: "aws-sdk-swift-core"),
-            ]
-        ),
         .testTarget(
             name: "AWSXRayRecorderTests",
             dependencies: ["AWSXRayRecorder"]

Sources/AWSXRayRecorder/Baggage.swift

@@ -0,0 +1,41 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the aws-xray-sdk-swift open source project
+//
+// Copyright (c) 2020 pokryfka and the aws-xray-sdk-swift project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import Baggage
+
+private enum XRayContextKey: BaggageContextKey {
+    typealias Value = XRayRecorder.TraceContext
+    var name: String { "XRayTraceContext" }
+}
+
+public extension BaggageContext {
+    /// XRay Trace Context..
+    var xRayContext: XRayRecorder.TraceContext? {
+        get {
+            self[XRayContextKey.self]
+        }
+        set {
+            self[XRayContextKey.self] = newValue
+        }
+    }
+}
+
+internal extension BaggageContext {
+    func withParent(_ parentId: XRayRecorder.Segment.ID) throws -> BaggageContext {
+        guard var context = xRayContext else { throw XRayRecorder.TraceError.missingContext }
+        context.parentId = parentId
+        var updated = self
+        updated.xRayContext = context
+        return updated
+    }
+}

Sources/AWSXRayRecorder/Recorder+Helpers.swift

@@ -11,14 +11,16 @@
 //
 //===----------------------------------------------------------------------===//
 
+import Baggage
+
 extension XRayRecorder {
     @inlinable
     @discardableResult
-    public func segment<T>(name: String, parentId: Segment.ID? = nil, metadata: XRayRecorder.Segment.Metadata? = nil,
+    public func segment<T>(name: String, context: TraceContext, metadata: XRayRecorder.Segment.Metadata? = nil,
                            body: (Segment) throws -> T)
         rethrows -> T
     {
-        let segment = beginSegment(name: name, parentId: parentId, metadata: metadata)
+        let segment = beginSegment(name: name, context: context, metadata: metadata)
         defer {
             segment.end()
         }
@@ -32,11 +34,11 @@ extension XRayRecorder {
 
     @inlinable
     @discardableResult
-    public func segment<T>(name: String, traceHeader: TraceContext, metadata: XRayRecorder.Segment.Metadata? = nil,
+    public func segment<T>(name: String, baggage: BaggageContext, metadata: XRayRecorder.Segment.Metadata? = nil,
                            body: (Segment) throws -> T)
         rethrows -> T
     {
-        let segment = beginSegment(name: name, context: traceHeader, metadata: metadata)
+        let segment = beginSegment(name: name, baggage: baggage, metadata: metadata)
         defer {
             segment.end()
         }

Sources/AWSXRayRecorder/Recorder+NIO.swift

@@ -15,8 +15,6 @@ import NIO
 
 // TODO: document
 
-// TODO: expose group provider intsead?
-
 extension XRayRecorder {
     public convenience init(config: Config = Config(), eventLoopGroup: EventLoopGroup? = nil) {
         if !config.enabled {
@@ -52,31 +50,16 @@ extension XRayRecorder {
 
 extension XRayRecorder {
     @inlinable
-    public func segment<T>(name: String, parentId: Segment.ID? = nil, metadata: Segment.Metadata? = nil,
+    public func segment<T>(name: String, context: TraceContext, metadata: Segment.Metadata? = nil,
                            body: () -> EventLoopFuture<T>) -> EventLoopFuture<T> {
-        let segment = beginSegment(name: name, parentId: parentId, metadata: metadata)
+        let segment = beginSegment(name: name, context: context, metadata: metadata)
         return body().always { result in
             if case Result<T, Error>.failure(let error) = result {
                 segment.setError(error)
             }
             segment.end()
         }
     }
-
-    // TODO: hopefully there will be a better way to pass the context, per https://github.com/slashmo/gsoc-swift-tracing/issues/48
-
-    @inlinable
-    public func beginSegment<T>(name: String, parentId: Segment.ID? = nil, metadata: Segment.Metadata? = nil,
-                                body: (Segment) -> EventLoopFuture<T>) -> EventLoopFuture<(Segment, T)> {
-        let segment = beginSegment(name: name, parentId: parentId, metadata: metadata)
-        return body(segment)
-            .always { result in
-                if case Result<T, Error>.failure(let error) = result {
-                    segment.setError(error)
-                }
-            }
-            .map { (segment, $0) }
-    }
 }
 
 extension XRayRecorder.Segment {

Sources/AWSXRayRecorder/Recorder.swift

@@ -11,6 +11,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+import Baggage
 import Dispatch
 import Logging
 
@@ -24,9 +25,6 @@ public class XRayRecorder {
 
     private lazy var logger = Logger(label: "xray.recorder.\(String.random32())")
 
-    @Synchronized private var traceId = TraceID()
-    @Synchronized private var sampled: Bool = true
-
     private let segmentsLock = ReadWriteLock()
     private var _segments = [Segment.ID: Segment]()
 
@@ -44,13 +42,14 @@ public class XRayRecorder {
         logger.logLevel = config.logLevel
     }
 
-    internal func beginSegment(name: String, parentId: Segment.ID? = nil, subsegment: Bool = false,
+    internal func beginSegment(name: String, context: TraceContext, baggage: BaggageContext,
                                aws: Segment.AWS? = nil, metadata: Segment.Metadata? = nil) -> Segment {
-        guard config.enabled, sampled else {
+        guard config.enabled, context.isSampled else {
+            // create dummy segment
             return Segment(
-                name: name, traceId: traceId, parentId: parentId, subsegment: subsegment,
+                id: .init(), name: name,
+                context: context, baggage: baggage,
                 aws: aws, metadata: metadata,
-                sampled: .notSampled,
                 callback: nil
             )
         }
@@ -65,41 +64,43 @@ public class XRayRecorder {
                 self.emitGroup.leave()
             }
         }
+        let segmentId = Segment.ID()
+        let subsegment = context.parentId != nil
         let newSegment = Segment(
-            name: name, traceId: traceId, parentId: parentId, subsegment: subsegment,
-            service: .init(version: config.serviceVersion),
+            id: segmentId, name: name,
+            context: context, baggage: baggage,
+            subsegment: subsegment,
             aws: aws, metadata: metadata,
             callback: callback
         )
-        let segmentId = newSegment.id
         segmentsLock.withWriterLockVoid { _segments[segmentId] = newSegment }
         return newSegment
     }
 
-    // TODO: pass Context including trace id and sampling decision?
-    public func beginSegment(name: String, parentId: Segment.ID? = nil,
-                             aws: Segment.AWS? = nil, metadata: Segment.Metadata? = nil) -> Segment {
-        beginSegment(name: name, parentId: parentId, subsegment: false, aws: aws, metadata: metadata)
-    }
-
-    // TODO: pass Context including trace id and sampling decision?
-    public func beginSubsegment(name: String, parentId: Segment.ID,
-                                aws: Segment.AWS? = nil, metadata: Segment.Metadata? = nil) -> Segment {
-        beginSegment(name: name, parentId: parentId, subsegment: true, aws: aws, metadata: metadata)
+    /// Creates new segment.
+    /// - Parameters:
+    ///   - name: segment name
+    ///   - context: the trace context
+    ///   - metadata: segment metadata
+    /// - Returns: new segment
+    public func beginSegment(name: String, context: TraceContext, metadata: Segment.Metadata? = nil) -> Segment {
+        var baggage = BaggageContext()
+        baggage.xRayContext = context
+        return beginSegment(name: name, context: context, baggage: baggage, metadata: metadata)
     }
 
-    public func beginSegment(name: String, context: TraceContext,
-                             aws: Segment.AWS? = nil, metadata: Segment.Metadata? = nil) -> Segment {
-        // TODO: do not store "current" or "active" traceId, segmentId or sampling decision
-        // we do it here only for the sake of XRayMiddleware which is conceptually broken and
-        // hopefully will be replaced
-        traceId = context.traceId
-        sampled = context.sampled.isSampled != false
-        if let parentId = context.parentId {
-            return beginSegment(name: name, parentId: parentId, subsegment: true, aws: aws, metadata: metadata)
-        } else {
-            return beginSegment(name: name, aws: aws, metadata: metadata)
-        }
+    /// Creates new segment.
+    /// Extracts the thre context from the baggage.
+    /// Creates new if the baggage does not contain a valid XRay Trace Context.
+    /// - Parameters:
+    ///   - name: segment name
+    ///   - baggage: baggage with the trace context
+    ///   - metadata: segment metadata
+    /// - Returns: new segment
+    public func beginSegment(name: String, baggage: BaggageContext, metadata: Segment.Metadata? = nil) -> XRayRecorder.Segment {
+        // TODO: Configure context missing strategy https://github.com/pokryfka/aws-xray-sdk-swift/issues/28
+        let context = baggage.xRayContext ?? TraceContext(sampled: .unknown)
+        return beginSegment(name: name, context: context, baggage: baggage, metadata: metadata)
     }
 
     internal func waitEmitting() {