From 2d5dbf787360655c836201bd691dfa204155fa74 Mon Sep 17 00:00:00 2001 From: Luke Howard Date: Mon, 12 Oct 2026 14:08:51 +1100 Subject: [PATCH 1/2] Add AsyncHandoff, an unbuffered handoff between two tasks An actor that passes elements from one task to another one at a time. send(_:) suspends until next() takes the element, so a sender gets no further ahead than its receiver and nothing accumulates between them. finish(with:) ends the handoff with a final element or error for the receiver and cancels any sender still waiting; cancelling a waiting sender or receiver removes its continuation. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01T8cahEd4GfsHQ89sNrmP4p --- .../Sources/WebSocket/AsyncHandoff.swift | 116 ++++++++++++++++ .../Tests/WebSocket/AsyncHandoffTests.swift | 131 ++++++++++++++++++ 2 files changed, 247 insertions(+) create mode 100644 FlyingFox/Sources/WebSocket/AsyncHandoff.swift create mode 100644 FlyingFox/Tests/WebSocket/AsyncHandoffTests.swift diff --git a/FlyingFox/Sources/WebSocket/AsyncHandoff.swift b/FlyingFox/Sources/WebSocket/AsyncHandoff.swift new file mode 100644 index 0000000..85f97f7 --- /dev/null +++ b/FlyingFox/Sources/WebSocket/AsyncHandoff.swift @@ -0,0 +1,116 @@ +// +// AsyncHandoff.swift +// FlyingFox +// +// Created by Luke Howard on 12/10/2026. +// Copyright © 2026 Simon Whitty. All rights reserved. +// +// Distributed under the permissive MIT license +// Get the latest version from here: +// +// https://github.com/swhitty/FlyingFox +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// + +import FlyingSocks + +/// Passes elements from one task to another, one at a time. `send(_:)` suspends +/// until `next()` takes the element, so a sender gets no further ahead than its +/// receiver and nothing is buffered between them. +actor AsyncHandoff { + + private typealias Sender = (element: Element, continuation: IdentifiableContinuation) + private typealias Receiver = IdentifiableContinuation + + private var senders = [Sender]() + private var receivers = [Receiver]() + private var isFinished = false + private var result: Result = .success(nil) + + /// Hands `element` to the receiver, suspending until it is taken. + /// Throws `CancellationError` if cancelled or finished before then. + func send(_ element: Element) async throws { + guard !isFinished else { + throw CancellationError() + } + if !receivers.isEmpty { + receivers.removeFirst().resume(returning: element) + return + } + try await withIdentifiableThrowingContinuation(isolation: self) { + senders.append((element, $0)) + } onCancel: { id in + Task { await self.cancelSender(id: id) } + } + } + + /// Takes the next element, suspending until one is sent. Once finished, + /// returns the final element or error, then nil. Returns nil if cancelled. + func next() async throws -> Element? { + if !senders.isEmpty { + let sender = senders.removeFirst() + sender.continuation.resume() + return sender.element + } + if isFinished { + defer { result = .success(nil) } + return try result.get() + } + return try await withIdentifiableThrowingContinuation(isolation: self) { + receivers.append($0) + } onCancel: { id in + Task { await self.cancelReceiver(id: id) } + } + } + + /// Ends the handoff. The receiver gets `result`, a last element or an error, + /// before nil; senders still waiting are cancelled. + func finish(with result: Result = .success(nil)) { + guard !isFinished else { return } + isFinished = true + self.result = result + for sender in senders { + sender.continuation.resume(throwing: CancellationError()) + } + senders = [] + if !receivers.isEmpty { + receivers.removeFirst().resume(with: result) + self.result = .success(nil) + } + for receiver in receivers { + receiver.resume(returning: nil) + } + receivers = [] + } + + func finish(throwing error: any Error) { + finish(with: .failure(error)) + } + + private func cancelSender(id: IdentifiableContinuation.ID) { + guard let index = senders.firstIndex(where: { $0.continuation.id == id }) else { return } + senders.remove(at: index).continuation.resume(throwing: CancellationError()) + } + + private func cancelReceiver(id: Receiver.ID) { + guard let index = receivers.firstIndex(where: { $0.id == id }) else { return } + receivers.remove(at: index).resume(returning: nil) + } +} diff --git a/FlyingFox/Tests/WebSocket/AsyncHandoffTests.swift b/FlyingFox/Tests/WebSocket/AsyncHandoffTests.swift new file mode 100644 index 0000000..13be583 --- /dev/null +++ b/FlyingFox/Tests/WebSocket/AsyncHandoffTests.swift @@ -0,0 +1,131 @@ +// +// AsyncHandoffTests.swift +// FlyingFox +// +// Created by Luke Howard on 12/10/2026. +// Copyright © 2026 Simon Whitty. All rights reserved. +// +// Distributed under the permissive MIT license +// Get the latest version from here: +// +// https://github.com/swhitty/FlyingFox +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// + +@testable import FlyingFox +@testable import FlyingSocks +import Foundation +import Testing + +struct AsyncHandoffTests { + + @Test + func send_Suspends_UntilElementIsTaken() async throws { + let handoff = AsyncHandoff() + let sent = Mutex([Int]()) + let sender = Task { + for element in 1...3 { + try await handoff.send(element) + sent.withLock { $0.append(element) } + } + } + + #expect(try await handoff.next() == 1) + // The second send cannot complete until the receiver takes it. + #expect(sent.withLock { !$0.contains(2) }) + #expect(try await handoff.next() == 2) + #expect(try await handoff.next() == 3) + try await sender.value + #expect(sent.withLock { $0 } == [1, 2, 3]) + } + + @Test + func next_Suspends_UntilElementIsSent() async throws { + let handoff = AsyncHandoff() + let receiver = Task { try await handoff.next() } + + try await handoff.send("Fish") + #expect(try await receiver.value == "Fish") + } + + @Test + func finish_EndsReceiver() async throws { + let handoff = AsyncHandoff() + await handoff.finish() + + #expect(try await handoff.next() == nil) + } + + @Test + func finish_DeliversFinalElement_ThenEnds() async throws { + let handoff = AsyncHandoff() + await handoff.finish(with: .success("Chips")) + + #expect(try await handoff.next() == "Chips") + #expect(try await handoff.next() == nil) + } + + @Test + func finish_DeliversError_ThenEnds() async throws { + let handoff = AsyncHandoff() + let receiver = Task { try await handoff.next() } + await handoff.finish(throwing: SocketError.disconnected) + + await #expect(throws: SocketError.disconnected) { + try await receiver.value + } + #expect(try await handoff.next() == nil) + } + + @Test + func finish_CancelsWaitingSenders() async throws { + let handoff = AsyncHandoff() + let sender = Task { try await handoff.send("Fish") } + await handoff.finish() + + await #expect(throws: CancellationError.self) { + try await sender.value + } + #expect(try await handoff.next() == nil) + } + + @Test + func send_ThrowsCancellationError_WhenSenderIsCancelled() async throws { + let handoff = AsyncHandoff() + let sender = Task { try await handoff.send("Fish") } + sender.cancel() + + await #expect(throws: CancellationError.self) { + try await sender.value + } + // The cancelled element is not delivered. + await handoff.finish() + #expect(try await handoff.next() == nil) + } + + @Test + func next_ReturnsNil_WhenReceiverIsCancelled() async throws { + let handoff = AsyncHandoff() + let receiver = Task { try await handoff.next() } + receiver.cancel() + + #expect(try await receiver.value == nil) + } +} From a9404f79cb79bb6de89cf6b24437af81e2d5f4fa Mon Sep 17 00:00:00 2001 From: Luke Howard Date: Mon, 12 Oct 2026 14:08:51 +1100 Subject: [PATCH 2/2] Hand WebSocket frames and messages off without buffering MessageFrameWSHandler moved frames and messages between the connection and the handler through two AsyncStream continuations with unbounded buffers. A client that floods a server whose handler is slow, or that stops reading while the handler keeps sending, grows those buffers until the process runs out of memory; and because yielding never waits, a handler could not tell when a message had reached the connection. Replace the continuations with AsyncHandoffs. A client frame is now read only when the handler has taken the previous message, and a message is taken from the handler's stream only when the connection is about to write the previous frame, so back-pressure reaches the client's socket in one direction and the handler's stream in the other. A handler that produces messages on demand, with AsyncStream(unfolding:) for instance, is paced by the client. The frames stream is now pull-based, so it has no continuation whose termination cancels the handler's tasks. protocolFrames gains an onTermination closure that runs when the stream ends or is released, and the handler cancels its tasks from there. On a client's close frame, the echo is sent and the frames finished before the .close message is handed to the handler as the final message of its stream, so the closing handshake does not wait on the handler. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01T8cahEd4GfsHQ89sNrmP4p --- .../WebSocket/AsyncStream+WSFrame.swift | 18 ++++- FlyingFox/Sources/WebSocket/WSHandler.swift | 67 +++++++++--------- .../WebSocket/AsyncStream+WSFrameTests.swift | 33 ++++++++- .../Tests/WebSocket/WSHandlerTests.swift | 69 ++++++++++++++++++- README.md | 2 + 5 files changed, 153 insertions(+), 36 deletions(-) diff --git a/FlyingFox/Sources/WebSocket/AsyncStream+WSFrame.swift b/FlyingFox/Sources/WebSocket/AsyncStream+WSFrame.swift index 910bc09..443835b 100644 --- a/FlyingFox/Sources/WebSocket/AsyncStream+WSFrame.swift +++ b/FlyingFox/Sources/WebSocket/AsyncStream+WSFrame.swift @@ -67,8 +67,14 @@ extension AsyncStream { typealias SendableMetatype = Any #endif - static func protocolFrames(from frames: S) -> Self where S.Element == WSFrame { - let iterator = Iterator(from: frames) + /// Converts an error from `frames` into a close frame that ends the stream. + /// `onTermination` runs when the stream ends, or when it is released without + /// being read to the end. + static func protocolFrames( + from frames: S, + onTermination: (@Sendable () -> Void)? = nil + ) -> Self where S.Element == WSFrame { + let iterator = Iterator(from: frames, onTermination: onTermination) return AsyncStream { do { return try await iterator.next() @@ -84,9 +90,15 @@ extension AsyncStream { private final class Iterator: @unchecked Sendable where S.Element == WSFrame { var iterator: S.AsyncIterator? + let onTermination: (@Sendable () -> Void)? - init(from frames: S){ + init(from frames: S, onTermination: (@Sendable () -> Void)?) { self.iterator = frames.makeAsyncIterator() + self.onTermination = onTermination + } + + deinit { + onTermination?() } func next() async throws -> WSFrame? { diff --git a/FlyingFox/Sources/WebSocket/WSHandler.swift b/FlyingFox/Sources/WebSocket/WSHandler.swift index f2f3660..bab09ea 100644 --- a/FlyingFox/Sources/WebSocket/WSHandler.swift +++ b/FlyingFox/Sources/WebSocket/WSHandler.swift @@ -74,66 +74,71 @@ public struct MessageFrameWSHandler: WSHandler { public func makeFrames(for client: AsyncThrowingStream) async throws -> AsyncStream { let framesIn = WSFrameValidator.validateFrames(from: client) - var messagesIn: AsyncStream.Continuation! - let messages = AsyncStream { - messagesIn = $0 + // Neither handoff buffers: a client frame is read only when the handler + // takes the previous message, and a message is taken from the handler + // only when the connection is ready to write the previous frame. + let messagesIn = AsyncHandoff() + let framesOut = AsyncHandoff() + + // messagesIn is never finished with an error. + let messagesOut = try await handler.makeMessages(for: AsyncStream { try? await messagesIn.next() }) + let task = Task { + await start(framesIn: framesIn, framesOut: framesOut, + messagesIn: messagesIn, messagesOut: messagesOut) } - let messagesOut = try await handler.makeMessages(for: messages) - let serverFrames = AsyncThrowingStream { [messagesIn] continuation in - let task = Task { - await start(framesIn: framesIn, framesOut: continuation, - messagesIn: messagesIn!, messagesOut: messagesOut) - } - continuation.onTermination = { @Sendable _ in task.cancel() } - } - - return AsyncStream.protocolFrames(from: serverFrames) + let serverFrames = AsyncThrowingStream { try await framesOut.next() } + return AsyncStream.protocolFrames(from: serverFrames, onTermination: { task.cancel() }) } func start( framesIn: S, - framesOut: AsyncThrowingStream.Continuation, - messagesIn: AsyncStream.Continuation, + framesOut: AsyncHandoff, + messagesIn: AsyncHandoff, messagesOut: AsyncStream ) async where S.Element == WSFrame { - defer { - messagesIn.finish() - framesOut.finish() - } await withTaskGroup(of: Void.self) { group in group.addTask { do { for try await frame in framesIn { - if let message = try makeMessage(for: frame) { - messagesIn.yield(message) - } else if let frame = try makeResponseFrames(for: frame) { - framesOut.yield(frame) - } if frame.opcode == .close { throw FrameError.closed(frame) + } else if let message = try makeMessage(for: frame) { + try await messagesIn.send(message) + } else if let frame = try makeResponseFrames(for: frame) { + try await framesOut.send(frame) } } } catch FrameError.closed(let frame) { - framesOut.yield(frame) + // Echo the close and end the frames before telling the + // handler, so the closing handshake does not wait on it. + try? await framesOut.send(frame) + await framesOut.finish() + await messagesIn.finish(with: Result { try makeMessage(for: frame) }) } catch is CancellationError { } catch { - framesOut.finish(throwing: error) + await framesOut.finish(throwing: error) } } group.addTask { - for await message in messagesOut { - for frame in makeFrames(for: message) { - framesOut.yield(frame) - if frame.opcode == .close { - return + do { + for await message in messagesOut { + for frame in makeFrames(for: message) { + try await framesOut.send(frame) + if frame.opcode == .close { + return + } } } + } catch { + // send throws only once framesOut has finished. } } await group.next() group.cancelAll() } + await messagesIn.finish() + await framesOut.finish() } func makeMessage(for frame: WSFrame) throws -> WSMessage? { diff --git a/FlyingFox/Tests/WebSocket/AsyncStream+WSFrameTests.swift b/FlyingFox/Tests/WebSocket/AsyncStream+WSFrameTests.swift index 55267c3..c014426 100644 --- a/FlyingFox/Tests/WebSocket/AsyncStream+WSFrameTests.swift +++ b/FlyingFox/Tests/WebSocket/AsyncStream+WSFrameTests.swift @@ -30,7 +30,7 @@ // @testable import FlyingFox -import FlyingSocks +@testable import FlyingSocks import Foundation import Testing @@ -105,6 +105,37 @@ struct WSFrameSequenceTests { ) } + @Test + func protocolFrames_CallOnTermination_WhenReleasedBeforeTheEnd() async throws { + let terminated = Mutex(false) + var frames: AsyncStream? = AsyncStream.protocolFrames(from: AsyncThrowingStream.make([.fish, .chips])) { + terminated.withLock { $0 = true } + } + var iterator = frames?.makeAsyncIterator() + + #expect( + try await iterator?.next() == .fish + ) + #expect(terminated.withLock { !$0 }) + + iterator = nil + frames = nil + #expect(terminated.withLock { $0 }) + } + + @Test + func protocolFrames_CallOnTermination_WhenReadToTheEnd() async throws { + let terminated = Mutex(false) + let frames = AsyncStream.protocolFrames(from: AsyncThrowingStream.make([.fish])) { + terminated.withLock { $0 = true } + } + + #expect( + try await frames.collectAll() == [.fish] + ) + #expect(terminated.withLock { $0 }) + } + @Test func protocolFrames_CatchErrors_AndCloseStream() async throws { let (stream, continuation) = AsyncThrowingStream.makeStream() diff --git a/FlyingFox/Tests/WebSocket/WSHandlerTests.swift b/FlyingFox/Tests/WebSocket/WSHandlerTests.swift index 0582210..d34122c 100644 --- a/FlyingFox/Tests/WebSocket/WSHandlerTests.swift +++ b/FlyingFox/Tests/WebSocket/WSHandlerTests.swift @@ -30,6 +30,7 @@ // @testable import FlyingFox +@testable import FlyingSocks import Foundation import Testing @@ -114,6 +115,7 @@ struct WSHandlerTests { let handler = MessageFrameWSHandler.make(handler: messages) let frames = try await handler.makeFrames(for: [.fish, .ping, .pong, .chips, .close]) + let collected = Task { try await frames.collectAll() } #expect( try await messages.input.takeNext() == .text("Fish") @@ -124,10 +126,64 @@ struct WSHandlerTests { ) #expect( - try await frames.collectAll() == [.pong, .close] + try await collected.value == [.pong, .close] ) } + @Test + func closeFrame_IsEchoed_AndDeliveredAsFinalMessage() async throws { + let messages = Messages() + let handler = MessageFrameWSHandler.make(handler: messages) + + let frames = try await handler.makeFrames(for: [.close]) + + #expect( + try await frames.collectAll() == [.close] + ) + #expect( + try await messages.input.collectAll() == [.close(WSCloseCode(1002, reason: "Bye"))] + ) + } + + @Test + func clientFrames_AreRead_OnlyAsMessagesAreTaken() async throws { + let read = Mutex(0) + let messages = Messages() + let handler = MessageFrameWSHandler.make(handler: messages) + + let frames = try await handler.makeFrames(for: AsyncThrowingStream { + read.withLock { $0 += 1 } + return .fish + }) + + #expect( + try await messages.input.takeNext() == .text("Fish") + ) + // The frame taken has been read, and at most one more waits to be taken. + #expect(read.withLock { $0 } <= 2) + withExtendedLifetime(frames) { } + } + + @Test + func messagesOut_ArePulled_OnlyAsFramesAreTaken() async throws { + let pulled = Mutex(0) + let handler = MessageFrameWSHandler.make(handler: UnfoldingMessages { + pulled.withLock { $0 += 1 } + return .text("Hello") + }) + let (clientFrames, clientContinuation) = AsyncThrowingStream.makeStream() + defer { clientContinuation.finish() } + + let frames = try await handler.makeFrames(for: clientFrames) + + var iterator = frames.makeAsyncIterator() + #expect( + try await iterator.next() == .make(fin: true, opcode: .text, payload: "Hello".data(using: .utf8)!) + ) + // The message taken has been pulled, and at most one more waits to be taken. + #expect(pulled.withLock { $0 } <= 2) + } + @Test func responseFramesEnds() async throws { let handler = MessageFrameWSHandler.make() @@ -186,6 +242,17 @@ extension MessageFrameWSHandler { } } +// Produces each outgoing message when asked for it, as a handler that must +// wait for a message to be taken before producing the next one would. +private struct UnfoldingMessages: WSMessageHandler { + + let next: @Sendable () async -> WSMessage? + + func makeMessages(for client: AsyncStream) async throws -> AsyncStream { + AsyncStream(unfolding: next) + } +} + final class Messages: WSMessageHandler, @unchecked Sendable { var input: AsyncStream! diff --git a/README.md b/README.md index f903414..f3f1aef 100644 --- a/README.md +++ b/README.md @@ -184,6 +184,8 @@ enum WSMessage { } ``` +Nothing is buffered between the streams and the connection: a frame is read from the client only when the handler takes the previous message, and a message is taken from the handler's stream only when the connection is ready to write the previous one. A handler that produces its messages on demand, for example with `AsyncStream(unfolding:)`, is therefore paced by the client. + Raw WebSocket frames can also be [provided](#websockets). ### RoutedHTTPHandler