improve test with a timeout

Author: sebsto

Sources/AWSLambdaRuntime/LambdaRuntimeClient.swift

@@ -162,7 +162,9 @@ final actor LambdaRuntimeClient: LambdaRuntimeClientProtocol {
     @usableFromInline
     func nextInvocation() async throws -> (Invocation, Writer) {
 
-        try await withTaskCancellationHandler {
+        try Task.checkCancellation()
+
+        return try await withTaskCancellationHandler {
             switch self.lambdaState {
             case .idle:
                 self.lambdaState = .waitingForNextInvocation

Tests/AWSLambdaRuntimeTests/LambdaRuntimeClientTests.swift

@@ -295,22 +295,30 @@ struct LambdaRuntimeClientTests {
                 logger: self.logger
             ) { runtimeClient in
                 do {
-                    let (_, writer) = try await runtimeClient.nextInvocation()
-                    try await writer.writeAndFinish(ByteBuffer(string: "hello"))
 
-                    // continue to simulate traffic until the server reports it has closed the connection
-                    for _ in 1...1000 {
-                        let (_, writer2) = try await runtimeClient.nextInvocation()
-                        try await writer2.writeAndFinish(ByteBuffer(string: "hello"))
+                    // simulate traffic until the server reports it has closed the connection
+                    // or a timeout, whichever comes first
+                    // result is ignored here, either there is a connection error or a timeout
+                    let _ = try await timeout(deadline: .seconds(1)) {
+                        while true {
+                            let (_, writer) = try await runtimeClient.nextInvocation()
+                            try await writer.writeAndFinish(ByteBuffer(string: "hello"))
+                        }
                     }
+                    // result is ignored here, we should never reach this line
+                    Issue.record("Connection reset test did not throw an error")
 
-                    Issue.record("Expected connection error but got successful invocation")
-
+                } catch is CancellationError {
+                    print("++CancellationError")
+                    Issue.record("Runtime client did not send connection closed error")
                 } catch let error as LambdaRuntimeError {
+                    print("++LambdaRuntimeError")
                     #expect(error.code == .connectionToControlPlaneLost)
                 } catch let error as ChannelError {
+                    print("++ChannelError")
                     #expect(error == .ioOnClosedChannel)
                 } catch let error as IOError {
+                    print("++IOError")
                     #expect(error.errnoCode == ECONNRESET || error.errnoCode == EPIPE)
                 } catch {
                     Issue.record("Unexpected error type: \(error)")

Tests/AWSLambdaRuntimeTests/Timeout.swift

@@ -0,0 +1,67 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftAWSLambdaRuntime open source project
+//
+// Copyright (c) 2025 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
+//
+//===----------------------------------------------------------------------===//
+
+// as suggested by https://github.com/vapor/postgres-nio/issues/489#issuecomment-2186509773
+func timeout<Success: Sendable>(
+    deadline: Duration,
+    _ closure: @escaping @Sendable () async throws -> Success
+) async throws -> Success {
+
+    let clock = ContinuousClock()
+
+    let result = await withTaskGroup(of: TimeoutResult<Success>.self, returning: Result<Success, any Error>.self) {
+        taskGroup in
+        taskGroup.addTask {
+            do {
+                try await clock.sleep(until: clock.now + deadline, tolerance: nil)
+                return .deadlineHit
+            } catch {
+                return .deadlineCancelled
+            }
+        }
+
+        taskGroup.addTask {
+            do {
+                let success = try await closure()
+                return .workFinished(.success(success))
+            } catch let error {
+                return .workFinished(.failure(error))
+            }
+        }
+
+        var r: Swift.Result<Success, any Error>?
+        while let taskResult = await taskGroup.next() {
+            switch taskResult {
+            case .deadlineCancelled:
+                continue  // loop
+
+            case .deadlineHit:
+                taskGroup.cancelAll()
+
+            case .workFinished(let result):
+                taskGroup.cancelAll()
+                r = result
+            }
+        }
+        return r!
+    }
+
+    return try result.get()
+}
+
+enum TimeoutResult<Success: Sendable> {
+    case deadlineHit
+    case deadlineCancelled
+    case workFinished(Result<Success, any Error>)
+}