From e02b9b6bab14b6fdf08261e51086e660dc2bc896 Mon Sep 17 00:00:00 2001 From: Alessio Buratti <9006089+Buratti@users.noreply.github.com> Date: Mon, 24 Jun 2024 17:36:06 +0200 Subject: [PATCH 01/13] First prototype --- .../AWSLambdaRuntimeCore/DetachedTasks.swift | 169 ++++++++++++++++++ .../AWSLambdaRuntimeCore/LambdaContext.swift | 25 ++- .../AWSLambdaRuntimeCore/LambdaRunner.swift | 12 +- 3 files changed, 200 insertions(+), 6 deletions(-) create mode 100644 Sources/AWSLambdaRuntimeCore/DetachedTasks.swift diff --git a/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift b/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift new file mode 100644 index 00000000..1c5896d5 --- /dev/null +++ b/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift @@ -0,0 +1,169 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the SwiftAWSLambdaRuntime open source project +// +// Copyright (c) 2022 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 Foundation +import NIOConcurrencyHelpers +import NIOCore +import Logging + +/// A container that allows tasks to finish after a synchronous invocation +/// has produced its response. +public final class DetachedTasksContainer { + + struct Context { + let eventLoop: EventLoop + let logger: Logger + } + + private var context: Context + private var storage: Storage + + init(context: Context) { + self.storage = Storage() + self.context = context + } + + /// Adds a detached task that runs on the given event loop. + /// + /// - Parameters: + /// - name: The name of the task. + /// - task: The task to execute. It receives an `EventLoop` and returns an `EventLoopFuture`. + /// - Returns: A `RegistrationKey` for the registered task. + @discardableResult + public func detached(name: String, task: @escaping (EventLoop) -> EventLoopFuture) -> RegistrationKey { + let key = RegistrationKey() + let task = task(context.eventLoop).always { _ in + self.storage.remove(key) + } + self.storage.add(key: key, name: name, task: task) + return key + } + + /// Adds a detached async task. + /// + /// - Parameters: + /// - name: The name of the task. + /// - task: The async task to execute. + /// - Returns: A `RegistrationKey` for the registered task. + @discardableResult + public func detached(name: String, task: @Sendable @escaping () async throws -> Void) -> RegistrationKey { + let key = RegistrationKey() + let promise = context.eventLoop.makePromise(of: Void.self) + promise.completeWithTask(task) + let task = promise.futureResult.always { result in + switch result { + case .success: + break + case .failure(let failure): + self.context.logger.warning( + "Execution of detached task failed with error.", + metadata: [ + "taskName": "\(name)", + "error": "\(failure)" + ] + ) + } + self.storage.remove(key) + } + self.storage.add(key: key, name: name, task: task) + return key + } + + /// Informs the runtime that the specified task should not be awaited anymore. + /// + /// - Warning: This method does not actually stop the execution of the + /// detached task, it only prevents the runtime from waiting for it before + /// `/next` is invoked. + /// + /// - Parameter key: The `RegistrationKey` of the task to cancel. + public func unsafeCancel(_ key: RegistrationKey) { + // To discuss: + // Canceling the execution doesn't seem to be an easy + // task https://github.com/apple/swift-nio/issues/2087 + // + // While removing the handler will allow the runtime + // to invoke `/next` without actually awaiting for the + // task to complete, it does not actually cancel + // the execution of the dispatched task. + // Since this is a bit counter-intuitive, we might not + // want this method to exist at all. + self.storage.remove(key) + } + + /// Awaits all registered tasks to complete. + /// + /// - Returns: An `EventLoopFuture` that completes when all tasks have finished. + internal func awaitAll() -> EventLoopFuture { + let tasks = storage.tasks + if tasks.isEmpty { + return context.eventLoop.makeSucceededVoidFuture() + } else { + return EventLoopFuture.andAllComplete(tasks.map(\.value.task), on: context.eventLoop).flatMap { + self.awaitAll() + } + } + } +} + +extension DetachedTasksContainer { + /// Lambda detached task registration key. + public struct RegistrationKey: Hashable, CustomStringConvertible { + var value: String + + init() { + // UUID basically + self.value = UUID().uuidString + } + + public var description: String { + self.value + } + } +} + +extension DetachedTasksContainer { + fileprivate final class Storage { + private let lock: NIOLock + + private var map: [RegistrationKey: (name: String, task: EventLoopFuture)] + + init() { + self.lock = .init() + self.map = [:] + } + + func add(key: RegistrationKey, name: String, task: EventLoopFuture) { + self.lock.withLock { + self.map[key] = (name: name, task: task) + } + } + + func remove(_ key: RegistrationKey) { + self.lock.withLock { + self.map[key] = nil + } + } + + var tasks: [RegistrationKey: (name: String, task: EventLoopFuture)] { + self.lock.withLock { + self.map + } + } + } +} + +// Ideally this would not be @unchecked Sendable, but Sendable checks do not understand locks +// We can transition this to an actor once we drop support for older Swift versions +extension DetachedTasksContainer: @unchecked Sendable {} +extension DetachedTasksContainer.Storage: @unchecked Sendable {} +extension DetachedTasksContainer.RegistrationKey: Sendable {} diff --git a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift index 24e960a4..adb10353 100644 --- a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift +++ b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift @@ -81,6 +81,7 @@ public struct LambdaContext: CustomDebugStringConvertible, Sendable { let logger: Logger let eventLoop: EventLoop let allocator: ByteBufferAllocator + let tasks: DetachedTasksContainer init( requestID: String, @@ -91,7 +92,8 @@ public struct LambdaContext: CustomDebugStringConvertible, Sendable { clientContext: String?, logger: Logger, eventLoop: EventLoop, - allocator: ByteBufferAllocator + allocator: ByteBufferAllocator, + tasks: DetachedTasksContainer ) { self.requestID = requestID self.traceID = traceID @@ -102,6 +104,7 @@ public struct LambdaContext: CustomDebugStringConvertible, Sendable { self.logger = logger self.eventLoop = eventLoop self.allocator = allocator + self.tasks = tasks } } @@ -158,6 +161,10 @@ public struct LambdaContext: CustomDebugStringConvertible, Sendable { public var allocator: ByteBufferAllocator { self.storage.allocator } + + public var tasks: DetachedTasksContainer { + self.storage.tasks + } init(requestID: String, traceID: String, @@ -177,7 +184,13 @@ public struct LambdaContext: CustomDebugStringConvertible, Sendable { clientContext: clientContext, logger: logger, eventLoop: eventLoop, - allocator: allocator + allocator: allocator, + tasks: DetachedTasksContainer( + context: DetachedTasksContainer.Context( + eventLoop: eventLoop, + logger: logger + ) + ) ) } @@ -209,7 +222,13 @@ public struct LambdaContext: CustomDebugStringConvertible, Sendable { deadline: .now() + timeout, logger: logger, eventLoop: eventLoop, - allocator: ByteBufferAllocator() + allocator: ByteBufferAllocator(), + tasks: DetachedTasksContainer( + context: DetachedTasksContainer.Context( + eventLoop: eventLoop, + logger: logger + ) + ) ) } } diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift index 23281b94..a9f82845 100644 --- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift +++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift @@ -95,13 +95,19 @@ internal final class LambdaRunner { if case .failure(let error) = result { logger.warning("lambda handler returned an error: \(error)") } - return (invocation, result) + return (invocation, result, context) } - }.flatMap { invocation, result in + }.flatMap { invocation, result, context in // 3. report results to runtime engine self.runtimeClient.reportResults(logger: logger, invocation: invocation, result: result).peekError { error in logger.error("could not report results to lambda runtime engine: \(error)") - } + // To discuss: + // Do we want to await the tasks in this case? + return context.tasks.awaitAll() + }.map { _ in context } + } + .flatMap { context in + context.tasks.awaitAll() } } From 31c78459e0c8f86882ec56b6857691616938341c Mon Sep 17 00:00:00 2001 From: Alessio Buratti <9006089+Buratti@users.noreply.github.com> Date: Mon, 24 Jun 2024 17:55:04 +0200 Subject: [PATCH 02/13] Fix build --- Sources/AWSLambdaRuntimeCore/LambdaContext.swift | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift index adb10353..5ee855f7 100644 --- a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift +++ b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift @@ -222,13 +222,7 @@ public struct LambdaContext: CustomDebugStringConvertible, Sendable { deadline: .now() + timeout, logger: logger, eventLoop: eventLoop, - allocator: ByteBufferAllocator(), - tasks: DetachedTasksContainer( - context: DetachedTasksContainer.Context( - eventLoop: eventLoop, - logger: logger - ) - ) + allocator: ByteBufferAllocator() ) } } From 22483e076ef228ad9621fa5b2e9f2d0eb546b19f Mon Sep 17 00:00:00 2001 From: Alessio Buratti <9006089+Buratti@users.noreply.github.com> Date: Thu, 1 Aug 2024 23:18:04 +0200 Subject: [PATCH 03/13] Removes task cancellation https://github.com/swift-server/swift-aws-lambda-runtime/pull/334#discussion_r1666713889 --- .../AWSLambdaRuntimeCore/DetachedTasks.swift | 21 ------------------- 1 file changed, 21 deletions(-) diff --git a/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift b/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift index 1c5896d5..f4ee1ecf 100644 --- a/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift +++ b/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift @@ -79,27 +79,6 @@ public final class DetachedTasksContainer { return key } - /// Informs the runtime that the specified task should not be awaited anymore. - /// - /// - Warning: This method does not actually stop the execution of the - /// detached task, it only prevents the runtime from waiting for it before - /// `/next` is invoked. - /// - /// - Parameter key: The `RegistrationKey` of the task to cancel. - public func unsafeCancel(_ key: RegistrationKey) { - // To discuss: - // Canceling the execution doesn't seem to be an easy - // task https://github.com/apple/swift-nio/issues/2087 - // - // While removing the handler will allow the runtime - // to invoke `/next` without actually awaiting for the - // task to complete, it does not actually cancel - // the execution of the dispatched task. - // Since this is a bit counter-intuitive, we might not - // want this method to exist at all. - self.storage.remove(key) - } - /// Awaits all registered tasks to complete. /// /// - Returns: An `EventLoopFuture` that completes when all tasks have finished. From 224039af1d5235cee4fcefaeb1a9299dba7615de Mon Sep 17 00:00:00 2001 From: Alessio Buratti <9006089+Buratti@users.noreply.github.com> Date: Thu, 1 Aug 2024 23:20:32 +0200 Subject: [PATCH 04/13] Force user to handle errors https://github.com/swift-server/swift-aws-lambda-runtime/pull/334#discussion_r1666712903 --- Sources/AWSLambdaRuntimeCore/DetachedTasks.swift | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift b/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift index f4ee1ecf..dc024f25 100644 --- a/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift +++ b/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift @@ -56,23 +56,11 @@ public final class DetachedTasksContainer { /// - task: The async task to execute. /// - Returns: A `RegistrationKey` for the registered task. @discardableResult - public func detached(name: String, task: @Sendable @escaping () async throws -> Void) -> RegistrationKey { + public func detached(name: String, task: @Sendable @escaping () async -> Void) -> RegistrationKey { let key = RegistrationKey() let promise = context.eventLoop.makePromise(of: Void.self) promise.completeWithTask(task) let task = promise.futureResult.always { result in - switch result { - case .success: - break - case .failure(let failure): - self.context.logger.warning( - "Execution of detached task failed with error.", - metadata: [ - "taskName": "\(name)", - "error": "\(failure)" - ] - ) - } self.storage.remove(key) } self.storage.add(key: key, name: name, task: task) From 7934a0f56f80cb4b19334e8f0a40cba2222ea819 Mon Sep 17 00:00:00 2001 From: Alessio Buratti <9006089+Buratti@users.noreply.github.com> Date: Thu, 1 Aug 2024 23:21:08 +0200 Subject: [PATCH 05/13] Remove EventLoop API https://github.com/swift-server/swift-aws-lambda-runtime/pull/334#discussion_r1666712244 --- Sources/AWSLambdaRuntimeCore/DetachedTasks.swift | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift b/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift index dc024f25..e3b9f5d4 100644 --- a/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift +++ b/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift @@ -33,22 +33,6 @@ public final class DetachedTasksContainer { self.context = context } - /// Adds a detached task that runs on the given event loop. - /// - /// - Parameters: - /// - name: The name of the task. - /// - task: The task to execute. It receives an `EventLoop` and returns an `EventLoopFuture`. - /// - Returns: A `RegistrationKey` for the registered task. - @discardableResult - public func detached(name: String, task: @escaping (EventLoop) -> EventLoopFuture) -> RegistrationKey { - let key = RegistrationKey() - let task = task(context.eventLoop).always { _ in - self.storage.remove(key) - } - self.storage.add(key: key, name: name, task: task) - return key - } - /// Adds a detached async task. /// /// - Parameters: From 2491f6fab7ae2aed782cd5bd27c6e5ab8faf9933 Mon Sep 17 00:00:00 2001 From: Alessio Buratti <9006089+Buratti@users.noreply.github.com> Date: Thu, 1 Aug 2024 23:31:49 +0200 Subject: [PATCH 06/13] Make DetachedTaskContainer internal https://github.com/swift-server/swift-aws-lambda-runtime/pull/334#discussion_r1666710596 https://github.com/swift-server/swift-aws-lambda-runtime/pull/334#discussion_r1666706576 --- .../AWSLambdaRuntimeCore/DetachedTasks.swift | 51 ++++--------------- .../AWSLambdaRuntimeCore/LambdaContext.swift | 12 +++-- 2 files changed, 17 insertions(+), 46 deletions(-) diff --git a/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift b/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift index e3b9f5d4..e21deea9 100644 --- a/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift +++ b/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift @@ -18,7 +18,7 @@ import Logging /// A container that allows tasks to finish after a synchronous invocation /// has produced its response. -public final class DetachedTasksContainer { +final class DetachedTasksContainer { struct Context { let eventLoop: EventLoop @@ -26,10 +26,9 @@ public final class DetachedTasksContainer { } private var context: Context - private var storage: Storage + private var storage: [RegistrationKey: EventLoopFuture] = [:] init(context: Context) { - self.storage = Storage() self.context = context } @@ -40,14 +39,14 @@ public final class DetachedTasksContainer { /// - task: The async task to execute. /// - Returns: A `RegistrationKey` for the registered task. @discardableResult - public func detached(name: String, task: @Sendable @escaping () async -> Void) -> RegistrationKey { + func detached(task: @Sendable @escaping () async -> Void) -> RegistrationKey { let key = RegistrationKey() let promise = context.eventLoop.makePromise(of: Void.self) promise.completeWithTask(task) let task = promise.futureResult.always { result in - self.storage.remove(key) + self.storage.removeValue(forKey: key) } - self.storage.add(key: key, name: name, task: task) + self.storage[key] = task return key } @@ -55,11 +54,11 @@ public final class DetachedTasksContainer { /// /// - Returns: An `EventLoopFuture` that completes when all tasks have finished. internal func awaitAll() -> EventLoopFuture { - let tasks = storage.tasks + let tasks = storage.values if tasks.isEmpty { return context.eventLoop.makeSucceededVoidFuture() } else { - return EventLoopFuture.andAllComplete(tasks.map(\.value.task), on: context.eventLoop).flatMap { + return EventLoopFuture.andAllComplete(Array(tasks), on: context.eventLoop).flatMap { self.awaitAll() } } @@ -68,7 +67,7 @@ public final class DetachedTasksContainer { extension DetachedTasksContainer { /// Lambda detached task registration key. - public struct RegistrationKey: Hashable, CustomStringConvertible { + struct RegistrationKey: Hashable, CustomStringConvertible { var value: String init() { @@ -76,45 +75,13 @@ extension DetachedTasksContainer { self.value = UUID().uuidString } - public var description: String { + var description: String { self.value } } } -extension DetachedTasksContainer { - fileprivate final class Storage { - private let lock: NIOLock - - private var map: [RegistrationKey: (name: String, task: EventLoopFuture)] - - init() { - self.lock = .init() - self.map = [:] - } - - func add(key: RegistrationKey, name: String, task: EventLoopFuture) { - self.lock.withLock { - self.map[key] = (name: name, task: task) - } - } - - func remove(_ key: RegistrationKey) { - self.lock.withLock { - self.map[key] = nil - } - } - - var tasks: [RegistrationKey: (name: String, task: EventLoopFuture)] { - self.lock.withLock { - self.map - } - } - } -} - // Ideally this would not be @unchecked Sendable, but Sendable checks do not understand locks // We can transition this to an actor once we drop support for older Swift versions extension DetachedTasksContainer: @unchecked Sendable {} -extension DetachedTasksContainer.Storage: @unchecked Sendable {} extension DetachedTasksContainer.RegistrationKey: Sendable {} diff --git a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift index 5ee855f7..c71f7913 100644 --- a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift +++ b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift @@ -161,10 +161,6 @@ public struct LambdaContext: CustomDebugStringConvertible, Sendable { public var allocator: ByteBufferAllocator { self.storage.allocator } - - public var tasks: DetachedTasksContainer { - self.storage.tasks - } init(requestID: String, traceID: String, @@ -201,6 +197,14 @@ public struct LambdaContext: CustomDebugStringConvertible, Sendable { let remaining = deadline - now return .milliseconds(remaining) } + + var tasks: DetachedTasksContainer { + self.storage.tasks + } + + func detachedBackgroundTask(_ body: @escaping @Sendable () async -> ()) { + self.tasks.detached(task: body) + } public var debugDescription: String { "\(Self.self)(requestID: \(self.requestID), traceID: \(self.traceID), invokedFunctionARN: \(self.invokedFunctionARN), cognitoIdentity: \(self.cognitoIdentity ?? "nil"), clientContext: \(self.clientContext ?? "nil"), deadline: \(self.deadline))" From a059ec5e20e92c2a738305772bdcb4dc676cc964 Mon Sep 17 00:00:00 2001 From: Alessio Buratti <9006089+Buratti@users.noreply.github.com> Date: Fri, 2 Aug 2024 00:14:40 +0200 Subject: [PATCH 07/13] Removes @unchecked Sendable https://github.com/swift-server/swift-aws-lambda-runtime/pull/334#discussion_r1666707646 --- .../AWSLambdaRuntimeCore/DetachedTasks.swift | 40 +++++++++++-------- .../AWSLambdaRuntimeCore/LambdaContext.swift | 4 +- 2 files changed, 27 insertions(+), 17 deletions(-) diff --git a/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift b/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift index e21deea9..f06750bf 100644 --- a/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift +++ b/Sources/AWSLambdaRuntimeCore/DetachedTasks.swift @@ -18,9 +18,9 @@ import Logging /// A container that allows tasks to finish after a synchronous invocation /// has produced its response. -final class DetachedTasksContainer { +actor DetachedTasksContainer: Sendable { - struct Context { + struct Context: Sendable { let eventLoop: EventLoop let logger: Logger } @@ -38,28 +38,41 @@ final class DetachedTasksContainer { /// - name: The name of the task. /// - task: The async task to execute. /// - Returns: A `RegistrationKey` for the registered task. - @discardableResult - func detached(task: @Sendable @escaping () async -> Void) -> RegistrationKey { + func detached(task: @Sendable @escaping () async -> Void) { let key = RegistrationKey() let promise = context.eventLoop.makePromise(of: Void.self) promise.completeWithTask(task) - let task = promise.futureResult.always { result in - self.storage.removeValue(forKey: key) + let task = promise.futureResult.always { [weak self] result in + guard let self else { return } + Task { + await self.removeTask(forKey: key) + } } self.storage[key] = task - return key + } + + func removeTask(forKey key: RegistrationKey) { + self.storage.removeValue(forKey: key) } /// Awaits all registered tasks to complete. /// /// - Returns: An `EventLoopFuture` that completes when all tasks have finished. - internal func awaitAll() -> EventLoopFuture { + func awaitAll() -> EventLoopFuture { let tasks = storage.values if tasks.isEmpty { return context.eventLoop.makeSucceededVoidFuture() } else { - return EventLoopFuture.andAllComplete(Array(tasks), on: context.eventLoop).flatMap { - self.awaitAll() + let context = context + return EventLoopFuture.andAllComplete(Array(tasks), on: context.eventLoop).flatMap { [weak self] in + guard let self else { + return context.eventLoop.makeSucceededFuture(()) + } + let promise = context.eventLoop.makePromise(of: Void.self) + promise.completeWithTask { + try await self.awaitAll().get() + } + return promise.futureResult } } } @@ -67,7 +80,7 @@ final class DetachedTasksContainer { extension DetachedTasksContainer { /// Lambda detached task registration key. - struct RegistrationKey: Hashable, CustomStringConvertible { + struct RegistrationKey: Hashable, CustomStringConvertible, Sendable { var value: String init() { @@ -80,8 +93,3 @@ extension DetachedTasksContainer { } } } - -// Ideally this would not be @unchecked Sendable, but Sendable checks do not understand locks -// We can transition this to an actor once we drop support for older Swift versions -extension DetachedTasksContainer: @unchecked Sendable {} -extension DetachedTasksContainer.RegistrationKey: Sendable {} diff --git a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift index c71f7913..99e06c48 100644 --- a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift +++ b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift @@ -203,7 +203,9 @@ public struct LambdaContext: CustomDebugStringConvertible, Sendable { } func detachedBackgroundTask(_ body: @escaping @Sendable () async -> ()) { - self.tasks.detached(task: body) + Task { + await self.tasks.detached(task: body) + } } public var debugDescription: String { From 9804c0438346fa27570fae4d4516126b6e11fafd Mon Sep 17 00:00:00 2001 From: Alessio Buratti <9006089+Buratti@users.noreply.github.com> Date: Fri, 2 Aug 2024 00:48:19 +0200 Subject: [PATCH 08/13] Invoke awaitAll() from async context --- Sources/AWSLambdaRuntimeCore/LambdaRunner.swift | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift index a9f82845..e3c68b43 100644 --- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift +++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift @@ -103,11 +103,19 @@ internal final class LambdaRunner { logger.error("could not report results to lambda runtime engine: \(error)") // To discuss: // Do we want to await the tasks in this case? - return context.tasks.awaitAll() + let promise = context.eventLoop.makePromise(of: Void.self) + promise.completeWithTask { + return try await context.tasks.awaitAll().get() + } + return promise.futureResult }.map { _ in context } } .flatMap { context in - context.tasks.awaitAll() + let promise = context.eventLoop.makePromise(of: Void.self) + promise.completeWithTask { + try await context.tasks.awaitAll().get() + } + return promise.futureResult } } From fb6e19ebfcea962729fbd48ecdcf6a3a4749d0e5 Mon Sep 17 00:00:00 2001 From: Alessio Buratti <9006089+Buratti@users.noreply.github.com> Date: Thu, 15 Aug 2024 03:15:51 +0200 Subject: [PATCH 09/13] Fix ambiguous expression type for swift 5.7 --- Sources/AWSLambdaRuntimeCore/LambdaRunner.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift index e3c68b43..87457f43 100644 --- a/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift +++ b/Sources/AWSLambdaRuntimeCore/LambdaRunner.swift @@ -110,7 +110,7 @@ internal final class LambdaRunner { return promise.futureResult }.map { _ in context } } - .flatMap { context in + .flatMap { (context: LambdaContext) -> EventLoopFuture in let promise = context.eventLoop.makePromise(of: Void.self) promise.completeWithTask { try await context.tasks.awaitAll().get() From b87bfeda82199fcfa9461cd330c833cd99a3fdfe Mon Sep 17 00:00:00 2001 From: Alessio Buratti <9006089+Buratti@users.noreply.github.com> Date: Thu, 15 Aug 2024 03:21:12 +0200 Subject: [PATCH 10/13] Fix visibility of detachedBackgroundTask --- Sources/AWSLambdaRuntimeCore/LambdaContext.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift index 99e06c48..d299534b 100644 --- a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift +++ b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift @@ -202,7 +202,7 @@ public struct LambdaContext: CustomDebugStringConvertible, Sendable { self.storage.tasks } - func detachedBackgroundTask(_ body: @escaping @Sendable () async -> ()) { + public func detachedBackgroundTask(_ body: @escaping @Sendable () async -> ()) { Task { await self.tasks.detached(task: body) } From 888af2981263e975492c1d40effec546c33b021e Mon Sep 17 00:00:00 2001 From: Alessio Buratti <9006089+Buratti@users.noreply.github.com> Date: Thu, 15 Aug 2024 03:31:23 +0200 Subject: [PATCH 11/13] Add swift-doc --- Sources/AWSLambdaRuntimeCore/LambdaContext.swift | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift index d299534b..943c4e5f 100644 --- a/Sources/AWSLambdaRuntimeCore/LambdaContext.swift +++ b/Sources/AWSLambdaRuntimeCore/LambdaContext.swift @@ -202,6 +202,13 @@ public struct LambdaContext: CustomDebugStringConvertible, Sendable { self.storage.tasks } + + /// Registers a background task that continues running after the synchronous invocation has completed. + /// This is useful for tasks like flushing metrics or performing clean-up operations without delaying the response. + /// + /// - Parameter body: An asynchronous closure that performs the background task. + /// - Warning: You will be billed for the milliseconds of Lambda execution time until the very last + /// background task is finished. public func detachedBackgroundTask(_ body: @escaping @Sendable () async -> ()) { Task { await self.tasks.detached(task: body) From c4aff5c0e457d2765800467e819a4419c50e27ce Mon Sep 17 00:00:00 2001 From: Alessio Buratti <9006089+Buratti@users.noreply.github.com> Date: Thu, 15 Aug 2024 03:32:21 +0200 Subject: [PATCH 12/13] Add example usage to readme --- readme.md | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/readme.md b/readme.md index 8a4111e6..a21a55eb 100644 --- a/readme.md +++ b/readme.md @@ -474,7 +474,7 @@ public protocol SimpleLambdaHandler { ### Context -When calling the user provided Lambda function, the library provides a `LambdaContext` class that provides metadata about the execution context, as well as utilities for logging and allocating buffers. +When calling the user provided Lambda function, the library provides a `LambdaContext` class that provides metadata about the execution context, as well as utilities for logging, allocating buffers and dispatch background tasks. ```swift public struct LambdaContext: CustomDebugStringConvertible, Sendable { @@ -555,6 +555,25 @@ public struct LambdaInitializationContext: Sendable { } ``` +### Background tasks + +The detachedBackgroundTask method allows you to register background tasks that continue running even after the Lambda runtime has reported the result of a synchronous invocation. This is particularly useful for integrations with services like API Gateway or CloudFront, where you can quickly return a response without waiting for non-essential tasks such as flushing metrics or performing non-critical clean-up operations. + +```swift + @main + struct MyLambda: SimpleLambdaHandler { + func handle(_ request: APIGatewayV2Request, context: LambdaContext) async throws -> APIGatewayV2Response { + let response = makeResponse() + context.detachedBackgroundTask { + try? await Task.sleep(for: .seconds(3)) + print("Background task completed") + } + print("Returning response") + return response + } + } +``` + ### Configuration The library’s behavior can be fine tuned using environment variables based configuration. The library supported the following environment variables: From 7e647e8130eb5b256e55ddc5edbf3dc3f4ea1949 Mon Sep 17 00:00:00 2001 From: Alessio Buratti <9006089+Buratti@users.noreply.github.com> Date: Thu, 15 Aug 2024 03:45:05 +0200 Subject: [PATCH 13/13] Add tests --- .../DetachedTasksTests.swift | 80 +++++++++++++++++++ 1 file changed, 80 insertions(+) create mode 100644 Tests/AWSLambdaRuntimeCoreTests/DetachedTasksTests.swift diff --git a/Tests/AWSLambdaRuntimeCoreTests/DetachedTasksTests.swift b/Tests/AWSLambdaRuntimeCoreTests/DetachedTasksTests.swift new file mode 100644 index 00000000..b0ec42bc --- /dev/null +++ b/Tests/AWSLambdaRuntimeCoreTests/DetachedTasksTests.swift @@ -0,0 +1,80 @@ +//===----------------------------------------------------------------------===// +// +// This source file is part of the SwiftAWSLambdaRuntime open source project +// +// Copyright (c) 2017-2018 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 +// +//===----------------------------------------------------------------------===// + +@testable import AWSLambdaRuntimeCore +import NIO +import XCTest +import Logging + +class DetachedTasksTest: XCTestCase { + + actor Expectation { + var isFulfilled = false + func fulfill() { + isFulfilled = true + } + } + + func testAwaitTasks() async throws { + let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) + defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } + + let context = DetachedTasksContainer.Context( + eventLoop: eventLoopGroup.next(), + logger: Logger(label: "test") + ) + let expectation = Expectation() + + let container = DetachedTasksContainer(context: context) + await container.detached { + try! await Task.sleep(for: .milliseconds(200)) + await expectation.fulfill() + } + + try await container.awaitAll().get() + let isFulfilled = await expectation.isFulfilled + XCTAssert(isFulfilled) + } + + func testAwaitChildrenTasks() async throws { + let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1) + defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) } + + let context = DetachedTasksContainer.Context( + eventLoop: eventLoopGroup.next(), + logger: Logger(label: "test") + ) + let expectation1 = Expectation() + let expectation2 = Expectation() + + let container = DetachedTasksContainer(context: context) + await container.detached { + await container.detached { + try! await Task.sleep(for: .milliseconds(300)) + await expectation1.fulfill() + } + try! await Task.sleep(for: .milliseconds(200)) + await container.detached { + try! await Task.sleep(for: .milliseconds(100)) + await expectation2.fulfill() + } + } + + try await container.awaitAll().get() + let isFulfilled1 = await expectation1.isFulfilled + let isFulfilled2 = await expectation2.isFulfilled + XCTAssert(isFulfilled1) + XCTAssert(isFulfilled2) + } +}