shutdown handler, lambda api adjustment, latest swift-nio (#5)

* shutdown handler, lambda api adjustment, latest swift-nio

motivation:
* shutdown: while not required, demonstrate better example for shutdown
* latest nio and api adjustments: fixes runtime issues observed when running for realz

changes:
* add signal based shutdown hook that terminates and cleans up resources
* pull latest swift-nio to get fix to thread renaming
* adjust aws lambda api prefix and headers
* reformat linux tests code

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", .upToNextMajor(from: "1.9.5")),
+        .package(url: "https://github.com/apple/swift-nio.git", .branch("master")),
     ],
     targets: targets
 )

Sources/SwiftAwsLambda/HttpClient.swift

@@ -25,16 +25,16 @@ internal class HTTPClient {
 
     func get(url: String) -> EventLoopFuture<HTTPResponse> {
         guard let request = HTTPRequest(url: url, method: .GET) else {
-            return eventLoop.newFailedFuture(error: HTTPClientError.invalidRequest)
+            return self.eventLoop.newFailedFuture(error: HTTPClientError.invalidRequest)
         }
-        return execute(request)
+        return self.execute(request)
     }
 
     func post(url: String, body: ByteBuffer? = nil) -> EventLoopFuture<HTTPResponse> {
         guard let request = HTTPRequest(url: url, method: .POST, body: body) else {
-            return eventLoop.newFailedFuture(error: HTTPClientError.invalidRequest)
+            return self.eventLoop.newFailedFuture(error: HTTPClientError.invalidRequest)
         }
-        return execute(request)
+        return self.execute(request)
     }
 
     func execute(_ request: HTTPRequest) -> EventLoopFuture<HTTPResponse> {
@@ -83,9 +83,9 @@ internal struct HTTPRequest: Equatable {
 
         self.version = version
         self.method = method
-        target = url.path + (url.query.map { "?" + $0 } ?? "")
+        self.target = url.path + (url.query.map { "?" + $0 } ?? "")
         self.host = host
-        port = url.port ?? 80
+        self.port = url.port ?? 80
         self.headers = headers
         self.body = body
     }
@@ -118,9 +118,9 @@ private struct HTTPResponseAccumulator {
     var state = State.idle
 
     mutating func handle(_ head: HTTPResponseHead) {
-        switch state {
+        switch self.state {
         case .idle:
-            state = .head(head)
+            self.state = .head(head)
         case .head:
             preconditionFailure("head already set")
         case .body:
@@ -131,11 +131,11 @@ private struct HTTPResponseAccumulator {
     }
 
     mutating func handle(_ part: ByteBuffer) {
-        switch state {
+        switch self.state {
         case .idle:
             preconditionFailure("no head received before body")
         case let .head(head):
-            state = .body(head, part)
+            self.state = .body(head, part)
         case .body(let head, var body):
             var part = part
             body.write(buffer: &part)
@@ -197,19 +197,19 @@ private class HTTPPartsHandler: ChannelInboundHandler, ChannelOutboundHandler {
 
         switch response {
         case let .head(head):
-            accumulator.handle(head)
+            self.accumulator.handle(head)
         case let .body(body):
-            accumulator.handle(body)
+            self.accumulator.handle(body)
         case .end:
-            switch accumulator.state {
+            switch self.accumulator.state {
             case .idle:
                 preconditionFailure("no head received before end")
             case let .head(head):
                 ctx.fireChannelRead(wrapInboundOut(HTTPResponse(status: head.status, headers: head.headers, body: nil)))
-                accumulator.state = .end
+                self.accumulator.state = .end
             case let .body(head, body):
                 ctx.fireChannelRead(wrapInboundOut(HTTPResponse(status: head.status, headers: head.headers, body: body)))
-                accumulator.state = .end
+                self.accumulator.state = .end
             case .end:
                 preconditionFailure("request already processed")
             }
@@ -244,17 +244,17 @@ private class UnaryHTTPHandler: ChannelInboundHandler, ChannelOutboundHandler {
         self.failAllPromises(error: error)
         ctx.close(promise: nil)
     }
-    
+
     func channelInactive(ctx: ChannelHandlerContext) {
         // Fail all promises
         self.failAllPromises(error: HTTPClientError.connectionClosed)
         ctx.fireChannelInactive()
     }
-    
+
     private func failAllPromises(error: Error) {
         while let promise = buffer.first {
             promise.fail(error: error)
-            buffer.removeFirst()
+            self.buffer.removeFirst()
         }
     }
 }

Sources/SwiftAwsLambda/Lambda+Codable.swift

@@ -16,21 +16,21 @@ import Foundation
 
 extension Lambda {
     public static func run<In: Decodable, Out: Encodable>(_ closure: @escaping LambdaCodableClosure<In, Out>) -> LambdaLifecycleResult {
-        return run(LambdaClosureWrapper(closure))
+        return self.run(LambdaClosureWrapper(closure))
     }
 
     public static func run<T>(_ handler: T) -> LambdaLifecycleResult where T: LambdaCodableHandler {
-        return run(handler as LambdaHandler)
+        return self.run(handler as LambdaHandler)
     }
 
     // for testing
     internal static func run<In: Decodable, Out: Encodable>(closure: @escaping LambdaCodableClosure<In, Out>, maxTimes: Int) -> LambdaLifecycleResult {
-        return run(handler: LambdaClosureWrapper(closure), maxTimes: maxTimes)
+        return self.run(handler: LambdaClosureWrapper(closure), maxTimes: maxTimes)
     }
 
     // for testing
     internal static func run<T>(handler: T, maxTimes: Int) -> LambdaLifecycleResult where T: LambdaCodableHandler {
-        return run(handler: handler as LambdaHandler, maxTimes: maxTimes)
+        return self.run(handler: handler as LambdaHandler, maxTimes: maxTimes)
     }
 }
 
@@ -66,7 +66,7 @@ public extension LambdaCodableHandler {
         guard let payloadAsCodable = codec.decode(payload) else {
             return callback(.failure("failed decoding payload (in)"))
         }
-        handle(context: context, payload: payloadAsCodable, callback: { result in
+        self.handle(context: context, payload: payloadAsCodable, callback: { result in
             switch result {
             case let .success(encodable):
                 guard let codableAsBytes = self.codec.encode(encodable) else {
@@ -80,16 +80,16 @@ public extension LambdaCodableHandler {
     }
 }
 
-// This is a class as encoder amd decoder are a class, which means its cheaper to hold a reference to both in a class then a struct. 
+// This is a class as encoder amd decoder are a class, which means its cheaper to hold a reference to both in a class then a struct.
 private class LambdaCodableJsonCodec<In: Decodable, Out: Encodable>: LambdaCodableCodec<In, Out> {
     private let encoder = JSONEncoder()
     private let decoder = JSONDecoder()
     public override func encode(_ value: Out) -> [UInt8]? {
-        return try? [UInt8](encoder.encode(value))
+        return try? [UInt8](self.encoder.encode(value))
     }
 
     public override func decode(_ data: [UInt8]) -> In? {
-        return try? decoder.decode(In.self, from: Data(data))
+        return try? self.decoder.decode(In.self, from: Data(data))
     }
 }
 
@@ -102,6 +102,6 @@ private struct LambdaClosureWrapper<In: Decodable, Out: Encodable>: LambdaCodabl
     }
 
     public func handle(context: LambdaContext, payload: In, callback: @escaping LambdaCodableCallback<Out>) {
-        closure(context, payload, callback)
+        self.closure(context, payload, callback)
     }
 }

Sources/SwiftAwsLambda/Lambda+String.swift

@@ -14,21 +14,21 @@
 
 extension Lambda {
     public static func run(_ closure: @escaping LambdaStringClosure) -> LambdaLifecycleResult {
-        return run(LambdaClosureWrapper(closure))
+        return self.run(LambdaClosureWrapper(closure))
     }
 
     public static func run(_ handler: LambdaStringHandler) -> LambdaLifecycleResult {
-        return run(handler as LambdaHandler)
+        return self.run(handler as LambdaHandler)
     }
 
     // for testing
     internal static func run(closure: @escaping LambdaStringClosure, maxTimes: Int) -> LambdaLifecycleResult {
-        return run(handler: LambdaClosureWrapper(closure), maxTimes: maxTimes)
+        return self.run(handler: LambdaClosureWrapper(closure), maxTimes: maxTimes)
     }
 
     // for testing
     internal static func run(handler: LambdaStringHandler, maxTimes: Int) -> LambdaLifecycleResult {
-        return run(handler: handler as LambdaHandler, maxTimes: maxTimes)
+        return self.run(handler: handler as LambdaHandler, maxTimes: maxTimes)
     }
 }
 
@@ -47,7 +47,7 @@ public extension LambdaStringHandler {
         guard let payloadAsString = String(bytes: payload, encoding: .utf8) else {
             return callback(.failure("failed casting payload to String"))
         }
-        handle(context: context, payload: payloadAsString, callback: { result in
+        self.handle(context: context, payload: payloadAsString, callback: { result in
             switch result {
             case let .success(string):
                 return callback(.success([UInt8](string.utf8)))
@@ -65,6 +65,6 @@ private struct LambdaClosureWrapper: LambdaStringHandler {
     }
 
     func handle(context: LambdaContext, payload: String, callback: @escaping LambdaStringCallback) {
-        closure(context, payload, callback)
+        self.closure(context, payload, callback)
     }
 }

Sources/SwiftAwsLambda/Lambda.swift

@@ -12,57 +12,127 @@
 //
 //===----------------------------------------------------------------------===//
 
+import Foundation
 import NIO
 
 public enum Lambda {
     public static func run(_ closure: @escaping LambdaClosure) -> LambdaLifecycleResult {
-        return run(LambdaClosureWrapper(closure))
+        return self._run(handler: LambdaClosureWrapper(closure))
     }
 
     public static func run(_ handler: LambdaHandler) -> LambdaLifecycleResult {
-        return Lifecycle(handler: handler).start()
+        return self._run(handler: handler)
     }
 
     // for testing
     internal static func run(closure: @escaping LambdaClosure, maxTimes: Int) -> LambdaLifecycleResult {
-        return run(handler: LambdaClosureWrapper(closure), maxTimes: maxTimes)
+        return self._run(handler: LambdaClosureWrapper(closure), maxTimes: maxTimes)
     }
 
     // for testing
     internal static func run(handler: LambdaHandler, maxTimes: Int) -> LambdaLifecycleResult {
-        return Lifecycle(handler: handler, maxTimes: maxTimes).start()
+        return self._run(handler: handler, maxTimes: maxTimes)
+    }
+
+    internal static func _run(handler: LambdaHandler, maxTimes: Int = 0, stopSignal: Signal = .INT) -> LambdaLifecycleResult {
+        do {
+            return try self._run(handler: handler, maxTimes: maxTimes, stopSignal: stopSignal).wait()
+        } catch {
+            return .failure(error)
+        }
+    }
+
+    internal static func _run(handler: LambdaHandler, maxTimes: Int = 0, stopSignal: Signal = .INT) -> EventLoopFuture<LambdaLifecycleResult> {
+        let lifecycle = Lifecycle(handler: handler, maxTimes: maxTimes)
+        let signalSource = trap(signal: stopSignal) { signal in
+            print("intercepted signal: \(signal)")
+            lifecycle.stop()
+        }
+        let future = lifecycle.start()
+        future.whenComplete {
+            signalSource.cancel()
+        }
+        return future
     }
 
     private class Lifecycle {
+        private let eventLoop = MultiThreadedEventLoopGroup(numberOfThreads: 1).next()
         private let handler: LambdaHandler
         private let max: Int
-        private var counter: Int = 0
 
-        init(handler: LambdaHandler, maxTimes: Int = 0) {
+        private var _state = LifecycleState.initialized
+        private let stateQueue = DispatchQueue(label: "LifecycleState")
+
+        init(handler: LambdaHandler, maxTimes: Int) {
+            print("lambda lifecycle init")
             self.handler = handler
-            max = maxTimes
-            assert(max >= 0)
+            self.max = maxTimes
+            assert(self.max >= 0)
+        }
+
+        deinit {
+            print("lambda lifecycle deinit")
+            assert(state == .shutdown)
+        }
+
+        private var state: LifecycleState {
+            get {
+                return self.stateQueue.sync {
+                    self._state
+                }
+            }
+            set {
+                self.stateQueue.sync {
+                    assert(newValue.rawValue > _state.rawValue, "invalid state \(newValue) after \(_state)")
+                    self._state = newValue
+                }
+            }
         }
 
-        func start() -> LambdaLifecycleResult {
-            var err: Error?
-            let runner = LambdaRunner(handler)
-            while nil == err && (0 == max || counter < max) {
-                do {
-                    // blocking! per aws lambda runtime spec the polling requets are to be done one at a time
-                    let result = try runner.run().wait()
-                    switch result {
-                    case .success:
-                        counter = counter + 1
-                    case let .failure(e):
-                        err = e
+        func start() -> EventLoopFuture<LambdaLifecycleResult> {
+            self.state = .active
+            let runner = LambdaRunner(eventLoop: eventLoop, lambdaHandler: handler)
+            let promise: EventLoopPromise<LambdaLifecycleResult> = eventLoop.newPromise()
+            print("lambda lifecycle statring")
+            DispatchQueue.global().async {
+                var err: Error?
+                var counter = 0
+                while .active == self.state && nil == err && (0 == self.max || counter < self.max) {
+                    do {
+                        // blocking! per aws lambda runtime spec the polling requets are to be done one at a time
+                        let result = try runner.run().wait()
+                        switch result {
+                        case .success:
+                            counter = counter + 1
+                        case let .failure(e):
+                            err = e
+                        }
+                    } catch {
+                        err = error
                     }
-                } catch {
-                    err = error
                 }
+                promise.succeed(result: err.map { _ in .failure(err!) } ?? .success(counter))
+                self.shutdown()
             }
-            return err.map { _ in .failure(err!) } ?? .success(counter)
+            return promise.futureResult
         }
+
+        func stop() {
+            print("lambda lifecycle stopping")
+            self.state = .stopping
+        }
+
+        private func shutdown() {
+            try! self.eventLoop.syncShutdownGracefully()
+            self.state = .shutdown
+        }
+    }
+
+    private enum LifecycleState: Int {
+        case initialized = 0
+        case active = 1
+        case stopping = 2
+        case shutdown = 3
     }
 }
 
@@ -84,7 +154,7 @@ public struct LambdaContext {
     public let invokedFunctionArn: String?
     public let cognitoIdentity: String?
     public let clientContext: String?
-    public let deadlineNs: String?
+    public let deadline: String?
 }
 
 private struct LambdaClosureWrapper: LambdaHandler {
@@ -94,6 +164,6 @@ private struct LambdaClosureWrapper: LambdaHandler {
     }
 
     func handle(context: LambdaContext, payload: [UInt8], callback: @escaping LambdaCallback) {
-        closure(context, payload, callback)
+        self.closure(context, payload, callback)
     }
 }