commonware_runtime/iouring/runtime.rs
1//! Worker execution and runner-wide lifecycle.
2//!
3//! A [pool](super::pool) of workers polls ordinary tasks, worker zero on the
4//! calling thread and the others on threads of their own. Dedicated and
5//! blocking tasks each get a separate thread and io_uring ring. All workers
6//! share configuration and services through [`Shared`], while execution state
7//! stays in each worker's [`Local`].
8//!
9//! ## State and ownership
10//!
11//! ```text
12//! Context
13//! +-- Shared (Arc) <-------------------------- Local.shared
14//! | +-- Config, metrics, stop/panic signals
15//! | +-- Network, Storage
16//! | +-- Pool (each pool worker's mailbox and run queue, global queue,
17//! | | idle state, shutdown barrier)
18//! | | ^-- task headers (Weak)
19//! | +-- Tasks (every live ordinary task, for teardown)
20//! | `-- Workers <------------------------ ActiveWorker
21//! `-- supervision tree
22//!
23//! Worker (one per thread)
24//! +-- Scope --------------------------------> CURRENT (TLS)
25//! | |
26//! +-- Rc<RefCell<Local>> <------------------------+
27//! | +-- Shared (Arc)
28//! | +-- Ready (owner's end of a pool worker's run queue)
29//! | +-- Timers (sleep registrations and deadlines)
30//! | +-- Driver
31//! | | +-- io_uring ring (SQ/CQ)
32//! | | +-- Waiters (requests and unconsumed results)
33//! | | `-- ready queue, deadlines, cancellations (IDs)
34//! | +-- Mailbox (messages and wake source)
35//! | `-- Deferred (callbacks and retired owners)
36//! +-- inbox (batch taken from Mailbox)
37//! +-- deferred (batch detached from Local)
38//! `-- panics (first failure retained through cleanup)
39//! ```
40//!
41//! [`Scope`] installs checked thread-local access. Task polls, callbacks, and
42//! user destructors run outside local borrows. [`Worker::callbacks`] swaps the
43//! local [`Deferred`] batch into the worker before running it, so callbacks can
44//! reenter [`Local`] and append another batch.
45//!
46//! The root is pinned separately and passed to [`Worker::drive`]. On a
47//! dedicated worker, the dedicated or blocking task it was started for runs as
48//! its root. Pool workers after worker zero run a root that ends when worker
49//! zero stops the pool.
50//!
51//! ## Task placement
52//!
53//! Factories run on the spawning caller. Registration allocates an ordinary
54//! task and retains it in [`Tasks`] on the caller's thread, then delivers its
55//! first runnable. A pool worker keeps it when nothing else is queued there or
56//! when it is the pool's only worker. Every other spawn, including one from a
57//! busy pool worker, goes to the pool's global queue, which wakes a parked
58//! worker.
59//!
60//! A wake queues the task in the run queue of the pool worker that wakes it. A
61//! wake from outside the pool or on a closing worker pushes it into the global
62//! queue, which any pool worker takes from. A wake during a poll requeues the
63//! task on its poller. A queued task can still move, to the global queue when
64//! its run queue overflows or to a worker that steals it, as the
65//! [pool](super::pool) describes.
66//!
67//! A push that leaves a stealable backlog offers it to a parked worker unless
68//! an existing search count covers finding or handing on that work. A count
69//! can belong to a searching worker or to a wake that has reserved it but has
70//! not yet claimed a worker.
71//!
72//! ```text
73//! ordinary spawn:
74//! check Tasks --> factory --> Tasks::register
75//! +-- pool worker, quiet or alone --> its run queue
76//! `-- otherwise --> global queue --> any pool worker
77//! +-- first: poll
78//! `-- rest: run queue
79//!
80//! dedicated / blocking spawn:
81//! Workers::reserve --> factory --> new thread
82//! |
83//! v
84//! Worker::run_task
85//! |
86//! v
87//! drop ActiveWorker
88//! ```
89//!
90//! [`ActiveWorker`] keeps the launch counted from before factory construction
91//! until worker cleanup and failure reporting finish. It retains only the
92//! [`Workers`] barrier, allowing that worker's [`Shared`] owners to be released first.
93//!
94//! I/O and sleeps register directly with the worker that first polls them. When
95//! one is polled later on another worker, the registering worker forwards its
96//! result through a channel. Mailboxes carry root wakes, observation transfers,
97//! and cancellation messages.
98//!
99//! Each turn polls a bounded batch of tasks and checks the root's wake flag,
100//! services I/O and timers, and applies a bounded batch of mailbox messages.
101//! Deferred callbacks run between these phases.
102//!
103//! A pool worker takes a share of the global queue every
104//! [`Config::with_global_queue_interval`] takes and whenever its own run queue
105//! is empty. Taking a share returns one runnable to poll next and queues the
106//! rest in the run queue, no more than it has room for and no more than would
107//! fill it to half. A pool worker that finds both queues empty after a service
108//! pass tries to steal before parking, or relies on workers already holding
109//! search counts when the search cap refuses it.
110//!
111//! Before parking, the worker checks readiness again because polls and
112//! callbacks can produce more work. Pool workers coordinate through
113//! [`Searcher`], and all workers wait through their mailbox's
114//! [`Waker`](super::waker::Waker). See [`super::idle`] for parking order.
115//!
116//! A held search count suppresses kicks, so a worker returns it before user
117//! code and before another turn without parking. Every drive-loop exit,
118//! including unwinding, withdraws any live idle publication and returns a held
119//! count before root destruction.
120//!
121//! ## Shutdown
122//!
123//! Once worker zero is running, runner exit follows this order:
124//!
125//! ```text
126//! root completes or fails
127//! |
128//! v
129//! destroy root (TLS and mailbox still available)
130//! |
131//! v
132//! close worker registry, task set, and global queue
133//! |
134//! v
135//! abort spawned tasks through the supervision tree, stop the other pool workers
136//! |
137//! v
138//! every pool worker discards its run queue and drains the task set
139//! |
140//! v
141//! every pool worker waits for the others to finish polling and draining
142//! |
143//! v
144//! pool workers close mailbox and io_uring ring --> join the pool workers
145//! |
146//! v
147//! collect a pool worker failure the root missed
148//! |
149//! v
150//! wait for dedicated workers to finish
151//! |
152//! v
153//! return result or resume panic
154//! ```
155//!
156//! No pool worker closes its mailbox or io_uring ring until every pool worker
157//! has drained the set and finished its last poll, since a task polled on one
158//! worker can hold registrations on another, whose mailbox forwards their
159//! results. A dedicated worker closes its mailbox once its root is gone.
160//!
161//! Work spawned by root destruction participates in shutdown. Each worker keeps
162//! TLS installed while destroying tasks, draining retained writes and syncs,
163//! retiring cancellations, and running deferred callbacks. Native thread-local
164//! destruction on dedicated threads can follow registration release.
165
166use super::{
167 driver::Driver,
168 idle::Searcher,
169 mailbox::{Cancel, Forward, Mailbox, Message},
170 pool::{MAX_WORKERS, Pool, Threads},
171 request::{RequestOutput, RetiredResources},
172 sleep::{Sleep, Timers},
173 spinner::{Config as SpinnerConfig, Spinner},
174 task::{AfterPoll, Ready, RootWaker, Runnable},
175 tasks::Tasks,
176 timeout::TimeoutWheel,
177 waker::SUBMISSION_SEQ_MASK,
178};
179#[cfg(feature = "external")]
180use crate::Pacer;
181use crate::{
182 BlobLayout, BlobVersion, Clock, Error, Execution, Handle, METRICS_PREFIX, Name, SinkOf,
183 Spawner as _, StreamOf, Supervisor as _, child_label,
184 network::{
185 iouring::{Config as NetworkConfig, Network},
186 metered::Network as MeteredNetwork,
187 },
188 prefixed_name,
189 process::metered::Metrics as ProcessMetrics,
190 signal::Signal,
191 storage::{
192 iouring::{Config as StorageConfig, Storage},
193 metered::Storage as MeteredStorage,
194 },
195 telemetry::metrics::{
196 CounterFamily, Gauge, GaugeFamily, Metric, Register, Registered, Registry, add_attribute,
197 raw, task::Label, validate_label,
198 },
199 utils::{self, FactoryGuard, MetricHandle, Panicker, signal::Stopper, supervision::Tree},
200};
201use commonware_macros::select;
202use commonware_parallel::{Rayon, rayon::ThreadPoolBuilder};
203use commonware_utils::{
204 channel::oneshot,
205 sync::{Condvar, Mutex},
206 sys_rng,
207};
208use governor::clock::{Clock as GClock, ReasonablyRealtime};
209use rand_core::{Rng, TryCryptoRng, TryRng};
210use std::{
211 any::Any,
212 cell::RefCell,
213 convert::Infallible,
214 env,
215 future::{Future, poll_fn},
216 mem,
217 net::{IpAddr, SocketAddr, ToSocketAddrs},
218 num::{NonZeroU64, NonZeroUsize},
219 ops::RangeInclusive,
220 panic::{AssertUnwindSafe, catch_unwind, resume_unwind},
221 path::PathBuf,
222 pin::{Pin, pin},
223 ptr,
224 rc::Rc,
225 sync::{Arc, Weak},
226 task::{Context as TaskContext, Poll, Waker},
227 time::{Duration, Instant, SystemTime},
228};
229
230/// Maximum task polls or mailbox messages processed before yielding to other work.
231const BATCH_SIZE: usize = 64;
232
233/// Pinned future polled as a dedicated worker's root, or the runner's service
234/// task before it becomes a [`Task`](super::task::Task).
235type BoxedTask = Pin<Box<dyn Future<Output = ()> + Send>>;
236
237/// Configuration of each worker's io_uring instance and operation timing wheel.
238///
239/// The runtime requires Linux 6.1 or newer and always uses single-issuer mode
240/// with deferred task work. The same configuration applies to pool and
241/// dedicated workers. The wheel horizon is derived from network timeout policy.
242#[derive(Clone, Debug)]
243pub struct RingConfig {
244 /// SQ size and maximum outstanding operation SQEs, rounded up to a power of two.
245 ///
246 /// Must be nonzero and round to at most 32,768. Defaults to 128. The runtime
247 /// chooses 1024 for its production default and 128 when built for tests.
248 ///
249 /// The CQ uses the kernel default of twice the rounded SQ size. Cancellation
250 /// and mailbox wake CQEs also use that space, and cancellation acknowledgements
251 /// can outlive their requests. The operation limit therefore does not bound
252 /// all pending CQEs. The runtime relies on the kernel's CQ overflow backlog
253 /// (`IORING_FEAT_NODROP`), drained by subsequent completion-service enters.
254 pub size: u32,
255 /// Nonzero operation deadline granularity, defaulting to 5 milliseconds.
256 ///
257 /// Smaller ticks improve precision but increase wheel storage and service
258 /// frequency. Configuration must fit within 1,048,576 wheel slots.
259 pub timeout_wheel_tick: Duration,
260}
261
262impl Default for RingConfig {
263 fn default() -> Self {
264 Self {
265 size: 128,
266 timeout_wheel_tick: Duration::from_millis(5),
267 }
268 }
269}
270
271/// Configuration for the native io_uring runtime.
272///
273/// A pool of [`Self::with_worker_threads`] workers polls ordinary tasks, the
274/// first on the thread calling [`crate::Runner::start`]. Tasks marked dedicated
275/// or blocking each receive a fresh thread and io_uring ring. Shutdown
276/// completes retained writes and syncs without a configured time limit.
277#[derive(Clone)]
278pub struct Config {
279 /// Pool workers, including the calling thread.
280 worker_threads: usize,
281 /// Every Nth poll of a busy pool worker looks at the global queue first,
282 /// defaulting to 31.
283 global_queue_interval: u32,
284 /// Per-worker io_uring ring capacity and operation wheel tick.
285 ring_config: RingConfig,
286 /// Idle spinning policy, shared by all workers.
287 idle_spinner: SpinnerConfig,
288 /// Stack size for pool, dedicated worker, and Rayon threads.
289 thread_stack_size: usize,
290 /// Whether spawned-task panics and dedicated worker failures are caught.
291 /// Task-disposal panics during worker execution are contained with either setting.
292 catch_panics: bool,
293 /// Base directory held while storage resources or requests remain alive.
294 storage_directory: PathBuf,
295 /// Accepted blob layouts, defaulting to the complete supported range.
296 storage_blob_layouts: RangeInclusive<BlobLayout>,
297 /// Optional TCP_NODELAY override, defaulting to Some(true).
298 tcp_nodelay: Option<bool>,
299 /// Request immediate reset on socket close, defaulting to true.
300 zero_linger: bool,
301 /// Whole-call outbound connection deadline, defaulting to 10 seconds.
302 connect_timeout: Duration,
303 /// Whole-call send and receive deadline, defaulting to 60 seconds.
304 read_write_timeout: Duration,
305 /// Network receive buffering capacity, defaulting to 64 KiB.
306 read_buffer_size: usize,
307}
308
309impl Config {
310 /// Return the default runtime configuration with a generated temporary directory.
311 pub fn new() -> Self {
312 let ring_config = RingConfig {
313 size: if cfg!(test) { 128 } else { 1024 },
314 ..RingConfig::default()
315 };
316 let suffix = sys_rng().next_u64();
317
318 Self {
319 worker_threads: 2,
320 global_queue_interval: 31,
321 ring_config,
322 idle_spinner: SpinnerConfig::default(),
323 thread_stack_size: utils::thread::system_thread_stack_size(),
324 catch_panics: false,
325 storage_directory: env::temp_dir().join(format!("commonware_iouring_runtime_{suffix}")),
326 storage_blob_layouts: BlobLayout::ALL,
327 tcp_nodelay: Some(true),
328 zero_linger: true,
329 connect_timeout: Duration::from_secs(10),
330 read_write_timeout: Duration::from_secs(60),
331 read_buffer_size: 64 * 1024,
332 }
333 }
334
335 /// Set the number of pool workers, including the calling thread. Defaults
336 /// to 2.
337 ///
338 /// Must be between 1 and 64. Each worker has its own io_uring ring.
339 /// Dedicated and blocking threads are additional to this count.
340 pub const fn with_worker_threads(mut self, workers: usize) -> Self {
341 self.worker_threads = workers;
342 self
343 }
344
345 /// Set how often a busy pool worker looks at the global queue first: every
346 /// `interval`th poll. Defaults to 31.
347 ///
348 /// Must be nonzero. A smaller interval starts tasks in the global queue
349 /// sooner at the cost of more synchronization. At 1, a busy worker checks
350 /// the global queue before every poll and polls its own queued tasks only
351 /// while the global queue is empty.
352 pub const fn with_global_queue_interval(mut self, interval: u32) -> Self {
353 self.global_queue_interval = interval;
354 self
355 }
356
357 /// Set io_uring ring capacity and deadline granularity. See [`RingConfig`].
358 pub const fn with_ring_config(mut self, config: RingConfig) -> Self {
359 self.ring_config = config;
360 self
361 }
362
363 /// Set idle spinning policy. Use [`SpinnerConfig::disabled`] to disable it.
364 pub const fn with_idle_spinner(mut self, config: SpinnerConfig) -> Self {
365 self.idle_spinner = config;
366 self
367 }
368
369 /// Set the stack size of pool, dedicated worker, and Rayon threads. Worker
370 /// zero runs on the calling thread and keeps its stack.
371 pub const fn with_thread_stack_size(mut self, size: usize) -> Self {
372 self.thread_stack_size = size;
373 self
374 }
375
376 /// Set whether spawned-task panics and dedicated worker failures are caught.
377 ///
378 /// Caught failures are logged without interrupting the root. Propagated
379 /// failures are observed only while the root is executing. A failure of a
380 /// pool worker itself, such as an io_uring error, fails the runner with
381 /// either setting. It interrupts the root while the root runs, and
382 /// otherwise the runner resumes it after shutdown.
383 ///
384 /// A worker that fails to start, for example because its io_uring ring
385 /// cannot be created, reports the error through the panic hook. With
386 /// `panic = "abort"`, such a failure aborts the process with either
387 /// setting, like a failure to spawn the worker's thread.
388 ///
389 /// During worker execution, cancellation-time and task-disposal panics
390 /// escaping the user-poll wrapper are contained with either setting. An
391 /// unpublished result may resolve to [`Error::Closed`], while an
392 /// already-published result remains available.
393 /// Task factories execute synchronously and propagate panics to their caller.
394 pub const fn with_catch_panics(mut self, catch: bool) -> Self {
395 self.catch_panics = catch;
396 self
397 }
398
399 /// Set the outbound connection timeout, including queueing and retries.
400 ///
401 /// Must be nonzero and no greater than 30 years. The maximum configured
402 /// network timeout must fit the wheel slot limit in [`RingConfig`].
403 pub const fn with_connect_timeout(mut self, timeout: Duration) -> Self {
404 self.connect_timeout = timeout;
405 self
406 }
407
408 /// Set the send and receive timeout, including queueing and partial progress.
409 ///
410 /// Must be nonzero and no greater than 30 years. The maximum configured
411 /// network timeout must fit the wheel slot limit in [`RingConfig`].
412 pub const fn with_read_write_timeout(mut self, timeout: Duration) -> Self {
413 self.read_write_timeout = timeout;
414 self
415 }
416
417 /// Set the best-effort TCP_NODELAY override. None leaves the system default.
418 pub const fn with_tcp_nodelay(mut self, enabled: Option<bool>) -> Self {
419 self.tcp_nodelay = enabled;
420 self
421 }
422
423 /// Set whether sockets request zero linger when configured.
424 pub const fn with_zero_linger(mut self, enabled: bool) -> Self {
425 self.zero_linger = enabled;
426 self
427 }
428
429 /// Set the network receive buffer size in bytes.
430 pub const fn with_read_buffer_size(mut self, size: usize) -> Self {
431 self.read_buffer_size = size;
432 self
433 }
434
435 /// Set the storage directory, created and held when the runner starts.
436 pub fn with_storage_directory(mut self, directory: impl Into<PathBuf>) -> Self {
437 self.storage_directory = directory.into();
438 self
439 }
440
441 /// Set the blob layouts accepted by storage.
442 ///
443 /// New blobs use the latest layout in this range. Existing blobs outside
444 /// it fail to open with [`crate::Error::BlobLayoutMismatch`]. Restrict the
445 /// range to what a rollback target can read before first opening storage.
446 ///
447 /// # Panics
448 ///
449 /// Panics if `layouts` is empty.
450 pub fn with_storage_blob_layouts(mut self, layouts: RangeInclusive<BlobLayout>) -> Self {
451 assert!(
452 !layouts.is_empty(),
453 "storage blob layouts must be non-empty"
454 );
455 self.storage_blob_layouts = layouts;
456 self
457 }
458
459 /// Return the number of pool workers, including the calling thread.
460 pub const fn worker_threads(&self) -> usize {
461 self.worker_threads
462 }
463
464 /// Return how often a busy pool worker looks at the global queue first:
465 /// every Nth poll.
466 pub const fn global_queue_interval(&self) -> u32 {
467 self.global_queue_interval
468 }
469
470 /// Return per-worker io_uring ring configuration.
471 pub const fn ring_config(&self) -> &RingConfig {
472 &self.ring_config
473 }
474
475 /// Return per-worker idle spinning configuration.
476 pub const fn idle_spinner(&self) -> &SpinnerConfig {
477 &self.idle_spinner
478 }
479
480 /// Return the configured pool, dedicated worker, and Rayon thread stack size.
481 pub const fn thread_stack_size(&self) -> usize {
482 self.thread_stack_size
483 }
484
485 /// Return whether spawned-task panics and dedicated worker failures are caught.
486 pub const fn catch_panics(&self) -> bool {
487 self.catch_panics
488 }
489
490 /// Return the whole-call outbound connection timeout.
491 pub const fn connect_timeout(&self) -> Duration {
492 self.connect_timeout
493 }
494
495 /// Return the whole-call send and receive timeout.
496 pub const fn read_write_timeout(&self) -> Duration {
497 self.read_write_timeout
498 }
499
500 /// Return the TCP_NODELAY override.
501 pub const fn tcp_nodelay(&self) -> Option<bool> {
502 self.tcp_nodelay
503 }
504
505 /// Return whether sockets request zero linger.
506 pub const fn zero_linger(&self) -> bool {
507 self.zero_linger
508 }
509
510 /// Return the network receive buffer size.
511 pub const fn read_buffer_size(&self) -> usize {
512 self.read_buffer_size
513 }
514
515 /// Return the storage directory retained by the runtime's storage resources.
516 pub const fn storage_directory(&self) -> &PathBuf {
517 &self.storage_directory
518 }
519
520 /// Return the accepted blob layout range.
521 pub const fn storage_blob_layouts(&self) -> &RangeInclusive<BlobLayout> {
522 &self.storage_blob_layouts
523 }
524
525 /// Validate and normalize settings before acquiring any startup resources.
526 fn validate(&mut self) {
527 assert!(
528 (1..=MAX_WORKERS).contains(&self.worker_threads),
529 "worker threads ({}) must be between 1 and {MAX_WORKERS}",
530 self.worker_threads,
531 );
532 assert!(
533 self.global_queue_interval != 0,
534 "global queue interval must be nonzero"
535 );
536 assert!(self.ring_config.size != 0, "ring size must be nonzero");
537 self.ring_config.size = self
538 .ring_config
539 .size
540 .checked_next_power_of_two()
541 .expect("ring size overflow");
542 assert!(self.ring_config.size <= 32_768, "ring size exceeds 32768");
543 assert!(
544 self.idle_spinner.budget_us <= self.idle_spinner.max_budget_us,
545 "spinner budget_us ({}) must not exceed max_budget_us ({})",
546 self.idle_spinner.budget_us,
547 self.idle_spinner.max_budget_us,
548 );
549 assert!(
550 !self.storage_blob_layouts.is_empty(),
551 "storage blob layouts must be non-empty"
552 );
553 assert!(
554 !self.connect_timeout.is_zero(),
555 "connect timeout must be nonzero"
556 );
557 assert!(
558 !self.read_write_timeout.is_zero(),
559 "read/write timeout must be nonzero"
560 );
561 TimeoutWheel::validate_layout(self.max_timeout(), self.ring_config.timeout_wheel_tick)
562 .expect("invalid io_uring operation deadline layout");
563 }
564
565 /// Largest network timeout supported by every worker's operation wheel.
566 fn max_timeout(&self) -> Duration {
567 self.connect_timeout.max(self.read_write_timeout)
568 }
569}
570
571impl Default for Config {
572 fn default() -> Self {
573 Self::new()
574 }
575}
576
577/// Task counters shared by all workers in one runner.
578struct TaskMetrics {
579 /// Number of tasks created, including rejected spawns.
580 tasks_spawned: CounterFamily<Label>,
581 /// Number of tasks that have not completed or been aborted.
582 tasks_running: GaugeFamily<Label>,
583}
584
585impl TaskMetrics {
586 /// Register task families beneath the runtime metrics namespace.
587 fn new(registry: &mut impl Register) -> Self {
588 Self {
589 tasks_spawned: registry.register(
590 "tasks_spawned",
591 "Total number of tasks spawned",
592 raw::Family::default(),
593 ),
594 tasks_running: registry.register(
595 "tasks_running",
596 "Number of tasks currently running",
597 raw::Family::default(),
598 ),
599 }
600 }
601}
602
603/// Registration and cleanup barrier for dedicated workers.
604#[derive(Default)]
605struct Workers {
606 /// Registration gate and count protected by the same lock.
607 state: Mutex<WorkerCount>,
608 /// Notify shutdown when the last registered worker finishes cleanup.
609 idle: Condvar,
610}
611
612/// Dedicated workers that shutdown must still wait for.
613#[derive(Default)]
614struct WorkerCount {
615 /// Whether shutdown has stopped accepting new workers.
616 closed: bool,
617 /// Accepted workers whose factories, execution, or cleanup have not finished.
618 active: usize,
619}
620
621impl Workers {
622 /// Reserve a worker before its factory runs, unless registration has closed.
623 fn reserve(self: &Arc<Self>) -> Option<ActiveWorker> {
624 let mut state = self.state.lock();
625 if state.closed {
626 return None;
627 }
628
629 // Closing registration under this same lock cannot miss an accepted launch.
630 state.active += 1;
631 Some(ActiveWorker(self.clone()))
632 }
633
634 /// Stop accepting workers while retaining every existing registration.
635 fn close(&self) {
636 self.state.lock().closed = true;
637 }
638
639 /// Wait until every accepted worker has finished runtime cleanup.
640 /// Call after [`Self::close`] so no new reservation can extend the count.
641 fn wait(&self) {
642 let mut state = self.state.lock();
643 while state.active != 0 {
644 self.idle.wait(&mut state);
645 }
646 }
647}
648
649/// Count an accepted worker until its runtime cleanup finishes.
650///
651/// Retains only [`Workers`], allowing Shared and its storage hold to be released
652/// before the worker count reaches zero.
653struct ActiveWorker(
654 /// Registry whose active count includes this reservation.
655 Arc<Workers>,
656);
657
658impl Drop for ActiveWorker {
659 fn drop(&mut self) {
660 let mut state = self.0.state.lock();
661 state.active -= 1;
662 if state.active == 0 {
663 self.0.idle.notify_all();
664 }
665 drop(state);
666
667 #[cfg(test)]
668 tests::after_release(&self.0);
669 }
670}
671
672/// Runner-wide services shared by pool and dedicated workers.
673pub struct Shared {
674 /// Validated configuration, immutable after startup.
675 pub cfg: Config,
676 /// Scheduling state of the pool workers.
677 pub pool: Arc<Pool>,
678 /// User-visible metrics registry.
679 registry: Registry,
680 /// Task counters and running gauges.
681 metrics: TaskMetrics,
682 /// Aggregate active requests, updated with each worker's own count delta.
683 pending_operations: Gauge,
684 /// Stop signal and acknowledgement state.
685 shutdown: Mutex<Stopper>,
686 /// User task panic policy and root notification.
687 panicker: Panicker,
688 /// Synchronized creation and closure of dedicated workers.
689 workers: Arc<Workers>,
690 /// Every live ordinary task, retained so the pool workers' cleanup can
691 /// drop its future.
692 tasks: Tasks,
693 /// Metered storage with its metadata lock and directory hold.
694 storage: MeteredStorage<Storage>,
695 /// Metered native socket adapter.
696 network: MeteredNetwork<Network>,
697}
698
699/// Task and runtime ownership transferred to a dedicated worker's thread.
700///
701/// Field order destroys a rejected task and Shared before the reservation is released.
702struct Launch {
703 /// Task to execute, or destroy if thread creation fails.
704 task: BoxedTask,
705 /// Runtime services retained until execution and cleanup finish.
706 shared: Arc<Shared>,
707 /// Registration released after the preceding owners are destroyed.
708 active: ActiveWorker,
709}
710
711impl Shared {
712 /// Transfer a reserved worker's task and its cleanup responsibility to a new thread.
713 fn launch(self: &Arc<Self>, task: BoxedTask, active: ActiveWorker) {
714 let payload = Launch {
715 task,
716 shared: self.clone(),
717 active,
718 };
719
720 #[cfg(test)]
721 let payload = tests::before_launch(payload);
722
723 utils::thread::spawn(self.cfg.thread_stack_size, move || {
724 Worker::run_task(payload)
725 });
726 }
727}
728
729/// Runtime capabilities and supervision context for the current task.
730///
731/// Ordinary children run on the pool, including those spawned by dedicated and
732/// blocking tasks. Task factories run on their caller, while ordinary futures
733/// run on the pool and can move between its workers, as can resources and
734/// pending I/O and sleep futures.
735///
736/// Registrations stay on their original worker. Closing that worker causes
737/// unresolved I/O futures to fail and unresolved sleeps to panic when polled.
738/// Pool workers close only at shutdown, after every pool task has been dropped,
739/// while dedicated and blocking tasks may still run. Detached sync completion
740/// handles can be awaited on any thread.
741pub struct Context {
742 /// User-facing task and metric namespace.
743 name: String,
744 /// Validated metric attributes inherited by children.
745 attributes: Vec<(String, String)>,
746 /// Shared services and runner-wide lifecycle state.
747 shared: Arc<Shared>,
748 /// This context's node in the mandatory supervision tree.
749 tree: Arc<Tree>,
750 /// Placement requested for the next consumed spawn.
751 execution: Execution,
752}
753
754impl Context {
755 /// Access shared task metric families for the common spawn helper.
756 fn metrics(&self) -> &TaskMetrics {
757 &self.shared.metrics
758 }
759}
760
761impl crate::Spawner for Context {
762 fn dedicated(mut self) -> Self {
763 self.execution = Execution::Dedicated;
764 self
765 }
766
767 fn shared(mut self, blocking: bool) -> Self {
768 self.execution = Execution::Shared(blocking);
769 self
770 }
771
772 fn spawn<F, Fut, T>(mut self, f: F) -> Handle<T>
773 where
774 F: FnOnce(Self) -> Fut + Send + 'static,
775 Fut: Future<Output = T> + Send + 'static,
776 T: Send + 'static,
777 {
778 let (_, metric) = spawn_metrics!(self);
779 let parent = self.tree.clone();
780 let execution = self.execution;
781
782 // Placement applies to this spawn. Children start with ordinary placement.
783 self.execution = Execution::default();
784 let (child, aborted) = Tree::child(&parent);
785 if aborted {
786 return Handle::closed(metric);
787 }
788
789 self.tree = child;
790 let shared = self.shared.clone();
791
792 // Dedicated and blocking spawns reserve a dedicated worker before the
793 // factory runs, so shutdown waits for them. Ordinary spawns only check
794 // that the task set still accepts tasks.
795 let reservation = if matches!(execution, Execution::Dedicated | Execution::Shared(true)) {
796 let Some(reservation) = shared.workers.reserve() else {
797 return Handle::closed(metric);
798 };
799 Some(reservation)
800 } else {
801 if shared.tasks.is_closed() {
802 return Handle::closed(metric);
803 }
804 None
805 };
806
807 // User construction runs on the caller with no runtime borrow or lock.
808 // A reserved dedicated worker remains counted through construction and
809 // launch, including when the factory unwinds or shutdown closes the
810 // registry.
811 let guard = FactoryGuard::new(&parent, metric);
812 let future = f(self);
813
814 // The execution wrapper takes over cleanup once the factory returns.
815 let (future, handle) = Handle::init(
816 future,
817 guard.disarm(),
818 shared.panicker.clone(),
819 parent.clone(),
820 );
821
822 // A dedicated worker polls the future as its root. An ordinary spawn
823 // becomes a task of the pool.
824 if let Some(reservation) = reservation {
825 shared.launch(Box::pin(future), reservation);
826 } else if let Err(task) = shared.tasks.register(future, Arc::downgrade(&shared.pool)) {
827 // Rejection after closure follows the caller's panic boundary.
828 // Cancel descendants and finish metrics before destroying captures.
829 parent.abort();
830 task.clear();
831 }
832 handle
833 }
834
835 async fn stop(self, value: i32, timeout: Option<Duration>) -> Result<(), Error> {
836 let resolved = self.shared.shutdown.lock().stop(value);
837 let timeout = timeout.map_or_else(
838 || futures::future::Either::Right(futures::future::pending()),
839 |duration| futures::future::Either::Left(self.sleep(duration)),
840 );
841 select! {
842 result = resolved => result.map_err(|_| Error::Closed),
843 _ = timeout => Err(Error::Timeout),
844 }
845 }
846
847 fn stopped(&self) -> Signal {
848 self.shared.shutdown.lock().stopped()
849 }
850}
851
852impl crate::Strategizer for Context {
853 fn strategy(&self, parallelism: NonZeroUsize) -> Rayon {
854 let pool = ThreadPoolBuilder::new()
855 .num_threads(parallelism.get())
856 .stack_size(self.shared.cfg.thread_stack_size)
857 .build()
858 .expect("failed to create io_uring Rayon thread pool");
859 Rayon::with_pool(Arc::new(pool))
860 }
861}
862
863impl crate::Supervisor for Context {
864 fn child(&self, label: &'static str) -> Self {
865 let (tree, _) = Tree::child(&self.tree);
866 Self {
867 name: child_label(&self.name, label),
868 attributes: self.attributes.clone(),
869 shared: self.shared.clone(),
870 tree,
871 execution: Execution::default(),
872 }
873 }
874
875 fn with_attribute(mut self, key: &'static str, value: impl std::fmt::Display) -> Self {
876 validate_label(key);
877 add_attribute(&mut self.attributes, key, value);
878 self
879 }
880
881 fn name(&self) -> Name {
882 Name {
883 label: self.name.clone(),
884 attributes: self.attributes.clone(),
885 }
886 }
887}
888
889impl crate::Metrics for Context {
890 fn register<N: Into<String>, H: Into<String>, M: Metric>(
891 &self,
892 name: N,
893 help: H,
894 metric: M,
895 ) -> Registered<M> {
896 self.shared.registry.register(
897 prefixed_name(&self.name, &name.into()),
898 help.into(),
899 self.attributes.clone(),
900 Arc::new(metric),
901 )
902 }
903
904 fn encode(&self) -> String {
905 self.shared.registry.encode()
906 }
907}
908
909impl Clock for Context {
910 fn current(&self) -> SystemTime {
911 SystemTime::now()
912 }
913
914 fn sleep(&self, duration: Duration) -> impl Future<Output = ()> + Send + 'static + use<> {
915 Sleep::new(duration)
916 }
917
918 fn sleep_until(
919 &self,
920 deadline: SystemTime,
921 ) -> impl Future<Output = ()> + Send + 'static + use<> {
922 Sleep::until(deadline)
923 }
924}
925
926#[cfg(feature = "external")]
927impl Pacer for Context {
928 fn pace<'a, F, T>(
929 &'a self,
930 _latency: Duration,
931 future: F,
932 ) -> impl Future<Output = T> + Send + 'a
933 where
934 F: Future<Output = T> + Send + 'a,
935 T: Send + 'a,
936 {
937 future
938 }
939}
940
941impl GClock for Context {
942 type Instant = SystemTime;
943
944 fn now(&self) -> Self::Instant {
945 self.current()
946 }
947}
948
949impl ReasonablyRealtime for Context {}
950
951impl crate::Network for Context {
952 type Listener = <MeteredNetwork<Network> as crate::Network>::Listener;
953
954 async fn bind(&self, socket: SocketAddr) -> Result<Self::Listener, Error> {
955 self.shared.network.bind(socket).await
956 }
957
958 async fn dial(&self, socket: SocketAddr) -> Result<(SinkOf<Self>, StreamOf<Self>), Error> {
959 self.shared.network.dial(socket).await
960 }
961}
962
963impl crate::Resolver for Context {
964 async fn resolve(&self, host: &str) -> Result<Vec<IpAddr>, Error> {
965 let host = host.to_owned();
966 self.child("resolver")
967 .shared(true)
968 .spawn(move |_| async move {
969 (host.as_str(), 0)
970 .to_socket_addrs()
971 .map(|addresses| addresses.map(|address| address.ip()).collect())
972 .map_err(|error| Error::ResolveFailed(error.to_string()))
973 })
974 .await
975 .map_err(|error| Error::ResolveFailed(error.to_string()))?
976 }
977}
978
979impl TryRng for Context {
980 type Error = Infallible;
981
982 fn try_next_u32(&mut self) -> Result<u32, Self::Error> {
983 Ok(sys_rng().next_u32())
984 }
985
986 fn try_next_u64(&mut self) -> Result<u64, Self::Error> {
987 Ok(sys_rng().next_u64())
988 }
989
990 fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Self::Error> {
991 sys_rng().fill_bytes(dest);
992 Ok(())
993 }
994}
995
996impl TryCryptoRng for Context {}
997
998impl crate::Storage for Context {
999 type Blob = <MeteredStorage<Storage> as crate::Storage>::Blob;
1000
1001 async fn open_versioned(
1002 &self,
1003 partition: &str,
1004 name: &[u8],
1005 versions: RangeInclusive<BlobVersion>,
1006 ) -> Result<(Self::Blob, u64, BlobVersion), Error> {
1007 self.shared
1008 .storage
1009 .open_versioned(partition, name, versions)
1010 .await
1011 }
1012
1013 async fn remove(&self, partition: &str, name: Option<&[u8]>) -> Result<(), Error> {
1014 self.shared.storage.remove(partition, name).await
1015 }
1016
1017 async fn scan(&self, partition: &str) -> Result<Vec<Vec<u8>>, Error> {
1018 self.shared.storage.scan(partition).await
1019 }
1020}
1021
1022/// Sync or timer completion notification published outside Local.
1023type Completion = (oneshot::Sender<Result<(), Error>>, Result<(), Error>);
1024
1025/// I/O result forwarded to another thread, including any returned buffers or descriptors.
1026type OperationResult = (
1027 oneshot::Sender<Result<RequestOutput, Error>>,
1028 Result<RequestOutput, Error>,
1029);
1030
1031/// Which work a worker thread runs.
1032#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1033pub enum Role {
1034 /// Pool worker at this index, polling ordinary tasks. Worker zero runs on
1035 /// the thread calling [`crate::Runner::start`] and shuts the runner down.
1036 Pool(u32),
1037 /// A worker with a thread and io_uring ring of its own, started for one
1038 /// dedicated or blocking task that it polls as its root.
1039 Dedicated,
1040}
1041
1042impl Role {
1043 /// Whether this worker polls the pool's ordinary tasks.
1044 pub const fn is_pool(self) -> bool {
1045 matches!(self, Self::Pool(_))
1046 }
1047
1048 /// Whether this is worker zero, which polls the runner's root and shuts the
1049 /// runner down.
1050 pub const fn owns_runner(self) -> bool {
1051 matches!(self, Self::Pool(0))
1052 }
1053}
1054
1055/// A pool worker's seat in the pool: its index, the task set whose tasks it
1056/// polls, and its run queue. Dedicated workers have no seat.
1057struct Seat {
1058 /// Index of this worker in the pool, as in [`Role::Pool`].
1059 index: u32,
1060 /// Identity of the pool's task set, which [`Local::serves`] compares with a
1061 /// task's owner.
1062 tasks: NonZeroU64,
1063 /// The owner's end of this worker's FIFO run queue, whose shared end the
1064 /// pool keeps for thieves.
1065 ready: Ready,
1066}
1067
1068/// Mutable execution state accessed only by its owning worker thread.
1069///
1070/// Polls and callbacks obtain this state through [`Self::current`] or
1071/// [`Self::with_current`]. They release each borrow before running user code,
1072/// including waker methods and destructors.
1073pub struct Local {
1074 /// This worker's seat, or `None` on a dedicated worker, which polls only
1075 /// its root.
1076 seat: Option<Seat>,
1077 /// The io_uring ring's owner, taken only after kernel retirement so it can
1078 /// be dropped unborrowed.
1079 pub driver: Option<Driver>,
1080 /// Sleeper registrations and deadlines.
1081 pub timers: Timers,
1082 /// Reject new registration and send wakes to the global queue, while
1083 /// allowing idempotent cancellation.
1084 pub closing: bool,
1085 /// Shared monotonic sample for the current service turn.
1086 pub now: Instant,
1087 /// Root wake flag, set by local wakes or applied mailbox messages and
1088 /// taken before each root poll so a wake during the poll is kept.
1089 pub root_ready: bool,
1090 /// Strong mailbox ownership retained through kernel retirement.
1091 pub mailbox: Arc<Mailbox>,
1092 /// Callbacks and retired owners collected while Local is borrowed.
1093 pub deferred: Deferred,
1094 /// Runner configuration, metrics, and shared adapters.
1095 shared: Arc<Shared>,
1096 /// This worker's contribution currently included in the aggregate gauge.
1097 reported_pending: usize,
1098}
1099
1100impl Local {
1101 /// Construct an io_uring ring on the thread that will own all its
1102 /// submissions. A pool worker takes the mailbox the pool created for it and
1103 /// the owner's end of its run queue, and panics if its index is out of
1104 /// range or that end was already taken. A dedicated worker creates its own
1105 /// mailbox and has no run queue. An I/O error at startup returns before the
1106 /// run queue's owner is taken.
1107 fn new(shared: Arc<Shared>, role: Role) -> std::io::Result<Self> {
1108 #[cfg(test)]
1109 tests::before_startup(&shared.workers)?;
1110
1111 let mailbox = match role {
1112 Role::Pool(index) => shared.pool.mailbox(index).clone(),
1113 Role::Dedicated => Arc::new(Mailbox::new()?),
1114 };
1115 let now = Instant::now();
1116 let driver = Driver::new(
1117 &shared.cfg.ring_config,
1118 shared.cfg.max_timeout(),
1119 mailbox.waker.clone(),
1120 now,
1121 )?;
1122
1123 // Only this worker pushes and pops, while the pool keeps the shared
1124 // end for thieves.
1125 let seat = match role {
1126 Role::Pool(index) => Some(Seat {
1127 index,
1128 tasks: shared.tasks.id(),
1129 ready: shared.pool.queue(index).owner(),
1130 }),
1131 Role::Dedicated => None,
1132 };
1133 Ok(Self {
1134 seat,
1135 driver: Some(driver),
1136 timers: Timers::default(),
1137 closing: false,
1138 now,
1139 root_ready: true,
1140 mailbox,
1141 deferred: Deferred::default(),
1142 shared,
1143 reported_pending: 0,
1144 })
1145 }
1146
1147 /// Return the worker installed on this thread, including during cleanup.
1148 /// Returns `None` when no worker is installed or its TLS key has been destroyed.
1149 pub fn current() -> Option<Rc<RefCell<Self>>> {
1150 // Escaped wakers can run from native TLS destructors after this key is gone.
1151 CURRENT
1152 .try_with(|current| current.borrow().clone())
1153 .ok()
1154 .flatten()
1155 }
1156
1157 /// Return the current worker if it owns `mailbox`, including during shutdown.
1158 #[inline(always)]
1159 pub fn owner(mailbox: &Weak<Mailbox>) -> Option<Rc<RefCell<Self>>> {
1160 let local = Self::current()?;
1161
1162 // The weak reference preserves allocation identity without retaining a worker.
1163 let matches = ptr::eq(Arc::as_ptr(&local.borrow().mailbox), mailbox.as_ptr());
1164 matches.then_some(local)
1165 }
1166
1167 /// Run `f` with the worker installed on this thread, or with `None` when
1168 /// no worker is installed or its TLS key has been destroyed. The worker is
1169 /// borrowed in place rather than through a clone of its handle, so the
1170 /// caller must hold no `Local` borrow, and `f` must run no user code while
1171 /// it holds the worker.
1172 #[inline]
1173 pub fn with_current<R>(f: impl FnOnce(Option<&mut Self>) -> R) -> R {
1174 let mut f = Some(f);
1175
1176 // Escaped wakers can run from native TLS destructors after this key is gone.
1177 let result = CURRENT.try_with(|current| {
1178 let current = current.borrow();
1179 let local = current.as_ref()?;
1180 let f = f.take().expect("the closure must be taken only once");
1181 Some(f(Some(&mut local.borrow_mut())))
1182 });
1183
1184 if let Ok(Some(result)) = result {
1185 return result;
1186 }
1187
1188 // Without a worker, `f` runs outside the TLS borrow.
1189 let f = f.take().expect("the closure must be taken only once");
1190 f(None)
1191 }
1192
1193 /// Whether this worker polls tasks retained by `owner` and is not closing.
1194 /// A closing worker polls nothing more, so it serves no task set.
1195 #[inline]
1196 pub fn serves(&self, owner: NonZeroU64) -> bool {
1197 self.seat.as_ref().is_some_and(|seat| seat.tasks == owner) && !self.closing
1198 }
1199
1200 /// Which work this worker runs.
1201 pub const fn role(&self) -> Role {
1202 match &self.seat {
1203 Some(seat) => Role::Pool(seat.index),
1204 None => Role::Dedicated,
1205 }
1206 }
1207
1208 /// Workers in this worker's pool.
1209 pub fn workers(&self) -> usize {
1210 self.shared.pool.workers()
1211 }
1212
1213 /// This pool worker's run queue. Panics on a dedicated worker.
1214 pub const fn ready(&self) -> &Ready {
1215 &self
1216 .seat
1217 .as_ref()
1218 .expect("the worker must be a pool worker")
1219 .ready
1220 }
1221
1222 /// This pool worker's run queue, mutably. Panics on a dedicated worker.
1223 const fn ready_mut(&mut self) -> &mut Ready {
1224 &mut self
1225 .seat
1226 .as_mut()
1227 .expect("the worker must be a pool worker")
1228 .ready
1229 }
1230
1231 /// Queue a runnable on this pool worker, offering stealable backlog or
1232 /// exporting overflow through [`Pool::enqueue`]. The runnable must belong
1233 /// to this pool. Panics on a dedicated worker.
1234 #[inline]
1235 pub fn enqueue(&mut self, runnable: Runnable) {
1236 let Self { seat, shared, .. } = self;
1237 let seat = seat.as_mut().expect("the worker must be a pool worker");
1238 shared.pool.enqueue(&mut seat.ready, runnable);
1239 }
1240
1241 /// Release an operation or timer on its worker, directly or through its mailbox.
1242 ///
1243 /// A worker whose mailbox is closed or gone has already taken responsibility
1244 /// for cleanup.
1245 pub fn cancel(mailbox: &Weak<Mailbox>, cancel: Cancel) {
1246 if let Some(local) = Self::owner(mailbox) {
1247 // Owner-local drops also reach this path after the mailbox has closed.
1248 local.borrow_mut().release(cancel);
1249 return;
1250 }
1251
1252 if let Some(mailbox) = mailbox.upgrade() {
1253 // If closure wins the race, worker cleanup will release the registration.
1254 let _ = mailbox.send(Message::Cancel(cancel));
1255 }
1256 }
1257
1258 /// Detach observation without running callbacks under the local borrow.
1259 fn release(&mut self, cancel: Cancel) {
1260 match cancel {
1261 Cancel::Waiter(id) => self.driver.as_mut().unwrap().orphan(id, &mut self.deferred),
1262 Cancel::Timer(id) => self.timers.cancel(id, &mut self.deferred),
1263 }
1264 }
1265
1266 /// Update aggregate pending-operation metrics using only this worker's delta.
1267 fn update_pending(&mut self) {
1268 // Count active requests, including queued ones. Completed results waiting
1269 // for their callers no longer contribute to the pending count.
1270 let pending = self.driver.as_ref().unwrap().len();
1271
1272 // Other workers update the same gauge. Apply only our change so their
1273 // contributions remain intact.
1274 if pending > self.reported_pending {
1275 self.shared
1276 .pending_operations
1277 .inc_by((pending - self.reported_pending) as _);
1278 } else if pending < self.reported_pending {
1279 self.shared
1280 .pending_operations
1281 .dec_by((self.reported_pending - pending) as _);
1282 }
1283 self.reported_pending = pending;
1284 }
1285
1286 /// Whether task polling or callbacks prevent the worker from parking. A
1287 /// pool worker also has the pool's global queue to take from.
1288 fn is_ready(&self) -> bool {
1289 let seat = self.seat.as_ref();
1290 seat.is_some_and(|seat| !seat.ready.is_empty())
1291 || self.root_ready
1292 || !self.deferred.is_empty()
1293 || (seat.is_some() && self.shared.pool.has_global())
1294 }
1295
1296 /// Earliest absolute deadline across driver requests and sleepers.
1297 fn next_deadline(&mut self) -> Option<Instant> {
1298 [
1299 self.driver.as_mut().unwrap().next_deadline(),
1300 self.timers.next_deadline(),
1301 ]
1302 .into_iter()
1303 .flatten()
1304 .min()
1305 }
1306}
1307
1308thread_local! {
1309 /// Checked access to the one worker allowed on the current thread.
1310 static CURRENT: RefCell<Option<Rc<RefCell<Local>>>> = const { RefCell::new(None) };
1311}
1312
1313/// Thread-local worker access, cleared after mandatory cleanup on every exit.
1314struct Scope;
1315
1316impl Scope {
1317 /// Reject nesting before any directory hold, thread, or io_uring ring is
1318 /// created.
1319 fn assert_vacant() {
1320 CURRENT.with(|current| {
1321 assert!(
1322 current.borrow().is_none(),
1323 "nested io_uring runtime entry is not supported"
1324 );
1325 });
1326 }
1327
1328 /// Install one worker, retaining the entry check as a defensive invariant.
1329 fn install(local: Rc<RefCell<Local>>) -> Self {
1330 CURRENT.with(|current| {
1331 let mut current = current.borrow_mut();
1332 assert!(current.is_none(), "nested io_uring worker scope");
1333 *current = Some(local);
1334 });
1335 Self
1336 }
1337}
1338
1339impl Drop for Scope {
1340 fn drop(&mut self) {
1341 CURRENT.with(|current| {
1342 // Worker still owns Local, so clearing TLS cannot destroy it under this borrow.
1343 current.borrow_mut().take();
1344 });
1345 }
1346}
1347
1348/// Panic payload reported after mandatory worker cleanup.
1349pub type Panic = Box<dyn Any + Send + 'static>;
1350
1351/// Retain the first panic while allowing the remaining cleanup to run.
1352///
1353/// Later payloads are leaked because their destructors may panic. An unclaimed
1354/// first payload is also leaked when this accumulator is dropped, including
1355/// during an existing unwind.
1356#[derive(Default)]
1357pub struct Panics {
1358 /// First failure, held until the worker can report it.
1359 first: Option<Panic>,
1360}
1361
1362impl Panics {
1363 /// Run a task poll or destructor, returning `None` if it panics.
1364 ///
1365 /// Cancellation can destroy the user future during polling, so both paths need
1366 /// this boundary. The caller must release worker borrows before invoking it.
1367 pub fn contain<T>(f: impl FnOnce() -> T) -> Option<T> {
1368 cfg_if::cfg_if! {
1369 if #[cfg(feature = "loom")] {
1370 // Loom reports a violation by panicking, and containing that
1371 // panic would hide it, so loom builds let every panic through.
1372 let result: std::thread::Result<T> = Ok(f());
1373 } else {
1374 let result = catch_unwind(AssertUnwindSafe(f));
1375 }
1376 }
1377 match result {
1378 Ok(output) => Some(output),
1379 Err(panic) => {
1380 // Payload destruction can also run user code. A secondary panic
1381 // cannot escape this task boundary or replace a worker failure.
1382 if let Err(secondary) = catch_unwind(AssertUnwindSafe(|| drop(panic))) {
1383 mem::forget(secondary);
1384 }
1385 None
1386 }
1387 }
1388 }
1389
1390 /// Run a callback without letting its panic interrupt cleanup.
1391 fn run(&mut self, f: impl FnOnce()) {
1392 if let Err(panic) = catch_unwind(AssertUnwindSafe(f)) {
1393 self.retain(panic);
1394 }
1395 }
1396
1397 /// Retain the first payload and leak any later ones.
1398 fn retain(&mut self, panic: Panic) {
1399 if self.first.is_none() {
1400 self.first = Some(panic);
1401 } else {
1402 mem::forget(panic);
1403 }
1404 }
1405
1406 /// Take the retained panic.
1407 fn take(&mut self) -> Option<Panic> {
1408 self.first.take()
1409 }
1410}
1411
1412impl Drop for Panics {
1413 fn drop(&mut self) {
1414 mem::forget(self.first.take());
1415 }
1416}
1417
1418/// Abort if cleanup exits before kernel-visible resources can be released.
1419/// Forget this guard only after all outstanding operations have retired.
1420struct RetirementGuard;
1421
1422impl Drop for RetirementGuard {
1423 fn drop(&mut self) {
1424 std::process::abort();
1425 }
1426}
1427
1428/// Callbacks and resources detached from Local before running user code.
1429///
1430/// A worker alternates two batches so callbacks can append to Local without
1431/// borrowing the batch currently being drained. Each vector retains its capacity.
1432#[derive(Default)]
1433pub struct Deferred {
1434 /// Notifications of I/O completion, timer expiry, or worker closure.
1435 pub wakes: Vec<Waker>,
1436 /// Wakers whose registrations were replaced or cancelled.
1437 pub drops: Vec<Waker>,
1438 /// Results whose observation has ended.
1439 pub outputs: Vec<RequestOutput>,
1440 /// Buffers and descriptors no longer used by the kernel.
1441 pub resources: Vec<RetiredResources>,
1442 /// Detached sync and forwarded timer notifications without I/O output.
1443 pub completions: Vec<Completion>,
1444 /// Forwarded I/O output or worker-closure errors.
1445 pub results: Vec<OperationResult>,
1446}
1447
1448impl Deferred {
1449 /// Whether the batch contains no callbacks or owners to release.
1450 pub const fn is_empty(&self) -> bool {
1451 self.wakes.is_empty()
1452 && self.drops.is_empty()
1453 && self.outputs.is_empty()
1454 && self.resources.is_empty()
1455 && self.completions.is_empty()
1456 && self.results.is_empty()
1457 }
1458
1459 /// Run each callback independently, retaining failures until cleanup finishes.
1460 /// The caller must release all Local borrows before invoking this method.
1461 fn run(&mut self, panics: &mut Panics) {
1462 for waker in self.wakes.drain(..) {
1463 panics.run(|| waker.wake());
1464 }
1465 for waker in self.drops.drain(..) {
1466 panics.run(|| drop(waker));
1467 }
1468 for output in self.outputs.drain(..) {
1469 panics.run(|| drop(output));
1470 }
1471 for resources in self.resources.drain(..) {
1472 panics.run(|| drop(resources));
1473 }
1474
1475 // Publish results after releasing the completed requests' owners.
1476 for (sender, output) in self.completions.drain(..) {
1477 panics.run(|| {
1478 let _ = sender.send(output);
1479 });
1480 }
1481 for (sender, output) in self.results.drain(..) {
1482 panics.run(|| {
1483 let _ = sender.send(output);
1484 });
1485 }
1486 }
1487}
1488
1489/// Execute one thread's tasks and own its state through shutdown.
1490///
1491/// Constructed before installing TLS so unwinding always has a cleanup owner.
1492/// Keeps callback execution outside [`Local`] borrows and retains failures until
1493/// all kernel-visible work has retired.
1494pub struct Worker {
1495 /// Local state retained until its driver has drained and been destroyed.
1496 local: Rc<RefCell<Local>>,
1497 /// Cleared only after local cleanup, including during an unexpected unwind.
1498 scope: Option<Scope>,
1499 /// Callback batch detached from Local and reused after draining.
1500 deferred: Deferred,
1501 /// First poll, infrastructure, or callback failure.
1502 panics: Panics,
1503 /// Retained mailbox batch, reversed once so bounded pops remain FIFO.
1504 inbox: Vec<Message>,
1505 /// Mailbox publication sequence acknowledged when whole batches enter the inbox.
1506 processed_seq: u32,
1507 /// Calls to [`Self::next_runnable`] left until the one that looks at the
1508 /// pool's global queue.
1509 until_global: u32,
1510 /// Configured interval between forced global checks. An empty run queue
1511 /// can trigger an earlier check.
1512 global_queue_interval: u32,
1513 /// False until kernel retirement and callback cleanup have finished.
1514 finished: bool,
1515}
1516
1517impl Worker {
1518 /// Install TLS only after constructing a cleanup owner. A pool worker
1519 /// counts itself in the pool's shutdown barrier, and its cleanup leaves it.
1520 /// Panics if this thread already has a worker installed.
1521 fn new(local: Local) -> Self {
1522 if local.role().is_pool() {
1523 local.shared.pool.enter();
1524 }
1525 let global_queue_interval = local.shared.cfg.global_queue_interval;
1526 let local = Rc::new(RefCell::new(local));
1527 let mut worker = Self {
1528 local: local.clone(),
1529 scope: None,
1530 deferred: Deferred::default(),
1531 panics: Panics::default(),
1532 inbox: Vec::new(),
1533 processed_seq: 0,
1534 until_global: global_queue_interval,
1535 global_queue_interval,
1536 finished: false,
1537 };
1538 worker.scope = Some(Scope::install(local));
1539 worker
1540 }
1541
1542 /// Create a worker and drive a stack-pinned root, retaining the worker for cleanup.
1543 ///
1544 /// The optional service task runs alongside the root. Only worker zero,
1545 /// which owns the runner, closes runner-wide reservations when its root
1546 /// exits.
1547 ///
1548 /// Root construction, execution, and destruction share a panic boundary.
1549 /// A startup error goes through the panic hook once and returns as that
1550 /// panic's payload. Later failures are retained in the worker so the caller
1551 /// can cancel accepted tasks and clean up with TLS still installed.
1552 pub fn run<F, Fut>(
1553 shared: Arc<Shared>,
1554 role: Role,
1555 build: F,
1556 service: Option<BoxedTask>,
1557 ) -> Result<(Self, Option<Fut::Output>), Panic>
1558 where
1559 F: FnOnce() -> Fut,
1560 Fut: Future,
1561 {
1562 // Raise a startup error as a panic and catch it at once, so the panic
1563 // hook reports it here. Callers propagate the payload with
1564 // resume_unwind, which skips the hook, or hand it to the panicker like
1565 // a task panic, so the hook reports it exactly once. Returning the
1566 // payload instead of unwinding drops `build` and `service` outside any
1567 // unwind, since their destructors may panic.
1568 let local = Local::new(shared.clone(), role).map_err(|error| {
1569 catch_unwind(AssertUnwindSafe(|| {
1570 panic!(
1571 "failed to create native io_uring worker (Linux 6.1 with SINGLE_ISSUER and DEFER_TASKRUN is required): {error}"
1572 )
1573 }))
1574 .expect_err("a startup error must panic")
1575 })?;
1576 let mut worker = Self::new(local);
1577 let mailbox = worker.local.borrow().mailbox.clone();
1578
1579 // Register the background service as an ordinary task before constructing
1580 // the separately polled root.
1581 if let Some(service) = service
1582 && let Err(service) = shared.tasks.register(service, Arc::downgrade(&shared.pool))
1583 {
1584 worker.panics.run(|| service.clear());
1585 }
1586 let root_waker = RootWaker::new(Arc::downgrade(&mailbox)).into();
1587
1588 // The catch owns the root, including when interrupted. Worker and TLS stay
1589 // outside it so root destruction can orphan operations or spawn more work.
1590 let result = catch_unwind(AssertUnwindSafe(|| {
1591 let root = build();
1592 worker.drive(pin!(root), &root_waker)
1593 }));
1594 let output = match result {
1595 Ok(output) => output,
1596 Err(panic) => {
1597 worker.panics.retain(panic);
1598 None
1599 }
1600 };
1601
1602 // Include work spawned by root destruction in the shutdown barrier.
1603 if role.owns_runner() {
1604 shared.workers.close();
1605 }
1606 worker.begin_close();
1607 Ok((worker, output))
1608 }
1609
1610 /// Execute a reserved task on this thread and report failures after cleanup.
1611 /// The caller's [`Launch`] owns the task, shared services, and worker reservation
1612 /// in drop order.
1613 fn run_task(payload: Launch) {
1614 // Bind the reservation first so unwinding releases it last.
1615 let Launch {
1616 active,
1617 shared,
1618 task,
1619 } = payload;
1620
1621 let result = catch_unwind(AssertUnwindSafe(|| {
1622 // Startup may reject the builder without invoking it. Its captured
1623 // root keeps the same disposal boundary as an executing spawned task.
1624 let root = TaskRoot { task: Some(task) };
1625 let (mut worker, output) = Self::run(shared.clone(), Role::Dedicated, || root, None)?;
1626 worker.cleanup();
1627 worker.result(output)
1628 }));
1629
1630 match result {
1631 Ok(Ok(())) => {}
1632 Ok(Err(panic)) | Err(panic) => shared.panicker.notify(panic),
1633 }
1634
1635 // Runtime cleanup, Shared destruction, and failure publication precede
1636 // counter release. Native TLS destruction remains outside this boundary.
1637 drop(shared);
1638 drop(active);
1639 }
1640
1641 /// Detach one callback batch and run it outside the local borrow.
1642 fn callbacks(&mut self, searcher: Option<&mut Searcher<'_>>) {
1643 // An empty batch runs no user code, so a searching worker keeps its
1644 // count rather than hand the search on.
1645 if self.local.borrow().deferred.is_empty() {
1646 return;
1647 }
1648
1649 // A held count stops kicks from waking parked workers. Callbacks can
1650 // wake work and then wait for a peer to run it, so return the count
1651 // before they run.
1652 if let Some(searcher) = searcher {
1653 searcher.leave_search();
1654 }
1655
1656 // Reuse the emptied batch while callbacks collect more work in Local.
1657 mem::swap(&mut self.deferred, &mut self.local.borrow_mut().deferred);
1658 self.deferred.run(&mut self.panics);
1659 }
1660
1661 /// Acknowledge one transferred mailbox batch.
1662 const fn acknowledge_batch(&mut self) {
1663 self.processed_seq = self.processed_seq.wrapping_add(1) & SUBMISSION_SEQ_MASK;
1664 }
1665
1666 /// Take the first failure retained so far.
1667 pub fn take_panic(&mut self) -> Option<Panic> {
1668 self.panics.take()
1669 }
1670
1671 /// Close local registration. Worker zero also closes the task set and the
1672 /// pool. A dedicated worker closes its mailbox, retaining queued messages,
1673 /// while a pool worker keeps its mailbox open until every pool worker has
1674 /// finished its last poll, and [`Self::cleanup`] then closes it. Repeated
1675 /// calls leave the worker closed and preserve its existing inbox.
1676 fn begin_close(&mut self) {
1677 let mailbox = {
1678 let mut local = self.local.borrow_mut();
1679 let role = local.role();
1680
1681 // Worker zero shuts the runner down. It first closes the task set,
1682 // so a spawn from then on is refused and disposed of on its caller.
1683 // It then closes the global queue, which releases workers waiting
1684 // for the pool to close. Because the set closed first, it retains
1685 // the task of every runnable the closed queue discards. Cleanup
1686 // drains the set. Closing again does nothing.
1687 if role.owns_runner() {
1688 local.shared.tasks.close();
1689 local.shared.pool.close();
1690 }
1691 local.closing = true;
1692 if role.is_pool() {
1693 return;
1694 }
1695 local.mailbox.clone()
1696 };
1697 self.close_mailbox(&mailbox);
1698 }
1699
1700 /// Close the mailbox, retaining its queued messages in the inbox.
1701 fn close_mailbox(&mut self, mailbox: &Mailbox) {
1702 let messages = mailbox.close();
1703
1704 if !messages.is_empty() {
1705 self.acknowledge_batch();
1706 }
1707
1708 // Keep queued messages until cleanup can dispose of them outside the
1709 // local borrow. Shutdown only disposes of these messages, so no FIFO
1710 // reversal is needed here.
1711 self.inbox.extend(messages);
1712 }
1713
1714 /// Cancel observers and finish all kernel-visible work before releasing TLS.
1715 /// Retains callback failures until cleanup finishes. An infrastructure
1716 /// failure before kernel retirement aborts the process through
1717 /// [`RetirementGuard`].
1718 ///
1719 /// A pool worker first waits for the pool to close, discards its run queue,
1720 /// and drains the task set, then waits for every other pool worker to
1721 /// finish polling and draining. Worker zero stops the other pool workers
1722 /// before it drains, so they reach that barrier too.
1723 pub fn cleanup(&mut self) {
1724 if self.finished {
1725 return;
1726 }
1727
1728 // Protect every retirement transition, including observer removal
1729 // before the first drain turn, against unexpected infrastructure unwind.
1730 let retirement = RetirementGuard;
1731 self.begin_close();
1732
1733 // A pool worker that stopped early, after a failure, waits for the
1734 // runner to close the pool, so the set it drains below is closed and
1735 // retains the task of every runnable it discards.
1736 let (role, shared) = {
1737 let local = self.local.borrow();
1738 (local.role(), local.shared.clone())
1739 };
1740 if role.is_pool() {
1741 // The drive loop withdrew this worker from the idle set and
1742 // returned its search count when it ended, unwinding included, so
1743 // no wake is spent on it and no kick waits for it here.
1744 shared.pool.wait_closed();
1745 }
1746
1747 // On the runner's normal path, worker zero gets here after the
1748 // supervision tree is aborted, so the workers it stops destroy only
1749 // cancelled tasks and the runner's metrics task. Stopping them in
1750 // cleanup also releases the barrier below when an unwind skips the
1751 // runner's shutdown.
1752 if role.owns_runner() {
1753 shared.pool.stop();
1754 }
1755
1756 if let Role::Pool(index) = role {
1757 // Each task is retained for teardown or already cleared, so
1758 // discarding its runnable runs no user code under this borrow. A
1759 // thief can retain runnables it took until it reaches the barrier
1760 // below.
1761 self.local.borrow_mut().ready_mut().discard();
1762
1763 // Destroy futures before closing I/O so their observers can detach.
1764 // Clear outside local borrows and set locks. A task still polling
1765 // on another worker leaves destruction to that poller.
1766 for task in shared.tasks.drain(index as usize) {
1767 Panics::contain(|| task.clear());
1768 }
1769
1770 // A task polled on one worker can hold registrations on another,
1771 // whose mailbox forwards their results. No pool worker closes its
1772 // mailbox or drops the forwards it took from it until every pool
1773 // worker has drained the set and finished its last poll, so no poll
1774 // finds a registration's worker closed. An empty set is no such
1775 // barrier, since another worker may still be dropping a future.
1776 shared.pool.finish();
1777
1778 // No pool task can run from here on, so the mailbox can close.
1779 let mailbox = self.local.borrow().mailbox.clone();
1780 self.close_mailbox(&mailbox);
1781 }
1782
1783 // Closing the driver and timer table below subsumes queued cancellations.
1784 // Discard queued forwarding messages so their receivers observe closure.
1785 // Senders can run callbacks on drop, so destroy these messages outside
1786 // the local borrow.
1787 for message in self.inbox.drain(..) {
1788 Panics::contain(|| drop(message));
1789 }
1790
1791 // Detach local and forwarded observers before running callbacks.
1792 // Admission is closed, so subsequent handle drops find no registration
1793 // to detach. Only driver service can produce further deferred work.
1794 {
1795 let mut local = self.local.borrow_mut();
1796 let Local {
1797 driver,
1798 timers,
1799 deferred,
1800 ..
1801 } = &mut *local;
1802 driver.as_mut().unwrap().close(deferred);
1803 timers.clear(deferred);
1804 local.update_pending();
1805 }
1806 self.callbacks(None);
1807
1808 // Retained writes and syncs still need to finish. Cancelled operations
1809 // retain their resources until their own completion arrives.
1810 loop {
1811 // A retained write may still be queued without an SQE in flight.
1812 // Service must stage that work before we consider waiting for a CQE.
1813 self.service(false);
1814 self.callbacks(None);
1815
1816 let mut local = self.local.borrow_mut();
1817
1818 if local.driver.as_ref().unwrap().is_empty() {
1819 assert!(
1820 local.deferred.is_empty(),
1821 "callbacks refilled deferred work after close"
1822 );
1823 break;
1824 }
1825
1826 let deadline = local.next_deadline();
1827
1828 // Parking only enters the kernel and invokes no user callbacks.
1829 // Keeping the driver in Local preserves the cleanup owner on unwind.
1830 let parked = local
1831 .driver
1832 .as_mut()
1833 .unwrap()
1834 .park(self.processed_seq, deadline);
1835 drop(local);
1836 parked.expect("io_uring shutdown wait failed");
1837 }
1838
1839 // No SQE can reference request resources now. Driver destruction can
1840 // unwind normally, and TLS remains available until that destruction ends.
1841 mem::forget(retirement);
1842 let driver = self.local.borrow_mut().driver.take();
1843 self.panics.run(|| drop(driver));
1844 self.finished = true;
1845 self.scope.take();
1846 }
1847
1848 /// Select the root result after cleanup.
1849 fn result<T>(&mut self, output: Option<T>) -> Result<T, Panic> {
1850 if let Some(panic) = self.panics.take() {
1851 // A failure may arrive after a successful root poll. Its output can
1852 // own arbitrary destructors, so preserve the first panic on disposal.
1853 self.panics.run(|| drop(output));
1854 return Err(panic);
1855 }
1856
1857 Ok(output.expect("worker root ended without an output or failure"))
1858 }
1859
1860 /// Apply a bounded batch of messages, taking another mailbox batch when needed.
1861 /// The publication sequence records accepted work independently of signaling.
1862 fn messages(&mut self, mailbox: &Arc<Mailbox>) {
1863 if self.inbox.is_empty()
1864 && mailbox.waker.pending(self.processed_seq)
1865 && mailbox.take(&mut self.inbox)
1866 {
1867 // Count transfer once, not each message application. Reversing
1868 // allows bounded FIFO processing without shifting the remainder.
1869 self.acknowledge_batch();
1870 self.inbox.reverse();
1871 }
1872
1873 for _ in 0..BATCH_SIZE {
1874 let Some(message) = self.inbox.pop() else {
1875 break;
1876 };
1877 match message {
1878 Message::WakeRoot => self.local.borrow_mut().root_ready = true,
1879 Message::Forward(forward) => {
1880 let mut local = self.local.borrow_mut();
1881 let Local {
1882 driver,
1883 timers,
1884 deferred,
1885 ..
1886 } = &mut *local;
1887 match forward {
1888 Forward::Waiter(id, sender) => {
1889 driver.as_mut().unwrap().forward(id, sender, deferred);
1890 }
1891 Forward::Timer(id, sender) => timers.forward(id, sender, deferred),
1892 }
1893 }
1894 Message::Cancel(cancel) => self.local.borrow_mut().release(cancel),
1895 }
1896 }
1897 }
1898
1899 /// Service the io_uring ring and timers using one time sample, collecting
1900 /// callbacks. Returns whether kernel service was deferred to an idle wait.
1901 fn service(&mut self, defer_kernel_service: bool) -> bool {
1902 let mut local = self.local.borrow_mut();
1903 local.now = Instant::now();
1904 let Local {
1905 now,
1906 driver,
1907 timers,
1908 deferred,
1909 ..
1910 } = &mut *local;
1911 let result = driver
1912 .as_mut()
1913 .unwrap()
1914 .service(*now, defer_kernel_service, deferred);
1915 let outcome = result.expect("io_uring driver service failed");
1916
1917 // Expiry defers observer notifications until after this borrow ends.
1918 timers.expire(*now, deferred);
1919 local.update_pending();
1920 outcome
1921 }
1922
1923 /// Take the next runnable to poll: a share of the pool's global queue every
1924 /// [`Config::with_global_queue_interval`] calls, otherwise the oldest
1925 /// runnable in this worker's run queue, otherwise a share of the global
1926 /// queue. An empty global take falls back to the run queue. `pool` must be
1927 /// this worker's pool, and the caller must hold no Local borrow.
1928 // The drive loop is instantiated in the caller's crate, so the common case
1929 // is inlined there and the global queue sits behind a call.
1930 #[inline]
1931 #[must_use]
1932 fn next_runnable(&mut self, pool: &Pool) -> Option<Runnable> {
1933 // The interval counts every take, so runnables queued here cannot
1934 // keep the global queue waiting however fast they requeue. A
1935 // countdown keeps a division off this path.
1936 self.until_global -= 1;
1937 if self.until_global == 0 {
1938 self.until_global = self.global_queue_interval;
1939 } else {
1940 if let Some(runnable) = self.local.borrow_mut().ready_mut().pop() {
1941 return Some(runnable);
1942 }
1943 if !pool.has_global() {
1944 return None;
1945 }
1946 }
1947 Self::take_global(&self.local, pool)
1948 }
1949
1950 /// Return the oldest global runnable and append the rest of its capped
1951 /// share to the run queue. Fall back to the oldest run queue runnable if
1952 /// the global queue was empty. `local` must be an unborrowed pool worker
1953 /// of `pool`.
1954 #[inline(never)]
1955 #[must_use]
1956 fn take_global(local: &RefCell<Local>, pool: &Pool) -> Option<Runnable> {
1957 let mut local = local.borrow_mut();
1958 let ready = local.ready_mut();
1959 pool.take(ready).or_else(|| ready.pop())
1960 }
1961
1962 /// Try to steal from the other run queues, then publish this worker idle
1963 /// and look once more. Returns whether the search found work. `searcher`
1964 /// must be this worker's, with its run queue empty and Local unborrowed. A
1965 /// true result leaves no publication but may hold a search count. A false
1966 /// result leaves a live publication for the caller to end through
1967 /// [`Searcher::finish_park`] once its wait or spin returns.
1968 fn search_before_park(&self, searcher: &mut Searcher<'_>) -> bool {
1969 if searcher.search(self.local.borrow_mut().ready_mut()) {
1970 return true;
1971 }
1972
1973 #[cfg(test)]
1974 tests::at_park(
1975 searcher.pool(),
1976 searcher.index(),
1977 tests::ParkPoint::BeforeIdle,
1978 );
1979
1980 if searcher.publish_and_look(self.local.borrow_mut().ready_mut()) {
1981 return true;
1982 }
1983
1984 #[cfg(test)]
1985 tests::at_park(
1986 searcher.pool(),
1987 searcher.index(),
1988 tests::ParkPoint::BeforeWait,
1989 );
1990
1991 false
1992 }
1993
1994 /// Drive tasks and one separately pinned root through bounded service turns.
1995 /// Returns the root output, or `None` with the first failure retained in the
1996 /// worker through root destruction. The caller then closes registration and
1997 /// cleans up the worker.
1998 fn drive<Fut: Future>(
1999 &mut self,
2000 mut root: Pin<&mut Fut>,
2001 root_waker: &Waker,
2002 ) -> Option<Fut::Output> {
2003 let (mailbox, shared, spinner_cfg, role) = {
2004 let local = self.local.borrow();
2005 (
2006 local.mailbox.clone(),
2007 local.shared.clone(),
2008 local.shared.cfg.idle_spinner.clone(),
2009 local.role(),
2010 )
2011 };
2012 let (pool, index) = match role {
2013 Role::Pool(index) => (Some(&*shared.pool), Some(index)),
2014 Role::Dedicated => (None, None),
2015 };
2016 let owner = shared.tasks.id();
2017
2018 // A pool worker's part in searching and waking. Dropped on every
2019 // return and unwind of this loop, it withdraws a live idle publication
2020 // and returns a held search count before the caller destroys the root.
2021 let mut searcher = index.map(|index| Searcher::new(&shared.pool, index));
2022
2023 // A pool claim signals this worker's waker without publishing a
2024 // message, so the spin watches the wake signal as well as the mailbox
2025 // sequence.
2026 let mut spinner = Spinner::new(&spinner_cfg, || {
2027 mailbox.waker.pending(self.processed_seq) || mailbox.waker.signalled()
2028 });
2029 let max_spin =
2030 Duration::from_micros(spinner_cfg.max_budget_us.try_into().unwrap_or(u64::MAX));
2031
2032 loop {
2033 if self.panics.first.is_some() {
2034 return None;
2035 }
2036
2037 // Bound task polling so a self-waking task cannot starve the root,
2038 // mailbox, or io_uring service. A dedicated worker polls only its
2039 // root.
2040 if let (Some(pool), Some(searcher)) = (pool, &mut searcher) {
2041 for _ in 0..BATCH_SIZE {
2042 let Some(runnable) = self.next_runnable(pool) else {
2043 break;
2044 };
2045
2046 // User code can wake a backlog and wait for a peer to run it.
2047 // Return the count so those wakes can kick a parked peer.
2048 searcher.leave_search();
2049
2050 // The inner wrapper handles user polling policy. The poll also
2051 // contains the destruction of a finished or cancelled future.
2052 match runnable.poll(owner) {
2053 AfterPoll::Done => {}
2054
2055 // Queue a self-wake at the tail of the run queue, behind
2056 // the runnables already queued there. Overflow or stealing
2057 // can still move it elsewhere.
2058 AfterPoll::Requeue(runnable) => self.local.borrow_mut().enqueue(runnable),
2059
2060 // The future is gone, so releasing the set's reference and
2061 // the runnable's runs no user code.
2062 AfterPoll::Retire(task) => drop(shared.tasks.remove(&task)),
2063 }
2064 }
2065 }
2066
2067 // The root has no runnable. Its flag also records notifications
2068 // forwarded from other workers through the mailbox.
2069 let poll_root = mem::take(&mut self.local.borrow_mut().root_ready);
2070
2071 if poll_root {
2072 // The root can wake work and block waiting for it just as a
2073 // task can, so return a held count before polling it as well.
2074 if let Some(searcher) = &mut searcher {
2075 searcher.leave_search();
2076 }
2077
2078 let mut cx = TaskContext::from_waker(root_waker);
2079
2080 // A local wake during this poll sets root_ready again, and a
2081 // foreign wake queues a message that sets it when applied.
2082 // Pending must not clear a wake caused by the root itself.
2083 if let Poll::Ready(output) = root.as_mut().poll(&mut cx) {
2084 return Some(output);
2085 }
2086 }
2087
2088 // Polls can leave cancelled observers behind. Run one batch before
2089 // deciding whether kernel service can wait until parking. Reentrant
2090 // work remains visible to the complete readiness checks below.
2091 self.callbacks(searcher.as_mut());
2092 if self.panics.first.is_some() {
2093 return None;
2094 }
2095
2096 // An apparently idle turn can combine kernel service with the wait.
2097 // The result records whether that enter was actually deferred.
2098 let defer = !self.local.borrow().is_ready()
2099 && self.inbox.is_empty()
2100 && !mailbox.waker.pending(self.processed_seq);
2101 let kernel_deferred = self.service(defer);
2102
2103 // Apply foreign messages and completion callbacks before considering
2104 // a wait. Any tasks they make ready will be polled on the next turn.
2105 // Messages run no user code, so a held search count stays held
2106 // through them. The callback batch returns it before running.
2107 self.messages(&mailbox);
2108 self.callbacks(searcher.as_mut());
2109 if self.panics.first.is_some() {
2110 return None;
2111 }
2112
2113 #[cfg(test)]
2114 if let (Some(pool), Some(index)) = (pool, index) {
2115 tests::at_park(pool, index, tests::ParkPoint::AfterMessages);
2116 }
2117
2118 // Callbacks and messages can submit I/O or make tasks runnable, so
2119 // the pre-service idle decision alone is not enough to permit parking.
2120 let (ready, needs_kernel, pending_submissions, deadline) = {
2121 let mut local = self.local.borrow_mut();
2122 (
2123 local.is_ready(),
2124 !local.driver.as_ref().unwrap().is_empty(),
2125 local.driver.as_ref().unwrap().has_pending_submissions(),
2126 local.next_deadline(),
2127 )
2128 };
2129
2130 // Queued submissions need another service turn even when no task
2131 // has a notification that would cause it to be polled. Going round
2132 // again without parking ends a search, and the catch-up's callbacks
2133 // run user code.
2134 if ready
2135 || pending_submissions
2136 || !self.inbox.is_empty()
2137 || mailbox.waker.pending(self.processed_seq)
2138 {
2139 if let Some(searcher) = &mut searcher {
2140 searcher.leave_search();
2141 }
2142
2143 if (kernel_deferred && needs_kernel) || (defer && pending_submissions) {
2144 // An idle turn must stage work introduced by its callbacks.
2145 // A wake alone needs catch-up only if GETEVENTS was deferred.
2146 self.service(false);
2147 self.callbacks(None);
2148 }
2149
2150 continue;
2151 }
2152
2153 // Busy turns use only service's time sample. Take a fresh sample
2154 // at this actual idle boundary after any elapsed callback time.
2155 // Waiting for an expired deadline would not block, so this ends a
2156 // search as well, and the catch-up's callbacks run user code.
2157 let now = Instant::now();
2158 if deadline.is_some_and(|deadline| deadline <= now) {
2159 if let Some(searcher) = &mut searcher {
2160 searcher.leave_search();
2161 }
2162
2163 if kernel_deferred && needs_kernel {
2164 self.service(false);
2165 self.callbacks(None);
2166 }
2167
2168 continue;
2169 }
2170
2171 // Steal, then publish idleness before the last look for work. A
2172 // push racing this look either shows up here or leaves some worker,
2173 // not necessarily this one, woken or running to take it.
2174 if let Some(searcher) = &mut searcher
2175 && self.search_before_park(searcher)
2176 {
2177 // Running again ends a search, and the catch-up's callbacks run
2178 // user code.
2179 searcher.leave_search();
2180
2181 if kernel_deferred && needs_kernel {
2182 self.service(false);
2183 self.callbacks(None);
2184 }
2185
2186 continue;
2187 }
2188
2189 #[cfg(test)]
2190 tests::before_park();
2191
2192 // Unfinished requests need the io_uring wait to run deferred kernel
2193 // work. A futex wake alone cannot advance their I/O.
2194 if needs_kernel {
2195 let result = self
2196 .local
2197 .borrow_mut()
2198 .driver
2199 .as_mut()
2200 .unwrap()
2201 .park(self.processed_seq, deadline);
2202
2203 // Withdraw before callbacks so their kicks cannot spend a wake
2204 // on this already running worker. A claim's count stays held
2205 // until the worker returns it before user code, when it goes
2206 // round again, or at its next park.
2207 if let Some(searcher) = &mut searcher {
2208 searcher.finish_park();
2209 }
2210
2211 if !result.expect("io_uring kernel wait failed") {
2212 // A rejected wait may not have entered the kernel. Complete
2213 // its deferred service obligation before another polling turn.
2214 self.service(false);
2215 self.callbacks(searcher.as_mut());
2216 }
2217 } else {
2218 // A future timer by itself uses a timed futex wait. Spinning
2219 // is skipped when it could consume the remaining deadline.
2220 let near_deadline = deadline
2221 .is_some_and(|deadline| deadline.saturating_duration_since(now) <= max_spin);
2222
2223 if !near_deadline
2224 && spinner.spin(|| {
2225 mailbox.waker.pending(self.processed_seq) || mailbox.waker.signalled()
2226 })
2227 {
2228 // The worker runs again, so the next wait needs a new wake.
2229 mailbox.waker.consume_signal();
2230 } else if let Some(duration) = mailbox.waker.park_idle(self.processed_seq, deadline)
2231 {
2232 spinner.on_wake(duration);
2233 }
2234
2235 // End this publication even when spinning or a latched wake
2236 // avoided blocking. The next park needs a fresh publication.
2237 if let Some(searcher) = &mut searcher {
2238 searcher.finish_park();
2239 }
2240 }
2241 }
2242 }
2243}
2244
2245impl Drop for Worker {
2246 fn drop(&mut self) {
2247 // Cleanup still runs during unwinding. Panics leaks any unclaimed
2248 // payload so its destruction cannot replace the original failure.
2249 self.cleanup();
2250 }
2251}
2252
2253/// Run a dedicated or blocking task as its worker's root, containing poll and
2254/// disposal panics.
2255struct TaskRoot {
2256 /// Task retained until the root is destroyed.
2257 task: Option<BoxedTask>,
2258}
2259
2260impl Future for TaskRoot {
2261 type Output = ();
2262
2263 fn poll(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<Self::Output> {
2264 // A contained cancellation or disposal panic ends the task so the worker
2265 // can proceed with I/O cleanup.
2266 Panics::contain(|| self.task.as_mut().unwrap().as_mut().poll(cx)).unwrap_or(Poll::Ready(()))
2267 }
2268}
2269
2270impl Drop for TaskRoot {
2271 fn drop(&mut self) {
2272 Panics::contain(|| drop(self.task.take()));
2273 }
2274}
2275
2276/// Native io_uring runner executing ordinary tasks on a pool of workers, the
2277/// first on its calling thread.
2278///
2279/// The root future doesn't have to be Send and always runs on the calling
2280/// thread. Ordinary spawned futures run on whichever pool worker takes them,
2281/// and dedicated and blocking ones on threads of their own. The runner waits
2282/// for every pool worker, dedicated worker cleanup, and retained writes and
2283/// syncs before returning or resuming a panic. Native thread-local destruction
2284/// may follow.
2285pub struct Runner {
2286 /// Settings validated before any runtime resources are created.
2287 cfg: Config,
2288}
2289
2290impl Runner {
2291 /// Construct a runner without acquiring storage or creating an io_uring ring.
2292 pub const fn new(cfg: Config) -> Self {
2293 Self { cfg }
2294 }
2295}
2296
2297impl Default for Runner {
2298 fn default() -> Self {
2299 Self::new(Config::default())
2300 }
2301}
2302
2303impl crate::Runner for Runner {
2304 type Context = Context;
2305
2306 fn start<F, Fut>(mut self, f: F) -> Fut::Output
2307 where
2308 F: FnOnce(Context) -> Fut,
2309 Fut: Future,
2310 {
2311 // Reject nesting and invalid settings before acquiring storage or threads.
2312 Scope::assert_vacant();
2313 self.cfg.validate();
2314
2315 let mut registry = Registry::new();
2316 let mut runtime_registry = registry.sub_registry(METRICS_PREFIX);
2317 let metrics = TaskMetrics::new(&mut runtime_registry);
2318 let pending_operations = runtime_registry.register(
2319 "pending_operations",
2320 "Number of active logical requests across io_uring workers",
2321 raw::Gauge::default(),
2322 );
2323 let process = ProcessMetrics::init(&mut runtime_registry);
2324
2325 let storage = Storage::new(StorageConfig {
2326 storage_directory: self.cfg.storage_directory.clone(),
2327 blob_layouts: self.cfg.storage_blob_layouts.clone(),
2328 });
2329
2330 // Storage construction acquires the directory hold first. This sync
2331 // therefore includes any straggling writes from a preceding runner.
2332 crate::storage::sync(&self.cfg.storage_directory).unwrap_or_else(|error| {
2333 panic!(
2334 "failed to sync storage filesystem at startup ({}): {error}",
2335 self.cfg.storage_directory.display()
2336 );
2337 });
2338
2339 let storage = MeteredStorage::new(storage, &mut runtime_registry);
2340 let network = MeteredNetwork::new(
2341 Network::new(NetworkConfig {
2342 tcp_nodelay: self.cfg.tcp_nodelay,
2343 zero_linger: self.cfg.zero_linger,
2344 connect_timeout: self.cfg.connect_timeout,
2345 read_write_timeout: self.cfg.read_write_timeout,
2346 read_buffer_size: self.cfg.read_buffer_size,
2347 }),
2348 &mut runtime_registry,
2349 );
2350 let (panicker, tasks) = Panicker::new(self.cfg.catch_panics);
2351
2352 // Every pool worker's mailbox exists before any worker starts, so a
2353 // push can wake any of them from the first spawn on.
2354 let mailboxes = (0..self.cfg.worker_threads)
2355 .map(|_| Mailbox::new().map(Arc::new))
2356 .collect::<Result<Vec<_>, _>>()
2357 .unwrap_or_else(|error| panic!("failed to create io_uring worker mailbox: {error}"));
2358 let workers = mailboxes.len();
2359
2360 let shared = Arc::new(Shared {
2361 cfg: self.cfg,
2362 pool: Arc::new(Pool::new(mailboxes)),
2363 registry,
2364 metrics,
2365 pending_operations,
2366 shutdown: Mutex::new(Stopper::default()),
2367 panicker,
2368 workers: Arc::new(Workers::default()),
2369 tasks: Tasks::new(workers),
2370 storage,
2371 network,
2372 });
2373
2374 let label = Label::root();
2375 shared.metrics.tasks_spawned.get_or_create(&label).inc();
2376
2377 let metric = MetricHandle::new(shared.metrics.tasks_running.get_or_create(&label).clone());
2378 let tree = Tree::root();
2379 let context_shared = shared.clone();
2380 let context_tree = tree.clone();
2381
2382 // A pool worker that fails interrupts the root even when task panics
2383 // are caught, since the pool cannot run without it. The worker wakes
2384 // the root, which takes the failure before each poll. The receiver
2385 // outlives the root, so a failure after the root completes still fails
2386 // the runner.
2387 let mut threads = Threads::default();
2388 let (pool_panicker, mut pool_failures) = Panicker::new(false);
2389
2390 // The root factory needs TLS installed so it can synchronously spawn or
2391 // register work, and it runs only once every pool worker has started.
2392 // Task panics interrupt the root inside the pool failure check, so they
2393 // are observed only until the root completes, before its state is
2394 // destroyed.
2395 let output = Worker::run(
2396 shared.clone(),
2397 Role::Pool(0),
2398 || {
2399 threads.start(&shared, &pool_panicker);
2400 let root = tasks.interrupt(f(Context {
2401 name: label.name().to_string(),
2402 attributes: Vec::new(),
2403 shared: context_shared,
2404 tree: context_tree,
2405 execution: Execution::default(),
2406 }));
2407 let failures = &mut pool_failures;
2408 async move {
2409 let mut root = pin!(root);
2410 poll_fn(|cx| {
2411 if let Some(panic) = failures.try_take() {
2412 resume_unwind(panic);
2413 }
2414 root.as_mut().poll(cx)
2415 })
2416 .await
2417 }
2418 },
2419 Some(Box::pin(process.collect(Sleep::new))),
2420 )
2421 .and_then(|(mut worker, output)| {
2422 #[cfg(test)]
2423 tests::before_abort();
2424
2425 // The root is gone and the pool is closed. Abort ordinary,
2426 // dedicated, and blocking tasks, then clean up this worker, which
2427 // stops the other pool workers and waits for them to drain the
2428 // task set.
2429 worker.panics.run(|| tree.abort());
2430 worker.cleanup();
2431 worker.panics.run(|| threads.join());
2432 if let Some(panic) = pool_failures.try_take() {
2433 worker.panics.retain(panic);
2434 }
2435
2436 // Dedicated worker cleanup may depend on resources released by pool
2437 // task disposal or I/O retirement, so drain the pool first.
2438 shared.workers.wait();
2439 worker.result(output)
2440 });
2441 metric.finish();
2442
2443 match output {
2444 Ok(output) => output,
2445 Err(panic) => resume_unwind(panic),
2446 }
2447 }
2448}
2449
2450#[cfg(test)]
2451#[path = "tests.rs"]
2452pub mod tests;