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#[repr(transparent)]
14pub struct Cpus(affinity::Cpus);
15
16pub struct Context {
17 affinity: binding::Binding,
18 seed: random::Seed,
19 shutdown: Shutdown,
20}
21
22pub struct Shutdown {
37 source: fd::OwnedFd,
38 ready: cell::Cell<Option<shutdown::Notify>>,
39}
40
41struct Trigger(shutdown::Notify);
42
43pub 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
64struct 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 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}