Swift: Using [weak self] as control flow: callbacks vs async/await

21 April 2024 · 516 words

software development

swift
async

Instances of types in Swift, particularly when doing UI programming, are sensitive to retain cycles: A holds on to B and B holds on to A — they're both prevented from being freed. [weak self] to the rescue?

Simple callback

When writing UIs in Swift, mainly UIKit from the point-of-view of this post, it usually goes something like this:

class SomeUI: UIViewController {
    func onButtonPress() {
        networkCall1() {
            [weak self] in
            guard let self else {
                return
            }
            self.thing1()
            networkCall2() {
                [weak self] in
                guard let self else {
                    return
                }
                self.thing2()
            }
        }
    }
}

Each time a network call has been performed, you get a chance to check if self has been freed, and decide what to do. This is butt ugly, boring to write, and a pain. It does have the nice property of being a reasonably easy way to avoid retain cycles.

It becomes a bit gnarlier when you want to write aesthetically pleasing async code, since we chose async / await to avoid indentation in the first place. There is no natural way to avoid retain cycles.

class SomeUI: UIViewController {
    func onButtonPress() {
        Task {
            @MainActor in
            await networkCall1()
            self.thing1()
            await networkCall1()
            self.thing2()
        }
    }
}

self will be retained until all the network calls are done. This might not be a big deal, so let's look at a more complicated example.

Checking in a loop

class SomeUI: UIViewController {
    func loop() {
        networkCall() {
            [weak self] in
            guard let self else {
                return
            }
            self.thing1()
            // Create a timer that fires after 5s
            delay(by: 5.0) {
                [weak self] in
                guard let self else {
                    return
                }
                self.loop()
            }
        }
    }
}

The loop above will be terminated as soon as the UI is freed, which is what we wanted.

What seems to be a nice way to write that using async?

class SomeUI: UIViewController {
    func loop() {
        Task {
            @MainActor in
            repeat {
                await self.networkCall()
                self.thing1()
                try? await Task.sleep(nanoseconds: 5_000_000_000)
            } while !Task.isCancelled
        }
    }
}

Really nice looking, but now we need to handle the fact that we can't stop this loop in any reasonable way.

Attempt 1: Task cancellation

Our first attempt might be storing the Task and cancelling it when needed.

class SomeUI: UIViewController {
    var loopTask: Task<(), Error>?
    
    func loop() {
        self.loopTask = Task {
            @MainActor in
            repeat {
                await networkCall()
                self.thing1()
                try? await Task.sleep(nanoseconds: 5_000_000_000)
            } while !Task.isCancelled
        }
    }
    
    func cancel() {
        self.loopTask?.cancel()
    }
}

This will work if we make sure to call cancel on every path that can leave our UIViewController. This might be feasible, but it's not very maintainable. Seven years from now the next programmer updating the UIViewController will need to do the right thing.

Attempt 2: Avoid the retain cycle

Our second attempt is to avoid a retain cycle, which might look like this:

class SomeUI: UIViewController {
    func loop() {
        Task {
            @MainActor [weak self] in
            await networkCall1()
            self?.thing1()
            await networkCall2()
            self?.thing2()
            try? await Task.sleep(nanoseconds: 5_000_000_000)
            self?.loop()
        }
    }
}

This is an uncomplicated example, but notice how we'll keep doing network calls even if self has been freed. The only thing we'll skip are the function calls self?.thing1() and self?.thing2(). This may be undesirable!

Another issue with this approach is that since we spawn new Tasks all the time, it's not immediately possible (without some work) to cancel a long chain or tree of these.

Attempt 3: checking self after every await

It's possible to check the state of self with continuous if-statements, but then we're writing icky code again:

class SomeUI: UIViewController {
    func loop() {
        Task {
            @MainActor [weak self] in
            await networkCall1()
            if let self {
                self.thing1()
            } else {
                return
            }
            await networkCall2()
            if let self {
                self.thing2()
            } else {
                return
            }
            try? await Task.sleep(nanoseconds: 5_000_000_000)
            self?.loop()
        }
    }
}

Final words

The lifecycle of UI in Swift in general, and UIKit in particular is substandard. There's a lot to be said about React but at least it's possible to be able to do cleanup using a useEffect.

This is harder than I wish it was in Swift. In the end, we're responsible for writing code that works well, but it should be easier to get the system to help out than it is in Swift / iOS.

There are probably more ways to solve the retain cycle problem than the ones listed above, but we have at least three tools with different tradeoffs.

Is there an even better way? Let me know.