The Task Cancellation Trap

Working with cancellation in tasks is a double-edged sword. It’s great when it works, but it’s so easy to shoot yourself in the foot.

Take the observerTask example from the previous post:

let observerTask = Task {
  let notifications = NotificationCenter.default.notifications(named: name)
  
  for await next in notifications {
    print(next)
  }
}

For the observation to stop when you no longer need it, something needs to cancel that task. Just letting go of the handle does not cancel—you just lose the opportunity to cancel it.

I’m quoting the Task documentation here because I feel this is really essential to understand (emphasis mine):

It’s not a programming error to discard a reference to a task without waiting for that task to finish or canceling it. 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.

If you’re handling notifications like in the example, it’s even easier to fall into that trap. Most likely, you won’t be printing or even needing them, but calling some update function on an object. To avoid retain cycles, you’d weak-reference self, lulling yourself into thinking this is fine:

Task { [weak self] in
  // ...
  
  for await _ in notifications {
    self?.updateSomething()
  }
}

This has happened to me and others on my team. So what’s the problem? No reference to the task is kept around. Even when self goes away, the task runs forever. The observer still fires but doesn’t call any function—just leaking resources.

What are potential fixes? One is to modify the iteration to break out of the loop as soon as self goes away:

Task { [weak self] in
  // ...
  
  for await _ in notifications {
    guard let self else { return }
    self.updateSomething()
  }
}

This approach is better, but has a significant downside. The loop and task only end after the next notification. Some notifications rarely occur but need to be observed by many clients. All these observation tasks would still leak with this solution. (Also, not all tasks have a loop like that.)

The other solution is to keep a reference to the task and cancel it when the client—self in this case—goes away. Store the task in an instance variable and cancel it on deinit:

@MainActor class MyObserver {
  var observerTask: Task<Void, Never>?
   
  init() {
    self.observerTask = Task { [weak self] in
      // same as before
    }
  }

  deinit {
    observerTask?.cancel()
  }
}

(Side note: observerTask needs to be optional because its closure references self, and all properties must be initialized before such a reference. Through the optional, the property first initializes to nil, then gets set to the actual task.)

While that’s better, we weren’t entirely happy. Two issues:

1. The Task initializer is declared as @discardableResult, which means you can create a task and let go of its handle without a warning, like Task {…}. This makes it too easy to unintentionally set up a task that’s forgotten and runs forever.
2. It’s too easy to forget to call the cancel method. Especially when the task is stored in an array or dictionary rather than a single instance variable. Every task handle would need to be cancelled individually when removed or replaced.

To tackle these, we did two things. To solve number 2, we wrote a little wrapper called ScopedTask, which handles cancellation automatically. Usage is essentially identical to Task, but the task gets cancelled when the last reference is released. Plus, its initializer is not @discardableResult—meaning you’ll get a warning when creating a ScopedTask without storing it in a variable. With it, the example above becomes:

@MainActor class MyObserver {
  var observerTask: ScopedTask<Void, Never>?
   
  init() {
    self.observerTask = ScopedTask { [weak self] in
      // same as before
    }
  }
}

To solve number 1, we plan to set up a custom linter rule for SwiftLint that warns about any unstructured Task use. If you want to do something similar, here’s a suggestion for a custom rule you can add to your .swiftlint.yml configuration:

custom_rules:
  unstructured_task:
    name: "Unstructured Task"
    message: "Prefer ScopedTask instead of plain Task."
    regex: '\bTask(?=\s*[<{(])'
    severity: warning
    match_kinds:
      - identifier

As with every linter rule, this isn’t a perfect solution. But it helps make contributors aware of the issue. Some places will require plain Task, like the ScopedTask implementation itself. There, we disable the rule while explicitly stating and documenting the intentional deviation from the norm. Our main goal is to ensure this doesn’t go unnoticed.

I published the source for ScopedTask on GitHub—feel free to use or adapt it to your needs. And let me know if you find these kinds of things interesting or useful. 😊