turbopack_trace_utils/
tokio_workers.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2
3static ACTIVE_WORKER_THREADS: AtomicUsize = AtomicUsize::new(0);
4
5/// Initialize the count of non-parked scheduler workers before building the
6/// process's Tokio runtime. Tokio's thread start/stop hooks also include the
7/// blocking pool, while park/unpark hooks apply only to scheduler workers.
8///
9/// This assumes a single traced Tokio runtime per process. We do not adjust
10/// the count when the runtime shuts down.
11pub fn set_worker_threads(worker_threads: usize) {
12    ACTIVE_WORKER_THREADS.store(worker_threads, Ordering::Relaxed);
13}
14
15pub fn park() {
16    ACTIVE_WORKER_THREADS.fetch_sub(1, Ordering::Relaxed);
17}
18
19pub fn unpark() {
20    ACTIVE_WORKER_THREADS.fetch_add(1, Ordering::Relaxed);
21}
22
23pub fn active_worker_threads() -> usize {
24    ACTIVE_WORKER_THREADS.load(Ordering::Relaxed)
25}
26
27/// Match Tokio's default worker count when a runtime does not set it explicitly.
28pub fn default_worker_threads() -> usize {
29    match std::env::var("TOKIO_WORKER_THREADS") {
30        Ok(value) => {
31            let count: usize = value.parse().expect("TOKIO_WORKER_THREADS must be usize");
32            assert!(count > 0, "TOKIO_WORKER_THREADS cannot be set to 0");
33            count
34        }
35        Err(std::env::VarError::NotPresent) => {
36            std::thread::available_parallelism().map_or(1, |n| n.get())
37        }
38        Err(error) => panic!("TOKIO_WORKER_THREADS must be valid unicode: {error}"),
39    }
40}
41
42#[cfg(test)]
43mod tests {
44    use std::{
45        sync::{Arc, Barrier, mpsc},
46        time::{Duration, Instant},
47    };
48
49    use super::*;
50
51    /// Long enough to survive a busy CI machine, short enough to report a
52    /// failure instead of hanging the test run.
53    const TIMEOUT: Duration = Duration::from_secs(10);
54
55    /// Tokio's workers change the count from their own threads whenever they
56    /// happen to be scheduled, so anything that depends on them has to be
57    /// awaited rather than asserted on right away.
58    #[track_caller]
59    fn wait_for_active_worker_threads(expected: usize) {
60        let deadline = Instant::now() + TIMEOUT;
61        loop {
62            let active = active_worker_threads();
63            if active == expected {
64                return;
65            }
66            assert!(
67                Instant::now() < deadline,
68                "expected {expected} active worker thread(s), found {active}"
69            );
70            // Sleeping rather than spinning: the worker threads we are waiting
71            // for need the CPU more than this thread does.
72            std::thread::sleep(Duration::from_millis(1));
73        }
74    }
75
76    /// This covers the whole counting contract in one test on purpose:
77    /// `ACTIVE_WORKER_THREADS` is process-global, so separate tests would race
78    /// with each other when the test binary runs them in parallel threads.
79    #[test]
80    fn counts_scheduler_workers_but_not_blocking_threads() {
81        // The seeded count is only observable while the runtime does not exist
82        // yet. `Builder::build` spawns the workers before it returns and each of
83        // them parks as soon as it finds no work, so the count already starts
84        // dropping at a moment this thread does not control.
85        set_worker_threads(2);
86        assert_eq!(active_worker_threads(), 2);
87        park();
88        assert_eq!(active_worker_threads(), 1);
89        unpark();
90        assert_eq!(active_worker_threads(), 2);
91
92        let mut builder = tokio::runtime::Builder::new_multi_thread();
93        builder
94            .worker_threads(2)
95            .on_thread_park(park)
96            .on_thread_unpark(unpark);
97        let rt = builder.build().unwrap();
98
99        // Both scheduler workers park once the runtime is idle.
100        wait_for_active_worker_threads(0);
101
102        // Blocking threads are not scheduler workers and must not be counted:
103        // the count stays at zero while a blocking task occupies one of them.
104        let (blocking_tx, blocking_rx) = mpsc::channel();
105        rt.spawn_blocking(move || blocking_tx.send(active_worker_threads()).unwrap());
106        assert_eq!(blocking_rx.recv_timeout(TIMEOUT).unwrap(), 0);
107
108        // A spawned task unparks a scheduler worker, which is counted. The task
109        // holds the worker on the barrier so that it cannot park again before
110        // the count is read.
111        let barrier = Arc::new(Barrier::new(2));
112        let worker_barrier = barrier.clone();
113        let (started_tx, started_rx) = mpsc::channel();
114        rt.spawn(async move {
115            started_tx.send(()).unwrap();
116            worker_barrier.wait();
117        });
118        started_rx.recv_timeout(TIMEOUT).unwrap();
119        assert!(active_worker_threads() > 0);
120        barrier.wait();
121
122        // Park and unpark stay balanced, so the count returns to zero once the
123        // runtime is idle again.
124        wait_for_active_worker_threads(0);
125    }
126}
turbopack_trace_utils/tokio_workers.rs

turbopack_trace_utils/
tokio_workers.rs
