Skip to main content

dope_core/driver/
ready.rs

1use std::cell::Cell;
2use std::io::{Error, ErrorKind, Result};
3use std::marker::PhantomData;
4use std::ptr::NonNull;
5
6use o3::collections::BatchSet;
7use o3::marker::ThreadBound;
8
9use crate::io::fd::FdSlot;
10
11use super::DriverRef;
12use super::token::{Epoch, ROUTE_FRAMEWORK, SlotIndex, Token};
13
14const NIL: u32 = u32::MAX;
15
16/// A generation-checked address in a driver's local ready queue.
17#[derive(Clone, Copy)]
18pub struct ReadyKey<'d> {
19    index: u32,
20    epoch: u32,
21    _arena: PhantomData<&'d Arena>,
22    _thread: ThreadBound,
23}
24
25impl ReadyKey<'static> {
26    pub const NONE: Self = Self {
27        index: NIL,
28        epoch: 0,
29        _arena: PhantomData,
30        _thread: ThreadBound::NEW,
31    };
32}
33
34impl PartialEq for ReadyKey<'_> {
35    fn eq(&self, other: &Self) -> bool {
36        self.index == other.index && self.epoch == other.epoch
37    }
38}
39
40impl Eq for ReadyKey<'_> {}
41
42/// A driver-scoped, type-erased wake target for an in-flight operation.
43///
44/// Unlike [`ReadyKey`], this preserves hierarchical task wakeups. A completion
45/// can therefore wake the exact child task that registered the operation,
46/// rather than only waking the root runtime task.
47#[derive(Clone, Copy)]
48pub struct CompletionWaker<'d> {
49    target: CompletionTarget<'d>,
50    _thread: ThreadBound,
51}
52
53#[derive(Clone, Copy)]
54enum CompletionTarget<'d> {
55    Ready(DriverRef<'d>, ReadyKey<'d>),
56    Callback(NonNull<()>, unsafe fn(NonNull<()>)),
57}
58
59impl<'d> CompletionWaker<'d> {
60    pub fn from_ready(driver: DriverRef<'d>, key: ReadyKey<'d>) -> Self {
61        Self {
62            target: CompletionTarget::Ready(driver, key),
63            _thread: ThreadBound::NEW,
64        }
65    }
66
67    /// # Safety
68    ///
69    /// `target` must remain valid for every call to `wake` while this handle is
70    /// live, and `callback` must accept that exact target.
71    pub unsafe fn from_callback(target: NonNull<()>, callback: unsafe fn(NonNull<()>)) -> Self {
72        Self {
73            target: CompletionTarget::Callback(target, callback),
74            _thread: ThreadBound::NEW,
75        }
76    }
77
78    pub fn wake(self) {
79        match self.target {
80            CompletionTarget::Ready(driver, key) => driver.activate_ready(key),
81            CompletionTarget::Callback(target, callback) => unsafe { callback(target) },
82        }
83    }
84}
85
86#[derive(Clone, Copy)]
87pub struct ReadyHandle<'d> {
88    arena: &'d Arena,
89    key: ReadyKey<'d>,
90}
91
92impl<'d> ReadyHandle<'d> {
93    pub fn set_target(self, target: Token) {
94        self.arena.set_target(self.key, target);
95    }
96
97    pub fn activate(self) {
98        self.arena.activate(self.key);
99    }
100
101    pub fn key(self) -> ReadyKey<'d> {
102        self.key
103    }
104}
105
106pub struct ReadySlot<'d> {
107    arena: &'d Arena,
108    key: ReadyKey<'d>,
109}
110
111impl<'d> ReadySlot<'d> {
112    fn new(arena: &'d Arena, index: u32) -> Self {
113        arena.live[index as usize].set(true);
114        Self {
115            arena,
116            key: ReadyKey {
117                index,
118                epoch: arena.epochs[index as usize].get(),
119                _arena: PhantomData,
120                _thread: ThreadBound::NEW,
121            },
122        }
123    }
124
125    pub fn set_target(&self, target: Token) {
126        self.arena.set_target(self.key, target);
127    }
128
129    pub fn activate(&self) {
130        self.arena.activate(self.key);
131    }
132
133    pub fn key(&self) -> ReadyKey<'d> {
134        self.key
135    }
136}
137
138impl Drop for ReadySlot<'_> {
139    fn drop(&mut self) {
140        self.arena.release(self.key);
141    }
142}
143
144pub(super) struct Arena {
145    fixed: usize,
146    ready: BatchSet,
147    targets: Box<[Cell<Token>]>,
148    epochs: Box<[Cell<u32>]>,
149    live: Box<[Cell<bool>]>,
150    next_free: Box<[Cell<u32>]>,
151    free: Cell<u32>,
152    free_len: Cell<usize>,
153}
154
155impl Arena {
156    pub(super) fn new(fixed: usize, dynamic: usize) -> Result<Box<Self>> {
157        let capacity = fixed
158            .checked_add(dynamic)
159            .ok_or_else(|| Error::new(ErrorKind::InvalidInput, "dope: ready capacity overflow"))?;
160        if capacity > u32::MAX as usize {
161            return Err(Error::new(
162                ErrorKind::InvalidInput,
163                "dope: ready capacity exceeds u32",
164            ));
165        }
166
167        let dummy = Token::new(ROUTE_FRAMEWORK, SlotIndex::new(0), Epoch::INITIAL);
168        Ok(Box::new(Self {
169            fixed,
170            ready: BatchSet::with_capacity(capacity),
171            targets: (0..capacity).map(|_| Cell::new(dummy)).collect(),
172            epochs: (0..capacity).map(|_| Cell::new(0)).collect(),
173            live: (0..capacity)
174                .map(|index| Cell::new(index < fixed))
175                .collect(),
176            next_free: (0..capacity)
177                .map(|index| {
178                    let next = if index + 1 < capacity {
179                        index as u32 + 1
180                    } else {
181                        NIL
182                    };
183                    Cell::new(next)
184                })
185                .collect(),
186            free: Cell::new(if fixed < capacity { fixed as u32 } else { NIL }),
187            free_len: Cell::new(capacity - fixed),
188        }))
189    }
190
191    fn valid(&self, key: ReadyKey<'_>) -> bool {
192        let index = key.index as usize;
193        self.live.get(index).is_some_and(Cell::get) && self.epochs[index].get() == key.epoch
194    }
195
196    pub(crate) fn fixed_slot(&self, slot: FdSlot) -> ReadyHandle<'_> {
197        let index = slot.raw();
198        debug_assert!((index as usize) < self.fixed);
199        ReadyHandle {
200            arena: self,
201            key: ReadyKey {
202                index,
203                epoch: 0,
204                _arena: PhantomData,
205                _thread: ThreadBound::NEW,
206            },
207        }
208    }
209
210    pub(crate) fn make_slot(&self, target: Token) -> Result<ReadySlot<'_>> {
211        self.make_slot_reserving(target, 0)
212    }
213
214    pub(crate) fn make_slots<I>(&self, targets: I) -> Result<Box<[ReadySlot<'_>]>>
215    where
216        I: IntoIterator<Item = Token>,
217        I::IntoIter: ExactSizeIterator,
218    {
219        let targets = targets.into_iter();
220        let requested = targets.len();
221        let available = self.free_len.get();
222        if requested > available {
223            return Err(Self::capacity_error(requested, available));
224        }
225        targets
226            .enumerate()
227            .map(|(index, target)| self.make_slot_reserving(target, requested - index - 1))
228            .collect()
229    }
230
231    pub(crate) fn make_slot_reserving(
232        &self,
233        target: Token,
234        reserve: usize,
235    ) -> Result<ReadySlot<'_>> {
236        if self.free_len.get() <= reserve {
237            return Err(Self::capacity_error(1, self.free_len.get()));
238        }
239        let index = self.free.get();
240        if index == NIL {
241            return Err(Self::capacity_error(1, 0));
242        }
243        self.free.set(self.next_free[index as usize].get());
244        self.free_len.set(self.free_len.get() - 1);
245        self.targets[index as usize].set(target);
246        Ok(ReadySlot::new(self, index))
247    }
248
249    fn capacity_error(requested: usize, available: usize) -> Error {
250        Error::new(
251            ErrorKind::WouldBlock,
252            format!(
253                "dope: dynamic ready capacity exhausted: requested {requested}, available {available}"
254            ),
255        )
256    }
257
258    fn set_target(&self, key: ReadyKey<'_>, target: Token) {
259        if self.valid(key) {
260            self.targets[key.index as usize].set(target);
261        }
262    }
263
264    pub(crate) fn activate(&self, key: ReadyKey<'_>) {
265        if self.valid(key) {
266            self.ready.insert(key.index as usize);
267        }
268    }
269
270    fn release(&self, key: ReadyKey<'_>) {
271        if !self.valid(key) {
272            return;
273        }
274        let index = key.index as usize;
275        self.live[index].set(false);
276        self.ready.remove(index);
277        let Some(epoch) = key.epoch.checked_add(1) else {
278            self.next_free[index].set(NIL);
279            return;
280        };
281        self.epochs[index].set(epoch);
282        self.next_free[index].set(self.free.get());
283        self.free.set(key.index);
284        self.free_len.set(self.free_len.get() + 1);
285    }
286
287    pub(crate) fn drain(&self, mut activate: impl FnMut(Token)) {
288        let Some(mut ready) = self.ready.drain_batch() else {
289            return;
290        };
291        for index in &mut ready {
292            if self.live[index].get() {
293                activate(self.targets[index].get());
294            }
295        }
296    }
297
298    pub(crate) fn has_ready(&self) -> bool {
299        !self.ready.is_empty()
300    }
301}