LambdaRuntime.swift

//===----------------------------------------------------------------------===//
//
// This source file is part of the SwiftAWSLambdaRuntime open source project
//
// Copyright (c) 2024 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
import Synchronization

#if canImport(FoundationEssentials)
import FoundationEssentials
#else
import Foundation
#endif

// This is our gardian to ensure only one LambdaRuntime is running at the time
// We use an Atomic here to ensure thread safety
private let _isRunning = Atomic<Bool>(false)

// We need `@unchecked` Sendable here, as `NIOLockedValueBox` does not understand `sending` today.
// We don't want to use `NIOLockedValueBox` here anyway. We would love to use Mutex here, but this
// sadly crashes the compiler today (on Linux).
public final class LambdaRuntime<Handler>: @unchecked Sendable where Handler: StreamingLambdaHandler {
    // TODO: We want to change this to Mutex as soon as this doesn't crash the Swift compiler on Linux anymore
    let handlerMutex: NIOLockedValueBox<Handler?>
    let logger: Logger
    let eventLoop: EventLoop

    public init(
        handler: sending Handler,
        eventLoop: EventLoop = Lambda.defaultEventLoop,
        logger: Logger = Logger(label: "LambdaRuntime")
    ) {
        self.handlerMutex = NIOLockedValueBox(handler)
        self.eventLoop = eventLoop

        // by setting the log level here, we understand it can not be changed dynamically at runtime
        // developers have to wait for AWS Lambda to dispose and recreate a runtime environment to pickup a change
        // this approach is less flexible but more performant than reading the value of the environment variable at each invocation
        var log = logger
        log.logLevel = Lambda.env("LOG_LEVEL").flatMap(Logger.Level.init) ?? .info
        self.logger = log
        self.logger.debug("LambdaRuntime initialized")
    }

    /// Make sure only one run() is called at a time
    public func run() async throws {

        // we use an atomic global variable to ensure only one LambdaRuntime is running at the time
        let (_, original) = _isRunning.compareExchange(expected: false, desired: true, ordering: .relaxed)

        // if the original value was already true, run() is already running
        if original {
            throw LambdaRuntimeError(code: .runtimeCanOnlyBeStartedOnce)
        }

        try await withTaskCancellationHandler {
            // call the internal _run() method
            do {
                try await self._run()
            } catch {
                // when we catch an error, flip back the global variable to false
                _isRunning.store(false, ordering: .relaxed)
                throw error
            }
        } onCancel: {
            // when task is cancelled, flip back the global variable to false
            _isRunning.store(false, ordering: .relaxed)
        }

        // when we're done without error and without cancellation, flip back the global variable to false
        _isRunning.store(false, ordering: .relaxed)
    }

    private func _run() async throws {
        let handler = self.handlerMutex.withLockedValue { handler in
            let result = handler
            handler = nil
            return result
        }

        guard let handler else {
            throw LambdaRuntimeError(code: .runtimeCanOnlyBeStartedOnce)
        }

        // are we running inside an AWS Lambda runtime environment ?
        // AWS_LAMBDA_RUNTIME_API is set when running on Lambda
        // https://docs.aws.amazon.com/lambda/latest/dg/runtimes-api.html
        if let runtimeEndpoint = Lambda.env("AWS_LAMBDA_RUNTIME_API") {

            let ipAndPort = runtimeEndpoint.split(separator: ":", maxSplits: 1)
            let ip = String(ipAndPort[0])
            guard let port = Int(ipAndPort[1]) else { throw LambdaRuntimeError(code: .invalidPort) }

            try await LambdaRuntimeClient.withRuntimeClient(
                configuration: .init(ip: ip, port: port),
                eventLoop: self.eventLoop,
                logger: self.logger
            ) { runtimeClient in
                try await Lambda.runLoop(
                    runtimeClient: runtimeClient,
                    handler: handler,
                    logger: self.logger
                )
            }

        } else {

            #if LocalServerSupport
            // we're not running on Lambda and we're compiled in DEBUG mode,
            // let's start a local server for testing
            try await Lambda.withLocalServer(invocationEndpoint: Lambda.env("LOCAL_LAMBDA_SERVER_INVOCATION_ENDPOINT"))
            {

                try await LambdaRuntimeClient.withRuntimeClient(
                    configuration: .init(ip: "127.0.0.1", port: 7000),
                    eventLoop: self.eventLoop,
                    logger: self.logger
                ) { runtimeClient in
                    try await Lambda.runLoop(
                        runtimeClient: runtimeClient,
                        handler: handler,
                        logger: self.logger
                    )
                }
            }
            #else
            // in release mode, we can't start a local server because the local server code is not compiled.
            throw LambdaRuntimeError(code: .missingLambdaRuntimeAPIEnvironmentVariable)
            #endif
        }
    }
}