Awaiting WebSockets in SwiftUI

Awaiting WebSockets in SwiftUI

ยท

6 min read

WebSockets have become the standard network protocol for chat and event-driven based applications given their ability to provide a persistent communication channel between clients and a server. Rather than poll a server for results at set intervals, WebSockets emit messages as they become available to connected participants resulting in reduced latency and bandwidth.

Apple provides in-built tools via the URLSession object class to connect to a WebSocket and receive events. However, despite recent updates leveraging Swift's newer concurrency features in iOS 15, the class does not make use of one of Swift's more powerful concurrency features; AsyncSequence

What is AsyncSequence?

Apple's documentation defines AsyncSequence as a type that list values you can step through one at a time while adding asynchroncity. Unlike a regular Sequence, an AsyncSequence may have all, some or none of its values when you first use it. As a result, you await to receive values as they become available.

The description above sounds oddly familiar to how a WebSocket emits messages, making it an ideal type for implementing a more concurrent-friendly api that can be consumed via the following syntax :

for try await message in socket {
    // do something with a message
}

Building an Async WebSocket Stream

import Foundation

class WebSocketStream: AsyncSequence {

    typealias Element = URLSessionWebSocketTask.Message
    typealias AsyncIterator = AsyncThrowingStream<URLSessionWebSocketTask.Message, Error>.Iterator

    //...
}

Above we define a WebSocketStream class that conforms to the AsyncSequence protocol. The protocol requires specifying the output Element of the sequence, which in this case is a URLSessionWebSocketTask.Message . We also need to specify associated type of the AsyncIterator, which will be an AsyncThrowingStream.Iterator as the stream can potentially error out while listening to the WebSocket.

class WebSocketStream: AsyncSequence {

    //..

    private var stream: AsyncThrowingStream<Element, Error>?
    private var continuation: AsyncThrowingStream<Element, Error>.Continuation?
    private let socket: URLSessionWebSocketTask

    init(url: String, session: URLSession = URLSession.shared) {
        socket = session.webSocketTask(with: URL(string: url)!)
        stream = AsyncThrowingStream { continuation in
            self.continuation = continuation
            self.continuation?.onTermination = { @Sendable [socket] _ in
                socket.cancel()
            }
        }
    }

    //..
}

Our initializer takes in a url string and a URLSession object which we use to build a URLSessionWebSocketTask to listen and process web socket messages.

We also create a local AsyncThrowingStream property which will be used to provide an AsyncIterator when a task awaits the values of this sequence.

Note that we also keep a reference to the stream's continuation, which will allow us to signal messages as they come in from the WebSocket server.

In order to conform to the AsyncSequence protocol we need to implement the makeAsyncIterator() method which is the stream we created earlier:

//..
func makeAsyncIterator() -> AsyncIterator {
    guard let stream = stream else {
        fatalError("stream was not initialized")
    }
    socket.resume()
    listenForMessages()
    return stream.makeAsyncIterator()
}
// ...

In the method above we start the socket task by calling its resume method. The makeAsyncIterator() method is a great place to start any related processes when an asynchronous work unit (i.e Task) is initiated. If you are familiar with Combine, it would be the equivalent of the handleEvents(receiveSubscription:) operator.

Prior to returning the steam's AsyncIterator we also call the listenForMessages() method defined below:

private func listenForMessages() {
    socket.receive { [unowned self] result in
        switch result {
        case .success(let message):
            continuation?.yield(message)
            listenForMessages()
        case .failure(let error):
            continuation?.finish(throwing: error)
        }
    }
}

When the socket receives a message, depending on the result we signal a new value to the iterator via the continuation's yield method or finish it with a failure in the event of an error.

Given that the socket's receive method registers a one-time callback, we need to recursively call this method in the event of a success to continue listening to messages as long as the asynchronous task is running.

Now that our WebSocketStream class is complete, let's see how we can go about listening and reacting to it's events in a SwiftUI View

Putting together a Mock WebSocket Server

To demonstrate our WebSocketStream, we will be leveraging the Overseed platform to set up our mock WebSocket server. Overseed is a data platform that enables you to easily generate streams of synthetic data according to a defined schema complete with frequency and probability specifications.

We will be connecting to a WebSocket server I put together that generates randomized data for battery percentages every 3 seconds across three devices: iPad Pro, Apple Watch and Air Pods. The message emitted is a json string in the following format:

{
    "record": {
        "device_name": "IPad Pro",
        "battery_level": "53"
    }
}

Reacting to WebSocket Messages in SwiftUI

import SwiftUI

struct ContentView: View {
    @State
    private var devices: [Device] = [
        Device(name: "iPad Pro",
               batteryLevel: 100),
        Device(name: "Apple Watch",
               batteryLevel: 100),
        Device(name: "Air Pods",
               batteryLevel: 100)
    ]
    private let stream = WebSocketStream(url: "wss://stream.overseed.io:443/v1/generators/org/68fb8641-e958-4c8a-84e0-bd24b199a2e7/generator/5f47628e-d92c-449b-8274-b690d2135204/stream/-1/?stream_token=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpYXQiOjE2MzYwNzIwNDIsImlkIjoiM2ZmNGFiNzctMTc0Zi00ZGYzLWIyNzgtZmYyMGUzYzM3MDVhIn0.yhxB50P-61y8hDosif8pYQEle7QBf-ZwLOEJwXYO4kc")
    var body: some View {
        NavigationView {
            List {
                ForEach(devices) { device in
                    HStack(spacing: 0) {
                        Image(systemName: device.icon)
                            .resizable()
                            .aspectRatio(contentMode: .fit)
                            .frame(width: 32)
                        VStack(alignment: .leading, spacing: 0) {
                            Text(device.name)
                                .font(.title3)
                                .fontWeight(.bold)
                            ProgressView.init(value: device.batteryLevel/100)
                                .padding(.top, 10)
                                .animation(.default, value: device.batteryLevel/100)
                        }
                        .padding(.leading, 20)
                    }
                    .padding()
                }
            }
            .listStyle(.insetGrouped)
            .task {
                do {
                    for try await message in stream {
                        let updateDevice = try message.device()
                        devices = devices.map({ device in
                            device.id == updateDevice.id ? updateDevice : device
                        })
                    }
                } catch {
                    debugPrint("Oops something didn't go right")
                }
            }
            .navigationTitle("My Devices")
        }
    }
}

Above we put together a view that performs an asynchronous task when it appears via the task(priority:_:) method. In the task we await and decode each message from our WebSocketStream into a Device and update our view's local @State property to reflect the changes in our UI.

.task {
  do {
      for try await message in stream {
          let updateDevice = try message.device()
          devices = devices.map({ device in
              device.id == updateDevice.id ? updateDevice : device
          })
      }
  } catch {
      debugPrint("Oops something didn't go right")
  }
}

As long as the task is running, the view will continue to process updates and display the battery info of our devices based on the last received message. Simulator Screen Recording - iPhone 13 - 2021-11-09 at 21.07.21.gif By attaching our task to the lifetime of the view, SwiftUI will cancel the task when it removes the view invoking the onTermination callback of our stream's continuation object, resulting in the cancelation of our socket task.

Conclusion

A WebSocket is a great solution for providing a persistent two-way communication layer between a client and a server. In this article we discovered how we can leverage the power of AsyncSequence to put together an idiomatic concurrency api for connecting and listening to WebSocket events in a Swift application.

You can find a sample project that demonstrates the concepts of this article here

I hope you enjoyed this article and found it useful. If you have any questions, comments or feedback i'd love to hear them. Contact me or follow me on Twitter

ย