Managing an app’s memory is something that tends to be especially tricky to do within the context of asynchronous code, as various objects and values often need to be captured and retained over time in order for our asynchronous calls to be performed and handled.

While Swift’s relatively new async/await syntax does make many kinds of asynchronous operations easier to write, it still requires us to be quite careful when it comes to managing the memory for the various tasks and objects that are involved in such asynchronous code.

Implicit captures

One interesting aspect of async/await (and the Task type that we need to use to wrap such code when calling it from a synchronous context) is how objects and values often end up being implicitly captured while our asynchronous code is being executed.

For example, let’s say that we’re working on a DocumentViewController, which downloads and displays a Document that was downloaded from a given URL. To make our download execute lazily when our view controller is about to be displayed to the user, we’re starting that operation within our view controller’s viewWillAppear method, and we’re then either rendering the downloaded document once available, or showing any error that was encountered — like this:

class DocumentViewController: UIViewController {
    private let documentURL: URL
    private let urlSession: URLSession

    ...

    override func viewWillAppear(_ animated: Bool) {
        super.viewWillAppear(animated)

        Task {
            do {
                let (data, _) = try await urlSession.data(from: documentURL)
                let decoder = JSONDecoder()
                let document = try decoder.decode(Document.self, from: data)
                renderDocument(document)
            } catch {
                showErrorView(for: error)
            }
        }
    }

    private func renderDocument(_ document: Document) {
        ...
    }

    private func showErrorView(for error: Error) {
        ...
    }
}

Now, if we just quickly look at the above code, it might not seem like there’s any object capturing going on whatsoever. After all, asynchronous capturing has traditionally only happened within escaping closures, which in turn require us to always explicitly refer to self whenever we’re accessing a local property or method within such a closure (when self refers to a class instance, that is).

So we might expect that if we start displaying our DocumentViewController, but then navigate away from it before its download has completed, that it’ll be successfully deallocated once no external code (such as its parent UINavigationController) maintains a strong reference to it. But that’s actually not the case.

That’s because of the aforementioned implicit capturing that happens whenever we create a Task, or use await to wait for the result of an asynchronous call. Any object used within a Task will automatically be retained until that task has finished (or failed), including self whenever we’re referencing any of its members, like we’re doing above.

In many cases, this behavior might not actually be a problem, and will likely not lead to any actual memory leaks, since all captured objects will eventually be released once their capturing task has completed. However, let’s say that we’re expecting the documents downloaded by our DocumentViewController to potentially be quite large, and that we wouldn’t want multiple view controllers (and their download operations) to remain in memory if the user quickly navigates between different screens.

The classic way to address this sort of problem would be to perform a weak self capture, which is often accompanied by a guard let self expression within the capturing closure itself — in order to turn that weak reference into a strong one that can then be used within the closure’s code:

class DocumentViewController: UIViewController {
    ...

    override func viewWillAppear(_ animated: Bool) {
        super.viewWillAppear(animated)

        Task { [weak self] in
    guard let self = self else { return }

            do {
                let (data, _) = try await self.urlSession.data(
                    from: self.documentURL
                )
                let decoder = JSONDecoder()
                let document = try decoder.decode(Document.self, from: data)
                self.renderDocument(document)
            } catch {
                self.showErrorView(for: error)
            }
        }
    }
    
    ...
}

Unfortunately, that’s not going to work in this case, as our local self reference will still be retained while our asynchronous URLSession call is suspended, and until all of our closure’s code has finished running (just like how a local variable within a function is retained until that scope has been exited).

So if we truly wanted to capture self weakly, then we’d have to consistently use that weak self reference throughout our closure. To make it somewhat simpler to use our urlSession and documentURL properties, we could capture those separately, as doing so won’t prevent our view controller itself from being deallocated:

class DocumentViewController: UIViewController {
    ...

    override func viewWillAppear(_ animated: Bool) {
        super.viewWillAppear(animated)

        Task { [weak self, urlSession, documentURL] in
            do {
                let (data, _) = try await urlSession.data(from: documentURL)
                let decoder = JSONDecoder()
                let document = try decoder.decode(Document.self, from: data)
                self?.renderDocument(document)
            } catch {
                self?.showErrorView(for: error)
            }
        }
    }
    
    ...
}

The good news is that, with the above in place, our view controller will now be successfully deallocated if it ends up being dismissed before its download has completed.

However, that doesn’t mean that its task will automatically be cancelled. That might not be a problem in this particular case, but if our network call resulted in some kind of side-effect (like a database update), then that code would still run even after our view controller would be deallocated, which could result in bugs or unexpected behavior.

Cancelling tasks

