Implementing Task timeout with Swift Concurrency

Swift Concurrency provides us with loads of cool and interesting capabilities. For example, Structured Concurrency allows us to write a hierarchy of tasks that always ensures all child tasks are completed before the parent task can complete. We also have features like cooperative cancellation in Swift Concurrency which means that whenever we want to cancel a task, that task must proactively check for cancellation, and exit when needed.

One API that Swift Concurrency doesn't provide out of the box is an API to have tasks that timeout when they take too long. More generally speaking, we don't have an API that allows us to "race" two or more tasks.

In this post, I'd like to explore how we can implement a feature like this using Swift's Task Group. If you're looking for a full-blown implementation of timeouts in Swift Concurrency, I've found this package to handle it well, and in a way that covers most (if not all edge cases).

Racing two tasks with a Task Group

At the core of implementing a timeout mechanism is the ability to race two tasks:

1. A task with the work you're looking to perform
2. A task that handles the timeout

whichever task completes first is the task that dictates the outcome of our operation. If the task with the work completes first, we return the result of that work. If the task with the timeout completes first, then we might throw an error or return some default value.

We could also say that we don't implement a timeout but we implement a race mechanism where we either take data from one source or the other, whichever one comes back fastest.

We could abstract this into a function that has a signature that looks a little bit like this:

func race<T>(
  _ lhs: sending @escaping () async throws -> T,
  _ rhs: sending @escaping () async throws -> T
) async throws -> T {
  // ...
}

Our race function take two asynchronous closures that are sending which means that these closures closely mimic the API provided by, for example, Task and TaskGroup. To learn more about sending, you can read my post where I compare sending and @Sendable.

The implementation of our race method can be relatively straightforward:

func race<T>(
  _ lhs: sending @escaping () async throws -> T,
  _ rhs: sending @escaping () async throws -> T
) async throws -> T {
  return try await withThrowingTaskGroup(of: T.self) { group in
    group.addTask { try await lhs() }
    group.addTask { try await rhs() }

    return try await group.next()!
  }
}

We're creating a TaskGroup and add both closures to it. This means that both closures will start making progress as soon as possible (usually immediately). Then, I wrote return try await group.next()!. This line will wait for the next result in our group. In other words, the first task to complete (either by returning something or throwing an error) is the task that "wins".

The other task, the one that's still running, will be me marked as cancelled and we ignore its result.

There are some caveats around cancellation that I'll get to in a moment. First, I'd like to show you how we can use this race function to implement a timeout.

Implementing timeout

Using our race function to implement a timeout means that we should pass two closures to race that do the following:

1. One closure should perform our work (for example load a URL)
2. The other closure should throw an error after a specified amount of time

We'll define our own TimeoutError for the second closure:

enum TimeoutError: Error {
  case timeout
}

Next, we can call race as follows:

let result = try await race({ () -> String in
  let url = URL(string: "https://www.donnywals.com")!
  let (data, _) = try await URLSession.shared.data(from: url)
  return String(data: data, encoding: .utf8)!
}, {
  try await Task.sleep(for: .seconds(0.3))
  throw TimeoutError.timeout
})

print(result)

In this case, we either load content from the web, or we throw a TimeoutError after 0.3 seconds.

This wait of implementing a timeout doesn't look very nice. We can define another function to wrap up our timeout pattern, and we can improve our Task.sleep by setting a deadline instead of duration. A deadline will ensure that our task never sleeps longer than we intended.

The key difference here is that if our timeout task starts running "late", it will still sleep for 0.3 seconds which means it might take a but longer than 0.3 second for the timeout to hit. When we specify a deadline, we will make sure that the timeout hits 0.3 seconds from now, which means the task might effectively sleep a bit shorter than 0.3 seconds if it started late.

It's a subtle difference, but it's one worth pointing out.

Let's wrap our call to race and update our timeout logic:

func performWithTimeout<T>(
  of timeout: Duration,
  _ work: sending @escaping () async throws -> T
) async throws -> T {
  return try await race(work, {
    try await Task.sleep(until: .now + timeout)
    throw TimeoutError.timeout
  })
}

We're now using Task.sleep(until:) to make sure we set a deadline for our timeout.

Running the same operation as before now looks as follows:

let result = try await performWithTimeout(of: .seconds(0.5)) {
  let url = URL(string: "https://www.donnywals.com")!
  let (data, _) = try await URLSession.shared.data(from: url)
  return String(data: data, encoding: .utf8)!
}

It's a little bit nicer this way since we don't have to pass two closures anymore.

There's one last thing to take into account here, and that's cancellation.

Respecting cancellation

Taks cancellation in Swift Concurrency is cooperative. This means that any task that gets cancelled must "accept" that cancellation by actively checking for cancellation, and then exiting early when cancellation has occured.

At the same time, TaskGroup leverages Structured Concurrency. This means that a TaskGroup cannot return until all of its child tasks have completed.

When we reach a timeout scenario in the code above, we make the closure that runs our timeout an error. In our race function, the TaskGroup receives this error on try await group.next() line. This means that the we want to throw an error from our TaskGroup closure which signals that our work is done. However, we can't do this until the other task has also ended.

As soon as we want our error to be thrown, the group cancels all its child tasks. Built in methods like URLSession's data and Task.sleep respect cancellation and exit early. However, let's say you've already loaded data from the network and the CPU is crunching a huge amount of JSON, that process will not be aborted automatically. This could mean that even though your work timed out, you won't receive a timeout until after your heavy processing has completed.

And at that point you might have still waited for a long time, and you're throwing out the result of that slow work. That would be pretty wasteful.

When you're implementing timeout behavior, you'll want to be aware of this. And if you're performing expensive processing in a loop, you might want to sprinkle some calls to try Task.checkCancellation() throughout your loop:

for item in veryLongList {
  await process(item)
  // stop doing the work if we're cancelled
  try Task.checkCancellation()
}

// no point in checking here, the work is already done...

Note that adding a check after the work is already done doesn't really do much. You've already paid the price and you might as well use the results.

In Summary

Swift Concurrency comes with a lot of built-in mechanisms but it's missing a timeout or task racing API.

In this post, we implemented a simple race function that we then used to implement a timeout mechanism. You saw how we can use Task.sleep to set a deadline for when our timeout should occur, and how we can use a task group to race two tasks.

We ended this post with a brief overview of task cancellation, and how not handling cancellation can lead to a less effective timeout mechanism. Cooperative cancellation is great but, in my opinion, it makes implementing features like task racing and timeouts a lot harder due to the guarantees made by Structured Concurrency.