Use Task with AsyncStream and avoid retain cycle
I’ve been playing with the Swift 5.5 concurrency feature as it finally compiles for iOS 13 in Xcode 13.2. It means it’s now viable to use for most production code & app that need to support more than only iOS 15.
When using Task to make async code works in non async code, you have to keep in mind that a Task guarantee the execution of the code inside. It means that the Task will retain everything within the Task. I was probably a bit to eager to play with async / await and at first I didn’t properly read Task documentation.
It provides some important information
A task runs regardless of whether you keep a reference to it. However, if you discard the reference to a task, you give up the ability to wait for that task’s result or cancel the task.
The UIKit world is a bit harsher than the SwiftUI world. Task created using the .task modifier in SwiftUI will automatically get cancelled once the view disapear. Not in UIKit, as we don’t have this tighly coupled language / view lifecycle mechanism.
In the following example, I’ll setup a long running job that return an AsyncStream that some UI code can work with. In our case we’ll use it directly in the UIViewController, it’s a good setup for the purpose of this example.
First, let’s setup a worker:
class Worker {
var task: Task<(), Never>?
func getSteps() -> AsyncStream<Int> {
return .init { continutation in
continutation.onTermination = { @Sendable _ in
self.task?.cancel()
}
task = Task {
var step: Int = 0
for _ in 0..<1000 {
continutation.yield(step)
try? await Task.sleep(nanoseconds: 100000000)
step += 1
}
}
}
}
}It have a getSteps() function returning an AsyncStream of Int. The init of AsyncStream make it very easy to build an AsyncStream from synchronous based code. As you can notice, we also span a Task, the purpose of this Task is to be able to use the new Task.sleep so we can wait some time between steps. We want to simulate a long running job, and we don’t want the AsyncStream iteration to complete instantly when you’ll start to iterate on it from the outside.
Let’s now setup our UIViewController which will use our Worker
class TaskViewController: UIViewController {
let worker = Worker()
var task: Task<(), Never>?
init() {
super.init(nibName: nil, bundle: nil)
doWork()
}
required init?(coder: NSCoder) {
fatalError("init(coder:) has not been implemented")
}
func doWork() {
print("Task began")
task = Task {
for await step in worker.getSteps() {
print(step)
}
}
}
deinit {
task?.cancel()
print("deinit called")
}
}We start the work in the init, we want to iterate over all the steps of AsyncStream returned by our worker. As it’s async code, we need to wrap it in a Task. As it’s a long runing task, we need to keep a reference to it, as Apple documentation is stating, we’re responsible for cancelling the Task. So it makes sense to keep a reference to it and cancel it int the deinit.
If you take another look at our Worker class, you’ll notice that we connect the onTermination callback of the AsyncStream continuation. We’ll receive this callback when the iteration ont the AsyncStream will be externally terminated or cancelled. So here we’re simply forwarding Task cancellation so we can cancel our inner Task.
Now as all this code run in an Xcode playground, let’s make it all run and see how does it works, here is the final piece of code, the App that make it all works together.
class App {
var controller: TaskViewController?
init() {
present()
DispatchQueue.main.async {
self.dismiss()
}
}
func present() {
print("Present controller")
controller = TaskViewController()
controller?.task
}
func dismiss() {
print("Dismiss controller")
controller = nil
}
}
let app = App()Our goal is to test if our task is actually getting cancelled and all our code deinit properly.
So now let’s run the playground and see what is printed in the console
As we can see, it’s not doing what we would except the code to do. Our App dismiss our UIViewController and set its reference to nil, but the deinit of our controller is not getting called and our tasks are never getting cancelled.
Obviously, we have a retain cycle, and it’s super obvious right? The closure of our Task init is retaining self as we call (self.)worker.getSteps() within the Task
So let’s add a weak self and try again:
func doWork() {
print("Task began")
task = Task { [weak self] in
guard let self = self else { return }
for await step in self.worker.getSteps() {
print(step)
}
}
}Let’s see the results printed in the console. Same….
At this point, we start to enter the not so obvious quirks of long running task in non async code. The truth is, wether it’s a long running Task or a one shot short running Task (just an await call to an async returning function), the release of owner of the Task will not be done until the Task is finished or cancelled (and so release all owned code). This is the case because we use self within the Task. So here, our controller would eventually get its deinit called once our Task end. But this is probably not what you want, right?
You could also have no reference to the task you create and don’t reference self within the task, then it would deinit our controller properly, but the Task would continue to execute until its end, as it doesn’t die with the controller.
We have two ways of fixing this, if you need a reference to self before entering the await loop, you can create your needed variable outside of the Task. Like so:
func doWork() {
print("Task began")
let steps = worker.getSteps()
task = Task {
for await step in steps {
print(step)
}
}
}Now the console will look like this:
The controller is properly deinit and the Task is properly cancelled as the console will not print steps anymore.
Now if you would need to use self within the loop, you need to do it a little bit different. Basically, you’ll need need to use [weak self], and if you want to use a non optional self, you’ll need to create your local self strong variable within the loop, not above it or you’ll end up with the same retain cycle.
func doWork() {
print("Task began")
let steps = worker.getSteps()
task = Task { [weak self] in
for await step in steps {
guard let self = self else { return }
self.currentStep = step
print(self.currentStep)
}
}
}Now why is that?
So in order to write this article, you have to understand the timeline I’m in. I wrote the code above and created a retain cycle, I quickly found the workaround by playing around, but I didn’t really understood why I needed do it like so. As Task init closure is @escaping, you actually don’t need an explicit self, but it’ll still be captured and retained. An again, the Task will actually release everything once it ends. So for most “fire & forget” Task that will wrap just one async function in your UIViewController to mutate self, you’ll not notice the problem. Sure, your controller will probably end up being dealocated a bit later if you dismiss it before the Task is ended. And you should be aware that the Task is not cancelled if you dismiss it, so it’s important to keep a reference and cancel it yourself if the Task is not relevant anymore once your UI is dismissed.
So I’ve been looking around, and found this thread which put me on the right track:
Then I’ve asked Twitter
and got a few interesting replies
So yes, the answer is that you need to keep a weak reference of self within the loop, as each new step would retain self as you use it within each loop.
I don’t know how you feel about all this? But it finally make sense in my head and I’m actually getting confortable with plugging long running task to non async code. The new concurency system is really good, and make a lot of code far more easier to read and much more natural, as it’s just Swift with some more keywords.
Maybe this was already obvious to you, maybe not, but to me it was definitely not something obvious from the get go. But now you know, if you’re using a long running AsyncStream wrapped in a Task, how to not retain the whole thing until the Task end. (And again, it could be what you want, but I bet that in most cases, it’s not).
Thanks for reading and happy coding! 🚀
