Hand WebSocket frames and messages off without buffering

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<Element: Sendable> {
+
+    private typealias Sender = (element: Element, continuation: IdentifiableContinuation<Void, any Error>)
+    private typealias Receiver = IdentifiableContinuation<Element?, any Error>
+
+    private var senders = [Sender]()
+    private var receivers = [Receiver]()
+    private var isFinished = false
+    private var result: Result<Element?, any Error> = .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<Element?, any Error> = .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<Void, any Error>.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)
+    }
+}

FlyingFox/Sources/WebSocket/AsyncStream+WSFrame.swift

@@ -67,8 +67,14 @@ extension AsyncStream<WSFrame> {
     typealias SendableMetatype = Any
 #endif
 
-    static func protocolFrames<S: AsyncSequence & SendableMetatype>(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<S: AsyncSequence & SendableMetatype>(
+        from frames: S,
+        onTermination: (@Sendable () -> Void)? = nil
+    ) -> Self where S.Element == WSFrame {
+        let iterator = Iterator(from: frames, onTermination: onTermination)
         return AsyncStream<WSFrame> {
             do {
                 return try await iterator.next()
@@ -84,9 +90,15 @@ extension AsyncStream<WSFrame> {
 
     private final class Iterator<S: AsyncSequence>: @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? {

FlyingFox/Sources/WebSocket/WSHandler.swift

@@ -74,66 +74,71 @@ public struct MessageFrameWSHandler: WSHandler {
     public func makeFrames(for client: AsyncThrowingStream<WSFrame, any Error>) async throws -> AsyncStream<WSFrame> {
         let framesIn = WSFrameValidator.validateFrames(from: client)
 
-        var messagesIn: AsyncStream<WSMessage>.Continuation!
-        let messages = AsyncStream<WSMessage> {
-            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<WSMessage>()
+        let framesOut = AsyncHandoff<WSFrame>()
+
+        // 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<WSFrame, any Error> { [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<WSFrame, any Error> { try await framesOut.next() }
+        return AsyncStream.protocolFrames(from: serverFrames, onTermination: { task.cancel() })
     }
 
     func start<S: AsyncSequence & Sendable>(
         framesIn: S,
-        framesOut: AsyncThrowingStream<WSFrame, any Error>.Continuation,
-        messagesIn: AsyncStream<WSMessage>.Continuation,
+        framesOut: AsyncHandoff<WSFrame>,
+        messagesIn: AsyncHandoff<WSMessage>,
         messagesOut: AsyncStream<WSMessage>
     ) 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? {

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<Int>()
+        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<String>()
+        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<String>()
+        await handoff.finish()
+
+        #expect(try await handoff.next() == nil)
+    }
+
+    @Test
+    func finish_DeliversFinalElement_ThenEnds() async throws {
+        let handoff = AsyncHandoff<String>()
+        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<String>()
+        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<String>()
+        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<String>()
+        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<String>()
+        let receiver = Task { try await handoff.next() }
+        receiver.cancel()
+
+        #expect(try await receiver.value == nil)
+    }
+}