fixup

Author: tomerd

Sources/AWSLambdaRuntime/Lambda+LocalServer.swift

@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#if DEBUG
 import Dispatch
 import Logging
 import NIO
@@ -26,8 +27,6 @@ import NIOHTTP1
 //         callback(.success("Hello, \(payload)!"))
 //     }
 // }
-
-#if DEBUG
 extension Lambda {
     /// Execute code in the context of a mock Lambda server.
     ///
@@ -91,10 +90,10 @@ private enum LocalLambda {
         public typealias InboundIn = HTTPServerRequestPart
         public typealias OutboundOut = HTTPServerResponsePart
 
-        private static let queueLock = Lock()
-        private static var queue = [String: Pending]()
+        private var pending = CircularBuffer<(head: HTTPRequestHead, body: ByteBuffer?)>()
 
-        private var processing = CircularBuffer<(head: HTTPRequestHead, body: ByteBuffer?)>()
+        private static var invocations = CircularBuffer<Invocation>()
+        private static var invocationState = InvocationState.waitingForLambdaRequest
 
         private let logger: Logger
         private let invocationEndpoint: String
@@ -109,72 +108,101 @@ private enum LocalLambda {
 
             switch requestPart {
             case .head(let head):
-                self.processing.append((head: head, body: nil))
+                self.pending.append((head: head, body: nil))
             case .body(var buffer):
-                var request = self.processing.removeFirst()
+                var request = self.pending.removeFirst()
                 if request.body == nil {
                     request.body = buffer
                 } else {
                     request.body!.writeBuffer(&buffer)
                 }
-                self.processing.prepend(request)
+                self.pending.prepend(request)
             case .end:
-                let request = self.processing.removeFirst()
+                let request = self.pending.removeFirst()
                 self.processRequest(context: context, request: request)
             }
         }
 
         func processRequest(context: ChannelHandlerContext, request: (head: HTTPRequestHead, body: ByteBuffer?)) {
-            if request.head.uri.hasSuffix(self.invocationEndpoint) {
-                if let work = request.body {
-                    let requestId = "\(DispatchTime.now().uptimeNanoseconds)" // FIXME:
-                    let promise = context.eventLoop.makePromise(of: Response.self)
+            switch (request.head.method, request.head.uri) {
+            // this endpoint is called by the client invoking the lambda
+            case (.POST, let url) where url.hasSuffix(self.invocationEndpoint):
+                guard let work = request.body else {
+                    return self.writeResponse(context: context, response: .init(status: .badRequest))
+                }
+                let requestID = "\(DispatchTime.now().uptimeNanoseconds)" // FIXME:
+                let promise = context.eventLoop.makePromise(of: Response.self)
+                promise.futureResult.whenComplete { result in
+                    switch result {
+                    case .failure(let error):
+                        self.logger.error("invocation error: \(error)")
+                        self.writeResponse(context: context, response: .init(status: .internalServerError))
+                    case .success(let response):
+                        self.writeResponse(context: context, response: response)
+                    }
+                }
+                let invocation = Invocation(requestID: requestID, request: work, responsePromise: promise)
+                switch Self.invocationState {
+                case .waitingForInvocation(let promise):
+                    promise.succeed(invocation)
+                case .waitingForLambdaRequest, .waitingForLambdaResponse:
+                    Self.invocations.append(invocation)
+                }
+            // /next endpoint is called by the lambda polling for work
+            case (.GET, let url) where url.hasSuffix(Consts.requestWorkURLSuffix):
+                // check if our server is in the correct state
+                guard case .waitingForLambdaRequest = Self.invocationState else {
+                    self.logger.error("invalid invocation state \(Self.invocationState)")
+                    self.writeResponse(context: context, response: .init(status: .unprocessableEntity))
+                    return
+                }
+
+                // pop the first task from the queue
+                switch Self.invocations.popFirst() {
+                case .none:
+                    // if there is nothing in the queue,
+                    // create a promise that we can fullfill when we get a new task
+                    let promise = context.eventLoop.makePromise(of: Invocation.self)
                     promise.futureResult.whenComplete { result in
                         switch result {
-                        case .success(let response):
-                            self.writeResponse(context: context, response: response)
-                        case .failure:
+                        case .failure(let error):
+                            self.logger.error("invocation error: \(error)")
                             self.writeResponse(context: context, response: .init(status: .internalServerError))
+                        case .success(let invocation):
+                            Self.invocationState = .waitingForLambdaResponse(invocation)
+                            self.writeResponse(context: context, response: invocation.makeResponse())
                         }
                     }
-                    Self.queueLock.withLock {
-                        Self.queue[requestId] = Pending(requestId: requestId, request: work, responsePromise: promise)
-                    }
+                    Self.invocationState = .waitingForInvocation(promise)
+                case .some(let invocation):
+                    // if there is a task pending, we can immediatly respond with it.
+                    Self.invocationState = .waitingForLambdaResponse(invocation)
+                    self.writeResponse(context: context, response: invocation.makeResponse())
                 }
-            } else if request.head.uri.hasSuffix("/next") {
-                switch (Self.queueLock.withLock { Self.queue.popFirst() }) {
-                case .none:
-                    self.writeResponse(context: context, response: .init(status: .noContent))
-                case .some(let pending):
-                    var response = Response()
-                    response.body = pending.value.request
-                    // required headers
-                    response.headers = [
-                        (AmazonHeaders.requestID, pending.key),
-                        (AmazonHeaders.invokedFunctionARN, "arn:aws:lambda:us-east-1:\(Int16.random(in: Int16.min ... Int16.max)):function:custom-runtime"),
-                        (AmazonHeaders.traceID, "Root=\(Int16.random(in: Int16.min ... Int16.max));Parent=\(Int16.random(in: Int16.min ... Int16.max));Sampled=1"),
-                        (AmazonHeaders.deadline, "\(DispatchWallTime.distantFuture.millisSinceEpoch)"),
-                    ]
-                    Self.queueLock.withLock {
-                        Self.queue[pending.key] = pending.value
-                    }
-                    self.writeResponse(context: context, response: response)
-                }
-
-            } else if request.head.uri.hasSuffix("/response") {
+            // :requestID/response endpoint is called by the lambda posting the response
+            case (.POST, let url) where url.hasSuffix(Consts.postResponseURLSuffix):
                 let parts = request.head.uri.split(separator: "/")
-                guard let requestId = parts.count > 2 ? String(parts[parts.count - 2]) : nil else {
+
+                guard let requestID = parts.count > 2 ? String(parts[parts.count - 2]) : nil else {
+                    // the request is malformed, since we were expecting a requestId in the path
                     return self.writeResponse(context: context, response: .init(status: .badRequest))
                 }
-                switch (Self.queueLock.withLock { Self.queue[requestId] }) {
-                case .none:
-                    self.writeResponse(context: context, response: .init(status: .badRequest))
-                case .some(let pending):
-                    pending.responsePromise.succeed(.init(status: .ok, body: request.body))
-                    self.writeResponse(context: context, response: .init(status: .accepted))
-                    Self.queueLock.withLock { Self.queue[requestId] = nil }
+                guard case .waitingForLambdaResponse(let invocation) = Self.invocationState else {
+                    // a response was send, but we did not expect to receive one
+                    self.logger.error("invalid invocation state \(Self.invocationState)")
+                    return self.writeResponse(context: context, response: .init(status: .unprocessableEntity))
+                }
+                guard requestID == invocation.requestID else {
+                    // the request's requestId is not matching the one we are expecting
+                    self.logger.error("invalid invocation state request ID \(requestID) does not match expected \(invocation.requestID)")
+                    return self.writeResponse(context: context, response: .init(status: .badRequest))
                 }
-            } else {
+
+                invocation.responsePromise.succeed(.init(status: .ok, body: request.body))
+                self.writeResponse(context: context, response: .init(status: .accepted))
+                Self.invocationState = .waitingForLambdaRequest
+            // unknown call
+            default:
                 self.writeResponse(context: context, response: .init(status: .notFound))
             }
         }
@@ -207,10 +235,29 @@ private enum LocalLambda {
             var body: ByteBuffer?
         }
 
-        struct Pending {
-            let requestId: String
+        struct Invocation {
+            let requestID: String
             let request: ByteBuffer
             let responsePromise: EventLoopPromise<Response>
+
+            func makeResponse() -> Response {
+                var response = Response()
+                response.body = self.request
+                // required headers
+                response.headers = [
+                    (AmazonHeaders.requestID, self.requestID),
+                    (AmazonHeaders.invokedFunctionARN, "arn:aws:lambda:us-east-1:\(Int16.random(in: Int16.min ... Int16.max)):function:custom-runtime"),
+                    (AmazonHeaders.traceID, "Root=\(Int16.random(in: Int16.min ... Int16.max));Parent=\(Int16.random(in: Int16.min ... Int16.max));Sampled=1"),
+                    (AmazonHeaders.deadline, "\(DispatchWallTime.distantFuture.millisSinceEpoch)"),
+                ]
+                return response
+            }
+        }
+
+        enum InvocationState {
+            case waitingForInvocation(EventLoopPromise<Invocation>)
+            case waitingForLambdaRequest
+            case waitingForLambdaResponse(Invocation)
         }
     }
 

Sources/AWSLambdaRuntimeCore/LambdaRunner.swift

@@ -46,7 +46,7 @@ extension Lambda {
         func run(logger: Logger, handler: Handler) -> EventLoopFuture<Void> {
             logger.debug("lambda invocation sequence starting")
             // 1. request work from lambda runtime engine
-            return self.runtimeClient.getNextInvocation(logger: logger).peekError { error in
+            return self.runtimeClient.getNextInvocation(logger: logger).peekError { (error) in
                 if case RuntimeError.badStatusCode(.noContent) = error {
                     return
                 }
@@ -67,13 +67,7 @@ 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)")
                 }
-            }.flatMapErrorThrowing { error in
-                if case RuntimeError.badStatusCode(.noContent) = error {
-                    return ()
-                }
-                throw error
-            }
-            .always { result in
+            }.always { result in
                 // we are done!
                 logger.log(level: result.successful ? .debug : .warning, "lambda invocation sequence completed \(result.successful ? "successfully" : "with failure")")
             }

Sources/StringSample/main.swift

@@ -26,9 +26,7 @@ struct Handler: EventLoopLambdaHandler {
     }
 }
 
-try Lambda.withLocalServer {
-    Lambda.run(Handler())
-}
+Lambda.run(Handler())
 
 // MARK: - this can also be expressed as a closure: