dope_core/driver/
ready.rs

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

dope_core/driver/
ready.rs
