Skip to content

RFC: instrument lambda handler #162

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 16 commits into from
Closed
Show file tree
Hide file tree
Changes from 14 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions Package.swift
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,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.7.1")),
],
targets: [
.target(name: "AWSLambdaRuntime", dependencies: [
Expand All @@ -29,6 +32,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"),
Expand All @@ -38,13 +42,15 @@ 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"]),
// testing helper
.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
Expand Down
25 changes: 16 additions & 9 deletions Sources/AWSLambdaRuntimeCore/HTTPClient.swift
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
//
//===----------------------------------------------------------------------===//

import Baggage
import NIO
import NIOConcurrencyHelpers
import NIOHTTP1
Expand All @@ -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
Expand All @@ -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
Expand All @@ -75,22 +80,24 @@ 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")
self.executing = false
}
let wrapper = HTTPRequestWrapper(request: request, promise: promise)
channel.writeAndFlush(wrapper).cascadeFailure(to: promise)
return promise.futureResult
return promise.futureResult.endSegment(segment)
}
}

Expand Down
3 changes: 2 additions & 1 deletion Sources/AWSLambdaRuntimeCore/LambdaConfiguration.swift
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,8 @@ extension Lambda {
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)
{
self.general = general ?? General()
self.lifecycle = lifecycle ?? Lifecycle()
self.runtimeEngine = runtimeEngine ?? RuntimeEngine()
Expand Down
27 changes: 26 additions & 1 deletion Sources/AWSLambdaRuntimeCore/LambdaContext.swift
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,12 @@
//
//===----------------------------------------------------------------------===//

import AWSXRayRecorder
import Baggage
import Dispatch
import Logging
import NIO
import NIOConcurrencyHelpers

// MARK: - InitializationContext

Expand Down Expand Up @@ -50,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()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

tbh we can start out with just a lock and change only if proven to matter a lot;

The https://blog.nelhage.com/post/rwlock-contention/ keeps being brought up when we reach for RWLocks recently

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

RWlock are great when you do single (or very very very little) write and all-reads. in mixed mode it can get tricky tp get good performance


/// The request ID, which identifies the request that triggered the function invocation.
public let requestID: String

Expand All @@ -68,11 +74,23 @@ extension Lambda {
/// For invocations from the AWS Mobile SDK, data about the client application and device.
public let clientContext: String?

// TODO: or should the Lambda "runtime" context and the Baggage context be separate?
private var _baggage: BaggageContext

/// Baggage context.
public var baggage: BaggageContext {
Copy link
Contributor

@tomerd tomerd Aug 17, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does this need to be a public var vs let? is it immutable or would the user ever write on this field?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's currently forced to this via the Carrier's requirement, but it may well be that requirement is quite wrong and should just be get rather than get/set tbh. We'll look into that. It loops into a few pieces fo feedback @pokryfka had here 👀

get { self.lock.withLock { _baggage } }
set { self.lock.withLockVoid { _baggage = newValue } }
}

/// `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.
///
Expand All @@ -91,8 +109,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
Expand All @@ -106,7 +126,12 @@ extension Lambda {
var logger = logger
logger[metadataKey: "awsRequestID"] = .string(requestID)
logger[metadataKey: "awsTraceID"] = .string(traceID)
var baggage = BaggageContext()
// TODO: use `swift-tracing` API, note that, regardless, we can ONLY extract X-Ray Context
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not really; "we" don't decide what we extract, the configuration of instruments decides that, and yes since xray would be configured it'd extract it's own context here.

Specifically:

tracer.extract(<from where to extract, could be a dictionary> / http headers etc, into: &baggage, using: extractor apropriate to the first parameter, so http headers or similar)

Copy link
Contributor Author

@pokryfka pokryfka Aug 17, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

XRayInstrument can use the Instrument.extract API (its implemented and tested)

the problem here is that only the X-Ray trace context is provided by Lambda Runtime API (in header), see https://docs.aws.amazon.com/lambda/latest/dg/runtimes-api.html#runtimes-api-next

Context for other instruments may be provided in invocation payload and needs to be extracted by user in lambda event handler they implement.
(Note that AWSLambdaRuntimeCore only requires and knows that events, provided in invocation payload, are Decodable).

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see, I missed that bit of the Lambda design. So in this integration style, even if one triggered the lambda via an http request, it would not get “the http request” but just the body, and the headers are the ones as listed on there, including the XRay trace header etc.

There AFAIR exists an integration mode though to get the entire request, right?

Pass through the entire request – A Lambda function can receive the entire HTTP request (instead of just the request body) and set the HTTP response (instead of just the response body) using the AWS_PROXY integration type.

Is this something that the runtime currently is able to handle? Seems more to be about how the API Gateway is configured right? Though I’ve not had the time to dig deeper into this yet.

You’d probably know more about this @fabianfett, we can catch up about this today maybe?

Copy link
Contributor Author

@pokryfka pokryfka Aug 18, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess you are referring to API Gateway Lambda proxy integration

This does NOT affect Lambda custom runtime API, it affects API Gateway v1 "REST API" routing which, if configured that way, does not try to route events based on a RESTful model, instead it forwards all events to lambda which needs to resolve HTTP method, path and arguments itself based on the content in the event payload (but it still remains in the event payload -> invocation payload):

ANY /{proxy+}: The client must choose a particular HTTP method, must set a particular resource path hierarchy, and can set any headers, query string parameters, and applicable payload to pass the data as input to the integrated Lambda function.

Note that:

  • events with HTTP requests created by API Gateway may have different syntax, specifically:
  • event types are not defined in AWSLambdaCore but in AWSLambdaEvents which, unlike AWSLambdaCore, does have dependency on Foundation
  • user can choose to define its own In and Out types as long as they are Decodable/Encodable (I do it that way).
  • AWS does not limit context propagation to API Gateway (but it will not propagate context for other instruments; it MAY copy headers from the original requests and includes them in event payload, which is true for API Gateway)

For reference Integrating AWS X-Ray with other AWS services

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for explaining this in more depth @pokryfka, I need to read up some more here about aws/lambda in general it seems.

baggage.xRayContext = try? XRayContext(tracingHeader: traceID)
Copy link
Contributor

@tomerd tomerd Aug 14, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

would we want to let the users decide what tracer to use, or since this is AWS oriented anyways just pin to x-ray?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe that's the main TODO in this PR - not using the xray explicitly.

Even if it is bound to xRay in reality, we should make sure to only use the abstract API, maybe maybe some day there would be some other tracer or maybe amazon decide to make their own or something, no idea, but let's keep the door open for future evolution.

Using the tracing API also means that while developing locally you could plug in the Instruments(.app) (naming gets confusing...) tracer: slashmo/gsoc-swift-tracing#97 and see spans in Instruments on the mac. Instruments does not really understand / visualize "traces" with parents etc well today... but it's something we can keep in mind, maybe it'll get better at displaying those and then when developing locally you get the same user experience with tracing as on prod :-)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ktoso I generally agree with you. I'm however a little concerned, that tracing will require manual adjustment (incl. adding another dependency) if we don't include the XRay tracing by default. Lambda tracing wouldn't work out of the box in this case.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah I get that -- though the lambda runtime "core" should not be instrumented using any specific tracer, regardless of "yes it'll be xray" but maybe some day down the road there's other impls, and you'd want to swap it.

We could absolutely though make some "batteries included" package, we should think how to pull that off, wdyt?

self._baggage = baggage
self.logger = logger
self.tracer = tracer
}

public func getRemainingTime() -> TimeAmount {
Expand Down
29 changes: 21 additions & 8 deletions Sources/AWSLambdaRuntimeCore/LambdaHandler.swift
Original file line number Diff line number Diff line change
Expand Up @@ -129,18 +129,31 @@ 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) {
let segment = context.tracer.beginSegment(name: "HandleEvent", baggage: context.baggage)
let decodedEvent = segment.subsegment(name: "DecodeIn") { _ in
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

btw, I've been wondering if we should offer tracer.withSpan() { ... } as built-in for wrapping synchronous blocks with a span in the swift-tracing API right away, or if we should stay away from adding any kind of sugar in the API package.

This is quite common I think so I think we could add it...

We could also do the same with a NIO extensions package then to handle Future returning blocks 🤔

WDYT @pokryfka @tomerd ?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ktoso +1

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I do use "helpers" both with closures and NIO futures as referenced here slashmo/gsoc-swift-tracing#125 (comment)
when API is changed to TracingInstrument the amount of sugar will depend on what swift-tracing provides

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`):
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
let subsegment = segment.beginSubsegment(name: "HandleIn")
context.baggage = subsegment.baggage
return self.handle(context: context, event: `in`)
.endSegment(subsegment)
.flatMapThrowing { out in
try context.tracer.segment(name: "EncodeOut", baggage: segment.baggage) { _ in
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Naming nitpick: I really would like to get all libs to be consistent with the use of the context parameter name, note that it can then accept a carrier and it becomes easier to just pass the context: segment / context: span / context: context (the lambda context); the only exception is context.baggage(like inLambda.Context`)

Copy link
Contributor Author

@pokryfka pokryfka Aug 17, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this will be "context" when using TracingInstrument API:

var span = tracer.startSpan(named: "EncodeOut", context: segment.baggage) 

I dont have strong opinion on that, but to avoid confusion in XRaySDK I use "context" for XRayContext type (which does have the X-Ray trace context, strongly typed) and "baggage" for BaggageContext which may have the X-Ray trace context

var baggage = BaggageContext() // empty
let context = XRayContext()
baggage.xRayContext =  context

let segment = tracer.beginSegment(name: "EncodeOut", baggage: baggage) // may report missing context
// or
let segment2 = tracer.beginSegment(name: "EncodeOut", context: context)

Copy link
Contributor

@ktoso ktoso Aug 17, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What would you say about the following spelling though:

var context = BaggageContext()
context.xRay =  XRayContext()
// since also:
// context.zipkin = ZipkinState() // I call it state in some things I worked on,
//
// because it aligns with  https://www.w3.org/TR/trace-context/#combined-header-value
// "tracestate" where each tracer may carry their own state by a vendor identified key
//
// Example: vendorname1=opaqueValue1,vendorname2=opaqueValue2

also because one can use the BaggageContextCarrier to assign through into the underlying baggage context;
This way all use sites, regardless if a carrier, raw baggage context, or any "my specific framework type" can use the same call site style:

func x(context: SomeFramework) { context.xRay ...
func x(context: BaggageContext) { context.xRay ...
func x(context: BaggageContextCarrier) { context.xRay ...

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see the point with

let segment = tracer.beginSegment(name: "EncodeOut", baggage: baggage)
// or
let segment2 = tracer.beginSegment(name: "EncodeOut", context: context)

though; but that's the segment API, you are free to do what you want there but still I would not recommend using baggage as parameter names, you'd want to accept a BaggageCarrierCarrier most likely, and I would suggest calling it context for consistency -- people don't need to overthink it. But that's your segment API - you're free to design that how you want, but just a suggestion to keep in mind.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I made a ticket to discuss naming once more: slashmo/gsoc-swift-baggage-context#23

switch self.encodeOut(allocator: context.allocator, value: out) {
case .failure(let error):
throw CodecError.responseEncoding(error)
case .success(let buffer):
return buffer
}
}
}
}
.endSegment(segment)
}
}

Expand Down
10 changes: 9 additions & 1 deletion Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
//
//===----------------------------------------------------------------------===//

import AWSXRaySDK
import Logging
import NIO
import NIOConcurrencyHelpers
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
32 changes: 25 additions & 7 deletions Sources/AWSLambdaRuntimeCore/LambdaRunner.swift
Original file line number Diff line number Diff line change
Expand Up @@ -12,22 +12,30 @@
//
//===----------------------------------------------------------------------===//

import AWSXRaySDK // TODO: use swift-tracing when available
import Baggage
import Dispatch
import Logging
import NIO

// type names defined in `TracingInstrument`, aliases will be removed
public typealias TracingInstrument = XRayRecorder
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah I see, that confused me for a moment why it's an instrument but used segment() ;)

Okey for the PoC, agreed on migrating once "released" though it would be good to use this PR to already migrate and see that the API supports everything one needs here.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this should make the API transition easier, temporary ...

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it be possible to PoC this out using the real swift-tracing API right away though?
We need real feedback on the API by using it in real use-cases, such as instrumenting lambda, http client and grpc; If there's things that show up as missing or sucky during that work now is the time to keep fixing the API. We should not wait for the API to move to its final destination to PoC things out and provide feedback. WDYT?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should be pretty much straight forward as far API is concerned.

I think the missing part is to freeze and release TracingInstrument API.
it will not be the final release but, if you consider existing API complete or a candidate for that, lets make a release and try to use it.

i am trying to keep XRayInstrument in sync but theres been many breaking changes recently in baggage and swift-tracing; I pin versions using git hashes, quite randomly.

public typealias NoOpTracingInstrument = XRayNoOpRecorder

extension Lambda {
/// LambdaRunner manages the Lambda runtime workflow, or business logic.
internal final class Runner {
private let runtimeClient: RuntimeClient
private let tracer: TracingInstrument
private let eventLoop: EventLoop
private let allocator: ByteBufferAllocator

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.runtimeClient = RuntimeClient(eventLoop: self.eventLoop, configuration: configuration.runtimeEngine, tracer: tracer)
self.tracer = tracer
self.allocator = ByteBufferAllocator()
}

Expand Down Expand Up @@ -60,15 +68,19 @@ extension Lambda {
logger.debug("lambda invocation sequence starting")
// 1. request invocation from lambda runtime engine
self.isGettingNextInvocation = true
// 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
// 2. send invocation to handler
self.isGettingNextInvocation = false
let context = Context(logger: logger,
tracer: self.tracer,
eventLoop: self.eventLoop,
allocator: self.allocator,
invocation: invocation)
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
Expand All @@ -80,14 +92,19 @@ extension Lambda {
if case .failure(let error) = result {
logger.warning("lambda handler returned an error: \(error)")
}
return (invocation, result)
return (invocation, result, context)
}
}.flatMap { invocation, result in
}.flatMap { (invocation, result, context: Context) 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: 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)")
}
}
}
// 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)
Expand All @@ -101,14 +118,15 @@ 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,
deadline: DispatchWallTime(millisSinceEpoch: invocation.deadlineInMillisSinceEpoch),
cognitoIdentity: invocation.cognitoIdentity,
clientContext: invocation.clientContext,
logger: logger,
tracer: tracer,
eventLoop: eventLoop,
allocator: allocator)
}
Expand Down
Loading