Tokio Gives Progress, Not Ordering: Scheduling 1M Tasks
This is a prequel to my previous article, Your Rust Service Isn't Leaking — It Could Be the Allocator.
In that post, I wrote about how a few memory allocators behaved differently under our workload. Before we figured out that the memory behavior was tied to allocators, we first tried to reduce the memory usage from the application side.
Our service was event-driven:
- Read events from a message queue (Kafka/Redis Streams/NATS)
- For each event, spawn a Tokio task to process it
Our Task Pattern
In our workload, each event had a maximum of 1000 user tokens. For every user token, we had to make an outbound call, wait on I/O and collect all responses belonging to the event.
Here is a simplified version of the code where we fan-out tokio tasks for user token, fan-in the responses, and generate a response event.
struct Event {
payload: Bytes, // ~4KB
user_tokens: Vec<String>, // up to 1000 tokens
// other fields
}
...
// in main
loop {
let event: Event = fetch_next_event().await;
tokio::spawn(async move {
let data = event.payload.clone();
let mut tasks = JoinSet::new();
for token in &event.user_tokens {
let token = token.clone();
let data = data.clone();
tasks.spawn(async move {
// hits an outbound API and returns response
process(token, data).await
});
}
let mut responses = Vec::with_capacity(event.user_tokens.len());
while let Some(res) = tasks.join_next().await {
responses.push(res)
}
generate_response_event(event, responses);
});
}Our primary focus here was throughput. While the fan-out above is unbounded, I assumed it practically wouldn't be a problem. These Tokio tasks are short-lived. They finish and drop once they receive responses, usually within a few milliseconds. So I assumed tasks spawned early would also finish early. Though individual outbound calls could finish out of order, I expected the earlier events overall to finish first, even while newer events were still coming in.
Our requirement was just to make each outbound call as fast as possible so that the burst could finish within the expected time.
What the Logs Looked Like
Here's how our logs looked during a burst of 1000 events, each containing ~1000 user tokens, spawning ~1M tasks in total:
started: event 1, user 5
started: event 1, user 8
started: event 1, user 2
started: event 2, user 6
started: event 2, user 4
started: event 3, user 8
...
finished: event 779
finished: event 976
started: event 900, user 42
started: event 900, user 261
started: event 1, user 974
started: event 1, user 831
...
finished: event 5
finished: event 3While the burst finished within the expected time, the logs showed that the token tasks from earlier events were being started much later.
I didn't expect strict ordering between tasks. They could complete in any order, and that was fine. What surprised me was how far some of the earlier tasks drifted from their submission order before receiving their first poll.
Inside Tokio's Scheduler
Tokio’s multi-threaded runtime has a fixed number of worker threads, a local queue per worker and a global queue shared across workers. Each worker has a local queue with a capacity of 256 tasks and on overflow it moves half of its tasks to the global queue. Workers prefer pulling from their own local queue first, check the global queue occasionally, and steal from other workers when they are idle.
When the tasks become independently schedulable, Tokio no longer knows which event created them. Now they are just runnable tasks which competing to be polled.
A simplified picture looks like this:
global queue
+------------------------------+
| * overflow from local queues |
| * remotely scheduled tasks |
| * mixed older/newer work |
+--------------+---------------+
|
+-------------------------+-------------------------+
| | |
v v v
worker 0 worker 1 worker 2
+-------------+ +-------------+ +-------------+
| local queue | | local queue | | local queue |
| up to 256 | | up to 256 | | up to 256 |
+------+------+ +------+------+ +------+------+
| | |
v v v
poll task poll task poll taskWith this architecture, once our event fans out into 1000 user token tasks, those tasks get mixed with the token tasks from other events, parent event tasks waiting on JoinSet, and tasks waking from I/O readiness.
With so many tasks being submitted at the same time, Tokio does not guarantee that earlier tasks get polled first. Tasks from different events compete for the available worker queues, and scheduling decisions like queue overflow and work stealing can change the order in which tasks are picked up. As a result, some tasks from earlier events received their first poll much later.
The core distinction is:
task created != task polled != task completedWhile Tokio continues making progress on every task, memory is affected by how many tasks are live at the same time.
Every Tokio task carries some state and though that state might not be huge individually, when thousands of these tasks are live at the same time, it adds up to the peak memory usage.
And since a few tasks from earlier events in our workload stayed alive till the end of the burst, their parent event tasks also remained alive, keeping the event state in memory, which increased our peak memory usage.
Task Creation Needs a Bound
Tokio guarantees fair scheduling under a bounded number of tasks and assuming no task blocks a worker thread.
Our code had no bound on task creation initially:
read events as fast as possible
└── spawn one event task per event
└── spawn up to 1000 token tasks per eventTokio takes as many tasks as you give it and keeps making progress on them. The boundedness assumed by Tokio for its fairness guarantee has to come from the application.
These bounds vary based on workload. In our case, what we wanted was event-level fairness. We wanted all the token tasks belonging to a single event to be polled and finish closer to their submission time.
To achieve this, we used a Semaphore to bound the number of events that can be processed at a time.
allow only N events into processing at once
└── spawn one event task per event
└── spawn up to 1000 token tasks per eventThe maximum number of tasks Tokio sees at a time, which affects how fair it can be, also depends on how long each task takes to return on poll. We did some trial-and-error to find the Semaphore count which worked for us.
An important detail here is that even with this bounded structure, the burst finished within the expected time while reducing the peak memory significantly. In general, it might seem that adding a Semaphore would impact the throughput of the service, but that wasn't the case for us.
Not a Tokio Problem
This wasn't a Tokio problem. Even with 1M live tasks, Tokio kept making progress until our burst finished on time. The difference was between my assumption (the event as a whole should finish early if started early) and what the Tokio runtime sees (a task is a task, whichever event it came from).
This was an easy miss for me because there was no correctness issue. There were no tasks which dropped midway or noticeable delays in the service, and our throughput requirements were also met. It just showed up as a higher-than-expected memory spike for reasons that didn't appear anywhere in the application logic.
Conclusion
Spawned early doesn't mean first-polled early. First-polled early doesn't mean completed early.
If there is an application-level unit of fairness you expect Tokio to honor, like the event in our case, the runtime doesn't know it exists. The bounds should be added by the application.
And before spawning a Tokio task, ask what the maximum number of live tasks at a time could be and see how that can impact memory. In some cases, it could also impact other resources which your tasks hold while they are alive.