nuclei/rt/runtime/
placement.rs

1//! Tasks placed on a worker the caller names: [`RuntimeHandle::workers`],
2//! [`spawn_on_worker`](RuntimeHandle::spawn_on_worker),
3//! [`spawn_local_on_worker`](RuntimeHandle::spawn_local_on_worker) and
4//! [`on_each_worker`](RuntimeHandle::on_each_worker).
5//!
6//! A task placed on worker `k` is a task of `k`'s shard, as one spawned
7//! there is, and stays on it. The handoff moves a shard with every task of
8//! it to another thread (on the same core, when workers are pinned): the
9//! task keeps running on worker `k`, and
10//! [`worker_index`](super::worker_index) keeps answering `k`. A worker that does not take tasks refuses the placement with
11//! `InvalidInput`, and so does one that begins to stop while the task is
12//! placed: the task is counted on the shard before the shard is asked again
13//! whether it takes tasks, so a stop either sees the task (and runs it to
14//! its end) or the spawn sees the stop (and drops the task unrun). No placed
15//! task is ever accepted and then dropped unrun (proofs/tla/placement).
16
17use core::future::Future;
18use std::io;
19use std::sync::Arc;
20
21use super::super::sched::Sched;
22use super::super::shard::Shared;
23use super::super::task::{self, Task};
24use super::{Inner, RuntimeHandle, spawned};
25use crate::sync::internal::{Ordering, fence};
26use crate::task::JoinSet;
27
28impl RuntimeHandle {
29    /// The workers the runtime has started, stopped ones included: the
30    /// indices [`spawn_on_worker`](Self::spawn_on_worker) takes are
31    /// `0..workers()`.
32    pub fn workers(&self) -> usize {
33        self.inner.workers.started()
34    }
35
36    /// Spawns `future` on worker `worker`, where it stays: a handoff of the
37    /// worker's shard moves the task with it, so it runs on that worker for
38    /// its whole life.
39    ///
40    /// # Errors
41    ///
42    /// `InvalidInput` when the runtime has no such worker, or the worker
43    /// stopped (or is stopping) and takes no new task.
44    pub fn spawn_on_worker<F>(&self, worker: usize, future: F) -> io::Result<Task<F::Output>>
45    where
46        F: Future + Send + 'static,
47        F::Output: Send + 'static,
48    {
49        spawn_on_worker(&self.inner, worker, future)
50    }
51
52    /// Runs `make` on worker `worker` and spawns the future it makes there,
53    /// which need not be `Send`: it stays on the thread that made it, as a
54    /// [`spawn_local`](crate::spawn_local) task does.
55    ///
56    /// A handoff of the worker's shard leaves such a future on the thread
57    /// that made it, off the worker: it finishes there, where
58    /// [`worker_index`](super::worker_index) answers `None` and shard-local
59    /// values are out of reach. The handle is that of the task that makes
60    /// the future on the worker and waits for it: its id is not the one
61    /// [`task::id`](crate::task::id) answers inside the future, and dropping
62    /// it cancels both.
63    ///
64    /// # Errors
65    ///
66    /// As for [`spawn_on_worker`](Self::spawn_on_worker).
67    pub fn spawn_local_on_worker<M, F>(&self, worker: usize, make: M) -> io::Result<Task<F::Output>>
68    where
69        M: FnOnce() -> F + Send + 'static,
70        F: Future + 'static,
71        F::Output: Send + 'static,
72    {
73        self.spawn_on_worker(worker, async move {
74            let task = crate::spawn_local(make());
75            task.await
76        })
77    }
78
79    /// Runs `make` on every worker that takes tasks and spawns the future it
80    /// makes there, as [`spawn_local_on_worker`](Self::spawn_local_on_worker)
81    /// does: the set holds one task per worker. A runtime with no worker
82    /// gives an empty set.
83    pub fn on_each_worker<M, F>(&self, make: M) -> JoinSet<F::Output>
84    where
85        M: Fn() -> F + Send + Sync + 'static,
86        F: Future + 'static,
87        F::Output: Send + 'static,
88    {
89        let make = Arc::new(make);
90        let mut set = JoinSet::new();
91        for worker in 0..self.workers() {
92            let make = Arc::clone(&make);
93            // A worker that stopped is no worker to run on.
94            if let Ok(task) = self.spawn_local_on_worker(worker, move || make()) {
95                set.insert(task);
96            }
97        }
98        set
99    }
100}
101
102/// The shard of worker `worker` of `inner`, if it takes tasks.
103fn worker_shard(inner: &Inner, worker: usize) -> io::Result<&Arc<Shared>> {
104    let started = inner.workers.started();
105    if worker >= started {
106        return Err(io::Error::new(
107            io::ErrorKind::InvalidInput,
108            format!("worker {worker} does not exist: the runtime has started {started} workers"),
109        ));
110    }
111    match inner.workers.get(worker) {
112        Some(shard) if shard.takes_tasks() => Ok(shard),
113        Some(_) => Err(stopped(worker)),
114        None => Err(io::Error::new(
115            io::ErrorKind::InvalidInput,
116            format!("worker {worker} never started: its shard could not be made"),
117        )),
118    }
119}
120
121fn stopped(worker: usize) -> io::Error {
122    io::Error::new(
123        io::ErrorKind::InvalidInput,
124        format!("worker {worker} has stopped and takes no new task"),
125    )
126}
127
128/// Spawns `future` on worker `worker` of `inner`: as `spawn_on` does, with
129/// the second look at the worker between the task's count and its schedule.
130pub(super) fn spawn_on_worker<F>(
131    inner: &Inner,
132    worker: usize,
133    future: F,
134) -> io::Result<Task<F::Output>>
135where
136    F: Future + Send + 'static,
137    F::Output: Send + 'static,
138{
139    let shard = worker_shard(inner, worker)?;
140    let hooks = inner.task_hooks.as_ref();
141    let (runnable, task) = task::spawn(
142        future,
143        Sched::counted(Arc::clone(shard), None, hooks.cloned()),
144    );
145    spawned(hooks.map(|hooks| &**hooks), &task);
146    // Pairs with the fence of `Shared::drain`: either the stop sees the
147    // task counted, or this sees the stop.
148    fence(Ordering::SeqCst);
149    if !shard.takes_tasks() {
150        // Never run: its future is dropped here, and its handle with it.
151        drop(runnable);
152        drop(task);
153        return Err(stopped(worker));
154    }
155    runnable.schedule();
156    Ok(task)
157}
nuclei/rt/runtime/placement.rs

nuclei/rt/runtime/
placement.rs