One way to ensure that any ongoing download task will indeed be cancelled once our DocumentViewController goes out of memory would be to store a reference to that task, and to then call its cancel method when our view controller is being deallocated:

class DocumentViewController: UIViewController {
    private let documentURL: URL
    private let urlSession: URLSession
    private var loadingTask: Task<Void, Never>?
    
    ...

    deinit {
    loadingTask?.cancel()
}

    override func viewWillAppear(_ animated: Bool) {
        super.viewWillAppear(animated)

        loadingTask = Task { [weak self, urlSession, documentURL] in
            ...
        }
    }
    
    ...
}

Now everything works as expected, and all of our view controller’s memory and asynchronous state will automatically be cleaned up once it has been dismissed — but our code has also become quite complicated in the process. Having to write all of that memory management code for every view controller that performs an asynchronous task would be quite tedious, and it could even make us question whether async/await actually gives us any real benefits over technologies like Combine, delegates, or closures.

Thankfully, there’s another way to implement the above pattern that doesn’t involve quite as much code and complexity. Since the convention is for long-running async methods to throw an error if they get cancelled (see this article about delaying asynchronous tasks for more info), we can simply cancel our loadingTask once our view controller is about to be dismissed — and that will make our task throw an error, exit, and release all of its captured objects (including self). That way, we no longer need to capture self weakly, or do any other kind of manual memory management work — giving us the following implementation:

class DocumentViewController: UIViewController {
    private let documentURL: URL
    private let urlSession: URLSession
    private var loadingTask: Task<Void, Never>?
    
    ...

    override func viewWillAppear(_ animated: Bool) {
        super.viewWillAppear(animated)

        loadingTask = Task {
            do {
                let (data, _) = try await urlSession.data(from: documentURL)
                let decoder = JSONDecoder()
                let document = try decoder.decode(Document.self, from: data)
                renderDocument(document)
            } catch {
                showErrorView(for: error)
            }
        }
    }

    override func viewWillDisappear(_ animated: Bool) {
    super.viewWillDisappear(animated)
    loadingTask?.cancel()
}

    ...
}

Note that when our task is cancelled, our showErrorView method will now still be called (since an error will be thrown, and self remains in memory at that point). However, that extra method call should be completely negligible in terms of performance.

Long-running observations

The above set of memory management techniques should become even more important once we start using async/await to set up long-running observations of some kind of async sequence or stream. For example, here we’re making a UserListViewController observe a UserList class in order to reload its table view data once an array of User models was changed:

class UserList: ObservableObject {
    @Published private(set) var users: [User]
    ...
}

class UserListViewController: UIViewController {
    private let list: UserList
    private lazy var tableView = UITableView()
    
    ...

    override func viewWillAppear(_ animated: Bool) {
        super.viewWillAppear(animated)

        Task {
            for await users in list.$users.values {
                updateTableView(withUsers: users)
            }
        }
    }

    private func updateTableView(withUsers users: [User]) {
        ...
    }
}

To learn more about Published properties, check out this article.

Note that the above implementation currently doesn’t include any of the task cancellation logic that we previously implemented within our DocumentViewController, which in this case will actually lead to a memory leak. The reason is that (unlike our previous Document-loading task) our UserList observation task will keep running indefinitely, as it’s iterating over a Publisher-based async sequence that can’t throw an error, or complete in any other way.

The good news is that we can easily fix the above memory leak using the exact same technique as we previously used to prevent our DocumentViewController from being retained in memory — that is, to cancel our observation task once our view controller is about to disappear:

class UserListViewController: UIViewController {
    private let list: UserList
    private lazy var tableView = UITableView()
    private var observationTask: Task<Void, Never>?

    ...

    override func viewWillAppear(_ animated: Bool) {
        super.viewWillAppear(animated)

        observationTask = Task {
            for await users in list.$users.values {
                updateTableView(withUsers: users)
            }
        }
    }

    override func viewWillDisappear(_ animated: Bool) {
        super.viewWillDisappear(animated)
        observationTask?.cancel()
    }

    ...
}

Note that performing the above cancellation within deinit wouldn’t work in this case, since we’re dealing with an actual memory leak — meaning that deinit will never be called unless we break our observation task’s endless loop.

Conclusion

At first, it might seem like technologies like Task and async/await make asynchronous, memory-related issues a thing of the past, but unfortunately we still have to be careful around how objects are captured and retained when performing various kinds of async-marked calls. While actual memory leaks and retain cycles are perhaps not as easily encountered as when using things like Combine or closures, we still have to ensure that our objects and tasks are managed in a way that makes our code robust and easy to maintain.

I hope that you found this article useful. If you did, please share it! If you have any questions, comments, or feedback, then feel free to reach out via either Twitter or email.

Thanks for reading!