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/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/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) + } +} 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