refactor (#12)

motivation: make code simpler to reason about, better use of swift-nio

changes:
* remove main loop on global queue + wait, replace with recursion
* make better use of EvenLoopFuture, instead of result types to signal errors
* inject lifecycleId so we can share offloading queue
* improve logging
* adjust and improve tests
* update sanity and generate linux tests script to work better with dates and swiftformat

Author: tomerd

Package.swift

@@ -19,7 +19,7 @@ let package = Package(
         .executable(name: "SwiftAwsLambdaCodableSample", targets: ["SwiftAwsLambdaCodableSample"]),
     ],
     dependencies: [
-        .package(url: "https://github.com/apple/swift-nio.git", from: "2.0.0"),
+        .package(url: "https://github.com/apple/swift-nio.git", from: "2.8.0"),
         .package(url: "https://github.com/apple/swift-log.git", from: "1.0.0"),
         .package(url: "https://github.com/ianpartridge/swift-backtrace.git", from: "1.1.0"),
     ],

Sources/SwiftAwsLambda/Lambda.swift

@@ -29,49 +29,49 @@ public enum Lambda {
     ///
     /// - note: This is a blocking operation that will run forever, as it's lifecycle is managed by the AWS Lambda Runtime Engine.
     public static func run(_ closure: @escaping LambdaClosure) {
-        let _: LambdaLifecycleResult = self.run(closure)
+        self.run(closure: closure)
     }
 
     /// Run a Lambda defined by implementing the `LambdaHandler` protocol.
     ///
     /// - note: This is a blocking operation that will run forever, as it's lifecycle is managed by the AWS Lambda Runtime Engine.
     public static func run(_ handler: LambdaHandler) {
-        let _: LambdaLifecycleResult = self.run(handler: handler)
+        self.run(handler: handler)
     }
 
-    // for testing
-    internal static func run(maxTimes: Int = 0, stopSignal: Signal = .TERM, _ closure: @escaping LambdaClosure) -> LambdaLifecycleResult {
+    // for testing and internal use
+    @discardableResult
+    internal static func run(maxTimes: Int = 0, stopSignal: Signal = .TERM, closure: @escaping LambdaClosure) -> LambdaLifecycleResult {
         return self.run(handler: LambdaClosureWrapper(closure), maxTimes: maxTimes, stopSignal: stopSignal)
     }
 
-    // for testing
+    // for testing and internal use
+    @discardableResult
     internal static func run(handler: LambdaHandler, maxTimes: Int = 0, stopSignal: Signal = .TERM) -> LambdaLifecycleResult {
         do {
-            return try self.runAsync(handler: handler, maxTimes: maxTimes, stopSignal: stopSignal).wait()
+            return try self.runAsync(handler: handler, maxTimes: maxTimes, stopSignal: stopSignal).map { .success($0) }.wait()
         } catch {
             return .failure(error)
         }
     }
 
-    internal static func runAsync(handler: LambdaHandler, maxTimes: Int = 0, stopSignal: Signal = .TERM) -> EventLoopFuture<LambdaLifecycleResult> {
+    internal static func runAsync(handler: LambdaHandler, maxTimes: Int = 0, stopSignal: Signal = .TERM) -> EventLoopFuture<Int> {
+        Backtrace.install()
         let logger = Logger(label: "Lambda")
         let lifecycle = Lifecycle(logger: logger, handler: handler, maxTimes: maxTimes)
-        Backtrace.install()
         let signalSource = trap(signal: stopSignal) { signal in
             logger.info("intercepted signal: \(signal)")
             lifecycle.stop()
         }
-        let future = lifecycle.start()
-        future.whenComplete { _ in
+        return lifecycle.start().always { _ in
             lifecycle.stop()
             signalSource.cancel()
         }
-        return future
     }
 
     private class Lifecycle {
         private let logger: Logger
-        private let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        private let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: System.coreCount)
         private let handler: LambdaHandler
         private let max: Int
 
@@ -88,7 +88,7 @@ public enum Lambda {
 
         deinit {
             self.logger.info("lambda lifecycle deinit")
-            assert(state == .shutdown)
+            assert(self.state == .shutdown, "invalid state, expected shutdown")
         }
 
         private var state: LifecycleState {
@@ -105,63 +105,60 @@ public enum Lambda {
             }
         }
 
-        func start() -> EventLoopFuture<LambdaLifecycleResult> {
+        func start() -> EventLoopFuture<Int> {
             self.state = .initializing
+            let lifecycleId = NSUUID().uuidString
+            let eventLoop = self.eventLoopGroup.next()
             var logger = self.logger
-            logger[metadataKey: "lifecycleId"] = .string(NSUUID().uuidString)
-            let runner = LambdaRunner(eventLoopGroup: eventLoopGroup, lambdaHandler: handler)
-            let promise = self.eventLoopGroup.next().makePromise(of: LambdaLifecycleResult.self)
-            self.logger.info("lambda lifecycle starting")
-
-            runner.initialize(logger: logger).whenComplete { result in
-                switch result {
-                case .failure(let error):
-                    promise.succeed(.failure(error))
-                case .success:
-                    DispatchQueue.global().async {
-                        var lastError: Error?
-                        var counter = 0
-                        self.state = .active // initializtion completed successfully, we're good to go!
-                        while self.state == .active, lastError == nil, self.max == 0 || counter < self.max {
-                            do {
-                                logger[metadataKey: "lifecycleIteration"] = "\(counter)"
-                                // blocking! per aws lambda runtime spec the polling requests are to be done one at a time
-                                switch try runner.run(logger: logger).wait() {
-                                case .success:
-                                    counter = counter + 1
-                                case .failure(let error):
-                                    if self.state == .active {
-                                        lastError = error
-                                    }
-                                }
-                            } catch {
-                                if self.state == .active {
-                                    lastError = error
-                                }
-                            }
-                            // flush the log streams so entries are printed with a single invocation
-                            // TODO: should we suuport a flush API in swift-log's default logger?
-                            fflush(stdout)
-                            fflush(stderr)
-                        }
-                        promise.succeed(lastError.map { .failure($0) } ?? .success(counter))
-                    }
-                }
+            logger[metadataKey: "lifecycleId"] = .string(lifecycleId)
+            logger.info("lambda lifecycle starting")
+
+            let runner = LambdaRunner(eventLoop: eventLoop, lambdaHandler: handler, lifecycleId: lifecycleId)
+            return runner.initialize(logger: logger).flatMap { _ in
+                self.state = .active
+                return self.run(logger: logger, eventLoop: eventLoop, runner: runner, count: 0)
             }
-            return promise.futureResult
         }
 
         func stop() {
-            if self.state == .stopping {
+            switch self.state {
+            case .stopping:
                 return self.logger.info("lambda lifecycle aready stopping")
-            }
-            if self.state == .shutdown {
+            case .shutdown:
                 return self.logger.info("lambda lifecycle aready shutdown")
+            default:
+                self.logger.info("lambda lifecycle stopping")
+                self.state = .stopping
+                try! self.eventLoopGroup.syncShutdownGracefully()
+                self.state = .shutdown
+            }
+        }
+
+        private func run(logger: Logger, eventLoop: EventLoop, runner: LambdaRunner, count: Int) -> EventLoopFuture<Int> {
+            var logger = logger
+            logger[metadataKey: "lifecycleIteration"] = "\(count)"
+            return runner.run(logger: logger).flatMap { _ in
+                switch self.state {
+                case .idle, .initializing:
+                    preconditionFailure("invalid run state: \(self.state)")
+                case .active:
+                    if self.max > 0, count >= self.max {
+                        return eventLoop.makeSucceededFuture(count)
+                    }
+                    // recursive! per aws lambda runtime spec the polling requests are to be done one at a time
+                    return self.run(logger: logger, eventLoop: eventLoop, runner: runner, count: count + 1)
+                case .stopping, .shutdown:
+                    return eventLoop.makeSucceededFuture(count)
+                }
+            }.flatMapErrorThrowing { error in
+                // if we run into errors while shutting down, we ignore them
+                switch self.state {
+                case .stopping, .shutdown:
+                    return count
+                default:
+                    throw error
+                }
             }
-            self.logger.info("lambda lifecycle stopping")
-            self.state = .stopping
-            try! self.eventLoopGroup.syncShutdownGracefully()
-            self.state = .shutdown
         }
     }
 
@@ -225,6 +222,12 @@ public struct LambdaContext {
         self.cognitoIdentity = cognitoIdentity
         self.clientContext = clientContext
         self.deadline = deadline
+        // mutate logger with context
+        var logger = logger
+        logger[metadataKey: "awsRequestId"] = .string(requestId)
+        if let traceId = traceId {
+            logger[metadataKey: "awsTraceId"] = .string(traceId)
+        }
         self.logger = logger
     }
 }

Sources/SwiftAwsLambda/LambdaRunner.swift

@@ -20,102 +20,83 @@ import NIO
 internal final class LambdaRunner {
     private let runtimeClient: LambdaRuntimeClient
     private let lambdaHandler: LambdaHandler
-    private let eventLoopGroup: EventLoopGroup
+    private let eventLoop: EventLoop
+    private let lifecycleId: String
 
-    init(eventLoopGroup: EventLoopGroup, lambdaHandler: LambdaHandler) {
-        self.eventLoopGroup = eventLoopGroup
-        self.runtimeClient = LambdaRuntimeClient(eventLoopGroup: self.eventLoopGroup)
+    init(eventLoop: EventLoop, lambdaHandler: LambdaHandler, lifecycleId: String) {
+        self.eventLoop = eventLoop
+        self.runtimeClient = LambdaRuntimeClient(eventLoop: self.eventLoop)
         self.lambdaHandler = lambdaHandler
+        self.lifecycleId = lifecycleId
     }
 
     /// Run the user provided initializer. This *must* only be called once.
     ///
     /// - Returns: An `EventLoopFuture<Void>` fulfilled with the outcome of the initialization.
     func initialize(logger: Logger) -> EventLoopFuture<Void> {
-        let initPromise = self.eventLoopGroup.next().makePromise(of: Void.self)
-        self.lambdaHandler.initialize(promise: initPromise)
-
+        logger.info("initializing lambda")
         // We need to use `flatMap` instead of `whenFailure` to ensure we complete reporting the result before stopping.
-        return initPromise.futureResult.flatMapError { error in
-            self.runtimeClient.reportInitError(logger: logger, error: error).flatMapResult { postResult -> Result<Void, Error> in
-                switch postResult {
-                case .failure(let postResultError):
+        return self.lambdaHandler.initialize(eventLoop: self.eventLoop, lifecycleId: self.lifecycleId).flatMapError { error in
+            self.runtimeClient.reportInitializationError(logger: logger, error: error).flatMapResult { result -> Result<Void, Error> in
+                if case .failure(let reportingError) = result {
                     // We're going to bail out because the init failed, so there's not a lot we can do other than log
                     // that we couldn't report this error back to the runtime.
-                    logger.error("could not report initialization error to lambda runtime engine: \(postResultError)")
-                case .success:
-                    logger.info("successfully reported initialization error")
+                    logger.error("failed reporting initialization error to lambda runtime engine: \(reportingError)")
                 }
                 // Always return the init error
                 return .failure(error)
             }
         }
     }
 
-    func run(logger: Logger) -> EventLoopFuture<LambdaRunResult> {
-        var logger = logger
+    func run(logger: Logger) -> EventLoopFuture<Void> {
         logger.info("lambda invocation sequence starting")
         // 1. request work from lambda runtime engine
         return self.runtimeClient.requestWork(logger: logger).flatMap { workRequestResult in
             switch workRequestResult {
             case .failure(let error):
                 logger.error("could not fetch work from lambda runtime engine: \(error)")
-                return self.makeSucceededFuture(result: .failure(error))
+                return self.eventLoop.makeFailedFuture(error)
             case .success(let context, let payload):
-                logger[metadataKey: "awsRequestId"] = .string(context.requestId)
-                if let traceId = context.traceId {
-                    logger[metadataKey: "awsTraceId"] = .string(traceId)
-                }
                 // 2. send work to handler
                 logger.info("sending work to lambda handler \(self.lambdaHandler)")
-                let promise = self.eventLoopGroup.next().makePromise(of: LambdaResult.self)
-                self.lambdaHandler.handle(context: context, payload: payload, promise: promise)
-                return promise.futureResult.flatMap { lambdaResult in
+                return self.lambdaHandler.handle(eventLoop: self.eventLoop, lifecycleId: self.lifecycleId, context: context, payload: payload).flatMap { lambdaResult in
                     // 3. report results to runtime engine
                     self.runtimeClient.reportResults(logger: logger, context: context, result: lambdaResult).flatMap { postResultsResult in
                         switch postResultsResult {
                         case .failure(let error):
-                            logger.error("could not report results to lambda runtime engine: \(error)")
-                            return self.makeSucceededFuture(result: .failure(error))
+                            logger.error("failed reporting results to lambda runtime engine: \(error)")
+                            return self.eventLoop.makeFailedFuture(error)
                         case .success():
                             // we are done!
                             logger.info("lambda invocation sequence completed successfully")
-                            return self.makeSucceededFuture(result: .success(()))
+                            return self.eventLoop.makeSucceededFuture(())
                         }
                     }
                 }
             }
         }
     }
-
-    private func makeSucceededFuture<T>(result: T) -> EventLoopFuture<T> {
-        return self.eventLoopGroup.next().makeSucceededFuture(result)
-    }
 }
 
-internal typealias LambdaRunResult = Result<Void, Error>
-
 private extension LambdaHandler {
-    func handle(context: LambdaContext, payload: [UInt8], promise: EventLoopPromise<LambdaResult>) {
+    func initialize(eventLoop: EventLoop, lifecycleId: String) -> EventLoopFuture<Void> {
         // offloading so user code never blocks the eventloop
-        DispatchQueue(label: "lambda-\(context.requestId)").async {
-            self.handle(context: context, payload: payload) { (result: LambdaResult) in
-                promise.succeed(result)
-            }
+        let promise = eventLoop.makePromise(of: Void.self)
+        DispatchQueue(label: "lambda-\(lifecycleId)").async {
+            self.initialize { promise.completeWith($0) }
         }
+        return promise.futureResult
     }
 
-    func initialize(promise: EventLoopPromise<Void>) {
+    func handle(eventLoop: EventLoop, lifecycleId: String, context: LambdaContext, payload: [UInt8]) -> EventLoopFuture<LambdaResult> {
         // offloading so user code never blocks the eventloop
-        DispatchQueue(label: "lambda-initialize").async {
-            self.initialize { result in
-                switch result {
-                case .failure(let error):
-                    return promise.fail(error)
-                case .success:
-                    return promise.succeed(())
-                }
+        let promise = eventLoop.makePromise(of: LambdaResult.self)
+        DispatchQueue(label: "lambda-\(lifecycleId)").async {
+            self.handle(context: context, payload: payload) { result in
+                promise.succeed(result)
             }
         }
+        return promise.futureResult
     }
 }