Remove AsyncHelpers usage

We don't want to use a type-erased sequence wrapper for many reasons. Performance, non-throwing-ness,…

Author: tkrajacic

Package.swift

@@ -30,15 +30,12 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-log.git", .upToNextMajor(from: "1.5.2")),
         // AsyncChannel with backpressure
         .package(url: "https://github.com/apple/swift-async-algorithms", branch: "main"),
-        // Helpers for async/await
-        .package(url: "https://github.com/stairtree/async-helpers.git", branch: "main"),
     ],
     targets: [
         .target(
             name: "StructuredWebSocketClient",
             dependencies: [
                 .product(name: "Logging", package: "swift-log"),
-                .product(name: "AsyncHelpers", package: "async-helpers"),
                 .product(name: "AsyncAlgorithms", package: "swift-async-algorithms"),
             ]),
         .target(

Sources/StructuredWebSocketClient/MessageTransport.swift

@@ -13,7 +13,6 @@
 
 import Foundation
 import Logging
-import AsyncHelpers
 import AsyncAlgorithms
 #if canImport(FoundationNetworking)
 import FoundationNetworking

Sources/StructuredWebSocketClient/URLSession+MessageTransport.swift

@@ -13,7 +13,6 @@
 
 import Foundation
 import Logging
-import AsyncHelpers
 import AsyncAlgorithms
 #if canImport(FoundationNetworking)
 import FoundationNetworking
@@ -52,6 +51,12 @@ public final class URLSessionWebSocketTransport: MessageTransport {
     }
 
     public func close(with closeCode: URLSessionWebSocketTask.CloseCode, reason: Data?) {
+        // If the task is already closed, we need to call onClose, as that is
+        // the only way the events channel is finished.
+        guard wsTask.closeCode == .invalid else {
+            Task { await self.onClose(closeCode: closeCode, reason: reason) }
+            return
+        }
         wsTask.cancel(with: closeCode, reason: reason)
     }
 
@@ -89,6 +94,11 @@ public final class URLSessionWebSocketTransport: MessageTransport {
             WebSocketClient did complete with error (code: \(nsError.code), reason: \(reason))
             """)
         await self.events.send(.failure(nsError))
+        // If the task is already closed, we need to call onClose, as that is
+        // the only way the events channel is finished.
+        if wsTask.closeCode != .invalid {
+            await self.onClose(closeCode: .abnormalClosure, reason: Data(reason.utf8))
+        }
     }
 
     private func readNextMessage() {
@@ -110,6 +120,7 @@ public final class URLSessionWebSocketTransport: MessageTransport {
                     self?.readNextMessage()
                 } catch {
                     await self?.events.send(.failure(error))
+                    await self?.onClose(closeCode: .abnormalClosure, reason: Data(error.localizedDescription.utf8))
                 }
             }
         }

Sources/StructuredWebSocketClient/WebSocketClient.swift

@@ -13,7 +13,7 @@
 
 import Foundation
 import Logging
-import AsyncHelpers
+import AsyncAlgorithms
 #if canImport(FoundationNetworking)
 import FoundationNetworking
 #endif
@@ -50,28 +50,34 @@ public final class WebSocketClient {
         self.logger.trace("♻️ Deinit of WebSocketClient")
     }
 
-    public func connect() -> AnyAsyncSequence<WebSocketEvent> {
+    public func connect() ->  AsyncCompactMapSequence<AsyncChannel<WebSocketEvent>, WebSocketEvent> {
         self.transport.connect().compactMap { e -> WebSocketEvent? in
             switch e {
             case let .state(s):
                 return .state(s)
             case let .message(m, metadata: meta):
                 // pipe message through middleware
                 if let middleware = self.inboundMiddleware {
-                    let handled = try await middleware.handle(m, metadata: meta)
-                    switch handled {
-                    case .handled:
-                        return nil
-                    case let .unhandled(unhandled):
-                        return.message(unhandled, metadata: meta)
+                    do {
+                        let handled = try await middleware.handle(m, metadata: meta)
+                        switch handled {
+                        case .handled:
+                            return nil
+                        case let .unhandled(unhandled):
+                            return.message(unhandled, metadata: meta)
+                        }
+                    } catch {
+                        // In case the middleware throws when handling the message
+                        // We could also just pretend it was unhandled.
+                        return .failure(error)
                     }
                 } else {
                     return .message(m, metadata: meta)
                 }
             case let .failure(error):
                 return .failure(error)
             }
-        }.eraseToAnyAsyncSequence()
+        }
     }
 
     public func disconnect(reason: String?) async {

Sources/StructuredWebSocketClientTestSupport/TestMessageTransport.swift

@@ -12,7 +12,6 @@
 //===----------------------------------------------------------------------===//
 
 import Foundation
-import AsyncHelpers
 import AsyncAlgorithms
 import StructuredWebSocketClient
 #if canImport(FoundationNetworking)
@@ -21,8 +20,8 @@ import FoundationNetworking
 
 public final class TestMessageTransport: MessageTransport {
     var messageNumber: Int = 0
+    private var _events: AsyncChannel<WebSocketEvent> = .init()
     private var events: AsyncChannel<WebSocketEvent> = .init()
-    private var awaiter: Awaiter = .init()
     private var _initialMessages: [WebSocketEvent]
     private var _onMessage: (URLSessionWebSocketTask.Message, TestMessageTransport) async throws -> Void
 
@@ -31,32 +30,27 @@ public final class TestMessageTransport: MessageTransport {
         _onMessage = onMessage
     }
 
-    // will wait until state is connected
     public func push(_ event: WebSocketEvent) async {
-        await awaiter.awaitUntilTriggered {
-            await self.events.send(event)
-        }
+        await self._events.send(event)
     }
 
     public func send(_ message: URLSessionWebSocketTask.Message) async throws {
         try await _onMessage(message, self)
-        // print("Sending: \(String(decoding: try message.data(), as: UTF8.self))")
     }
 
     public func close(with closeCode: URLSessionWebSocketTask.CloseCode, reason: Data?) {
         Task {
-            await events.send(.state(.disconnected(closeCode: closeCode, reason: reason)))
-            self.events.finish()
+            await _events.send(.state(.disconnected(closeCode: closeCode, reason: reason)))
+            self._events.finish()
         }
     }
 
     public func connect() -> AsyncChannel<WebSocketEvent> {
-        // if push is called before connect this will only send the messages after the return
-        Task { await awaiter.trigger() }
-        Task {
-            for await event in chain([.state(.connected)].async, _initialMessages.async, events) {
-                await events.send(event)
+        Task { [unowned self] in
+            for await event in chain([.state(.connected)].async, self._initialMessages.async, self._events) {
+                await self.events.send(event)
             }
+            self.events.finish()
         }
         return events
     }