Skip to main content

dope_runtime/process/
layout.rs

1use std::{convert, mem, num::NonZeroUsize, vec};
2
3use dope_core::platform::affinity;
4use o3::alloc::{self, VecExt as _};
5
6use crate::process::{
7    self,
8    failure::{config, create, input},
9};
10
11/// Validated, ordered CPU topology used to map one exact startup
12/// resource per worker.
13pub struct Layout {
14    head: affinity::Cpu,
15    tail: Vec<affinity::Cpu>,
16}
17
18pub(in crate::process) struct Assignment<I> {
19    pub(in crate::process) cpu: affinity::Cpu,
20    pub(in crate::process) input: I,
21}
22
23pub(in crate::process) struct Set<I> {
24    prior: Vec<Assignment<I>>,
25    last: Assignment<I>,
26}
27
28impl Layout {
29    pub fn pinned(cpus: impl IntoIterator<Item = affinity::Cpu>) -> Result<Self, create::Error> {
30        let mut source = cpus.into_iter();
31        let Some(head) = source.next() else {
32            return Err(create::Error::Config(config::Error::Empty));
33        };
34        let mut tail = Vec::new();
35        alloc::VecExt::try_reserve(&mut tail, source.size_hint().0)?;
36        for cpu in source {
37            if cpu == head || tail.contains(&cpu) {
38                return Err(create::Error::Config(config::Error::DuplicateCpu(cpu)));
39            }
40            alloc::Vacant::try_new(&mut tail)?.insert(cpu);
41        }
42        Ok(Self { head, tail })
43    }
44
45    pub fn count(&self) -> NonZeroUsize {
46        NonZeroUsize::MIN.saturating_add(self.tail.len())
47    }
48
49    /// Assigns exactly one affine input to each CPU, in topology order.
50    pub fn try_assign<I, R>(self, items: R) -> Result<process::Runtime<I>, create::Error>
51    where
52        R: IntoIterator<Item = I>,
53        R::IntoIter: ExactSizeIterator,
54    {
55        self.try_map(items, convert::identity)
56    }
57
58    /// Consumes exactly one affine resource for each CPU, in topology order.
59    pub fn try_map<T, I, R>(
60        self,
61        items: R,
62        issue: impl Fn(T) -> I,
63    ) -> Result<process::Runtime<I>, create::Error>
64    where
65        R: IntoIterator<Item = T>,
66        R::IntoIter: ExactSizeIterator,
67    {
68        let mut items = Input::new(items.into_iter(), self.count())?;
69        let mut prior = Vec::try_new(self.tail.len())?;
70        let item = items.next(0)?;
71        let mut last = Assignment {
72            cpu: self.head,
73            input: issue(item),
74        };
75        for cpu in self.tail {
76            let item = items.next(prior.len() + 1)?;
77            let previous = mem::replace(
78                &mut last,
79                Assignment {
80                    cpu,
81                    input: issue(item),
82                },
83            );
84            prior.push(previous);
85        }
86        items.finish()?;
87        Ok(process::Runtime {
88            workers: Set { prior, last },
89        })
90    }
91
92    /// Replicates prior inputs and moves the original owner into the final CPU.
93    pub fn try_fork<O, T, I, E, R>(
94        self,
95        owner: O,
96        items: R,
97        issue: impl Fn(&O, T) -> Result<I, E>,
98        finish: impl FnOnce(O, T) -> I,
99    ) -> Result<process::Runtime<I>, input::Error<E>>
100    where
101        R: IntoIterator<Item = T>,
102        R::IntoIter: ExactSizeIterator,
103    {
104        let mut items =
105            Input::new(items.into_iter(), self.count()).map_err(input::Error::Layout)?;
106        let mut prior = Vec::try_new(self.tail.len())
107            .map_err(create::Error::Allocation)
108            .map_err(input::Error::Layout)?;
109        let item = items.next(0).map_err(input::Error::Layout)?;
110        let mut last = Assignment {
111            cpu: self.head,
112            input: item,
113        };
114        for cpu in self.tail {
115            let item = items.next(prior.len() + 1).map_err(input::Error::Layout)?;
116            let input = issue(&owner, last.input).map_err(|cause| input::Error::Replica {
117                cpu: last.cpu,
118                cause,
119            })?;
120            prior.push(Assignment {
121                cpu: last.cpu,
122                input,
123            });
124            last = Assignment { cpu, input: item };
125        }
126        items.finish().map_err(input::Error::Layout)?;
127        let last = Assignment {
128            cpu: last.cpu,
129            input: finish(owner, last.input),
130        };
131        Ok(process::Runtime {
132            workers: Set { prior, last },
133        })
134    }
135}
136
137struct Input<I> {
138    source: I,
139    required: NonZeroUsize,
140}
141
142impl<I: ExactSizeIterator> Input<I> {
143    fn new(source: I, required: NonZeroUsize) -> Result<Self, create::Error> {
144        let actual = source.len();
145        if actual != required.get() {
146            return Err(create::Error::Config(config::Error::ResourceCount {
147                actual,
148                required: required.get(),
149            }));
150        }
151        Ok(Self { source, required })
152    }
153
154    fn next(&mut self, yielded: usize) -> Result<I::Item, create::Error> {
155        self.source
156            .next()
157            .ok_or(create::Error::Config(config::Error::ResourceIterator {
158                yielded,
159                required: self.required.get(),
160            }))
161    }
162
163    fn finish(&mut self) -> Result<(), create::Error> {
164        if self.source.next().is_some() {
165            return Err(create::Error::Config(config::Error::ResourceIterator {
166                yielded: self.required.get() + 1,
167                required: self.required.get(),
168            }));
169        }
170        Ok(())
171    }
172}
173
174impl<I> Set<I> {
175    pub(in crate::process) fn len(&self) -> usize {
176        self.prior.len() + 1
177    }
178
179    pub(in crate::process) fn into_parts(self) -> (vec::IntoIter<Assignment<I>>, Assignment<I>) {
180        (self.prior.into_iter(), self.last)
181    }
182}
183
184const _: () = {
185    assert!(mem::size_of::<Layout>() == 4 * mem::size_of::<usize>());
186    assert!(mem::size_of::<Set<()>>() == 4 * mem::size_of::<usize>());
187};