From 3bfc54f347ee731e803c8d08a5862efbd72d6dcc Mon Sep 17 00:00:00 2001
From: Michal A <dev@pokryfka.net>
Date: Thu, 6 Aug 2020 12:11:46 +0800
Subject: [PATCH 01/15] Instrument Lambda Runner

---
 Package.swift                                 |  5 +++-
 .../AWSLambdaRuntimeCore/LambdaContext.swift  | 13 ++++++++-
 .../AWSLambdaRuntimeCore/LambdaHandler.swift  | 28 ++++++++++++-------
 .../AWSLambdaRuntimeCore/LambdaRunner.swift   |  9 +++++-
 4 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/Package.swift b/Package.swift
index 1d56806e..d1b6f6df 100644
--- a/Package.swift
+++ b/Package.swift
@@ -5,7 +5,7 @@ import PackageDescription
 let package = Package(
     name: "swift-aws-lambda-runtime",
     platforms: [
-        .macOS(.v10_13),
+        .macOS(.v10_14), // TODO: fix in aws-xray-sdk-swift
     ],
     products: [
         // this library exports `AWSLambdaRuntimeCore` and adds Foundation convenience methods
@@ -21,6 +21,8 @@ let package = Package(
         .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/swift-server/swift-backtrace.git", .upToNextMajor(from: "1.1.0")),
+        // TODO: use swift-tracing when available
+        .package(url: "https://github.com/pokryfka/aws-xray-sdk-swift.git", .upToNextMinor(from: "0.6.0")),
     ],
     targets: [
         .target(name: "AWSLambdaRuntime", dependencies: [
@@ -32,6 +34,7 @@ let package = Package(
             .product(name: "Logging", package: "swift-log"),
             .product(name: "Backtrace", package: "swift-backtrace"),
             .product(name: "NIOHTTP1", package: "swift-nio"),
+            .product(name: "AWSXRaySDK", package: "aws-xray-sdk-swift"),
         ]),
         .testTarget(name: "AWSLambdaRuntimeCoreTests", dependencies: [
             .byName(name: "AWSLambdaRuntimeCore"),
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
index ab30dd7b..098debb2 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+import AWSXRayRecorder
 import Dispatch
 import Logging
 import NIO
@@ -68,11 +69,16 @@ extension Lambda {
         /// For invocations from the AWS Mobile SDK, data about the client application and device.
         public let clientContext: String?
 
+        public let baggage: XRayContext // TODO: use BaggageContext when swift-tracign is ready
+
         /// `Logger` to log with
         ///
         /// - note: The `LogLevel` can be configured using the `LOG_LEVEL` environment variable.
         public let logger: Logger
 
+        /// Tracing instrument.
+        public let tracer: TracingInstrument
+
         /// The `EventLoop` the Lambda is executed on. Use this to schedule work with.
         /// This is useful when implementing the `EventLoopLambdaHandler` protocol.
         ///
@@ -91,8 +97,10 @@ extension Lambda {
                       cognitoIdentity: String? = nil,
                       clientContext: String? = nil,
                       logger: Logger,
+                      tracer: TracingInstrument,
                       eventLoop: EventLoop,
-                      allocator: ByteBufferAllocator) {
+                      allocator: ByteBufferAllocator)
+        {
             self.requestID = requestID
             self.traceID = traceID
             self.invokedFunctionARN = invokedFunctionARN
@@ -106,7 +114,10 @@ extension Lambda {
             var logger = logger
             logger[metadataKey: "awsRequestID"] = .string(requestID)
             logger[metadataKey: "awsTraceID"] = .string(traceID)
+            // TODO: handle error in better way
+            self.baggage = try! XRayContext(tracingHeader: traceID)
             self.logger = logger
+            self.tracer = tracer
         }
 
         public func getRemainingTime() -> TimeAmount {
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
index 16eba1cb..3534f03c 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
@@ -129,16 +129,24 @@ public protocol EventLoopLambdaHandler: ByteBufferLambdaHandler {
 public extension EventLoopLambdaHandler {
     /// Driver for `ByteBuffer` -> `In` decoding and `Out` -> `ByteBuffer` encoding
     func handle(context: Lambda.Context, event: ByteBuffer) -> EventLoopFuture<ByteBuffer?> {
-        switch self.decodeIn(buffer: event) {
-        case .failure(let error):
-            return context.eventLoop.makeFailedFuture(CodecError.requestDecoding(error))
-        case .success(let `in`):
-            return self.handle(context: context, event: `in`).flatMapThrowing { out in
-                switch self.encodeOut(allocator: context.allocator, value: out) {
-                case .failure(let error):
-                    throw CodecError.responseEncoding(error)
-                case .success(let buffer):
-                    return buffer
+        // TODO: creating subsegments with NIO which will be easier if the baggage is passes in channel
+        // see https://github.com/slashmo/gsoc-swift-tracing/issues/48
+        context.tracer.segment(name: "HandleEvent", context: context.baggage) {
+            // TODO: create helper to record errors in result types
+            let decodedEvent = context.tracer.segment(name: "DecodeIn", context: context.baggage) { _ in self.decodeIn(buffer: event) }
+            switch decodedEvent {
+            case .failure(let error):
+                return context.eventLoop.makeFailedFuture(CodecError.requestDecoding(error))
+            case .success(let `in`):
+                return self.handle(context: context, event: `in`).flatMapThrowing { out in
+                    try context.tracer.segment(name: "encodeOut", context: context.baggage) { _ in
+                        switch self.encodeOut(allocator: context.allocator, value: out) {
+                        case .failure(let error):
+                            throw CodecError.responseEncoding(error)
+                        case .success(let buffer):
+                            return buffer
+                        }
+                    }
                 }
             }
         }
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
index 8fc22de3..78e59af0 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
@@ -12,14 +12,18 @@
 //
 //===----------------------------------------------------------------------===//
 
+import AWSXRaySDK // TODO: use swift-tracing when available
 import Dispatch
 import Logging
 import NIO
 
+public typealias TracingInstrument = XRayRecorder
+
 extension Lambda {
     /// LambdaRunner manages the Lambda runtime workflow, or business logic.
     internal final class Runner {
         private let runtimeClient: RuntimeClient
+        internal let tracer: TracingInstrument
         private let eventLoop: EventLoop
         private let allocator: ByteBufferAllocator
 
@@ -28,6 +32,7 @@ extension Lambda {
         init(eventLoop: EventLoop, configuration: Configuration) {
             self.eventLoop = eventLoop
             self.runtimeClient = RuntimeClient(eventLoop: self.eventLoop, configuration: configuration.runtimeEngine)
+            self.tracer = XRayRecorder(eventLoopGroupProvider: .shared(eventLoop))
             self.allocator = ByteBufferAllocator()
         }
 
@@ -66,6 +71,7 @@ extension Lambda {
                 // 2. send invocation to handler
                 self.isGettingNextInvocation = false
                 let context = Context(logger: logger,
+                                      tracer: self.tracer,
                                       eventLoop: self.eventLoop,
                                       allocator: self.allocator,
                                       invocation: invocation)
@@ -101,7 +107,7 @@ extension Lambda {
 }
 
 private extension Lambda.Context {
-    convenience init(logger: Logger, eventLoop: EventLoop, allocator: ByteBufferAllocator, invocation: Lambda.Invocation) {
+    convenience init(logger: Logger, tracer: TracingInstrument, eventLoop: EventLoop, allocator: ByteBufferAllocator, invocation: Lambda.Invocation) {
         self.init(requestID: invocation.requestID,
                   traceID: invocation.traceID,
                   invokedFunctionARN: invocation.invokedFunctionARN,
@@ -109,6 +115,7 @@ private extension Lambda.Context {
                   cognitoIdentity: invocation.cognitoIdentity,
                   clientContext: invocation.clientContext,
                   logger: logger,
+                  tracer: tracer,
                   eventLoop: eventLoop,
                   allocator: allocator)
     }

From 4fa2e4ceba86175165fc9e10eac927c60173fa76 Mon Sep 17 00:00:00 2001
From: Michal A <dev@pokryfka.net>
Date: Thu, 6 Aug 2020 14:44:54 +0800
Subject: [PATCH 02/15] Fix compilation of AWSLambdaTesting and tests

---
 Package.swift                                         | 2 ++
 Sources/AWSLambdaTesting/Lambda+Testing.swift         | 4 ++++
 Tests/AWSLambdaRuntimeCoreTests/LambdaTest.swift      | 6 ++++++
 Tests/AWSLambdaRuntimeTests/Lambda+CodeableTest.swift | 4 ++++
 4 files changed, 16 insertions(+)

diff --git a/Package.swift b/Package.swift
index d1b6f6df..3b8322f7 100644
--- a/Package.swift
+++ b/Package.swift
@@ -44,6 +44,7 @@ let package = Package(
         .testTarget(name: "AWSLambdaRuntimeTests", dependencies: [
             .byName(name: "AWSLambdaRuntimeCore"),
             .byName(name: "AWSLambdaRuntime"),
+            .product(name: "AWSXRayRecorder", package: "aws-xray-sdk-swift"),
         ]),
         .target(name: "AWSLambdaEvents", dependencies: []),
         .testTarget(name: "AWSLambdaEventsTests", dependencies: ["AWSLambdaEvents"]),
@@ -51,6 +52,7 @@ let package = Package(
         .target(name: "AWSLambdaTesting", dependencies: [
             .byName(name: "AWSLambdaRuntime"),
             .product(name: "NIO", package: "swift-nio"),
+            .product(name: "AWSXRayRecorder", package: "aws-xray-sdk-swift"),
         ]),
         .testTarget(name: "AWSLambdaTestingTests", dependencies: ["AWSLambdaTesting"]),
         // for perf testing
diff --git a/Sources/AWSLambdaTesting/Lambda+Testing.swift b/Sources/AWSLambdaTesting/Lambda+Testing.swift
index 981ca736..38716ef9 100644
--- a/Sources/AWSLambdaTesting/Lambda+Testing.swift
+++ b/Sources/AWSLambdaTesting/Lambda+Testing.swift
@@ -36,10 +36,13 @@
 #if DEBUG
 @testable import AWSLambdaRuntime
 @testable import AWSLambdaRuntimeCore
+import AWSXRayRecorder
 import Dispatch
 import Logging
 import NIO
 
+private let noOpTracer = XRayRecorder(emitter: XRayNoOpEmitter())
+
 extension Lambda {
     public struct TestConfig {
         public var requestID: String
@@ -102,6 +105,7 @@ extension Lambda {
                               invokedFunctionARN: config.invokedFunctionARN,
                               deadline: .now() + config.timeout,
                               logger: logger,
+                              tracer: noOpTracer,
                               eventLoop: eventLoop,
                               allocator: ByteBufferAllocator())
 
diff --git a/Tests/AWSLambdaRuntimeCoreTests/LambdaTest.swift b/Tests/AWSLambdaRuntimeCoreTests/LambdaTest.swift
index 30fed618..5d5d035e 100644
--- a/Tests/AWSLambdaRuntimeCoreTests/LambdaTest.swift
+++ b/Tests/AWSLambdaRuntimeCoreTests/LambdaTest.swift
@@ -13,10 +13,13 @@
 //===----------------------------------------------------------------------===//
 
 @testable import AWSLambdaRuntimeCore
+import AWSXRayRecorder
 import Logging
 import NIO
 import XCTest
 
+private let noOpTracer = XRayRecorder(emitter: XRayNoOpEmitter())
+
 class LambdaTest: XCTestCase {
     func testSuccess() {
         let server = MockLambdaServer(behavior: Behavior())
@@ -263,6 +266,7 @@ class LambdaTest: XCTestCase {
                                      cognitoIdentity: nil,
                                      clientContext: nil,
                                      logger: Logger(label: "test"),
+                                     tracer: noOpTracer,
                                      eventLoop: MultiThreadedEventLoopGroup(numberOfThreads: 1).next(),
                                      allocator: ByteBufferAllocator())
         XCTAssertGreaterThan(context.deadline, .now())
@@ -274,6 +278,7 @@ class LambdaTest: XCTestCase {
                                             cognitoIdentity: context.cognitoIdentity,
                                             clientContext: context.clientContext,
                                             logger: context.logger,
+                                            tracer: noOpTracer,
                                             eventLoop: context.eventLoop,
                                             allocator: context.allocator)
         XCTAssertLessThan(expiredContext.deadline, .now())
@@ -287,6 +292,7 @@ class LambdaTest: XCTestCase {
                                      cognitoIdentity: nil,
                                      clientContext: nil,
                                      logger: Logger(label: "test"),
+                                     tracer: noOpTracer,
                                      eventLoop: MultiThreadedEventLoopGroup(numberOfThreads: 1).next(),
                                      allocator: ByteBufferAllocator())
         XCTAssertLessThanOrEqual(context.getRemainingTime(), .seconds(1))
diff --git a/Tests/AWSLambdaRuntimeTests/Lambda+CodeableTest.swift b/Tests/AWSLambdaRuntimeTests/Lambda+CodeableTest.swift
index 9aa3f72a..d1d1cc05 100644
--- a/Tests/AWSLambdaRuntimeTests/Lambda+CodeableTest.swift
+++ b/Tests/AWSLambdaRuntimeTests/Lambda+CodeableTest.swift
@@ -14,11 +14,14 @@
 
 @testable import AWSLambdaRuntime
 @testable import AWSLambdaRuntimeCore
+import AWSXRayRecorder
 import Logging
 import NIO
 import NIOFoundationCompat
 import XCTest
 
+private let noOpTracer = XRayRecorder(emitter: XRayNoOpEmitter())
+
 class CodableLambdaTest: XCTestCase {
     var eventLoopGroup: EventLoopGroup!
     let allocator = ByteBufferAllocator()
@@ -72,6 +75,7 @@ class CodableLambdaTest: XCTestCase {
                        cognitoIdentity: nil,
                        clientContext: nil,
                        logger: Logger(label: "test"),
+                       tracer: noOpTracer,
                        eventLoop: self.eventLoopGroup.next(),
                        allocator: ByteBufferAllocator())
     }

From 99e2d366800b80c15085d6e2243621345e407df1 Mon Sep 17 00:00:00 2001
From: Michal A <dev@pokryfka.net>
Date: Thu, 6 Aug 2020 14:48:53 +0800
Subject: [PATCH 03/15] Flush the tracer in each invocation

---
 Sources/AWSLambdaRuntimeCore/LambdaHandler.swift | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
index 3534f03c..26a53e5c 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
@@ -129,17 +129,17 @@ public protocol EventLoopLambdaHandler: ByteBufferLambdaHandler {
 public extension EventLoopLambdaHandler {
     /// Driver for `ByteBuffer` -> `In` decoding and `Out` -> `ByteBuffer` encoding
     func handle(context: Lambda.Context, event: ByteBuffer) -> EventLoopFuture<ByteBuffer?> {
-        // TODO: creating subsegments with NIO which will be easier if the baggage is passes in channel
-        // see https://github.com/slashmo/gsoc-swift-tracing/issues/48
         context.tracer.segment(name: "HandleEvent", context: context.baggage) {
-            // TODO: create helper to record errors in result types
-            let decodedEvent = context.tracer.segment(name: "DecodeIn", context: context.baggage) { _ in self.decodeIn(buffer: event) }
+            // TODO: create helper to record errors in passed in result types
+            let decodedEvent = context.tracer.segment(name: "DecodeIn", context: context.baggage) { _ in
+                self.decodeIn(buffer: event)
+            }
             switch decodedEvent {
             case .failure(let error):
                 return context.eventLoop.makeFailedFuture(CodecError.requestDecoding(error))
             case .success(let `in`):
                 return self.handle(context: context, event: `in`).flatMapThrowing { out in
-                    try context.tracer.segment(name: "encodeOut", context: context.baggage) { _ in
+                    try context.tracer.segment(name: "EncodeOut", context: context.baggage) { _ in
                         switch self.encodeOut(allocator: context.allocator, value: out) {
                         case .failure(let error):
                             throw CodecError.responseEncoding(error)
@@ -149,6 +149,9 @@ public extension EventLoopLambdaHandler {
                     }
                 }
             }
+        }.flatMap { result in
+            // flush the tracer after each invocation and return the invocation result
+            context.tracer.flush(on: context.eventLoop).map { result }
         }
     }
 

From b63990137a26beef71b4784fa7a8fce48ba15b05 Mon Sep 17 00:00:00 2001
From: Michal A <dev@pokryfka.net>
Date: Fri, 7 Aug 2020 10:45:33 +0800
Subject: [PATCH 04/15] Flush the tracer in LambdaRunner, handler is a
 subsegment of HandleIn

---
 Package.swift                                 |  5 ++-
 .../AWSLambdaRuntimeCore/LambdaContext.swift  | 10 +++--
 .../AWSLambdaRuntimeCore/LambdaHandler.swift  | 43 ++++++++++++-------
 .../AWSLambdaRuntimeCore/LambdaRunner.swift   |  4 ++
 4 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/Package.swift b/Package.swift
index 3b8322f7..aa63ba9f 100644
--- a/Package.swift
+++ b/Package.swift
@@ -5,7 +5,7 @@ import PackageDescription
 let package = Package(
     name: "swift-aws-lambda-runtime",
     platforms: [
-        .macOS(.v10_14), // TODO: fix in aws-xray-sdk-swift
+        .macOS(.v10_14), // TODO: should not be needed soon
     ],
     products: [
         // this library exports `AWSLambdaRuntimeCore` and adds Foundation convenience methods
@@ -19,10 +19,11 @@ let package = Package(
     ],
     dependencies: [
         .package(url: "https://github.com/apple/swift-nio.git", .upToNextMajor(from: "2.17.0")),
+//        .package(url: "https://github.com/pokryfka/swift-nio.git", .branch("feature/tracing")),
         .package(url: "https://github.com/apple/swift-log.git", .upToNextMajor(from: "1.0.0")),
         .package(url: "https://github.com/swift-server/swift-backtrace.git", .upToNextMajor(from: "1.1.0")),
         // TODO: use swift-tracing when available
-        .package(url: "https://github.com/pokryfka/aws-xray-sdk-swift.git", .upToNextMinor(from: "0.6.0")),
+        .package(url: "https://github.com/pokryfka/aws-xray-sdk-swift.git", .upToNextMinor(from: "0.6.1")),
     ],
     targets: [
         .target(name: "AWSLambdaRuntime", dependencies: [
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
index 098debb2..3b404778 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 import AWSXRayRecorder
+import Baggage
 import Dispatch
 import Logging
 import NIO
@@ -69,7 +70,8 @@ extension Lambda {
         /// For invocations from the AWS Mobile SDK, data about the client application and device.
         public let clientContext: String?
 
-        public let baggage: XRayContext // TODO: use BaggageContext when swift-tracign is ready
+        /// Context baggage.
+        public var baggage: BaggageContext
 
         /// `Logger` to log with
         ///
@@ -114,8 +116,10 @@ extension Lambda {
             var logger = logger
             logger[metadataKey: "awsRequestID"] = .string(requestID)
             logger[metadataKey: "awsTraceID"] = .string(traceID)
-            // TODO: handle error in better way
-            self.baggage = try! XRayContext(tracingHeader: traceID)
+            var baggage = BaggageContext()
+            // TODO: handle error
+            baggage.xRayContext = try? XRayContext(tracingHeader: traceID)
+            self.baggage = baggage
             self.logger = logger
             self.tracer = tracer
         }
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
index 26a53e5c..43e50e32 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
@@ -129,17 +129,29 @@ public protocol EventLoopLambdaHandler: ByteBufferLambdaHandler {
 public extension EventLoopLambdaHandler {
     /// Driver for `ByteBuffer` -> `In` decoding and `Out` -> `ByteBuffer` encoding
     func handle(context: Lambda.Context, event: ByteBuffer) -> EventLoopFuture<ByteBuffer?> {
-        context.tracer.segment(name: "HandleEvent", context: context.baggage) {
-            // TODO: create helper to record errors in passed in result types
-            let decodedEvent = context.tracer.segment(name: "DecodeIn", context: context.baggage) { _ in
-                self.decodeIn(buffer: event)
-            }
-            switch decodedEvent {
-            case .failure(let error):
-                return context.eventLoop.makeFailedFuture(CodecError.requestDecoding(error))
-            case .success(let `in`):
-                return self.handle(context: context, event: `in`).flatMapThrowing { out in
-                    try context.tracer.segment(name: "EncodeOut", context: context.baggage) { _ in
+        let segment = context.tracer.beginSegment(name: "HandleEvent", baggage: context.baggage)
+        // TODO: record errors propagated in result types?
+        let decodedEvent = segment.subsegment(name: "DecodeIn") { _ in
+            self.decodeIn(buffer: event)
+        }
+        switch decodedEvent {
+        case .failure(let error):
+            segment.addError(error)
+            segment.end()
+            return context.eventLoop.makeFailedFuture(CodecError.requestDecoding(error))
+        case .success(let `in`):
+            // TODO: use NIO helpers?
+            let subsegment = segment.beginSubsegment(name: "HandleIn")
+            context.baggage = subsegment.baggage
+            return self.handle(context: context, event: `in`)
+                .always { result in
+                    if case .failure(let error) = result {
+                        subsegment.addError(error)
+                    }
+                    subsegment.end()
+                }
+                .flatMapThrowing { out in
+                    try context.tracer.segment(name: "EncodeOut", baggage: segment.baggage) { _ in
                         switch self.encodeOut(allocator: context.allocator, value: out) {
                         case .failure(let error):
                             throw CodecError.responseEncoding(error)
@@ -147,11 +159,12 @@ public extension EventLoopLambdaHandler {
                             return buffer
                         }
                     }
+                }.always { result in
+                    if case .failure(let error) = result {
+                        segment.addError(error)
+                    }
+                    segment.end()
                 }
-            }
-        }.flatMap { result in
-            // flush the tracer after each invocation and return the invocation result
-            context.tracer.flush(on: context.eventLoop).map { result }
         }
     }
 
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
index 78e59af0..938e3868 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
@@ -61,6 +61,7 @@ extension Lambda {
                 }
         }
 
+        // TODO: instrument custom runtime API
         func run(logger: Logger, handler: Handler) -> EventLoopFuture<Void> {
             logger.debug("lambda invocation sequence starting")
             // 1. request invocation from lambda runtime engine
@@ -93,6 +94,9 @@ extension Lambda {
                 self.runtimeClient.reportResults(logger: logger, invocation: invocation, result: result).peekError { error in
                     logger.error("could not report results to lambda runtime engine: \(error)")
                 }
+            }.flatMap {
+                // flush the tracer after each invocation
+                self.tracer.flush(on: self.eventLoop)
             }
         }
 

From 6609e519b1ea04a213727e357fe0daea76537f8b Mon Sep 17 00:00:00 2001
From: Michal A <dev@pokryfka.net>
Date: Fri, 7 Aug 2020 11:09:28 +0800
Subject: [PATCH 05/15] Record ReportResults

---
 Sources/AWSLambdaRuntimeCore/LambdaRunner.swift | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
index 938e3868..9223cb1f 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 import AWSXRaySDK // TODO: use swift-tracing when available
+import Baggage
 import Dispatch
 import Logging
 import NIO
@@ -61,11 +62,11 @@ extension Lambda {
                 }
         }
 
-        // TODO: instrument custom runtime API
         func run(logger: Logger, handler: Handler) -> EventLoopFuture<Void> {
             logger.debug("lambda invocation sequence starting")
             // 1. request invocation from lambda runtime engine
             self.isGettingNextInvocation = true
+            // TODO: add API to explicitly set startTime, after all
             return self.runtimeClient.getNextInvocation(logger: logger).peekError { error in
                 logger.error("could not fetch work from lambda runtime engine: \(error)")
             }.flatMap { invocation, event in
@@ -76,6 +77,7 @@ extension Lambda {
                                       eventLoop: self.eventLoop,
                                       allocator: self.allocator,
                                       invocation: invocation)
+                let baggage = context.baggage
                 logger.debug("sending invocation to lambda handler \(handler)")
                 return handler.handle(context: context, event: event)
                     // Hopping back to "our" EventLoop is importnant in case the handler returns a future that
@@ -87,12 +89,14 @@ extension Lambda {
                         if case .failure(let error) = result {
                             logger.warning("lambda handler returned an error: \(error)")
                         }
-                        return (invocation, result)
+                        return (invocation, result, baggage)
                     }
-            }.flatMap { invocation, result in
+            }.flatMap { (invocation, result, baggage: BaggageContext) in
                 // 3. report results to runtime engine
-                self.runtimeClient.reportResults(logger: logger, invocation: invocation, result: result).peekError { error in
-                    logger.error("could not report results to lambda runtime engine: \(error)")
+                self.tracer.segment(name: "ReportResults", baggage: baggage) { _ in
+                    self.runtimeClient.reportResults(logger: logger, invocation: invocation, result: result).peekError { error in
+                        logger.error("could not report results to lambda runtime engine: \(error)")
+                    }
                 }
             }.flatMap {
                 // flush the tracer after each invocation

From 7b42510c286e7d878a964e911ea4f76f8242a3e5 Mon Sep 17 00:00:00 2001
From: Michal A <dev@pokryfka.net>
Date: Fri, 7 Aug 2020 16:59:33 +0800
Subject: [PATCH 06/15] Remove dependency on Foundation

---
 Package.swift | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/Package.swift b/Package.swift
index aa63ba9f..e0fb2dc7 100644
--- a/Package.swift
+++ b/Package.swift
@@ -23,7 +23,8 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-log.git", .upToNextMajor(from: "1.0.0")),
         .package(url: "https://github.com/swift-server/swift-backtrace.git", .upToNextMajor(from: "1.1.0")),
         // TODO: use swift-tracing when available
-        .package(url: "https://github.com/pokryfka/aws-xray-sdk-swift.git", .upToNextMinor(from: "0.6.1")),
+        //.package(url: "https://github.com/pokryfka/aws-xray-sdk-swift.git", .upToNextMinor(from: "0.6.1")),
+        .package(url: "https://github.com/pokryfka/aws-xray-sdk-swift.git", .branch("feature/foundation")),
     ],
     targets: [
         .target(name: "AWSLambdaRuntime", dependencies: [

From 50336b62987ddd0ec5722fe69fdd293ab0bc268a Mon Sep 17 00:00:00 2001
From: Michal A <dev@pokryfka.net>
Date: Wed, 12 Aug 2020 13:49:04 +0800
Subject: [PATCH 07/15] Use X-Ray SDK 0.7.0

---
 Package.swift | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/Package.swift b/Package.swift
index 0f9d7f16..966cfd60 100644
--- a/Package.swift
+++ b/Package.swift
@@ -20,8 +20,7 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-log.git", .upToNextMajor(from: "1.0.0")),
         .package(url: "https://github.com/swift-server/swift-backtrace.git", .upToNextMajor(from: "1.1.0")),
         // TODO: use swift-tracing when available
-        //.package(url: "https://github.com/pokryfka/aws-xray-sdk-swift.git", .upToNextMinor(from: "0.6.1")),
-        .package(url: "https://github.com/pokryfka/aws-xray-sdk-swift.git", .branch("feature/foundation")),
+        .package(url: "https://github.com/pokryfka/aws-xray-sdk-swift.git", .upToNextMinor(from: "0.7.0")),
     ],
     targets: [
         .target(name: "AWSLambdaRuntime", dependencies: [

From 5ee3cfe78af452f94cf800b8677f0a89448a36ab Mon Sep 17 00:00:00 2001
From: Michal A <dev@pokryfka.net>
Date: Mon, 17 Aug 2020 12:38:48 +0800
Subject: [PATCH 08/15] Use X-Ray SDK 0.7.1, use convenience endSegment and
 flush

---
 Package.swift                                    |  2 +-
 Sources/AWSLambdaRuntimeCore/LambdaHandler.swift | 15 ++-------------
 Sources/AWSLambdaRuntimeCore/LambdaRunner.swift  |  5 ++---
 3 files changed, 5 insertions(+), 17 deletions(-)

diff --git a/Package.swift b/Package.swift
index 966cfd60..cacb867b 100644
--- a/Package.swift
+++ b/Package.swift
@@ -20,7 +20,7 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-log.git", .upToNextMajor(from: "1.0.0")),
         .package(url: "https://github.com/swift-server/swift-backtrace.git", .upToNextMajor(from: "1.1.0")),
         // TODO: use swift-tracing when available
-        .package(url: "https://github.com/pokryfka/aws-xray-sdk-swift.git", .upToNextMinor(from: "0.7.0")),
+        .package(url: "https://github.com/pokryfka/aws-xray-sdk-swift.git", .upToNextMinor(from: "0.7.1")),
     ],
     targets: [
         .target(name: "AWSLambdaRuntime", dependencies: [
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
index 43e50e32..0889eccf 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
@@ -130,7 +130,6 @@ public extension EventLoopLambdaHandler {
     /// Driver for `ByteBuffer` -> `In` decoding and `Out` -> `ByteBuffer` encoding
     func handle(context: Lambda.Context, event: ByteBuffer) -> EventLoopFuture<ByteBuffer?> {
         let segment = context.tracer.beginSegment(name: "HandleEvent", baggage: context.baggage)
-        // TODO: record errors propagated in result types?
         let decodedEvent = segment.subsegment(name: "DecodeIn") { _ in
             self.decodeIn(buffer: event)
         }
@@ -140,16 +139,10 @@ public extension EventLoopLambdaHandler {
             segment.end()
             return context.eventLoop.makeFailedFuture(CodecError.requestDecoding(error))
         case .success(let `in`):
-            // TODO: use NIO helpers?
             let subsegment = segment.beginSubsegment(name: "HandleIn")
             context.baggage = subsegment.baggage
             return self.handle(context: context, event: `in`)
-                .always { result in
-                    if case .failure(let error) = result {
-                        subsegment.addError(error)
-                    }
-                    subsegment.end()
-                }
+                .endSegment(subsegment)
                 .flatMapThrowing { out in
                     try context.tracer.segment(name: "EncodeOut", baggage: segment.baggage) { _ in
                         switch self.encodeOut(allocator: context.allocator, value: out) {
@@ -159,12 +152,8 @@ public extension EventLoopLambdaHandler {
                             return buffer
                         }
                     }
-                }.always { result in
-                    if case .failure(let error) = result {
-                        segment.addError(error)
-                    }
-                    segment.end()
                 }
+                .endSegment(segment)
         }
     }
 
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
index 9223cb1f..c0fd11c0 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
@@ -98,10 +98,9 @@ extension Lambda {
                         logger.error("could not report results to lambda runtime engine: \(error)")
                     }
                 }
-            }.flatMap {
-                // flush the tracer after each invocation
-                self.tracer.flush(on: self.eventLoop)
             }
+            // flush the tracer after each invocation
+            .flush(self.tracer, recover: false)
         }
 
         /// cancels the current run, if we are waiting for next invocation (long poll from Lambda control plane)

From 5f9773b24ba191e41449ccd4db6bd075e4f9723b Mon Sep 17 00:00:00 2001
From: Michal A <dev@pokryfka.net>
Date: Mon, 17 Aug 2020 13:20:56 +0800
Subject: [PATCH 09/15] Bootstrap tracer implementation

---
 Sources/AWSLambdaRuntimeCore/Lambda.swift      | 18 ++++++++++++++++++
 .../LambdaConfiguration.swift                  |  6 +++++-
 .../AWSLambdaRuntimeCore/LambdaRunner.swift    |  6 ++++--
 Sources/AWSLambdaTesting/Lambda+Testing.swift  |  5 +----
 .../AWSLambdaRuntimeCoreTests/LambdaTest.swift |  9 +++------
 .../Lambda+CodeableTest.swift                  |  5 +----
 6 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/Sources/AWSLambdaRuntimeCore/Lambda.swift b/Sources/AWSLambdaRuntimeCore/Lambda.swift
index 5dc27648..80ef7a6f 100644
--- a/Sources/AWSLambdaRuntimeCore/Lambda.swift
+++ b/Sources/AWSLambdaRuntimeCore/Lambda.swift
@@ -21,6 +21,7 @@ import Darwin.C
 import Backtrace
 import Logging
 import NIO
+import NIOConcurrencyHelpers
 
 public enum Lambda {
     public typealias Handler = ByteBufferLambdaHandler
@@ -30,6 +31,23 @@ public enum Lambda {
     /// A function that takes a `InitializationContext` and returns an `EventLoopFuture` of a `ByteBufferLambdaHandler`
     public typealias HandlerFactory = (InitializationContext) -> EventLoopFuture<Handler>
 
+    internal typealias TracerFactory = (EventLoop) -> TracingInstrument
+
+    private static let lock = Lock()
+    private static var tracerFactory: TracerFactory = { _ in NoOpTracingInstrument() }
+
+    /// Select the desired `TracingInstrument` implementation.
+    ///
+    /// The caller is responsible for the tracer lifecycle.
+    ///
+    /// - parameters:
+    ///     -  tracerFactory: creates `TracingInstrument` implementation
+    public static func bootstrap(_ tracerFactory: @escaping (EventLoop) -> TracingInstrument) {
+        self.lock.withLockVoid {
+            self.tracerFactory = tracerFactory
+        }
+    }
+
     /// Run a Lambda defined by implementing the `LambdaHandler` protocol.
     ///
     /// - parameters:
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaConfiguration.swift b/Sources/AWSLambdaRuntimeCore/LambdaConfiguration.swift
index 9b9ec8fb..38d7b4d1 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaConfiguration.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaConfiguration.swift
@@ -21,15 +21,19 @@ extension Lambda {
         let general: General
         let lifecycle: Lifecycle
         let runtimeEngine: RuntimeEngine
+        let tracerFactory: TracerFactory
 
         init() {
             self.init(general: .init(), lifecycle: .init(), runtimeEngine: .init())
         }
 
-        init(general: General? = nil, lifecycle: Lifecycle? = nil, runtimeEngine: RuntimeEngine? = nil) {
+        init(general: General? = nil, lifecycle: Lifecycle? = nil, runtimeEngine: RuntimeEngine? = nil,
+             tracerFactory: TracerFactory? = nil)
+        {
             self.general = general ?? General()
             self.lifecycle = lifecycle ?? Lifecycle()
             self.runtimeEngine = runtimeEngine ?? RuntimeEngine()
+            self.tracerFactory = tracerFactory ?? { _ in NoOpTracingInstrument() }
         }
 
         struct General: CustomStringConvertible {
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
index c0fd11c0..c84b89a3 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
@@ -18,13 +18,15 @@ import Dispatch
 import Logging
 import NIO
 
+// type names defined in `TracingInstrument`, aliases will be removed
 public typealias TracingInstrument = XRayRecorder
+public typealias NoOpTracingInstrument = XRayNoOpRecorder
 
 extension Lambda {
     /// LambdaRunner manages the Lambda runtime workflow, or business logic.
     internal final class Runner {
         private let runtimeClient: RuntimeClient
-        internal let tracer: TracingInstrument
+        private let tracer: TracingInstrument
         private let eventLoop: EventLoop
         private let allocator: ByteBufferAllocator
 
@@ -33,7 +35,7 @@ extension Lambda {
         init(eventLoop: EventLoop, configuration: Configuration) {
             self.eventLoop = eventLoop
             self.runtimeClient = RuntimeClient(eventLoop: self.eventLoop, configuration: configuration.runtimeEngine)
-            self.tracer = XRayRecorder(eventLoopGroupProvider: .shared(eventLoop))
+            self.tracer = configuration.tracerFactory(eventLoop)
             self.allocator = ByteBufferAllocator()
         }
 
diff --git a/Sources/AWSLambdaTesting/Lambda+Testing.swift b/Sources/AWSLambdaTesting/Lambda+Testing.swift
index 38716ef9..2dadddc4 100644
--- a/Sources/AWSLambdaTesting/Lambda+Testing.swift
+++ b/Sources/AWSLambdaTesting/Lambda+Testing.swift
@@ -36,13 +36,10 @@
 #if DEBUG
 @testable import AWSLambdaRuntime
 @testable import AWSLambdaRuntimeCore
-import AWSXRayRecorder
 import Dispatch
 import Logging
 import NIO
 
-private let noOpTracer = XRayRecorder(emitter: XRayNoOpEmitter())
-
 extension Lambda {
     public struct TestConfig {
         public var requestID: String
@@ -105,7 +102,7 @@ extension Lambda {
                               invokedFunctionARN: config.invokedFunctionARN,
                               deadline: .now() + config.timeout,
                               logger: logger,
-                              tracer: noOpTracer,
+                              tracer: NoOpTracingInstrument(),
                               eventLoop: eventLoop,
                               allocator: ByteBufferAllocator())
 
diff --git a/Tests/AWSLambdaRuntimeCoreTests/LambdaTest.swift b/Tests/AWSLambdaRuntimeCoreTests/LambdaTest.swift
index 5d5d035e..5f67e395 100644
--- a/Tests/AWSLambdaRuntimeCoreTests/LambdaTest.swift
+++ b/Tests/AWSLambdaRuntimeCoreTests/LambdaTest.swift
@@ -13,13 +13,10 @@
 //===----------------------------------------------------------------------===//
 
 @testable import AWSLambdaRuntimeCore
-import AWSXRayRecorder
 import Logging
 import NIO
 import XCTest
 
-private let noOpTracer = XRayRecorder(emitter: XRayNoOpEmitter())
-
 class LambdaTest: XCTestCase {
     func testSuccess() {
         let server = MockLambdaServer(behavior: Behavior())
@@ -266,7 +263,7 @@ class LambdaTest: XCTestCase {
                                      cognitoIdentity: nil,
                                      clientContext: nil,
                                      logger: Logger(label: "test"),
-                                     tracer: noOpTracer,
+                                     tracer: NoOpTracingInstrument(),
                                      eventLoop: MultiThreadedEventLoopGroup(numberOfThreads: 1).next(),
                                      allocator: ByteBufferAllocator())
         XCTAssertGreaterThan(context.deadline, .now())
@@ -278,7 +275,7 @@ class LambdaTest: XCTestCase {
                                             cognitoIdentity: context.cognitoIdentity,
                                             clientContext: context.clientContext,
                                             logger: context.logger,
-                                            tracer: noOpTracer,
+                                            tracer: NoOpTracingInstrument(),
                                             eventLoop: context.eventLoop,
                                             allocator: context.allocator)
         XCTAssertLessThan(expiredContext.deadline, .now())
@@ -292,7 +289,7 @@ class LambdaTest: XCTestCase {
                                      cognitoIdentity: nil,
                                      clientContext: nil,
                                      logger: Logger(label: "test"),
-                                     tracer: noOpTracer,
+                                     tracer: NoOpTracingInstrument(),
                                      eventLoop: MultiThreadedEventLoopGroup(numberOfThreads: 1).next(),
                                      allocator: ByteBufferAllocator())
         XCTAssertLessThanOrEqual(context.getRemainingTime(), .seconds(1))
diff --git a/Tests/AWSLambdaRuntimeTests/Lambda+CodeableTest.swift b/Tests/AWSLambdaRuntimeTests/Lambda+CodeableTest.swift
index d1d1cc05..4175fd11 100644
--- a/Tests/AWSLambdaRuntimeTests/Lambda+CodeableTest.swift
+++ b/Tests/AWSLambdaRuntimeTests/Lambda+CodeableTest.swift
@@ -14,14 +14,11 @@
 
 @testable import AWSLambdaRuntime
 @testable import AWSLambdaRuntimeCore
-import AWSXRayRecorder
 import Logging
 import NIO
 import NIOFoundationCompat
 import XCTest
 
-private let noOpTracer = XRayRecorder(emitter: XRayNoOpEmitter())
-
 class CodableLambdaTest: XCTestCase {
     var eventLoopGroup: EventLoopGroup!
     let allocator = ByteBufferAllocator()
@@ -75,7 +72,7 @@ class CodableLambdaTest: XCTestCase {
                        cognitoIdentity: nil,
                        clientContext: nil,
                        logger: Logger(label: "test"),
-                       tracer: noOpTracer,
+                       tracer: NoOpTracingInstrument(),
                        eventLoop: self.eventLoopGroup.next(),
                        allocator: ByteBufferAllocator())
     }

From ccb11daac2823f087a5d8d9785f8fa002d888081 Mon Sep 17 00:00:00 2001
From: Michal A <dev@pokryfka.net>
Date: Mon, 17 Aug 2020 13:27:00 +0800
Subject: [PATCH 10/15] Do not change baggage in Lambda.Context (which is not
 thread safe)

---
 Sources/AWSLambdaRuntimeCore/LambdaContext.swift | 2 +-
 Sources/AWSLambdaRuntimeCore/LambdaHandler.swift | 3 ++-
 Sources/AWSLambdaRuntimeCore/LambdaRunner.swift  | 2 +-
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
index 3b404778..b78573bc 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
@@ -71,7 +71,7 @@ extension Lambda {
         public let clientContext: String?
 
         /// Context baggage.
-        public var baggage: BaggageContext
+        public let baggage: BaggageContext
 
         /// `Logger` to log with
         ///
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
index 0889eccf..308aee75 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
@@ -140,7 +140,8 @@ public extension EventLoopLambdaHandler {
             return context.eventLoop.makeFailedFuture(CodecError.requestDecoding(error))
         case .success(let `in`):
             let subsegment = segment.beginSubsegment(name: "HandleIn")
-            context.baggage = subsegment.baggage
+            // context is not thread safe, do not change it
+            // context.baggage = subsegment.baggage
             return self.handle(context: context, event: `in`)
                 .endSegment(subsegment)
                 .flatMapThrowing { out in
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
index c84b89a3..a0149200 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
@@ -93,7 +93,7 @@ extension Lambda {
                         }
                         return (invocation, result, baggage)
                     }
-            }.flatMap { (invocation, result, baggage: BaggageContext) in
+            }.flatMap { invocation, result, baggage in
                 // 3. report results to runtime engine
                 self.tracer.segment(name: "ReportResults", baggage: baggage) { _ in
                     self.runtimeClient.reportResults(logger: logger, invocation: invocation, result: result).peekError { error in

From bda00143d308cd8b0365aa104522cac551223c52 Mon Sep 17 00:00:00 2001
From: Michal A <dev@pokryfka.net>
Date: Mon, 17 Aug 2020 13:44:33 +0800
Subject: [PATCH 11/15] Shutdown tracer properly, remove tracer bootstrap (for
 now)

---
 Sources/AWSLambdaRuntimeCore/Lambda.swift      | 18 ------------------
 .../LambdaConfiguration.swift                  |  5 +----
 .../AWSLambdaRuntimeCore/LambdaLifecycle.swift | 10 +++++++++-
 .../AWSLambdaRuntimeCore/LambdaRunner.swift    |  4 ++--
 4 files changed, 12 insertions(+), 25 deletions(-)

diff --git a/Sources/AWSLambdaRuntimeCore/Lambda.swift b/Sources/AWSLambdaRuntimeCore/Lambda.swift
index 80ef7a6f..5dc27648 100644
--- a/Sources/AWSLambdaRuntimeCore/Lambda.swift
+++ b/Sources/AWSLambdaRuntimeCore/Lambda.swift
@@ -21,7 +21,6 @@ import Darwin.C
 import Backtrace
 import Logging
 import NIO
-import NIOConcurrencyHelpers
 
 public enum Lambda {
     public typealias Handler = ByteBufferLambdaHandler
@@ -31,23 +30,6 @@ public enum Lambda {
     /// A function that takes a `InitializationContext` and returns an `EventLoopFuture` of a `ByteBufferLambdaHandler`
     public typealias HandlerFactory = (InitializationContext) -> EventLoopFuture<Handler>
 
-    internal typealias TracerFactory = (EventLoop) -> TracingInstrument
-
-    private static let lock = Lock()
-    private static var tracerFactory: TracerFactory = { _ in NoOpTracingInstrument() }
-
-    /// Select the desired `TracingInstrument` implementation.
-    ///
-    /// The caller is responsible for the tracer lifecycle.
-    ///
-    /// - parameters:
-    ///     -  tracerFactory: creates `TracingInstrument` implementation
-    public static func bootstrap(_ tracerFactory: @escaping (EventLoop) -> TracingInstrument) {
-        self.lock.withLockVoid {
-            self.tracerFactory = tracerFactory
-        }
-    }
-
     /// Run a Lambda defined by implementing the `LambdaHandler` protocol.
     ///
     /// - parameters:
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaConfiguration.swift b/Sources/AWSLambdaRuntimeCore/LambdaConfiguration.swift
index 38d7b4d1..f8243256 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaConfiguration.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaConfiguration.swift
@@ -21,19 +21,16 @@ extension Lambda {
         let general: General
         let lifecycle: Lifecycle
         let runtimeEngine: RuntimeEngine
-        let tracerFactory: TracerFactory
 
         init() {
             self.init(general: .init(), lifecycle: .init(), runtimeEngine: .init())
         }
 
-        init(general: General? = nil, lifecycle: Lifecycle? = nil, runtimeEngine: RuntimeEngine? = nil,
-             tracerFactory: TracerFactory? = nil)
+        init(general: General? = nil, lifecycle: Lifecycle? = nil, runtimeEngine: RuntimeEngine? = nil)
         {
             self.general = general ?? General()
             self.lifecycle = lifecycle ?? Lifecycle()
             self.runtimeEngine = runtimeEngine ?? RuntimeEngine()
-            self.tracerFactory = tracerFactory ?? { _ in NoOpTracingInstrument() }
         }
 
         struct General: CustomStringConvertible {
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift b/Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift
index ec609901..4dab20b0 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift
@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+import AWSXRaySDK
 import Logging
 import NIO
 import NIOConcurrencyHelpers
@@ -78,7 +79,9 @@ extension Lambda {
 
             var logger = self.logger
             logger[metadataKey: "lifecycleId"] = .string(self.configuration.lifecycle.id)
-            let runner = Runner(eventLoop: self.eventLoop, configuration: self.configuration)
+
+            let tracer = XRayRecorder(eventLoopGroupProvider: .shared(eventLoop))
+            let runner = Runner(eventLoop: self.eventLoop, configuration: self.configuration, tracer: tracer)
 
             let startupFuture = runner.initialize(logger: logger, factory: self.factory)
             startupFuture.flatMap { handler -> EventLoopFuture<(ByteBufferLambdaHandler, Result<Int, Error>)> in
@@ -92,6 +95,11 @@ extension Lambda {
             .flatMap { (handler, runnerResult) -> EventLoopFuture<Int> in
                 // after the lambda finishPromise has succeeded or failed we need to
                 // shutdown the handler
+                tracer.shutdown { error in
+                    if let error = error {
+                        logger.error("Failed to shutdown tracer: \(error)")
+                    }
+                }
                 let shutdownContext = ShutdownContext(logger: logger, eventLoop: self.eventLoop)
                 return handler.shutdown(context: shutdownContext).flatMapErrorThrowing { error in
                     // if, we had an error shuting down the lambda, we want to concatenate it with
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
index a0149200..b17a2fb8 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
@@ -32,10 +32,10 @@ extension Lambda {
 
         private var isGettingNextInvocation = false
 
-        init(eventLoop: EventLoop, configuration: Configuration) {
+        init(eventLoop: EventLoop, configuration: Configuration, tracer: TracingInstrument) {
             self.eventLoop = eventLoop
             self.runtimeClient = RuntimeClient(eventLoop: self.eventLoop, configuration: configuration.runtimeEngine)
-            self.tracer = configuration.tracerFactory(eventLoop)
+            self.tracer = tracer
             self.allocator = ByteBufferAllocator()
         }
 

From f3b4beb3e4806e3673f728deb7c9f4ee80790bc8 Mon Sep 17 00:00:00 2001
From: Michal A <dev@pokryfka.net>
Date: Mon, 17 Aug 2020 19:08:37 +0800
Subject: [PATCH 12/15] Instrument getNextInvocation and HTTClient

---
 Sources/AWSLambdaRuntimeCore/HTTPClient.swift | 25 ++++++++++++-------
 .../AWSLambdaRuntimeCore/LambdaContext.swift  |  2 +-
 .../AWSLambdaRuntimeCore/LambdaHandler.swift  |  2 +-
 .../AWSLambdaRuntimeCore/LambdaRunner.swift   | 11 +++++---
 .../LambdaRuntimeClient.swift                 | 15 ++++++-----
 5 files changed, 34 insertions(+), 21 deletions(-)

diff --git a/Sources/AWSLambdaRuntimeCore/HTTPClient.swift b/Sources/AWSLambdaRuntimeCore/HTTPClient.swift
index fcd2a450..9b894890 100644
--- a/Sources/AWSLambdaRuntimeCore/HTTPClient.swift
+++ b/Sources/AWSLambdaRuntimeCore/HTTPClient.swift
@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+import Baggage
 import NIO
 import NIOConcurrencyHelpers
 import NIOHTTP1
@@ -23,31 +24,35 @@ internal final class HTTPClient {
     private let eventLoop: EventLoop
     private let configuration: Lambda.Configuration.RuntimeEngine
     private let targetHost: String
+    private let tracer: TracingInstrument
 
     private var state = State.disconnected
     private var executing = false
 
-    init(eventLoop: EventLoop, configuration: Lambda.Configuration.RuntimeEngine) {
+    init(eventLoop: EventLoop, configuration: Lambda.Configuration.RuntimeEngine, tracer: TracingInstrument) {
         self.eventLoop = eventLoop
         self.configuration = configuration
         self.targetHost = "\(self.configuration.ip):\(self.configuration.port)"
+        self.tracer = tracer
     }
 
-    func get(url: String, headers: HTTPHeaders, timeout: TimeAmount? = nil) -> EventLoopFuture<Response> {
+    func get(url: String, headers: HTTPHeaders, timeout: TimeAmount? = nil, context: BaggageContext) -> EventLoopFuture<Response> {
         self.execute(Request(targetHost: self.targetHost,
                              url: url,
                              method: .GET,
                              headers: headers,
-                             timeout: timeout ?? self.configuration.requestTimeout))
+                             timeout: timeout ?? self.configuration.requestTimeout),
+                     context: context)
     }
 
-    func post(url: String, headers: HTTPHeaders, body: ByteBuffer?, timeout: TimeAmount? = nil) -> EventLoopFuture<Response> {
+    func post(url: String, headers: HTTPHeaders, body: ByteBuffer?, timeout: TimeAmount? = nil, context: BaggageContext) -> EventLoopFuture<Response> {
         self.execute(Request(targetHost: self.targetHost,
                              url: url,
                              method: .POST,
                              headers: headers,
                              body: body,
-                             timeout: timeout ?? self.configuration.requestTimeout))
+                             timeout: timeout ?? self.configuration.requestTimeout),
+                     context: context)
     }
 
     /// cancels the current request if there is one
@@ -65,7 +70,7 @@ internal final class HTTPClient {
     }
 
     // TODO: cap reconnect attempt
-    private func execute(_ request: Request, validate: Bool = true) -> EventLoopFuture<Response> {
+    private func execute(_ request: Request, validate: Bool = true, context: BaggageContext) -> EventLoopFuture<Response> {
         if validate {
             precondition(self.executing == false, "expecting single request at a time")
             self.executing = true
@@ -75,14 +80,16 @@ internal final class HTTPClient {
         case .disconnected:
             return self.connect().flatMap { channel -> EventLoopFuture<Response> in
                 self.state = .connected(channel)
-                return self.execute(request, validate: false)
+                return self.execute(request, validate: false, context: context)
             }
         case .connected(let channel):
             guard channel.isActive else {
                 self.state = .disconnected
-                return self.execute(request, validate: false)
+                return self.execute(request, validate: false, context: context)
             }
 
+            let segment = self.tracer.beginSegment(name: "HTTPClient", baggage: context)
+            segment.setHTTPRequest(method: request.method.rawValue, url: request.url)
             let promise = channel.eventLoop.makePromise(of: Response.self)
             promise.futureResult.whenComplete { _ in
                 precondition(self.executing == true, "invalid execution state")
@@ -90,7 +97,7 @@ internal final class HTTPClient {
             }
             let wrapper = HTTPRequestWrapper(request: request, promise: promise)
             channel.writeAndFlush(wrapper).cascadeFailure(to: promise)
-            return promise.futureResult
+            return promise.futureResult.endSegment(segment)
         }
     }
 
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
index b78573bc..0cbd2962 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
@@ -117,7 +117,7 @@ extension Lambda {
             logger[metadataKey: "awsRequestID"] = .string(requestID)
             logger[metadataKey: "awsTraceID"] = .string(traceID)
             var baggage = BaggageContext()
-            // TODO: handle error
+            // TODO: use `swift-tracing` API, note that we can ONLY extract X-Ray Context from the invocation data
             baggage.xRayContext = try? XRayContext(tracingHeader: traceID)
             self.baggage = baggage
             self.logger = logger
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
index 308aee75..daeb3b62 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
@@ -140,7 +140,7 @@ public extension EventLoopLambdaHandler {
             return context.eventLoop.makeFailedFuture(CodecError.requestDecoding(error))
         case .success(let `in`):
             let subsegment = segment.beginSubsegment(name: "HandleIn")
-            // context is not thread safe, do not change it
+            // TODO: context is not thread safe, do not change it
             // context.baggage = subsegment.baggage
             return self.handle(context: context, event: `in`)
                 .endSegment(subsegment)
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
index b17a2fb8..beb6efd3 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
@@ -34,7 +34,7 @@ extension Lambda {
 
         init(eventLoop: EventLoop, configuration: Configuration, tracer: TracingInstrument) {
             self.eventLoop = eventLoop
-            self.runtimeClient = RuntimeClient(eventLoop: self.eventLoop, configuration: configuration.runtimeEngine)
+            self.runtimeClient = RuntimeClient(eventLoop: self.eventLoop, configuration: configuration.runtimeEngine, tracer: tracer)
             self.tracer = tracer
             self.allocator = ByteBufferAllocator()
         }
@@ -68,7 +68,8 @@ extension Lambda {
             logger.debug("lambda invocation sequence starting")
             // 1. request invocation from lambda runtime engine
             self.isGettingNextInvocation = true
-            // TODO: add API to explicitly set startTime, after all
+            // we will get the trace context in the invocation
+            let startTime = XRayRecorder.Timestamp.now()
             return self.runtimeClient.getNextInvocation(logger: logger).peekError { error in
                 logger.error("could not fetch work from lambda runtime engine: \(error)")
             }.flatMap { invocation, event in
@@ -80,6 +81,7 @@ extension Lambda {
                                       allocator: self.allocator,
                                       invocation: invocation)
                 let baggage = context.baggage
+                self.tracer.beginSegment(name: "getNextInvocation", baggage: baggage, startTime: startTime).end()
                 logger.debug("sending invocation to lambda handler \(handler)")
                 return handler.handle(context: context, event: event)
                     // Hopping back to "our" EventLoop is importnant in case the handler returns a future that
@@ -95,8 +97,9 @@ extension Lambda {
                     }
             }.flatMap { invocation, result, baggage in
                 // 3. report results to runtime engine
-                self.tracer.segment(name: "ReportResults", baggage: baggage) { _ in
-                    self.runtimeClient.reportResults(logger: logger, invocation: invocation, result: result).peekError { error in
+                self.tracer.segment(name: "ReportResults", baggage: baggage) { segment in
+                    self.runtimeClient.reportResults(logger: logger, invocation: invocation, result: result,
+                                                     context: segment.baggage).peekError { error in
                         logger.error("could not report results to lambda runtime engine: \(error)")
                     }
                 }
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift b/Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift
index 5e9e6aea..da3ceab7 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift
@@ -12,6 +12,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+import AWSXRayRecorder
+import Baggage
 import Logging
 import NIO
 import NIOHTTP1
@@ -27,16 +29,17 @@ extension Lambda {
         private let allocator = ByteBufferAllocator()
         private let httpClient: HTTPClient
 
-        init(eventLoop: EventLoop, configuration: Configuration.RuntimeEngine) {
+        init(eventLoop: EventLoop, configuration: Configuration.RuntimeEngine, tracer: TracingInstrument) {
             self.eventLoop = eventLoop
-            self.httpClient = HTTPClient(eventLoop: eventLoop, configuration: configuration)
+            self.httpClient = HTTPClient(eventLoop: eventLoop, configuration: configuration, tracer: tracer)
         }
 
         /// Requests invocation from the control plane.
         func getNextInvocation(logger: Logger) -> EventLoopFuture<(Invocation, ByteBuffer)> {
             let url = Consts.invocationURLPrefix + Consts.getNextInvocationURLSuffix
             logger.debug("requesting work from lambda runtime engine using \(url)")
-            return self.httpClient.get(url: url, headers: RuntimeClient.defaultHeaders).flatMapThrowing { response in
+            // we will get the trace context in the invocation
+            return self.httpClient.get(url: url, headers: RuntimeClient.defaultHeaders, context: .init()).flatMapThrowing { response in
                 guard response.status == .ok else {
                     throw RuntimeError.badStatusCode(response.status)
                 }
@@ -58,7 +61,7 @@ extension Lambda {
         }
 
         /// Reports a result to the Runtime Engine.
-        func reportResults(logger: Logger, invocation: Invocation, result: Result<ByteBuffer?, Error>) -> EventLoopFuture<Void> {
+        func reportResults(logger: Logger, invocation: Invocation, result: Result<ByteBuffer?, Error>, context: BaggageContext) -> EventLoopFuture<Void> {
             var url = Consts.invocationURLPrefix + "/" + invocation.requestID
             var body: ByteBuffer?
             let headers: HTTPHeaders
@@ -77,7 +80,7 @@ extension Lambda {
                 headers = RuntimeClient.errorHeaders
             }
             logger.debug("reporting results to lambda runtime engine using \(url)")
-            return self.httpClient.post(url: url, headers: headers, body: body).flatMapThrowing { response in
+            return self.httpClient.post(url: url, headers: headers, body: body, context: context).flatMapThrowing { response in
                 guard response.status == .accepted else {
                     throw RuntimeError.badStatusCode(response.status)
                 }
@@ -102,7 +105,7 @@ extension Lambda {
             var body = self.allocator.buffer(capacity: bytes.count)
             body.writeBytes(bytes)
             logger.warning("reporting initialization error to lambda runtime engine using \(url)")
-            return self.httpClient.post(url: url, headers: RuntimeClient.errorHeaders, body: body).flatMapThrowing { response in
+            return self.httpClient.post(url: url, headers: RuntimeClient.errorHeaders, body: body, context: .init()).flatMapThrowing { response in
                 guard response.status == .accepted else {
                     throw RuntimeError.badStatusCode(response.status)
                 }

From 8cc1072dc5bdcc03a08ccf5ceb6bbb61683fb79f Mon Sep 17 00:00:00 2001
From: Michal A <dev@pokryfka.net>
Date: Mon, 17 Aug 2020 19:59:04 +0800
Subject: [PATCH 13/15] Make Lambda.Context.BaggageContext thread safe

---
 .../AWSLambdaRuntimeCore/LambdaContext.swift   | 18 ++++++++++++++----
 .../AWSLambdaRuntimeCore/LambdaHandler.swift   |  3 +--
 .../AWSLambdaRuntimeCore/LambdaRunner.swift    |  9 ++++-----
 .../LambdaRuntimeClient.swift                  |  2 +-
 4 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
index 0cbd2962..cd7a04ca 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift
@@ -17,6 +17,7 @@ import Baggage
 import Dispatch
 import Logging
 import NIO
+import NIOConcurrencyHelpers
 
 // MARK: - InitializationContext
 
@@ -52,6 +53,9 @@ extension Lambda {
     /// Lambda runtime context.
     /// The Lambda runtime generates and passes the `Context` to the Lambda handler as an argument.
     public final class Context: CustomDebugStringConvertible {
+        // TODO: use RWLock (separate PR)
+        private let lock = Lock()
+
         /// The request ID, which identifies the request that triggered the function invocation.
         public let requestID: String
 
@@ -70,8 +74,14 @@ extension Lambda {
         /// For invocations from the AWS Mobile SDK, data about the client application and device.
         public let clientContext: String?
 
-        /// Context baggage.
-        public let baggage: BaggageContext
+        // TODO: or should the Lambda "runtime" context and the Baggage context be separate?
+        private var _baggage: BaggageContext
+
+        /// Baggage context.
+        public var baggage: BaggageContext {
+            get { self.lock.withLock { _baggage } }
+            set { self.lock.withLockVoid { _baggage = newValue } }
+        }
 
         /// `Logger` to log with
         ///
@@ -117,9 +127,9 @@ extension Lambda {
             logger[metadataKey: "awsRequestID"] = .string(requestID)
             logger[metadataKey: "awsTraceID"] = .string(traceID)
             var baggage = BaggageContext()
-            // TODO: use `swift-tracing` API, note that we can ONLY extract X-Ray Context from the invocation data
+            // TODO: use `swift-tracing` API, note that, regardless, we can ONLY extract X-Ray Context
             baggage.xRayContext = try? XRayContext(tracingHeader: traceID)
-            self.baggage = baggage
+            self._baggage = baggage
             self.logger = logger
             self.tracer = tracer
         }
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
index daeb3b62..0889eccf 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
@@ -140,8 +140,7 @@ public extension EventLoopLambdaHandler {
             return context.eventLoop.makeFailedFuture(CodecError.requestDecoding(error))
         case .success(let `in`):
             let subsegment = segment.beginSubsegment(name: "HandleIn")
-            // TODO: context is not thread safe, do not change it
-            // context.baggage = subsegment.baggage
+            context.baggage = subsegment.baggage
             return self.handle(context: context, event: `in`)
                 .endSegment(subsegment)
                 .flatMapThrowing { out in
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
index beb6efd3..e85cee7b 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
@@ -80,8 +80,7 @@ extension Lambda {
                                       eventLoop: self.eventLoop,
                                       allocator: self.allocator,
                                       invocation: invocation)
-                let baggage = context.baggage
-                self.tracer.beginSegment(name: "getNextInvocation", baggage: baggage, startTime: startTime).end()
+                self.tracer.beginSegment(name: "getNextInvocation", baggage: context.baggage, startTime: startTime).end()
                 logger.debug("sending invocation to lambda handler \(handler)")
                 return handler.handle(context: context, event: event)
                     // Hopping back to "our" EventLoop is importnant in case the handler returns a future that
@@ -93,11 +92,11 @@ extension Lambda {
                         if case .failure(let error) = result {
                             logger.warning("lambda handler returned an error: \(error)")
                         }
-                        return (invocation, result, baggage)
+                        return (invocation, result, context)
                     }
-            }.flatMap { invocation, result, baggage in
+            }.flatMap { (invocation, result, context: Context) in
                 // 3. report results to runtime engine
-                self.tracer.segment(name: "ReportResults", baggage: baggage) { segment in
+                self.tracer.segment(name: "ReportResults", baggage: context.baggage) { segment in
                     self.runtimeClient.reportResults(logger: logger, invocation: invocation, result: result,
                                                      context: segment.baggage).peekError { error in
                         logger.error("could not report results to lambda runtime engine: \(error)")
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift b/Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift
index da3ceab7..fb073c06 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift
@@ -38,7 +38,7 @@ extension Lambda {
         func getNextInvocation(logger: Logger) -> EventLoopFuture<(Invocation, ByteBuffer)> {
             let url = Consts.invocationURLPrefix + Consts.getNextInvocationURLSuffix
             logger.debug("requesting work from lambda runtime engine using \(url)")
-            // we will get the trace context in the invocation
+            // we will get the trace context in the invocation response
             return self.httpClient.get(url: url, headers: RuntimeClient.defaultHeaders, context: .init()).flatMapThrowing { response in
                 guard response.status == .ok else {
                     throw RuntimeError.badStatusCode(response.status)

From 03b745f2ba95f6aa6581b482e675fe88a7a14da9 Mon Sep 17 00:00:00 2001
From: Michal A <dev@pokryfka.net>
Date: Mon, 17 Aug 2020 20:49:16 +0800
Subject: [PATCH 14/15] Fix tests, change tracer lifecycle WIP

---
 Sources/AWSLambdaRuntimeCore/Lambda.swift       | 15 ++++++++++++++-
 .../AWSLambdaRuntimeCore/LambdaLifecycle.swift  | 17 +++++++----------
 Sources/AWSLambdaRuntimeCore/LambdaRunner.swift |  5 +++--
 .../LambdaRuntimeClient.swift                   |  3 ++-
 .../LambdaRuntimeClientTest.swift               |  5 +++--
 5 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/Sources/AWSLambdaRuntimeCore/Lambda.swift b/Sources/AWSLambdaRuntimeCore/Lambda.swift
index 5dc27648..9287489b 100644
--- a/Sources/AWSLambdaRuntimeCore/Lambda.swift
+++ b/Sources/AWSLambdaRuntimeCore/Lambda.swift
@@ -18,6 +18,7 @@ import Glibc
 import Darwin.C
 #endif
 
+import AWSXRaySDK
 import Backtrace
 import Logging
 import NIO
@@ -106,8 +107,12 @@ public enum Lambda {
             logger.logLevel = configuration.general.logLevel
 
             var result: Result<Int, Error>!
+            // TODO: bootstrap/configure the trace lifecycle
+            let tracer: XRayRecorder = XRayRecorder(eventLoopGroupProvider: .createNew)
             MultiThreadedEventLoopGroup.withCurrentThreadAsEventLoop { eventLoop in
-                let lifecycle = Lifecycle(eventLoop: eventLoop, logger: logger, configuration: configuration, factory: factory)
+//                let tracer = XRayRecorder(eventLoopGroupProvider: .shared(eventLoop))
+                let lifecycle = Lifecycle(eventLoop: eventLoop, logger: logger, tracer: tracer,
+                                          configuration: configuration, factory: factory)
                 #if DEBUG
                 let signalSource = trap(signal: configuration.lifecycle.stopSignal) { signal in
                     logger.info("intercepted signal: \(signal)")
@@ -130,6 +135,14 @@ public enum Lambda {
                 }
             }
 
+            // TODO: this is broken, the tracer tries to eat cake and have cake: flash on eventLoop then syncShutdown
+            // fix in XRayRecorder
+            tracer.shutdown { error in
+                if let error = error {
+                    preconditionFailure("Failed to shutdown tracer: \(error)")
+                }
+            }
+
             logger.info("shutdown completed")
             return result
         }
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift b/Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift
index 4dab20b0..028f1dc3 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift
@@ -27,6 +27,7 @@ extension Lambda {
         private let logger: Logger
         private let configuration: Configuration
         private let factory: HandlerFactory
+        private let tracer: TracingInstrument
 
         private var state = State.idle {
             willSet {
@@ -42,13 +43,16 @@ extension Lambda {
         ///     - logger: A `Logger` to log the Lambda events.
         ///     - factory: A `LambdaHandlerFactory` to create the concrete  Lambda handler.
         public convenience init(eventLoop: EventLoop, logger: Logger, factory: @escaping HandlerFactory) {
-            self.init(eventLoop: eventLoop, logger: logger, configuration: .init(), factory: factory)
+            self.init(eventLoop: eventLoop, logger: logger, tracer: NoOpTracingInstrument(), configuration: .init(),
+                      factory: factory)
         }
 
-        init(eventLoop: EventLoop, logger: Logger, configuration: Configuration, factory: @escaping HandlerFactory) {
+        init(eventLoop: EventLoop, logger: Logger, tracer: TracingInstrument, configuration: Configuration,
+             factory: @escaping HandlerFactory) {
             self.eventLoop = eventLoop
             self.shutdownPromise = eventLoop.makePromise(of: Int.self)
             self.logger = logger
+            self.tracer = tracer
             self.configuration = configuration
             self.factory = factory
         }
@@ -79,9 +83,7 @@ extension Lambda {
 
             var logger = self.logger
             logger[metadataKey: "lifecycleId"] = .string(self.configuration.lifecycle.id)
-
-            let tracer = XRayRecorder(eventLoopGroupProvider: .shared(eventLoop))
-            let runner = Runner(eventLoop: self.eventLoop, configuration: self.configuration, tracer: tracer)
+            let runner = Runner(eventLoop: self.eventLoop, configuration: self.configuration, tracer: self.tracer)
 
             let startupFuture = runner.initialize(logger: logger, factory: self.factory)
             startupFuture.flatMap { handler -> EventLoopFuture<(ByteBufferLambdaHandler, Result<Int, Error>)> in
@@ -95,11 +97,6 @@ extension Lambda {
             .flatMap { (handler, runnerResult) -> EventLoopFuture<Int> in
                 // after the lambda finishPromise has succeeded or failed we need to
                 // shutdown the handler
-                tracer.shutdown { error in
-                    if let error = error {
-                        logger.error("Failed to shutdown tracer: \(error)")
-                    }
-                }
                 let shutdownContext = ShutdownContext(logger: logger, eventLoop: self.eventLoop)
                 return handler.shutdown(context: shutdownContext).flatMapErrorThrowing { error in
                     // if, we had an error shuting down the lambda, we want to concatenate it with
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
index e85cee7b..d8a751ee 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
@@ -32,9 +32,10 @@ extension Lambda {
 
         private var isGettingNextInvocation = false
 
-        init(eventLoop: EventLoop, configuration: Configuration, tracer: TracingInstrument) {
+        init(eventLoop: EventLoop, configuration: Configuration, tracer: TracingInstrument = NoOpTracingInstrument()) {
             self.eventLoop = eventLoop
-            self.runtimeClient = RuntimeClient(eventLoop: self.eventLoop, configuration: configuration.runtimeEngine, tracer: tracer)
+            self.runtimeClient = RuntimeClient(eventLoop: self.eventLoop, configuration: configuration.runtimeEngine,
+                                               tracer: tracer)
             self.tracer = tracer
             self.allocator = ByteBufferAllocator()
         }
diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift b/Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift
index fb073c06..e6b9cd0b 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift
@@ -29,7 +29,8 @@ extension Lambda {
         private let allocator = ByteBufferAllocator()
         private let httpClient: HTTPClient
 
-        init(eventLoop: EventLoop, configuration: Configuration.RuntimeEngine, tracer: TracingInstrument) {
+        init(eventLoop: EventLoop, configuration: Configuration.RuntimeEngine,
+             tracer: TracingInstrument = NoOpTracingInstrument()) {
             self.eventLoop = eventLoop
             self.httpClient = HTTPClient(eventLoop: eventLoop, configuration: configuration, tracer: tracer)
         }
diff --git a/Tests/AWSLambdaRuntimeCoreTests/LambdaRuntimeClientTest.swift b/Tests/AWSLambdaRuntimeCoreTests/LambdaRuntimeClientTest.swift
index 94c8ac62..00a2e4ce 100644
--- a/Tests/AWSLambdaRuntimeCoreTests/LambdaRuntimeClientTest.swift
+++ b/Tests/AWSLambdaRuntimeCoreTests/LambdaRuntimeClientTest.swift
@@ -263,7 +263,7 @@ class LambdaRuntimeClientTest: XCTestCase {
         XCTAssertNoThrow(inv = try Lambda.Invocation(headers: header))
         guard let invocation = inv else { return }
 
-        let result = client.reportResults(logger: logger, invocation: invocation, result: Result.failure(TestError("boom")))
+        let result = client.reportResults(logger: logger, invocation: invocation, result: Result.failure(TestError("boom")), context: .init())
 
         var inboundHeader: HTTPServerRequestPart?
         XCTAssertNoThrow(inboundHeader = try server.readInbound())
@@ -303,7 +303,8 @@ class LambdaRuntimeClientTest: XCTestCase {
         XCTAssertNoThrow(inv = try Lambda.Invocation(headers: header))
         guard let invocation = inv else { return }
 
-        let result = client.reportResults(logger: logger, invocation: invocation, result: Result.success(nil))
+        let result = client.reportResults(logger: logger, invocation: invocation, result: Result.success(nil),
+                                          context: .init())
 
         var inboundHeader: HTTPServerRequestPart?
         XCTAssertNoThrow(inboundHeader = try server.readInbound())

From dc5c7b98af7e833df3252c6ce5da24e899a03734 Mon Sep 17 00:00:00 2001
From: Michal A <dev@pokryfka.net>
Date: Mon, 17 Aug 2020 22:08:53 +0800
Subject: [PATCH 15/15] Pass the baggage to the result segment (not context
 with baggage)

---
 Sources/AWSLambdaRuntimeCore/LambdaRunner.swift | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
index d8a751ee..8ca19b4c 100644
--- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
+++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
@@ -81,6 +81,7 @@ extension Lambda {
                                       eventLoop: self.eventLoop,
                                       allocator: self.allocator,
                                       invocation: invocation)
+                let baggage = context.baggage
                 self.tracer.beginSegment(name: "getNextInvocation", baggage: context.baggage, startTime: startTime).end()
                 logger.debug("sending invocation to lambda handler \(handler)")
                 return handler.handle(context: context, event: event)
@@ -93,11 +94,11 @@ extension Lambda {
                         if case .failure(let error) = result {
                             logger.warning("lambda handler returned an error: \(error)")
                         }
-                        return (invocation, result, context)
+                        return (invocation, result, baggage)
                     }
-            }.flatMap { (invocation, result, context: Context) in
+            }.flatMap { invocation, result, baggage in
                 // 3. report results to runtime engine
-                self.tracer.segment(name: "ReportResults", baggage: context.baggage) { segment in
+                self.tracer.segment(name: "ReportResults", baggage: baggage) { segment in
                     self.runtimeClient.reportResults(logger: logger, invocation: invocation, result: result,
                                                      context: segment.baggage).peekError { error in
                         logger.error("could not report results to lambda runtime engine: \(error)")