@@ -23,13 +23,7 @@ protocol WebSocketClient: Sendable {
2323 func send( _ message: RealtimeMessageV2 ) async throws
2424 func receive( ) -> AsyncThrowingStream < RealtimeMessageV2 , any Error >
2525 func connect( ) -> AsyncStream < ConnectionStatus >
26- func disconnect( closeCode: URLSessionWebSocketTask . CloseCode )
27- }
28-
29- extension WebSocketClient {
30- func disconnect( ) {
31- disconnect ( closeCode: . normalClosure)
32- }
26+ func disconnect( )
3327}
3428
3529final class WebSocket : NSObject , URLSessionWebSocketDelegate , WebSocketClient , @unchecked Sendable {
@@ -39,7 +33,7 @@ final class WebSocket: NSObject, URLSessionWebSocketDelegate, WebSocketClient, @
3933
4034 struct MutableState {
4135 var continuation : AsyncStream < ConnectionStatus > . Continuation ?
42- var stream : SocketStream ?
36+ var connection : WebSocketConnection < RealtimeMessageV2 , RealtimeMessageV2 > ?
4337 }
4438
4539 private let mutableState = LockIsolated ( MutableState ( ) )
@@ -57,7 +51,7 @@ final class WebSocket: NSObject, URLSessionWebSocketDelegate, WebSocketClient, @
5751 mutableState. withValue { state in
5852 let session = URLSession ( configuration: configuration, delegate: self , delegateQueue: nil )
5953 let task = session. webSocketTask ( with: realtimeURL)
60- state. stream = SocketStream ( task: task)
54+ state. connection = WebSocketConnection ( task: task)
6155 task. resume ( )
6256
6357 let ( stream, continuation) = AsyncStream< ConnectionStatus> . makeStream( )
@@ -66,51 +60,27 @@ final class WebSocket: NSObject, URLSessionWebSocketDelegate, WebSocketClient, @
6660 }
6761 }
6862
69- func disconnect( closeCode : URLSessionWebSocketTask . CloseCode ) {
63+ func disconnect( ) {
7064 mutableState. withValue { state in
71- state. stream ? . cancel ( with : closeCode )
65+ state. connection ? . close ( )
7266 }
7367 }
7468
7569 func receive( ) -> AsyncThrowingStream < RealtimeMessageV2 , any Error > {
76- mutableState. withValue { mutableState in
77- guard let stream = mutableState. stream else {
78- return . finished(
79- throwing: RealtimeError (
80- " receive() called before connect(). Make sure to call `connect()` before calling `receive()`. "
81- )
70+ guard let connection = mutableState. connection else {
71+ return . finished(
72+ throwing: RealtimeError (
73+ " receive() called before connect(). Make sure to call `connect()` before calling `receive()`. "
8274 )
83- }
84-
85- return stream. map { message in
86- switch message {
87- case let . string( stringMessage) :
88- self . logger? . verbose ( " Received message: \( stringMessage) " )
89-
90- guard let data = stringMessage. data ( using: . utf8) else {
91- throw RealtimeError ( " Expected a UTF8 encoded message. " )
92- }
93-
94- let message = try JSONDecoder ( ) . decode ( RealtimeMessageV2 . self, from: data)
95- return message
96-
97- case . data:
98- fallthrough
99-
100- default :
101- throw RealtimeError ( " Unsupported message type. " )
102- }
103- }
104- . eraseToThrowingStream ( )
75+ )
10576 }
77+
78+ return connection. receive ( )
10679 }
10780
10881 func send( _ message: RealtimeMessageV2 ) async throws {
109- let data = try JSONEncoder ( ) . encode ( message)
110- let string = String ( decoding: data, as: UTF8 . self)
111-
112- logger? . verbose ( " Sending message: \( string) " )
113- try await mutableState. stream? . send ( . string( string) )
82+ logger? . verbose ( " Sending message: \( message) " )
83+ try await mutableState. connection? . send ( message)
11484 }
11585
11686 // MARK: - URLSessionWebSocketDelegate
@@ -145,85 +115,3 @@ final class WebSocket: NSObject, URLSessionWebSocketDelegate, WebSocketClient, @
145115 mutableState. continuation? . yield ( . error( error) )
146116 }
147117}
148-
149- typealias WebSocketStream = AsyncThrowingStream < URLSessionWebSocketTask . Message , any Error >
150-
151- final class SocketStream : AsyncSequence , Sendable {
152- typealias AsyncIterator = WebSocketStream . Iterator
153- typealias Element = URLSessionWebSocketTask . Message
154-
155- struct MutableState {
156- var continuation : WebSocketStream . Continuation ?
157- var stream : WebSocketStream ?
158- }
159-
160- private let task : URLSessionWebSocketTask
161- private let mutableState = LockIsolated ( MutableState ( ) )
162-
163- private func makeStreamIfNeeded( ) -> WebSocketStream {
164- mutableState. withValue { state in
165- if let stream = state. stream {
166- return stream
167- }
168-
169- let stream = WebSocketStream { continuation in
170- state. continuation = continuation
171- waitForNextValue ( )
172- }
173-
174- state. stream = stream
175- return stream
176- }
177- }
178-
179- private func waitForNextValue( ) {
180- guard task. closeCode == . invalid else {
181- mutableState. continuation? . finish ( )
182- return
183- }
184-
185- task. receive { [ weak self] result in
186- guard let continuation = self ? . mutableState. continuation else { return }
187-
188- do {
189- let message = try result. get ( )
190- continuation. yield ( message)
191- self ? . waitForNextValue ( )
192- } catch {
193- continuation. finish ( throwing: error)
194- }
195- }
196- }
197-
198- init ( task: URLSessionWebSocketTask ) {
199- self . task = task
200- }
201-
202- deinit {
203- mutableState. continuation? . finish ( )
204- }
205-
206- func makeAsyncIterator( ) -> WebSocketStream . Iterator {
207- makeStreamIfNeeded ( ) . makeAsyncIterator ( )
208- }
209-
210- func cancel( with closeCode: URLSessionWebSocketTask . CloseCode = . goingAway) {
211- task. cancel ( with: closeCode, reason: nil )
212- mutableState. continuation? . finish ( )
213- }
214-
215- func send( _ message: URLSessionWebSocketTask . Message ) async throws {
216- try await task. send ( message)
217- }
218- }
219-
220- #if os(Linux) || os(Windows)
221- extension URLSessionWebSocketTask {
222- func receive( completionHandler: @Sendable @escaping ( Result < URLSessionWebSocketTask . Message , any Error > ) -> Void ) {
223- Task {
224- let result = await Result ( catching: { try await self . receive ( ) } )
225- completionHandler ( result)
226- }
227- }
228- }
229- #endif
0 commit comments