How to use [weak self] in Swift Concurrency Tasks?

As a developer who uses Swift regularly, [weak self] should be something that's almost muscle memory to you. I've written about using [weak self] before in the context of when you should generally capture self weakly in your closures to avoid retain cycles. The bottom line of that post is that closures that aren't @escaping will usually not need a [weak self] because the closures aren't retained beyond the scope of the function you're passing them to. In other words, closures that aren't @escaping don't usually cause memory leaks. I'm sure there are exceptions but generally speaking I've found this rule of thumb to hold up.

This idea of not needing [weak self] for all closures is reinforced by the introduction of SE-0269 which allows us to leverage implicit self captures in situations where closures aren’t retained, making memory leaks unlikely.

Later, I also wrote about how Task instances that iterate async sequences are fairly likely to have memory leaks due to this implicit usage of self.

So how do we use [weak self] on Task? And if we shouldn't, how do we avoid memory leaks?

In this post, I aim to answer these questions.

The basics of using [weak self] in completion handlers

As Swift developers, our first instinct is to do a weak -> strong dance in pretty much every closure. For example:

loadData { [weak self] data in 
  guard let self else { return }

  // use data
}

This approach makes a lot of sense. We start the call to loadData, and once the data is loaded our closure is called. Because we don't need to run the closure if self has been deallocated during our loadData call, we use guard let self to make sure self is still there before we proceed.

This becomes increasingly important when we stack work:

loadData { [weak self] data in 
  guard let self else { return }

  processData(data) { [weak self] models in 
    // use models
  }
}

Notice that we use [weak self] in both closures. Once we grab self with guard let self our reference is strong again. This means that for the rest of our closure, self is held on to as a strong reference. Due to SE-0269 we can call processData without writing self.processData if we have a strong reference to self.

The closure we pass to processData also captures self weakly. That's because we don't want that closure to capture our strong reference. We need a new [weak self] to prevent the closure that we passed to processData from creating a (shortly lived) memory leak.

When we take all this knowledge and we transfer it to Task, things get interesting...

Using [weak self] and unwrapping it immediately in a Task

Let's say that we want to write an equivalent of our loadData and processData chain, but they're now async functions that don't take a completion handler.

A common first approach would be to do the following:

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

  let data = await loadData()
  let models = await processData(data)
}

Unfortunately, this code does not solve the memory leak that we solved in our original example.

An unstructured Task you create will start running as soon as possible. This means that if we have a function like below, the task will run as soon as the function reaches the end of its body:

func loadModels() {
  // 1
  Task { [weak self] in
    // 3: _immediately_ after the function ends
    guard  let self else { return }

    let data = await loadData()
    let models = await processData(data)
  }
  // 2
}

More complex call stacks might push the start of our task back by a bit, but generally speaking, the task will run pretty much immediately.

The problem with guard let self at the start of your Task

Because Task in Swift starts running as soon as possible, the chance of self getting deallocated in the time between creating and starting the task is very small. It's not impossible, but by the time your Task starts, it's likely self is still around no matter what.

After we make our reference to self strong, the Task holds on to self until the Task completes. In our call that means that we retain self until our call to processData completes. If we translate this back to our old code, here's what the equivalent would look like in callback based code:

loadData { data in 
  self.processData(data) { models in 
    // for example, self.useModels
  }
}

We don't have [weak self] anywhere. This means that self is retained until the closure we pass to processData has run.

The exact same thing is happening in our Task above.

Generally speaking, this isn't a problem. Your work will finish and self is released. Maybe it sticks around a bit longer than you'd like but it's not a big deal in the grand scheme of things.

But how would we prevent kicking off processData if self has been deallocated in this case?

Preventing a strong self inside of your Task

We could make sure that we never make our reference to self into a strong one. For example, by checking if self is still around through a nil check or by guarding the result of processData. I'm using both techniques in the snippet above but the guard self != nil could be omitted in this case:

Task { [weak self] in
  let data = await loadData()
  guard self != nil else { return }

  guard let models = await self?.processData(data) else {
    return
  }

  // use models
}

The code isn't pretty, but it would achieve our goal.

Let's take a look at a slightly more complex issue that involves repeatedly fetching data in an unstructured Task.

Using [weak self] in a longer running Task

Our original example featured two async calls that, based on their names, probably wouldn't take all that long to complete. In other words, we were solving a memory leak that would typically solve itself within a matter of seconds and you could argue that's not actually a memory leak worth solving.

A more complex and interesting example could look as follows:

func loadAllPages() {
  // only fetch pages once
  guard fetchPagesTask == nil else { return }

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

    var hasMorePages = true
    while hasMorePages && !Task.isCancelled {
      let page = await fetchNextPage()
      hasMorePages = !page.isLastPage
    }

    // we're done, we could call loadAllPages again to restart the loading process
    fetchPagesTask = nil
  }
}

Let's remove some noise from this function so we can see the bits that are actually relevant to whether or not we have a memory leak. I wanted to show you the full example to help you understand the bigger picture of this code sample...

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

  var hasMorePages = true
  while hasMorePages {
    let page = await fetchNextPage()
    hasMorePages = !page.isLastPage
  }
}

There. That's much easier to look at, isn't it?

So in our Task we have a [weak self] capture and immediately we unwrap with a guard self. You already know this won't do what we want it to. The Task will start running immediately, and self will be held on to strongly until our task ends. That said, we do want our Task to end if self is deallocated.

To achieve this, we can actually move our guard let self into the while loop:

Task { [weak self] in
  var hasMorePages = true

  while hasMorePages {
    guard let self else { break }
    let page = await fetchNextPage()
    hasMorePages = !page.isLastPage
  }
}

Now, every iteration of the while loop gets its own strong self that's released at the end of the iteration. The next one attempts to capture its own strong copy. If that fails because self is now gone, we break out of the loop.

We fixed our problem by capturing a strong reference to self only when we need it, and by making it as short-lived as possible.

In Summary

Most Task closures in Swift don't strictly need [weak self] because the Task generally only exists for a relatively short amount of time. If you find that you do want to make sure that the Task doesn't cause memory leaks, you should make sure that the first line in your Task isn't guard let self else { return }. If that's the first line in your Task, you're capturing a strong reference to self as soon as the Task starts running which usually is almost immediately.

Instead, unwrap self only when you need it and make sure you only keep the unwrapped self around as short as possible (for example in a loop's body). You could also use self? to avoid unwrapping altogether, that way you never grab a strong reference to self. Lastly, you could consider not capturing self at all. If you can, capture only the properties you need so that you don't rely on all of self to stick around when you only need parts of self.