Cancelling works.

Author: fabianfett

Sources/AWSLambdaRuntimeCore/HTTPClient.swift

@@ -48,6 +48,20 @@ internal final class HTTPClient {
                              timeout: timeout ?? self.configuration.requestTimeout))
     }
 
+    /// cancels the current request if there is one
+    func cancel() {
+        guard self.executing.exchange(with: true) else {
+            // there is no request running. nothing to cancel
+            return
+        }
+
+        guard case .connected(let channel) = self.state else {
+            preconditionFailure("if we are executing, we expect to have an open channel")
+        }
+
+        channel.triggerUserOutboundEvent(RequestCancelEvent(), promise: nil)
+    }
+
     // TODO: cap reconnect attempt
     private func execute(_ request: Request, validate: Bool = true) -> EventLoopFuture<Response> {
         precondition(!validate || self.executing.compareAndExchange(expected: false, desired: true), "expecting single request at a time")
@@ -120,6 +134,7 @@ internal final class HTTPClient {
     internal enum Errors: Error {
         case connectionResetByPeer
         case timeout
+        case cancelled
     }
 
     private enum State {
@@ -284,6 +299,18 @@ private final class UnaryHandler: ChannelDuplexHandler {
         context.fireChannelInactive()
     }
 
+    func triggerUserOutboundEvent(context: ChannelHandlerContext, event: Any, promise: EventLoopPromise<Void>?) {
+        switch event {
+        case is RequestCancelEvent:
+            guard self.pending != nil else {
+                return
+            }
+            self.completeWith(.failure(HTTPClient.Errors.cancelled))
+        default:
+            context.triggerUserOutboundEvent(event, promise: promise)
+        }
+    }
+
     private func completeWith(_ result: Result<HTTPClient.Response, Error>) {
         guard let pending = self.pending else {
             preconditionFailure("invalid state, no pending request")
@@ -299,3 +326,5 @@ private struct HTTPRequestWrapper {
     let request: HTTPClient.Request
     let promise: EventLoopPromise<HTTPClient.Response>
 }
+
+private struct RequestCancelEvent {}

Sources/AWSLambdaRuntimeCore/Lambda.swift

@@ -107,10 +107,10 @@ public enum Lambda {
         var logger = Logger(label: "Lambda")
         logger.logLevel = configuration.general.logLevel
         let eventLoop = eventLoopGroup.next()
-        
+
         return eventLoop.flatSubmit {
             let lifecycle = Lifecycle(eventLoop: eventLoopGroup.next(), logger: logger, configuration: configuration, factory: factory)
-            
+
             #if DEBUG
             // We only offer coordinated shutdown in DEBUG, since Lambda functions
             // running on AWS are not shutdown in a coordinated way.
@@ -119,13 +119,13 @@ public enum Lambda {
                 lifecycle.shutdown()
             }
             #endif
-            
+
             return lifecycle.start().flatMap {
                 let future = lifecycle.shutdownFuture
                 #if DEBUG
-                    return future.always { _ in signalSource.cancel() }
+                return future.always { _ in signalSource.cancel() }
                 #else
-                    return future
+                return future
                 #endif
             }
         }

Sources/AWSLambdaRuntimeCore/LambdaLifecycle.swift

@@ -81,10 +81,20 @@ extension Lambda {
 
         // MARK: -  Private
 
+        #if DEBUG
         /// Begin the `Lifecycle` shutdown.
+        /// Only needed for debugging purposes.
         public func shutdown() {
-            self.state = .shuttingdown
+            // make this method thread safe by dispatching onto the eventloop
+            self.eventLoop.execute {
+                guard case .active(let runner, _) = self.state else {
+                    return
+                }
+                self.state = .shuttingdown
+                runner.cancelWaitingForNextInvocation()
+            }
         }
+        #endif
 
         private func markShutdown() {
             self.state = .shutdown

Sources/AWSLambdaRuntimeCore/LambdaRunner.swift

@@ -18,10 +18,12 @@ import NIO
 
 extension Lambda {
     /// LambdaRunner manages the Lambda runtime workflow, or business logic.
-    internal struct Runner {
+    internal final class Runner {
         private let runtimeClient: RuntimeClient
         private let eventLoop: EventLoop
 
+        private var isGettingNextInvocation = false
+
         init(eventLoop: EventLoop, configuration: Configuration) {
             self.eventLoop = eventLoop
             self.runtimeClient = RuntimeClient(eventLoop: self.eventLoop, configuration: configuration.runtimeEngine)
@@ -46,13 +48,12 @@ 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
-                if case RuntimeError.badStatusCode(.noContent) = error {
-                    return
-                }
+            self.isGettingNextInvocation = true
+            return self.runtimeClient.getNextInvocation(logger: logger).peekError { error in
                 logger.error("could not fetch work from lambda runtime engine: \(error)")
             }.flatMap { invocation, payload in
-                // 2. send invocation to handler
+                // 2. send work to handler
+                self.isGettingNextInvocation = false
                 let context = Context(logger: logger, eventLoop: self.eventLoop, invocation: invocation)
                 logger.debug("sending invocation to lambda handler \(handler)")
                 return handler.handle(context: context, payload: payload)
@@ -72,6 +73,15 @@ extension Lambda {
                 logger.log(level: result.successful ? .debug : .warning, "lambda invocation sequence completed \(result.successful ? "successfully" : "with failure")")
             }
         }
+
+        #if DEBUG
+        /// cancels the current run, if we are waiting for next invocation (long poll from Lambda control plane)
+        /// only needed for debugging purposes.
+        func cancelWaitingForNextInvocation() {
+            guard self.isGettingNextInvocation else { return }
+            self.runtimeClient.cancel()
+        }
+        #endif
     }
 }
 

Sources/AWSLambdaRuntimeCore/LambdaRuntimeClient.swift

@@ -114,6 +114,13 @@ extension Lambda {
                 }
             }
         }
+
+        #if DEBUG
+        /// Cancels the current request, if one is running. Only needed for debugging purposes
+        func cancel() {
+            self.httpClient.cancel()
+        }
+        #endif
     }
 }