feat: create XRayInstrument

Examples/Package.swift

@@ -9,14 +9,17 @@ let package = Package(
     ],
     products: [
         .executable(name: "AWSXRayRecorderExample", targets: ["AWSXRayRecorderExample"]),
+        .executable(name: "AWSXRayInstrumentExample", targets: ["AWSXRayInstrumentExample"]),
         .executable(name: "AWSXRayRecorderExampleSDK", targets: ["AWSXRayRecorderExampleSDK"]),
         .executable(name: "AWSXRayRecorderExampleLambda", targets: ["AWSXRayRecorderExampleLambda"]),
     ],
     dependencies: [
         .package(name: "aws-xray-sdk-swift", path: ".."),
+        .package(url: "https://github.com/slashmo/gsoc-swift-tracing.git", .revision("0d96630f614bda1bd88c9422cf05b077cf034886")),
         .package(url: "https://github.com/swift-aws/aws-sdk-swift.git", .upToNextMinor(from: "5.0.0-alpha.5")),
         .package(url: "https://github.com/swift-server/swift-aws-lambda-runtime.git", .upToNextMajor(from: "0.2.0")),
         .package(url: "https://github.com/swift-server/async-http-client.git", .upToNextMajor(from: "1.0.0")),
+        .package(url: "https://github.com/apple/swift-nio.git", .upToNextMajor(from: "2.17.0")),
     ],
     targets: [
         .target(
@@ -25,6 +28,16 @@ let package = Package(
                 .product(name: "AWSXRayRecorder", package: "aws-xray-sdk-swift"),
             ]
         ),
+        .target(
+            name: "AWSXRayInstrumentExample",
+            dependencies: [
+                .product(name: "AWSXRayRecorder", package: "aws-xray-sdk-swift"),
+                .product(name: "AWSXRayInstrument", package: "aws-xray-sdk-swift"),
+                .product(name: "NIOInstrumentation", package: "gsoc-swift-tracing"),
+                .product(name: "NIO", package: "swift-nio"),
+                .product(name: "NIOHTTP1", package: "swift-nio"),
+            ]
+        ),
         .target(
             name: "AWSXRayRecorderExampleSDK",
             dependencies: [

Examples/Sources/AWSXRayInstrumentExample/main.swift

@@ -0,0 +1,35 @@
+import AWSXRayInstrument
+import AWSXRayRecorder
+import Baggage // BaggageContext
+import Instrumentation // InstrumentationSystem
+import NIOHTTP1 // HTTPHeaders
+import NIOInstrumentation // HTTPHeadersExtractor
+
+// create and boostrap the instrument
+let instrument = XRayRecorder(emitter: XRayLogEmitter(), config: .init(logLevel: .debug))
+InstrumentationSystem.bootstrap(instrument)
+
+let tracer = InstrumentationSystem.tracer // the instrument
+
+// extract the context from HTTP headers
+let headers = HTTPHeaders([
+    ("X-Amzn-Trace-Id", "Root=1-5759e988-bd862e3fe1be46a994272793;Sampled=1"),
+])
+var baggage = BaggageContext()
+tracer.extract(headers, into: &baggage, using: HTTPHeadersExtractor())
+
+// TODO: create instrumented HTTP client
+
+// trace
+var span = tracer.startSpan(named: "Span 1", context: baggage)
+span.setAttribute("Attribute 1", forKey: "key1")
+span.addEvent(.init(name: "Event"))
+
+var span2 = tracer.startSpan(named: "Span 2", context: baggage)
+span.setAttribute("Attribute 2", forKey: "key2")
+
+span.end()
+span2.end()
+
+// (tracer as? XRayRecorder)?.wait()
+instrument.wait()

Package.swift

@@ -9,13 +9,18 @@ let package = Package(
     ],
     products: [
         .library(name: "AWSXRayRecorder", targets: ["AWSXRayRecorder"]),
+        .library(name: "AWSXRayInstrument", targets: ["AWSXRayInstrument"]),
         .library(name: "AWSXRayRecorderLambda", targets: ["AWSXRayRecorderLambda"]),
-        .library(name: "AWSXRayRecorderSDK", targets: ["AWSXRayRecorderSDK"]),
+        .library(name: "AWSXRayRecorderSDK", targets: ["AWSXRayRecorderSDK"]), // TODO: remove
     ],
     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(name: "swift-baggage-context", url: "https://github.com/slashmo/gsoc-swift-baggage-context.git", .upToNextMinor(from: "0.1.0")),
+//        .package(url: "https://github.com/slashmo/gsoc-swift-tracing.git", .branch("main")),
+        // "bleeding edge"
+        .package(url: "https://github.com/slashmo/gsoc-swift-tracing.git", .revision("0d96630f614bda1bd88c9422cf05b077cf034886")),
         .package(url: "https://github.com/swift-aws/aws-sdk-swift-core.git", .upToNextMinor(from: "5.0.0-alpha.5")),
         .package(url: "https://github.com/swift-server/swift-aws-lambda-runtime.git", .upToNextMajor(from: "0.2.0")),
     ],
@@ -26,25 +31,43 @@ 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"),
+            ]
+        ),
+        .testTarget(
+            name: "AWSXRayRecorderTests",
+            dependencies: [.target(name: "AWSXRayRecorder")]
+        ),
+        .target(
+            name: "AWSXRayInstrument",
+            dependencies: [
+                .target(name: "AWSXRayRecorder"),
+                .product(name: "Instrumentation", package: "gsoc-swift-tracing"),
+                .product(name: "Baggage", package: "swift-baggage-context"),
+            ]
+        ),
+        .testTarget(
+            name: "AWSXRayInstrumentTests",
+            dependencies: [
+                .target(name: "AWSXRayInstrument"),
+                .product(name: "NIOInstrumentation", package: "gsoc-swift-tracing"),
+                .product(name: "NIO", package: "swift-nio"),
+                .product(name: "NIOHTTP1", package: "swift-nio"),
             ]
         ),
         .target(
             name: "AWSXRayRecorderLambda",
             dependencies: [
-                .byName(name: "AWSXRayRecorder"),
+                .target(name: "AWSXRayRecorder"),
                 .product(name: "AWSLambdaRuntime", package: "swift-aws-lambda-runtime"),
             ]
         ),
         .target(
             name: "AWSXRayRecorderSDK",
             dependencies: [
-                .byName(name: "AWSXRayRecorder"),
+                .target(name: "AWSXRayRecorder"),
                 .product(name: "AWSSDKSwiftCore", package: "aws-sdk-swift-core"),
             ]
         ),
-        .testTarget(
-            name: "AWSXRayRecorderTests",
-            dependencies: ["AWSXRayRecorder"]
-        ),
     ]
 )

Sources/AWSXRayInstrument/Context.swift

@@ -0,0 +1,38 @@
+// TODO: move to XRayRecorder at some point
+
+import AWSXRayRecorder
+import Baggage
+
+// TODO: consider renaming the type
+internal typealias XRayContext = XRayRecorder.TraceContext
+
+internal enum AmazonHeaders {
+    static let traceId = "X-Amzn-Trace-Id"
+}
+
+private enum XRayContextKey: BaggageContextKey {
+    typealias Value = XRayContext
+
+    var name: String { "XRayContext" }
+}
+
+internal extension BaggageContext {
+    var xRayContext: XRayContext? {
+        get {
+            self[XRayContextKey.self]
+        }
+        set {
+            self[XRayContextKey.self] = newValue
+        }
+    }
+}
+
+public extension XRayRecorder {
+    func beginSegment(name: String, context: BaggageContext,
+                      aws: Segment.AWS? = nil, metadata: Segment.Metadata? = nil) -> XRayRecorder.Segment {
+        // TODO: use `AWS_XRAY_CONTEXT_MISSING` to configure how to handle missing context
+        // TODO: log error if context is missing, create new trace
+        let context = context.xRayContext ?? TraceContext(sampled: .unknown)
+        return beginSegment(name: name, context: context, aws: aws, metadata: metadata)
+    }
+}

Sources/AWSXRayInstrument/Instrument.swift

@@ -0,0 +1,24 @@
+import AWSXRayRecorder
+import Baggage
+import Dispatch // TODO: remove if/when not needed
+import Instrumentation
+
+extension XRayRecorder: TracingInstrument {
+    public func extract<Carrier, Extractor>(_ carrier: Carrier, into baggage: inout BaggageContext, using extractor: Extractor) where Carrier == Extractor.Carrier, Extractor: ExtractorProtocol {
+        guard let tracingHeader = extractor.extract(key: AmazonHeaders.traceId, from: carrier) else { return }
+        if let context = try? XRayContext(tracingHeader: tracingHeader) {
+            baggage.xRayContext = context
+        }
+    }
+
+    public func inject<Carrier, Injector>(_ baggage: BaggageContext, into carrier: inout Carrier, using injector: Injector) where Carrier == Injector.Carrier, Injector: InjectorProtocol {
+        guard let context = baggage.xRayContext else { return }
+        injector.inject(context.tracingHeader, forKey: AmazonHeaders.traceId, into: &carrier)
+    }
+
+    public func startSpan(named operationName: String, context: BaggageContext, ofKind kind: SpanKind, at timestamp: DispatchTime?) -> Span {
+        // TODO: map time type, see https://github.com/slashmo/gsoc-swift-tracing/pull/82#issuecomment-661868753
+        // TODO: does kind map anyhow to Subsegment?
+        beginSegment(name: operationName, context: context)
+    }
+}

Sources/AWSXRayInstrument/Span.swift

