Method listen URLSessionWebSocketTask without recursion

Question

When called, the method of URLSessionWebSocketTask.receive(completionHandler:), receives only one message.

So in order to continue receiving messages I need to call this method again and again. This is easily implemented using recursion:

private func listen(completion: @escaping ((Error) -> Void)) {
    guard let task, !task.isCancelled else { return }
    task.receive { [weak self] result in
        switch result {
        case let .success(message):
            self?.handle(message: message)
            self?.listen(completion: completion)
        case let .failure(error):
            completion(error)
        }
    }
}

But the question is, will this recursion cause a stack overflow, and is it possible to implement this logic without recursion?

Answer 1

You asked:

will this recursion cause a stack overflow?

No, this would not cause a stack overflow because the recursive call is in the asynchronous completion handler. The listen method has already returned by the time receive’s completion handler is called. This pattern is not uncommon with legacy asynchronous patterns such as delegates or completion handler closures.

In short, your code is fine.

is it possible to implement this logic without recursion?

The above notwithstanding, you can eliminate the recursion with Swift concurrency, if you would like. To do this, you can use the async rendition of receive, and then you can have a simple loop without recursion:

func receiveMessages(for socket: URLSessionWebSocketTask) async throws {
    while !Task.isCancelled {
        let message = try await socket.receive()
        …
    }
}

While your code works, this is arguably easier to reason about at a glance. But it entails using Swift concurrency, so you might only do this if you are embracing async-await throughout your codebase. One would generally use Swift concurrency or legacy patterns, but generally not both at the same time.

---

Personally, I might wrap this loop in an asynchronous sequence of messages:

extension URLSessionWebSocketTask {
    func messages() -> AsyncThrowingStream<URLSessionWebSocketTask.Message, Error> {
        AsyncThrowingStream { continuation in
            let task = Task {
                do {
                    while !Task.isCancelled {
                        let message = try await receive()
                        continuation.yield(message)
                    }
                    continuation.finish()
                } catch {
                    continuation.finish(throwing: error)
                }
            }
            continuation.onTermination = { state in
                if case .cancelled = state {
                    task.cancel()
                }
            }
        }
    }
}

Then you can just for try await that sequence. For example:

let url = URL(string: "wss://echo.websocket.org")!

let socket = URLSession.shared.webSocketTask(with: url)
socket.resume()

let receiveTask = Task {
    for try await message in socket.messages() {
        print("received", message)
    }
}

defer {
    receiveTask.cancel()
    socket.cancel(with: .normalClosure, reason: nil)
}

// this demo web socket service echos what I send, so
// let’s slowly send some messages so the service will
// respond with messages, received above.

for i in 0 ..< 10 {
    try await socket.send(.string("\(i)"))
    try await Task.sleep(for: .seconds(1))
}