LambdaRuntime.swift

//===----------------------------------------------------------------------===//
//
// This source file is part of the SwiftAWSLambdaRuntime open source project
//
// Copyright (c) 2017-2020 Apple Inc. and the SwiftAWSLambdaRuntime project authors
// Licensed under Apache License v2.0
//
// See LICENSE.txt for license information
// See CONTRIBUTORS.txt for the list of SwiftAWSLambdaRuntime project authors
//
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//

import Logging
import NIOConcurrencyHelpers
import NIOCore

/// `LambdaRuntime` manages the Lambda process lifecycle.
///
/// - note: It is intended to be used within a single `EventLoop`. For this reason this class is not thread safe.
public final class LambdaRuntime<Handler: ByteBufferLambdaHandler> {
    private let eventLoop: EventLoop
    private let shutdownPromise: EventLoopPromise<Int>
    private let logger: Logger
    private let configuration: Lambda.Configuration

    private var state = State.idle {
        willSet {
            self.eventLoop.assertInEventLoop()
            precondition(newValue.order > self.state.order, "invalid state \(newValue) after \(self.state.order)")
        }
    }

    /// Create a new `LambdaRuntime`.
    ///
    /// - parameters:
    ///     - eventLoop: An `EventLoop` to run the Lambda on.
    ///     - logger: A `Logger` to log the Lambda events.
    public convenience init(eventLoop: EventLoop, logger: Logger) {
        self.init(eventLoop: eventLoop, logger: logger, configuration: .init())
    }

    init(eventLoop: EventLoop, logger: Logger, configuration: Lambda.Configuration) {
        self.eventLoop = eventLoop
        self.shutdownPromise = eventLoop.makePromise(of: Int.self)
        self.logger = logger
        self.configuration = configuration
    }

    deinit {
        guard case .shutdown = self.state else {
            preconditionFailure("invalid state \(self.state)")
        }
    }

    /// The `Lifecycle` shutdown future.
    ///
    /// - Returns: An `EventLoopFuture` that is fulfilled after the Lambda lifecycle has fully shutdown.
    public var shutdownFuture: EventLoopFuture<Int> {
        self.shutdownPromise.futureResult
    }

    /// Start the `LambdaRuntime`.
    ///
    /// - Returns: An `EventLoopFuture` that is fulfilled after the Lambda hander has been created and initialized, and a first run has been scheduled.
    public func start() -> EventLoopFuture<Void> {
        if self.eventLoop.inEventLoop {
            return self._start()
        } else {
            return self.eventLoop.flatSubmit { self._start() }
        }
    }

    private func _start() -> EventLoopFuture<Void> {
        // This method must be called on the `EventLoop` the `LambdaRuntime` has been initialized with.
        self.eventLoop.assertInEventLoop()

        logger.info("lambda runtime starting with \(self.configuration)")
        self.state = .initializing

        var logger = self.logger
        logger[metadataKey: "lifecycleId"] = .string(self.configuration.lifecycle.id)
        let terminator = LambdaTerminator()
        let runner = Lambda.Runner(eventLoop: self.eventLoop, configuration: self.configuration)

        let startupFuture = runner.initialize(logger: logger, terminator: terminator, handlerType: Handler.self)
        startupFuture.flatMap { handler -> EventLoopFuture<Result<Int, Error>> in
            // after the startup future has succeeded, we have a handler that we can use
            // to `run` the lambda.
            let finishedPromise = self.eventLoop.makePromise(of: Int.self)
            self.state = .active(runner, handler)
            self.run(promise: finishedPromise)
            return finishedPromise.futureResult.mapResult { $0 }
        }.flatMap { runnerResult -> EventLoopFuture<Int> in
            // after the lambda finishPromise has succeeded or failed we need to
            // shutdown the handler
            terminator.terminate(eventLoop: self.eventLoop).flatMapErrorThrowing { error in
                // if, we had an error shutting down the handler, we want to concatenate it with
                // the runner result
                logger.error("Error shutting down handler: \(error)")
                throw Lambda.RuntimeError.shutdownError(shutdownError: error, runnerResult: runnerResult)
            }.flatMapResult { _ -> Result<Int, Error> in
                // we had no error shutting down the lambda. let's return the runner's result
                runnerResult
            }
        }.always { _ in
            // triggered when the Lambda has finished its last run or has a startup failure.
            self.markShutdown()
        }.cascade(to: self.shutdownPromise)

        return startupFuture.map { _ in }
    }

    // MARK: -  Private

    #if DEBUG
    /// Begin the `LambdaRuntime` shutdown. Only needed for debugging purposes, hence behind a `DEBUG` flag.
    public func shutdown() {
        // make this method thread safe by dispatching onto the eventloop
        self.eventLoop.execute {
            let oldState = self.state
            self.state = .shuttingdown
            if case .active(let runner, _) = oldState {
                runner.cancelWaitingForNextInvocation()
            }
        }
    }
    #endif

    private func markShutdown() {
        self.state = .shutdown
    }

    @inline(__always)
    private func run(promise: EventLoopPromise<Int>) {
        func _run(_ count: Int) {
            switch self.state {
            case .active(let runner, let handler):
                if self.configuration.lifecycle.maxTimes > 0, count >= self.configuration.lifecycle.maxTimes {
                    return promise.succeed(count)
                }
                var logger = self.logger
                logger[metadataKey: "lifecycleIteration"] = "\(count)"
                runner.run(logger: logger, handler: handler).whenComplete { result in
                    switch result {
                    case .success:
                        logger.log(level: .debug, "lambda invocation sequence completed successfully")
                        // recursive! per aws lambda runtime spec the polling requests are to be done one at a time
                        _run(count + 1)
                    case .failure(HTTPClient.Errors.cancelled):
                        if case .shuttingdown = self.state {
                            // if we ware shutting down, we expect to that the get next
                            // invocation request might have been cancelled. For this reason we
                            // succeed the promise here.
                            logger.log(level: .info, "lambda invocation sequence has been cancelled for shutdown")
                            return promise.succeed(count)
                        }
                        logger.log(level: .error, "lambda invocation sequence has been cancelled unexpectedly")
                        promise.fail(HTTPClient.Errors.cancelled)
                    case .failure(let error):
                        logger.log(level: .error, "lambda invocation sequence completed with error: \(error)")
                        promise.fail(error)
                    }
                }
            case .shuttingdown:
                promise.succeed(count)
            default:
                preconditionFailure("invalid run state: \(self.state)")
            }
        }

        _run(0)
    }

    private enum State {
        case idle
        case initializing
        case active(Lambda.Runner, Handler)
        case shuttingdown
        case shutdown

        internal var order: Int {
            switch self {
            case .idle:
                return 0
            case .initializing:
                return 1
            case .active:
                return 2
            case .shuttingdown:
                return 3
            case .shutdown:
                return 4
            }
        }
    }
}

/// This is safe since lambda runtime synchronizes by dispatching all methods to a single `EventLoop`
#if compiler(>=5.5) && canImport(_Concurrency)
extension LambdaRuntime: @unchecked Sendable {}
#endif