@@ -0,0 +1,136 @@
+import AnyCodable
+import AWSXRayRecorder
+import Baggage
+import Dispatch
+import Instrumentation
+
+extension XRayRecorder.Segment: Instrumentation.Span {
+    public var operationName: String { name }
+
+    public var kind: SpanKind { .internal }
+
+    public var status: SpanStatus? {
+        get { nil } // TODO: getter to be removed
+        set(newValue) {
+            if let status = newValue {
+                setStatus(status)
+            }
+        }
+    }
+
+    public func setStatus(_ status: SpanStatus) {
+        // TODO: should the status be set just once?
+        guard status.cannonicalCode != .ok else { return }
+        // note that contrary to what the name may suggest, exceptions are added not set
+        setException(message: status.message ?? "\(status.cannonicalCode)", type: "\(status.cannonicalCode)")
+    }
+
+    public var startTimestamp: DispatchTime { DispatchTime.now() } // TODO: getter to be removed
+
+    public var endTimestamp: DispatchTime? { nil } // TODO: getter to be removed
+
+    public func end(at timestamp: DispatchTime) {
+        end()
+    }
+
+    public var baggage: BaggageContext {
+        // TODO: make a Segment attribute
+        var baggage = BaggageContext()
+        baggage.xRayContext = context
+        return baggage
+    }
+
+    public var events: [SpanEvent] { [SpanEvent]() } // TODO: getter to be removed
+
+    public func addEvent(_ event: SpanEvent) {
+        // XRay segment does not have direct Span Event equivalent
+        // Arguably the closest match is a subsegment with startTime == endTime (?)
+        // TODO: test different approaches, make it configurable
+        beginSubsegment(name: event.name, metadata: nil).end()
+        // TODO: set Event attributes once interface is refined
+        // we can also store it as metadata as in https://github.com/awslabs/aws-xray-sdk-with-opentelemetry/commit/89f941af2b32844652c190b79328f9f783fe60f8
+        appendMetadata(AnyEncodable(event), forKey: MetadataKeys.events.rawValue)
+    }
+
+    public func setAttribute(_ value: String, forKey key: String) {
+        setAnnotation(value, forKey: key)
+    }
+
+    public func setAttribute(_ value: [String], forKey key: String) {
+        setMetadata(AnyEncodable(value), forKey: "attr_\(key)")
+    }
+
+    public func setAttribute(_ value: Int, forKey key: String) {
+        setAnnotation(value, forKey: key)
+    }
+
+    public func setAttribute(_ value: [Int], forKey key: String) {
+        setMetadata(AnyEncodable(value), forKey: "attr_\(key)")
+    }
+
+    public func setAttribute(_ value: Double, forKey key: String) {
+        setAnnotation(value, forKey: key)
+    }
+
+    public func setAttribute(_ value: [Double], forKey key: String) {
+        setMetadata(AnyEncodable(value), forKey: "attr_\(key)")
+    }
+
+    public func setAttribute(_ value: Bool, forKey key: String) {
+        setAnnotation(value, forKey: key)
+    }
+
+    public func setAttribute(_ value: [Bool], forKey key: String) {
+        setMetadata(AnyEncodable(value), forKey: "attr_\(key)")
+    }
+
+    public var isRecording: Bool {
+        context.sampled == .sampled
+    }
+
+    public var links: [SpanLink] { [SpanLink]() } // TODO: getter to be removed
+
+    public func addLink(_ link: SpanLink) {
+        appendMetadata(AnyEncodable(link), forKey: MetadataKeys.links.rawValue)
+    }
+}
+
+// MARK: -
+
+private enum MetadataKeys: String {
+    case events
+    case links
+}
+
+extension BaggageContext: Encodable {
+    public func encode(to encoder: Encoder) throws {
+        guard let context = xRayContext else { return }
+        var container = encoder.singleValueContainer()
+        // TODO: not sure if it makes sense to encode sampling decision
+        try container.encode(context.tracingHeader)
+    }
+}
+
+extension Instrumentation.SpanEvent: Encodable {
+    enum CodingKeys: String, CodingKey {
+        case name
+        // TODO: add attributes and timestamp after their types are updated
+    }
+
+    public func encode(to encoder: Encoder) throws {
+        var container = encoder.container(keyedBy: CodingKeys.self)
+        try container.encode(name, forKey: .name)
+    }
+}
+
+extension Instrumentation.SpanLink: Encodable {
+    enum CodingKeys: String, CodingKey {
+        case context
+        // TODO: add attributes
+    }
+
+    public func encode(to encoder: Encoder) throws {
+        var container = encoder.container(keyedBy: CodingKeys.self)
+        try container.encode(context, forKey: .context)
+    }
+}

Tests/AWSXRayInstrumentTests/ContextFactory.swift

@@ -0,0 +1,27 @@
+import Baggage
+
+@testable import AWSXRayInstrument
+
+extension BaggageContext {
+    static func empty() -> BaggageContext {
+        BaggageContext()
+    }
+
+    static func withTracingHeader(_ tracingHeader: String) -> BaggageContext {
+        var context = BaggageContext()
+        context.xRayContext = try? XRayContext(tracingHeader: tracingHeader)
+        return context
+    }
+
+    static func withoutParentSampled() -> BaggageContext {
+        var context = BaggageContext()
+        context.xRayContext = XRayContext(traceId: .init(), parentId: nil, sampled: .sampled)
+        return context
+    }
+
+    static func withoutParentNotSampled() -> BaggageContext {
+        var context = BaggageContext()
+        context.xRayContext = XRayContext(traceId: .init(), parentId: nil, sampled: .notSampled)
+        return context
+    }
+}