Handle the `Result` returned by tasks in `tokio_util`'s `TaskTracker`

I was initially spawning tasks on a JoinSet, and hooked onto the ctrl_c signal in the main thread to allow the user to shut down the server. I was aborting tasks when shutting down, which was not ideal since the tasks were handling websocket connections. Then I learned about CancellationToken and TaskTracker, which was a nice way to gracefully tell tasks to clean up and return. Unfortunately there's one regression, JoinSet has a join_next() method which I was using to handle any Result::Errs my tasks were returning. TaskTracker doesn't have such a method, and I couldn't find an equivalent way to do it. Any ideas? Maybe there is a solution and I missed it.

Edit: I thought about manually adding the JoinHandle returned by TaskTracker#spawn to a JoinSet, but that's not possible (I saw on GH Issues that the feature was not accepted). Anyways there's got to be a better way

what's the reason you choose TaskTracker instead of JoinSet? unless you need the special features of TaskTracker that JoinSet does not provide, why not just use CancellationToken with JoinSet like what you did before?

The entire reason that TaskTracker exists is: (from the docs)

  1. When tasks exit, a TaskTracker will allow the task to immediately free its memory.
  2. Some other stuff that's less important.

The first point is the most important one. A JoinSet keeps track of the return value of every inserted task. This means that if the caller keeps inserting tasks and never calls join_next, then their return values will keep building up and consuming memory, even if most of the tasks have already exited. This can cause the process to run out of memory. With a TaskTracker , this does not happen. Once tasks exit, they are immediately removed from the TaskTracker .

So if you want the return values of the tasks, you need to capture them using a strategy other than the TaskTracker. You could use a JoinSet?

You could even use both, see the examples on track_future:

use tokio::task::JoinSet;
use tokio_util::task::TaskTracker;

let tracker = TaskTracker::new();
let mut join_set = JoinSet::new();

join_set.spawn(tracker.track_future(my_async_fn()));
3 Likes

Ok I've realised I don't actually need TaskTracker at all, I was able to just use JoinSet#join_all to replace TaskTracker#wait and the rest of the program works like it did before.

1 Like