-
Notifications
You must be signed in to change notification settings - Fork 118
Accept multiple POST /invoke
requests to allow parallel testing
#585
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
sebsto
wants to merge
36
commits into
awslabs:main
Choose a base branch
from
sebsto:sebsto/fix_584
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+734
−118
Open
Changes from 31 commits
Commits
Show all changes
36 commits
Select commit
Hold shift + click to select a range
1de4b19
Gently decline subsequent POST /invoke request while the Lambda handl…
f35685c
Update Sources/AWSLambdaRuntime/Lambda+LocalServer.swift
sebsto 44068af
Merge branch 'main' into sebsto/fix_584
af43436
Merge branch 'sebsto/fix_584' of github.com:sebsto/swift-aws-lambda-r…
8aa9a46
fix typo and language in comments
6d02848
fix test
0566ef5
remove async constraint on `push()`
b341bb4
swift-format
d73dce6
remove the fatal error to make testinge easier
0cd73da
fix comment
5cc8ed9
Merge branch 'main' into sebsto/fix_584
b3576fe
use Result instead of Tupe
ae18d7b
swift format
605e86c
onCancel: resume continuation outside of the lock
2d6842a
swift-format
009b5c6
add logging trace details
0ed15f9
fix parallel invocation for non streaming lambda functions
a2ce3be
move pool to a separate file
5529cae
add unit test for the pool
9abbe8e
remove forced trace level for logging
6311b46
add license header
18e09c6
skip some tests on Swift 6.0
c07555a
swift-format
bff7507
replace a do {} catch {} with #expect(throwing:)
86591a7
swift-format
df989af
Generate time independent requestID + add a test for rapid fire
585a414
fix test on Linux
8845467
Also check that the test responses are in the correct order
6138790
simplify the test #expect
669ad43
use String(decoding:) instead of String(data:)
9c85560
State : use a struct instead of enum
54744c5
Update Sources/AWSLambdaRuntime/Lambda+LocalServer.swift
sebsto 503f2da
Update Sources/AWSLambdaRuntime/Lambda+LocalServer.swift
sebsto 30f19ad
Merge branch 'main' into sebsto/fix_584
b662024
Merge branch 'main' into sebsto/fix_584
ab77e5a
return an HTTP 5xx error in case we mess up something in the Pool
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,192 @@ | ||
//===----------------------------------------------------------------------===// | ||
// | ||
// 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 | ||
// | ||
//===----------------------------------------------------------------------===// | ||
|
||
#if LocalServerSupport | ||
import DequeModule | ||
import Synchronization | ||
|
||
@available(LambdaSwift 2.0, *) | ||
extension LambdaHTTPServer { | ||
|
||
/// A shared data structure to store the current invocation or response requests and the continuation objects. | ||
/// This data structure is shared between instances of the HTTPHandler | ||
/// (one instance to serve requests from the Lambda function and one instance to serve requests from the client invoking the lambda function). | ||
internal final class Pool<T>: AsyncSequence, AsyncIteratorProtocol, Sendable where T: Sendable { | ||
private let poolName: String | ||
internal init(name: String = "Pool") { self.poolName = name } | ||
|
||
typealias Element = T | ||
|
||
struct State { | ||
var buffer: Deque<T> = [] | ||
var waitingForAny: CheckedContinuation<T, any Error>? | ||
var waitingForSpecific: [String: CheckedContinuation<T, any Error>] = [:] | ||
} | ||
|
||
private let lock = Mutex<State>(State()) | ||
|
||
/// enqueue an element, or give it back immediately to the iterator if it is waiting for an element | ||
public func push(_ item: T) { | ||
let continuationToResume = self.lock.withLock { state -> CheckedContinuation<T, any Error>? in | ||
// First check if there's a waiting continuation that can handle this item | ||
|
||
// Check for FIFO waiter first | ||
if let continuation = state.waitingForAny { | ||
state.waitingForAny = nil | ||
return continuation | ||
} | ||
|
||
// Check for specific waiter | ||
if let response = item as? LocalServerResponse, | ||
let requestId = response.requestId, | ||
let continuation = state.waitingForSpecific.removeValue(forKey: requestId) | ||
{ | ||
return continuation | ||
} | ||
|
||
// No waiting continuation, add to buffer | ||
state.buffer.append(item) | ||
return nil | ||
} | ||
|
||
// Resume continuation outside the lock to prevent potential deadlocks | ||
continuationToResume?.resume(returning: item) | ||
} | ||
|
||
/// Unified next() method that handles both FIFO and requestId-specific waiting | ||
private func _next(for requestId: String?) async throws -> T { | ||
// exit if the task is cancelled | ||
guard !Task.isCancelled else { | ||
throw CancellationError() | ||
} | ||
|
||
return try await withTaskCancellationHandler { | ||
try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<T, any Error>) in | ||
let nextAction: Result<T, PoolError>? = self.lock.withLock { state -> Result<T, PoolError>? in | ||
if let requestId = requestId { | ||
// Look for oldest (first) item for this requestId in buffer | ||
if let index = state.buffer.firstIndex(where: { item in | ||
if let response = item as? LocalServerResponse { | ||
return response.requestId == requestId | ||
} | ||
return false | ||
}) { | ||
let item = state.buffer.remove(at: index) | ||
return .success(item) | ||
} else { | ||
// Check for conflicting waiters | ||
if state.waitingForAny != nil { | ||
return .failure(PoolError(cause: .mixedWaitingModes)) | ||
} | ||
if state.waitingForSpecific[requestId] != nil { | ||
return .failure(PoolError(cause: .duplicateRequestIdWait(requestId))) | ||
} | ||
|
||
// No matching item, wait for it | ||
state.waitingForSpecific[requestId] = continuation | ||
return nil | ||
} | ||
} else { | ||
// FIFO mode - take first item | ||
if let first = state.buffer.popFirst() { | ||
return .success(first) | ||
} else { | ||
// Check for conflicting waiters | ||
if !state.waitingForSpecific.isEmpty { | ||
return .failure(PoolError(cause: .mixedWaitingModes)) | ||
} | ||
if state.waitingForAny != nil { | ||
return .failure(PoolError(cause: .nextCalledTwice(state.waitingForAny!))) | ||
} | ||
|
||
state.waitingForAny = continuation | ||
return nil | ||
} | ||
} | ||
} | ||
|
||
switch nextAction { | ||
case .success(let item): | ||
continuation.resume(returning: item) | ||
case .failure(let error): | ||
if case let .nextCalledTwice(prevContinuation) = error.cause { | ||
prevContinuation.resume(throwing: error) | ||
} | ||
continuation.resume(throwing: error) | ||
case .none: | ||
// do nothing - continuation is stored in state | ||
break | ||
} | ||
} | ||
} onCancel: { | ||
// Ensure we properly handle cancellation by removing stored continuation | ||
let continuationsToCancel = self.lock.withLock { state -> [CheckedContinuation<T, any Error>] in | ||
var toCancel: [CheckedContinuation<T, any Error>] = [] | ||
|
||
if let continuation = state.waitingForAny { | ||
toCancel.append(continuation) | ||
state.waitingForAny = nil | ||
} | ||
|
||
for continuation in state.waitingForSpecific.values { | ||
toCancel.append(continuation) | ||
} | ||
state.waitingForSpecific.removeAll() | ||
|
||
return toCancel | ||
} | ||
|
||
// Resume all continuations outside the lock to avoid potential deadlocks | ||
for continuation in continuationsToCancel { | ||
continuation.resume(throwing: CancellationError()) | ||
sebsto marked this conversation as resolved.
Show resolved
Hide resolved
|
||
} | ||
} | ||
} | ||
|
||
/// Simple FIFO next() method - used by AsyncIteratorProtocol | ||
func next() async throws -> T? { | ||
try await _next(for: nil) | ||
} | ||
|
||
/// RequestId-specific next() method for LocalServerResponse - NOT part of AsyncIteratorProtocol | ||
func next(for requestId: String) async throws -> T { | ||
try await _next(for: requestId) | ||
} | ||
|
||
func makeAsyncIterator() -> Pool { | ||
self | ||
} | ||
|
||
struct PoolError: Error { | ||
let cause: Cause | ||
var message: String { | ||
switch self.cause { | ||
case .nextCalledTwice: | ||
return "Concurrent invocations to next(). This is not allowed." | ||
case .duplicateRequestIdWait(let requestId): | ||
return "Already waiting for requestId: \(requestId)" | ||
case .mixedWaitingModes: | ||
return "Cannot mix FIFO waiting (next()) with specific waiting (next(for:))" | ||
} | ||
} | ||
|
||
enum Cause { | ||
case nextCalledTwice(CheckedContinuation<T, any Error>) | ||
case duplicateRequestIdWait(String) | ||
case mixedWaitingModes | ||
} | ||
} | ||
} | ||
} | ||
#endif |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.