@@ -23,13 +23,7 @@ protocol WebSocketClient: Sendable {
23
23
func send( _ message: RealtimeMessageV2 ) async throws
24
24
func receive( ) -> AsyncThrowingStream < RealtimeMessageV2 , any Error >
25
25
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( )
33
27
}
34
28
35
29
final class WebSocket : NSObject , URLSessionWebSocketDelegate , WebSocketClient , @unchecked Sendable {
@@ -39,7 +33,7 @@ final class WebSocket: NSObject, URLSessionWebSocketDelegate, WebSocketClient, @
39
33
40
34
struct MutableState {
41
35
var continuation : AsyncStream < ConnectionStatus > . Continuation ?
42
- var stream : SocketStream ?
36
+ var connection : WebSocketConnection < RealtimeMessageV2 , RealtimeMessageV2 > ?
43
37
}
44
38
45
39
private let mutableState = LockIsolated ( MutableState ( ) )
@@ -57,7 +51,7 @@ final class WebSocket: NSObject, URLSessionWebSocketDelegate, WebSocketClient, @
57
51
mutableState. withValue { state in
58
52
let session = URLSession ( configuration: configuration, delegate: self , delegateQueue: nil )
59
53
let task = session. webSocketTask ( with: realtimeURL)
60
- state. stream = SocketStream ( task: task)
54
+ state. connection = WebSocketConnection ( task: task)
61
55
task. resume ( )
62
56
63
57
let ( stream, continuation) = AsyncStream< ConnectionStatus> . makeStream( )
@@ -66,51 +60,27 @@ final class WebSocket: NSObject, URLSessionWebSocketDelegate, WebSocketClient, @
66
60
}
67
61
}
68
62
69
- func disconnect( closeCode : URLSessionWebSocketTask . CloseCode ) {
63
+ func disconnect( ) {
70
64
mutableState. withValue { state in
71
- state. stream ? . cancel ( with : closeCode )
65
+ state. connection ? . close ( )
72
66
}
73
67
}
74
68
75
69
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()`. "
82
74
)
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
+ )
105
76
}
77
+
78
+ return connection. receive ( )
106
79
}
107
80
108
81
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)
114
84
}
115
85
116
86
// MARK: - URLSessionWebSocketDelegate
@@ -145,85 +115,3 @@ final class WebSocket: NSObject, URLSessionWebSocketDelegate, WebSocketClient, @
145
115
mutableState. continuation? . yield ( . error( error) )
146
116
}
147
117
}
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