4 Commits
Author SHA1 Message Date
Manoj MahapatraandSébastien Stormacq eccd045d80 fix: Local Server should pass HTTP headers down to the Lambda Runtime (#643)
<!--- Provide a general summary of your changes in the Title above -->

## Issue \#
<!--- If it fixes an issue, please link to the issue here -->
https://github.com/awslabs/swift-aws-lambda-runtime/issues/607

## Description of changes
<!--- Why is this change required? What problem does it solve? -->
The local HTTP server was not forwarding user‑provided headers to the
runtime’s response. It passes all headers through to the runtime. This
it makes local behavior match the Lambda runtime API contract and allows
developers to opt into metadata by sending the appropriate runtime
headers.

## New/existing dependencies impact assessment, if applicable
<!--- No new dependencies were added to this change. -->
<!--- If any dependency was added / modified / removed,
THIRD-PARTY-LICENSES must be updated accordingly. -->
N/A

## Conventional Commits
<!--- Please use conventional commits to let us know what kind of change
this is.-->
<!--- More info can be found here:
https://www.conventionalcommits.org/en/v1.0.0/-->

By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 license.

---------

Co-authored-by: Sébastien Stormacq <sebastien.stormacq@gmail.com>
2026-02-18 15:15:04 +01:00
Sébastien StormacqandSebastien Stormacq 4815273dc3 Fix race condition in Lambda+LocalServer causing NIOAsyncWriter fatal error (Bug #635) (#636)
On fast machines, the local Lambda server crashes with:
```
Fatal error: Deinited NIOAsyncWriter without calling finish()
```

This occurs in `NIOAsyncChannelHandler.channelActive()` when child
connection channels are created.

## Root Cause

This is a known issue with NIO's async server channel API (see
[swift-nio#2637](https://github.com/apple/swift-nio/issues/2637)).

**The fundamental problem:**

1. The async `bind()` API creates `NIOAsyncChannel` instances for
incoming connections
2. These channels are yielded through an async stream to the server loop
3. When the serving task is cancelled (or completes), the async stream
iteration stops
4. Any channels that were accepted but not yet read from the stream are
dropped
5. These unread channels never have `executeThenClose()` called on them
6. Their `NIOAsyncWriter` is deallocated without `finish()` being called
→ fatal error

**Why graceful shutdown doesn't help:**

Even closing the server channel gracefully doesn't eliminate the race -
there's a timing window where:
- A connection is accepted and queued in the async stream
- The server task is cancelled or completes
- The queued channel is never read and gets dropped

IMHO, this is an inherent limitation of the `async bind()` API when
combined with task cancellation.

## Solution

I stopped using the `async bind()` API entirely. Instead, I use the
traditional callback-based `childChannelInitializer`:

1. Create `NIOAsyncChannel` directly in `childChannelInitializer`
(synchronous context)
2. Immediately spawn a `Task.detached` to handle the connection
3. Each connection is handled independently, not through a cancellable
async stream
4. Detached tasks are not affected by task group cancellation
5. Every channel has `executeThenClose()` called immediately, preventing
the writer from being dropped

This approach avoids the async stream entirely, eliminating the race
condition.

## Changes

- Replaced `async bind()` with traditional `childChannelInitializer`
- Each connection spawns a `Task.detached` that immediately calls
`executeThenClose()`
- Removed the connection iteration loop (no longer needed)
- Server task now simply waits for the channel to close
- Simplified shutdown logic since there's no async stream to drain

## Trade-offs

- Uses `Task.detached` (unstructured concurrency) to bridge NIO's
event-loop world with Swift concurrency
- This is necessary until NIO provides a new bootstrap API that properly
handles cancellation
- Each connection is handled independently rather than through
structured concurrency

## Testing

Tested on fast machines where the race condition was reliably
reproducible. The crash no longer occurs.

## References

- [swift-nio#2637](https://github.com/apple/swift-nio/issues/2637) -
Known issue with async server channels and cancellation
- [Comment from NIO
maintainer](https://github.com/apple/swift-nio/issues/2637#issuecomment-1921317577)
- Recommends avoiding cancellation or using callback-based API

Fixes #635

---------

Co-authored-by: Sebastien Stormacq <stormacq@amazon.lu>
2026-01-27 09:10:58 +00:00
Sébastien StormacqandSebastien Stormacq 34e89b4027 Fix Test hangs in Lambda+LocalServer (#630) (#631)
# Fix test hangs caused by Pool cancellation race conditions

## Summary

This PR fixes two related race conditions in
`Lambda+LocalServer+Pool.swift` that were causing the test suite to hang
approximately 10% of the time.

## Problem

The test suite exhibited intermittent hangs (~10% frequency) due to two
bugs in the Pool implementation:

1. **Individual task cancellation bug**: When one task waiting for a
specific `requestId` was cancelled, the cancellation handler would
incorrectly cancel ALL waiting tasks instead of just the cancelled one.

2. **Server shutdown hang**: When the server shut down, waiting
continuations in the pools were never cancelled, causing handlers to
wait indefinitely for responses that would never arrive.

## Root Causes

### Root Cause #1: Cancellation Handler Removes ALL Continuations

The `onCancel` handler in `Pool._next()` was removing all continuations
from the `waitingForSpecific` dictionary when any single task was
cancelled:

```swift
onCancel: {
    // BUG: Removes ALL continuations, not just the cancelled task's
    for continuation in state.waitingForSpecific.values {
        toCancel.append(continuation)
    }
    state.waitingForSpecific.removeAll()
}
```

This caused unrelated concurrent invocations to fail with
`CancellationError` when one client cancelled their request.

### Root Cause #2: No Pool Cleanup During Server Shutdown

When the server shut down (e.g., test completes), the task group was
cancelled but the pools' waiting continuations were never notified. The
`/invoke` endpoint handlers would continue waiting for responses that
would never arrive because the Lambda function had stopped.

## Solution

### Fix #1: Only Remove Specific Continuation on Cancellation

Modified the cancellation handler to only remove the continuation for
the specific cancelled task:

```swift
onCancel: {
    // Only remove THIS task's continuation
    let continuationToCancel = self.lock.withLock { state -> CheckedContinuation<T, any Error>? in
        if let requestId = requestId {
            return state.waitingForSpecific.removeValue(forKey: requestId)
        } else {
            let cont = state.waitingForAny
            state.waitingForAny = nil
            return cont
        }
    }
    
    continuationToCancel?.resume(throwing: CancellationError())
}
```

### Fix #2: Add Pool Cleanup During Server Shutdown

Added `cancelAll()` method to the Pool class and call it during server
shutdown:

```swift
func cancelAll() {
    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
    }
    
    for continuation in continuationsToCancel {
        continuation.resume(throwing: CancellationError())
    }
}
```

Called during server shutdown:

```swift
let serverOrHandlerResult1 = await group.next()!
group.cancelAll()

// Cancel all waiting continuations in the pools to prevent hangs
server.invocationPool.cancelAll()
server.responsePool.cancelAll()
```

## Changes

### Modified Files

- **Sources/AWSLambdaRuntime/HTTPServer/Lambda+LocalServer+Pool.swift**
- Fixed cancellation handler in `_next()` to only remove specific
continuation
  - Added `cancelAll()` method for server shutdown cleanup

- **Sources/AWSLambdaRuntime/HTTPServer/Lambda+LocalServer.swift**
  - Call `cancelAll()` on both pools during server shutdown

### New Files

- **Tests/AWSLambdaRuntimeTests/LocalServerPoolCancellationTests.swift**
  - Added comprehensive test suite with 3 tests
- `testCancellationOnlyAffectsOwnTask`: Verifies only the cancelled task
receives CancellationError
- `testConcurrentInvocationsWithCancellation`: Tests real-world scenario
with 5 concurrent invocations
- `testFIFOModeCancellation`: Ensures FIFO mode cancellation works
correctly

## Testing

### Before Fix
- Test suite hung ~10% of the time
- When 1 task was cancelled, all 5 concurrent tasks received
`CancellationError`
- Streaming tests would occasionally hang during shutdown

### After Fix
- All 91 tests pass consistently without hangs
- When 1 task is cancelled, only that specific task receives
`CancellationError`
- Other tasks continue waiting normally
- Server shutdown properly cleans up all waiting continuations
- Multiple consecutive test runs confirm stability

### Test Coverage

The new test suite reproduces both bugs and verifies the fixes:

1. **testCancellationOnlyAffectsOwnTask**: Creates 3 tasks waiting for
different requestIds, cancels only one, and verifies the others are not
affected
2. **testConcurrentInvocationsWithCancellation**: Simulates 5 concurrent
invocations with one cancellation
3. **testFIFOModeCancellation**: Tests FIFO mode to ensure it still
works correctly

---------

Co-authored-by: Sebastien Stormacq <stormacq@amazon.lu>
2026-01-15 20:58:03 +01:00
Sébastien Stormacq e0f064a93e Refactor project directories (#621)
This PR refactors the project's directories.
As the number of source files grows, I created subdirectories to
separate the runtime itself, from its HTTP Client (`RuntimeClient`) and
local HTTP Server (`Lambda+LocalServer`).

The new layout looks like this:

```text
Sources
├── AWSLambdaRuntime
│   ├── FoundationSupport
│   │   ├── Context+Foundation.swift
│   │   ├── Lambda+JSON.swift
│   │   └── Vendored
│   │       ├── ByteBuffer-foundation.swift
│   │       └── JSON+ByteBuffer.swift
│   ├── HTTPClient
│   │   ├── ControlPlaneRequest.swift
│   │   ├── ControlPlaneRequestEncoder.swift
│   │   ├── LambdaRuntimeClient+ChannelHandler.swift
│   │   ├── LambdaRuntimeClient.swift
│   │   └── LambdaRuntimeClientProtocol.swift
│   ├── HTTPServer
│   │   ├── Lambda+LocalServer+Pool.swift
│   │   └── Lambda+LocalServer.swift
│   ├── Lambda.swift
│   ├── LambdaClock.swift
│   ├── LambdaContext.swift
│   ├── LambdaRequestID.swift
│   ├── LambdaResponseStreamWriter+Headers.swift
│   ├── LambdaRuntimeError.swift
│   ├── Runtime
│   │   ├── LambdaHandlers.swift
│   │   ├── LambdaRuntime+Codable.swift
│   │   ├── LambdaRuntime+Handler.swift
│   │   ├── LambdaRuntime+ServiceLifecycle.swift
│   │   └── LambdaRuntime.swift
│   ├── SendableMetatype.swift
│   ├── Utils.swift
│   └── Version.swift
└── MockServer
    └── MockHTTPServer.swift
```
2026-01-01 11:19:13 -05:00