nuclei/rt/runtime/knobs.rs
1//! The builder's hooks and scheduler knobs: what the runtime calls on its
2//! threads and around its tasks, how its workers run their loops, what it
3//! does when a task panics, and the seed of its random choices.
4
5use std::sync::Arc;
6
7use super::{Builder, RngSeed, TaskMeta, UnhandledPanic};
8
9impl Builder {
10 /// Called on each thread the runtime starts (workers and blocking-pool
11 /// threads), on that thread, before it does anything else.
12 pub fn on_thread_start(mut self, f: impl Fn() + Send + Sync + 'static) -> Self {
13 self.tuning.hooks.start = Some(Arc::new(f));
14 self
15 }
16
17 /// Called on each thread the runtime starts, on that thread, as it ends
18 /// (after a panic too).
19 pub fn on_thread_stop(mut self, f: impl Fn() + Send + Sync + 'static) -> Self {
20 self.tuning.hooks.stop = Some(Arc::new(f));
21 self
22 }
23
24 /// Called on a worker thread just before it sleeps with nothing to run.
25 pub fn on_thread_park(mut self, f: impl Fn() + Send + Sync + 'static) -> Self {
26 self.tuning.hooks.park = Some(Arc::new(f));
27 self
28 }
29
30 /// Called on a worker thread just after it wakes from a sleep.
31 pub fn on_thread_unpark(mut self, f: impl Fn() + Send + Sync + 'static) -> Self {
32 self.tuning.hooks.unpark = Some(Arc::new(f));
33 self
34 }
35
36 /// Called with each task this runtime spawns (not the blocking pool's
37 /// jobs), on the spawning thread, once the task exists.
38 pub fn on_task_spawn(mut self, f: impl Fn(&TaskMeta<'_>) + Send + Sync + 'static) -> Self {
39 self.task_hooks.spawn = Some(Arc::new(f));
40 self
41 }
42
43 /// Called just before each poll of a task of this runtime, on the thread
44 /// that polls it.
45 pub fn on_before_task_poll(
46 mut self,
47 f: impl Fn(&TaskMeta<'_>) + Send + Sync + 'static,
48 ) -> Self {
49 self.task_hooks.before_poll = Some(Arc::new(f));
50 self
51 }
52
53 /// Called just after each poll of a task of this runtime, on the thread
54 /// that polled it.
55 pub fn on_after_task_poll(mut self, f: impl Fn(&TaskMeta<'_>) + Send + Sync + 'static) -> Self {
56 self.task_hooks.after_poll = Some(Arc::new(f));
57 self
58 }
59
60 /// Called when a task of this runtime is done with: its future completed
61 /// or was dropped by a cancel, on the thread that dropped it.
62 pub fn on_task_terminate(mut self, f: impl Fn(&TaskMeta<'_>) + Send + Sync + 'static) -> Self {
63 self.task_hooks.terminate = Some(Arc::new(f));
64 self
65 }
66
67 /// Called with each poll of a task of this runtime that its worker's
68 /// watch sees run past [`handoff_after`](Builder::handoff_after), on the
69 /// watch's own thread: the task's id and name say which one holds its
70 /// worker up. The task is the innermost one whose poll was running when
71 /// the watch looked; the poll may have ended by the time the hook runs.
72 /// A poll is told once, unless a `block_on` inside it ran a long poll of
73 /// its own and the shard stayed on that thread: the outer poll may then
74 /// be told again as it goes on. The shard is handed on first, so a slow
75 /// hook delays no rescue, but every other worker waits for its next look
76 /// while the hook runs: keep it short. A panic in it is caught (after the
77 /// panic hook reported it) and the watch goes on. Only workers are
78 /// watched; with `handoff_after(None)` [`build`](Builder::build) fails
79 /// with `InvalidInput`.
80 pub fn on_long_poll(mut self, f: impl Fn(&TaskMeta<'_>) + Send + Sync + 'static) -> Self {
81 self.long_poll = Some(Arc::new(f));
82 self
83 }
84
85 /// Task runs between a worker's looks at its timers and IO when it is
86 /// busy (default 61): lower answers IO sooner, higher runs tasks with
87 /// fewer interruptions. At least 1; [`build`](Builder::build) fails with
88 /// `InvalidInput` at 0.
89 pub fn event_interval(mut self, runs: u32) -> Self {
90 self.tuning.driver_interval = runs;
91 self
92 }
93
94 /// Task runs between the times a worker takes from its inbox (tasks
95 /// woken by other threads) before its own queue (default 31). At least 1;
96 /// [`build`](Builder::build) fails with `InvalidInput` at 0.
97 pub fn global_queue_interval(mut self, runs: u32) -> Self {
98 self.tuning.inbox_interval = runs;
99 self
100 }
101
102 /// The most IO events a worker takes from its driver in one look
103 /// (default: all that are ready). At least 1; [`build`](Builder::build)
104 /// fails with `InvalidInput` at 0.
105 pub fn max_io_events_per_tick(mut self, events: usize) -> Self {
106 self.tuning.max_io_events = events;
107 self
108 }
109
110 /// Turns off each worker's LIFO slot, where a task woken by the task
111 /// running on the worker goes to run next. Without it the woken task
112 /// goes to the back of the worker's queue: fairer when the queue is
113 /// long, slower for a task that hands work to another and waits for it.
114 pub fn disable_lifo_slot(mut self) -> Self {
115 self.tuning.lifo = false;
116 self
117 }
118
119 /// Seeds the runtime's random choices, so that they repeat run after
120 /// run; see [`RngSeed`].
121 pub fn rng_seed(mut self, seed: RngSeed) -> Self {
122 self.rng_seed = Some(seed);
123 self
124 }
125
126 /// What the runtime does when one of its tasks panics (default
127 /// [`UnhandledPanic::Ignore`]: the panic goes to the task's handle).
128 ///
129 /// ```
130 /// use nuclei::runtime::{Runtime, UnhandledPanic};
131 ///
132 /// let rt = Runtime::builder()
133 /// .unhandled_panic(UnhandledPanic::ShutdownRuntime)
134 /// .build()
135 /// .unwrap();
136 /// let ran = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
137 /// rt.block_on(async {
138 /// let _task = nuclei::spawn(async { panic!("a task fails") });
139 /// core::future::pending::<()>().await
140 /// })
141 /// }));
142 /// assert!(ran.is_err(), "the block_on panics once the task did");
143 /// ```
144 pub fn unhandled_panic(mut self, policy: UnhandledPanic) -> Self {
145 self.unhandled_panic = policy;
146 self
147 }
148}