WWDC Review: How to use URLSessionWebSocketTask in Swift

With the latest updates from WWDC 2019, WebSocket is now a first-class citizen in iOS, macOS, tvOS, and watchOS. With iOS 13, Apple introduced URLSessionWebSocketTask, designed explicitly for WebSockets. Intrigued, we decided to take it for a spin and find out what's cool about this new API.

URLSessionWebSocketTask

URLSession has a new set of methods for WebSocket operations:

@available(iOS 13.0, *)
open func webSocketTask(with url: URL) -> URLSessionWebSocketTask


@available(iOS 13.0, *)
open func webSocketTask(with url: URL, protocols: [String]) -> URLSessionWebSocketTask


@available(iOS 13.0, *)
open func webSocketTask(with request: URLRequest) -> URLSessionWebSocketTask

URLSessionWebSocketTask extends URLSessionTask, so you have the same base set of APIs that you use for regular URLSessionTask. It may seem reasonable initially, but in reality, 2-way real-time communication via WebSocket differs from typical HTTP request workflow. It would be much better to have a separate API explicitly designed for WebSocket, but we have what we have.

Let's create our first WebSocket connection. I will use the WebSocket echo server running on my machine for testing. This server sends back every message it receives.

let urlSession = URLSession(configuration: .default)
let webSocketTask = urlSession.webSocketTask(with: "ws://0.0.0.0:8080/echo")

To open a connection, we need to resume the webSocketTask

webSocketTask.resume()

Sending messages

`URLSessionWebSocketTask#send` method accepts messages of type `URLSessionWebSocketTask.Message`. It is an enum with 2 cases:

public enum Message {
    case data(Data)
    case string(String)
}

We can only send a text or binary data. We can't send low-level WebSocket frames or control messages, which is fine for most cases. You rarely need to work with custom WebSocket frames.

let message = URLSessionWebSocketTask.Message.string("Hello Socket)
webSocketTask.send(message) { error in
    if let error = error {                
        print("WebSocket sending error: \(error)")
    }
}

The Echo server received the message "Hello Socket" and sent it back.

Receiving messages

We need to use the URLSessionWebSocketTask#receive method to receive data from the server. This method accepts a completion handler that receives a value of the new Resulttype introduced in Swift 5.1

webSocketTask.receive { result in

    switch result {
    case .failure(let error):
        print("Failed to receive message: \(error)")
    case .success(let message):
        switch message {
        case .string(let text):
            print("Received text message: \(text)")
        case .data(let data):
            print("Received binary message: \(data)")
        @unknown default:
            fatalError()
        }                
    }
}

The value of the success case is the same URLSessionWebSocketTask.Message that we used to send messages.

The tricky part is this: receive method only calls once. If you want to receive another message, you must call receive again.

func readMessage()  {
    webSocketTask.receive { result in
        switch result {
        case .failure(let error):
            print("Failed to receive message: \(error)")
        case .success(let message):
            switch message {
            case .string(let text):
                print("Received text message: \(text)")
            case .data(let data):
                print("Received binary message: \(data)")
            @unknown default:
                fatalError()
            }
            
            self.readMessage()
        }
    }
}

I have used dozens of WebSocket libraries across multiple programming languages, and this implementation is the strangest I have ever seen. Typically, you would set a callback on WebSocket, and it delivers all future messages to this callback. It is what you'd expect from a WebSocket API, and why do I have to create a new closure every time?

Ping Pongs

The server may drop your connection due to inactivity if your app is not sending messages over WebSocket with "acceptable" frequency. The system uses particular ping-pong messages to solve this problem. You will need to send them periodically (approximately every 10 seconds). This ensures that the server won't kill the connection (using the webSocketTask.sendPing).

func ping() {
    webSocketTask.sendPing { (error) in
        if let error = error {
            print("Ping failed: \(error)")
        }
        self.scheduleNextPing()
    }
}

This method is elementary. It sends default ping messages and notifies about its results, and the API won't allow customizing the ping message payload. You can use ping messages with custom payloads to deliver additional information to the server (statistics, for example).

Closing connection

When we no longer need a WebSocket connection, we must close it properly. URLSessionWebSocketTask has a separate WebSocket-specific cancel method.

open func cancel(with closeCode: URLSessionWebSocketTask.CloseCode, reason: Data?)
...
func disconnect() {
    webSocketTask.cancel(with: .goingAway, reason: nil)
}

It accepts close code and optional reason payload. Use close code to tell the server why this connection is disconnecting. For example, it may indicate errors related to the WebSocket or application protocols.

Observing connection state changes

We need to know when our WebSocket connection is connected or disconnected and why. To do that, we need to use URLSessionWebSocketDelegate.

@available(iOS 13.0, *)
public protocol URLSessionWebSocketDelegate : URLSessionTaskDelegate {
    
    optional func urlSession(_ session: URLSession, webSocketTask: URLSessionWebSocketTask, didOpenWithProtocol protocol: String?)
   
    optional func urlSession(_ session: URLSession, webSocketTask: URLSessionWebSocketTask, didCloseWith closeCode: URLSessionWebSocketTask.CloseCode, reason: Data?)
}

The API is pretty simple - it notifies you when the connection was opened or closed with a close code and a reason from the server.

We need to pass this delegate to the URLSession constructor.

urlSession = URLSession(configuration: sessionConfiguration, delegate: delegate, delegateQueue: OperationQueue())

Again, Apple is trying to follow the existing URLSession architecture. It's weird why I need to share the same delegate for multiple WebSocket connections. Of course, we can still create an individual URLSession for each connection.

Is it ready for production?

It's only available in iOS 13, and, to be honest, URLSessionWebSocketTask is not the best public WebSocket library. You can use it if you need basic WebSocket functionality and if you don't need to support old iOS versions. In other cases, I recommend using something more mature, like SwiftNIO or one of many open-source frameworks (Starscream, for example).

Sample project

I have created a sample project with a wrapper class around the URLSessionWebSocketTask. It removes part of the boilerplate code and gives you an API similar to the standard WebSocket library. This project has a SwiftNIO-based echo server that you can use for testing. You can find it here.

Summary

It is fascinating to see that the iOS networking stack is evolving. It's moving in the right direction. It could be better, but at least you have your basic required functionality out of the box.

In our next article, I will dig deeper into low-level networking with Network.framework. Oh, the latest iOS update also brought WebSocket support to NVConnection.