Skip to main content

dope_runtime/process/
mod.rs

1use std::{cell, io, iter, mem, os::fd, process, thread};
2
3use dope_core::{driver::settings, platform::affinity};
4
5use crate::{executor, random, shutdown};
6
7mod binding;
8mod sealed;
9
10pub(crate) use sealed::{Blocked, Limit, Set};
11
12/// Allocation-free snapshot of the logical CPUs available to this process.
13#[repr(transparent)]
14pub struct Cpus(affinity::Cpus);
15
16pub struct Context {
17    affinity: binding::Binding,
18    seed: random::Seed,
19    shutdown: Shutdown,
20}
21
22/// Shutdown source owned by exactly one production worker runtime.
23///
24/// A proof issued by the explicit single-runtime shutdown API cannot be
25/// retagged as a production process proof:
26///
27/// ```compile_fail,E0308
28/// use dope_runtime::{process, shutdown};
29///
30/// fn retag(
31///     requested: shutdown::Requested,
32/// ) -> shutdown::Requested<process::Shutdown> {
33///     requested
34/// }
35/// ```
36pub struct Shutdown {
37    source: fd::OwnedFd,
38    ready: cell::Cell<Option<shutdown::Notify>>,
39}
40
41struct Trigger(shutdown::Notify);
42
43/// Explicit stop authority for a process runtime started with [`Runtime::controlled`].
44pub struct Control(shutdown::Notify);
45
46enum Wait {
47    Signal,
48    Control(shutdown::Wait),
49}
50
51pub struct Runtime<I> {
52    workers: Vec<Worker<I>>,
53    triggers: Vec<Trigger>,
54    ready: Vec<shutdown::Wait>,
55    wait: Wait,
56}
57
58struct Worker<I> {
59    cpu: u16,
60    input: I,
61    shutdown: Shutdown,
62}
63
64/// Armed proof that a worker may only return through its shutdown request.
65struct Completion;
66
67const STACK_SIZE: usize = 8 * 1024 * 1024;
68const _: () = assert!(mem::size_of::<Completion>() == 0);
69
70impl Cpus {
71    pub fn current() -> io::Result<Self> {
72        affinity::Cpus::current().map(Self)
73    }
74}
75
76impl Iterator for Cpus {
77    type Item = u16;
78
79    fn next(&mut self) -> Option<Self::Item> {
80        self.0.next()
81    }
82
83    fn size_hint(&self) -> (usize, Option<usize>) {
84        self.0.size_hint()
85    }
86}
87
88impl DoubleEndedIterator for Cpus {
89    fn next_back(&mut self) -> Option<Self::Item> {
90        self.0.next_back()
91    }
92}
93
94impl ExactSizeIterator for Cpus {
95    fn len(&self) -> usize {
96        self.0.len()
97    }
98}
99
100impl iter::FusedIterator for Cpus {}
101
102const _: () = {
103    assert!(std::mem::size_of::<Cpus>() == std::mem::size_of::<affinity::Cpus>());
104    assert!(std::mem::align_of::<Cpus>() == std::mem::align_of::<affinity::Cpus>());
105};
106
107impl Context {
108    pub fn cpu(&self) -> u16 {
109        self.affinity.cpu()
110    }
111}
112
113impl executor::Factory for Context {
114    type Shutdown = Shutdown;
115
116    fn executor(
117        self,
118        config: settings::Config,
119    ) -> io::Result<executor::Executor<(), Self::Shutdown>> {
120        executor::Executor::with_seed(config, self.seed)?
121            .with_source(self.shutdown, |shutdown| fd::AsFd::as_fd(&shutdown.source))
122    }
123}
124
125impl<I> Runtime<I> {
126    pub fn pinned(workers: impl IntoIterator<Item = (u16, I)>) -> io::Result<Self> {
127        Self::build(workers, Wait::Signal)
128    }
129
130    /// Builds a runtime whose process wait can be completed by the returned
131    /// single-use control authority instead of a process signal.
132    pub fn controlled(workers: impl IntoIterator<Item = (u16, I)>) -> io::Result<(Self, Control)> {
133        let runtime = Self::build(workers, Wait::Signal)?;
134        let (wait, notify) = shutdown::Ends::blocking()?.split();
135        Ok((
136            Self {
137                wait: Wait::Control(wait),
138                ..runtime
139            },
140            Control(notify),
141        ))
142    }
143
144    fn build(workers: impl IntoIterator<Item = (u16, I)>, wait: Wait) -> io::Result<Self> {
145        use std::collections::HashSet;
146
147        let mut source = workers.into_iter();
148        let Some(first) = source.next() else {
149            return Err(io::Error::new(
150                io::ErrorKind::InvalidInput,
151                "runtime requires at least one CPU",
152            ));
153        };
154        Limit::current()?.raise()?;
155        let (lower, _) = source.size_hint();
156        let capacity = lower.saturating_add(1);
157        let mut unique = HashSet::with_capacity(capacity);
158        let mut values = Vec::with_capacity(capacity);
159        let mut triggers = Vec::with_capacity(capacity);
160        let mut ready = Vec::with_capacity(capacity);
161        for (cpu, input) in iter::once(first).chain(source) {
162            if !unique.insert(cpu) {
163                return Err(io::Error::new(
164                    io::ErrorKind::InvalidInput,
165                    "runtime requires distinct CPUs",
166                ));
167            }
168            let (source, notify) = shutdown::Ends::event()?.split();
169            let (ready_wait, ready_notify) = shutdown::Ends::blocking()?.split();
170            values.push(Worker {
171                cpu,
172                input,
173                shutdown: Shutdown {
174                    source,
175                    ready: cell::Cell::new(Some(ready_notify)),
176                },
177            });
178            triggers.push(Trigger(notify));
179            ready.push(ready_wait);
180        }
181        Ok(Self {
182            workers: values,
183            triggers,
184            ready,
185            wait,
186        })
187    }
188
189    pub fn run<N>(
190        self,
191        entry: fn(I, Context) -> io::Result<shutdown::Requested<Shutdown>>,
192        notify: N,
193    ) -> io::Result<()>
194    where
195        I: Send,
196        N: FnOnce() -> io::Result<()>,
197    {
198        let Self {
199            workers,
200            triggers,
201            ready,
202            wait,
203        } = self;
204        match wait {
205            Wait::Signal => {
206                let blocked = Set::termination()?.block()?;
207                run_waiting(workers, triggers, ready, entry, notify, || blocked.wait())
208            }
209            Wait::Control(reader) => {
210                run_waiting(workers, triggers, ready, entry, notify, || reader.wait())
211            }
212        }
213    }
214}
215
216fn run_waiting<I>(
217    workers: Vec<Worker<I>>,
218    triggers: Vec<Trigger>,
219    ready: Vec<shutdown::Wait>,
220    entry: fn(I, Context) -> io::Result<shutdown::Requested<Shutdown>>,
221    notify: impl FnOnce() -> io::Result<()>,
222    wait: impl FnOnce() -> io::Result<()>,
223) -> io::Result<()>
224where
225    I: Send,
226{
227    let seed = random::Seed::random()?;
228    thread::scope(|scope| {
229        let mut handles = Vec::with_capacity(workers.len());
230        for (index, worker) in workers.into_iter().enumerate() {
231            let seed_index = u64::try_from(index).map_err(io::Error::other)?;
232            let worker_seed = seed.derive(seed_index);
233            let spawn = thread::Builder::new()
234                .name(format!("dope-core-{index}"))
235                .stack_size(STACK_SIZE)
236                .spawn_scoped(scope, move || execute(worker, worker_seed, entry));
237            match spawn {
238                Ok(handle) => handles.push(handle),
239                Err(error) => {
240                    stop(triggers);
241                    join(handles);
242                    return Err(error);
243                }
244            }
245        }
246
247        for worker in ready {
248            if let Err(error) = worker.wait() {
249                stop(triggers);
250                join(handles);
251                return Err(error);
252            }
253        }
254        if let Err(error) = notify() {
255            stop(triggers);
256            join(handles);
257            return Err(error);
258        }
259        let waited = wait();
260        stop(triggers);
261        join(handles);
262        waited
263    })
264}
265
266impl Shutdown {
267    pub(crate) fn installed(&self) {
268        let Some(ready) = self.ready.take() else {
269            return;
270        };
271        if ready.notify().is_err() {
272            process::abort();
273        }
274    }
275}
276
277impl Control {
278    pub fn fire(self) -> io::Result<()> {
279        self.0.notify()
280    }
281}
282
283impl Completion {
284    fn accept(self, _requested: shutdown::Requested<Shutdown>) {
285        mem::forget(self);
286    }
287}
288
289impl Drop for Completion {
290    fn drop(&mut self) {
291        eprintln!("dope-core: worker exited without shutdown completion");
292        process::abort();
293    }
294}
295
296fn execute<I>(
297    worker: Worker<I>,
298    seed: random::Seed,
299    entry: fn(I, Context) -> io::Result<shutdown::Requested<Shutdown>>,
300) where
301    I: Send,
302{
303    let completion = Completion;
304    let Worker {
305        cpu,
306        input,
307        shutdown,
308    } = worker;
309    let result = binding::Binding::bind(cpu).and_then(|affinity| {
310        entry(
311            input,
312            Context {
313                affinity,
314                seed,
315                shutdown,
316            },
317        )
318    });
319    match result {
320        Ok(requested) => completion.accept(requested),
321        Err(error) => {
322            eprintln!("dope-core cpu={cpu}: {error}");
323            process::exit(1);
324        }
325    }
326}
327
328fn stop(triggers: Vec<Trigger>) {
329    let mut failed = false;
330    for trigger in triggers {
331        if trigger.0.notify().is_err() {
332            failed = true;
333        }
334    }
335    if failed {
336        process::exit(1);
337    }
338}
339
340fn join(handles: Vec<thread::ScopedJoinHandle<'_, ()>>) {
341    for handle in handles {
342        if handle.join().is_err() {
343            process::abort();
344        }
345    }
346}