dope_core/driver/
ready.rs

1use std::cell::Cell;
2use std::fmt;
3use std::io::{Error, ErrorKind, Result};
4use std::marker::PhantomData;
5use std::ptr::NonNull;
6
7use o3::cell::RegionToken;
8use o3::collections::BatchSet;
9use o3::marker::ThreadBound;
10
11use super::DriverRef;
12use super::token::{Epoch, ROUTE_FRAMEWORK, SlotIndex, Token};
13use crate::io::fd::FdSlot;
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///
49/// This is a linear registration capability: it cannot be cloned or copied,
50/// and `wake` consumes it. A rejected registration must return the value to
51/// its caller instead of silently retaining a duplicate.
52#[repr(transparent)]
53pub struct CompletionWaker<'d>(RawCompletionWaker<'d>);
54
55#[doc(hidden)]
56#[derive(Clone, Copy)]
57#[repr(transparent)]
58pub struct CompletionKey<'d>(RawCompletionWaker<'d>);
59
60#[derive(Clone, Copy)]
61struct RawCompletionWaker<'d> {
62    target: CompletionTarget<'d>,
63    _region: PhantomData<fn(&'d RegionToken<'d>) -> &'d RegionToken<'d>>,
64    _thread: ThreadBound,
65}
66
67impl fmt::Debug for CompletionWaker<'_> {
68    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
69        f.debug_struct("CompletionWaker").finish_non_exhaustive()
70    }
71}
72
73/// An owning slot for a linear [`CompletionWaker`].
74///
75/// The slot keeps the copyable representation private so owners can wake
76/// without clearing the slot on the hot path. The surrounding operation state
77/// remains responsible for deciding whether a wake is due.
78#[doc(hidden)]
79#[repr(transparent)]
80pub struct CompletionSlot<'d> {
81    wake: Cell<Option<RawCompletionWaker<'d>>>,
82}
83
84/// A proof-bearing sink for a completion handle retained after the current
85/// operation.
86///
87/// # Safety
88///
89/// `register` must not expose the supplied handle through its output. If it
90/// retains the handle, the implementing owner must ensure it cannot be invoked
91/// after its source is invalidated. Deferred removal is valid only when every
92/// path that could observe or invoke the handle first completes that removal.
93/// Re-registration and every owner drop path must preserve this invariant.
94#[doc(hidden)]
95pub unsafe trait CompletionRegistrar<'d> {
96    type Output;
97
98    fn register(self, wake: CompletionWaker<'d>) -> Self::Output;
99}
100
101/// A completion registrar that also needs mutable executor-region access.
102///
103/// # Safety
104///
105/// The same requirements as [`CompletionRegistrar`] apply. The supplied
106/// region token must not escape `register`.
107#[doc(hidden)]
108pub unsafe trait CompletionRegistrarWithRegion<'d> {
109    type Output;
110
111    fn register(self, wake: CompletionWaker<'d>, region: &mut RegionToken<'d>) -> Self::Output;
112}
113
114/// A callback target whose raw representation is valid for completion wakeups.
115///
116/// # Safety
117///
118/// `into_raw_parts` must return a callback that accepts exactly the returned
119/// target. The target must remain valid until every registrar that receives
120/// its completion handle has made that handle unobservable.
121#[doc(hidden)]
122pub unsafe trait CompletionCallback<'d> {
123    fn into_raw_parts(self) -> (NonNull<()>, unsafe fn(NonNull<()>));
124}
125
126#[derive(Clone, Copy)]
127enum CompletionTarget<'d> {
128    Ready(DriverRef<'d>, ReadyKey<'d>),
129    Callback(NonNull<()>, unsafe fn(NonNull<()>)),
130}
131
132impl<'d> CompletionWaker<'d> {
133    pub fn from_ready(driver: DriverRef<'d>, key: ReadyKey<'d>) -> Self {
134        Self(RawCompletionWaker {
135            target: CompletionTarget::Ready(driver, key),
136            _region: PhantomData,
137            _thread: ThreadBound::NEW,
138        })
139    }
140
141    /// Creates and delivers a callback-backed handle without exposing it.
142    #[doc(hidden)]
143    #[inline(always)]
144    pub fn register_callback<C, R>(source: C, region: &RegionToken<'d>, registrar: R) -> R::Output
145    where
146        C: CompletionCallback<'d>,
147        R: CompletionRegistrar<'d>,
148    {
149        let _ = region;
150        let (target, callback) = source.into_raw_parts();
151        registrar.register(Self(RawCompletionWaker {
152            target: CompletionTarget::Callback(target, callback),
153            _region: PhantomData,
154            _thread: ThreadBound::NEW,
155        }))
156    }
157
158    /// Creates and delivers a callback-backed handle with mutable region
159    /// access, without exposing either capability.
160    #[doc(hidden)]
161    #[inline(always)]
162    pub fn register_callback_with_region<C, R>(
163        source: C,
164        region: &mut RegionToken<'d>,
165        registrar: R,
166    ) -> R::Output
167    where
168        C: CompletionCallback<'d>,
169        R: CompletionRegistrarWithRegion<'d>,
170    {
171        let (target, callback) = source.into_raw_parts();
172        registrar.register(
173            Self(RawCompletionWaker {
174                target: CompletionTarget::Callback(target, callback),
175                _region: PhantomData,
176                _thread: ThreadBound::NEW,
177            }),
178            region,
179        )
180    }
181
182    pub fn wake(self) {
183        self.0.wake();
184    }
185
186    #[doc(hidden)]
187    pub fn key(&self) -> CompletionKey<'d> {
188        CompletionKey(self.0)
189    }
190}
191
192impl RawCompletionWaker<'_> {
193    fn same_target(&self, other: &Self) -> bool {
194        match (self.target, other.target) {
195            (
196                CompletionTarget::Ready(left_driver, left_key),
197                CompletionTarget::Ready(right_driver, right_key),
198            ) => left_driver == right_driver && left_key == right_key,
199            (
200                CompletionTarget::Callback(left_target, left_callback),
201                CompletionTarget::Callback(right_target, right_callback),
202            ) => left_target == right_target && std::ptr::fn_addr_eq(left_callback, right_callback),
203            _ => false,
204        }
205    }
206
207    fn wake(self) {
208        match self.target {
209            CompletionTarget::Ready(driver, key) => driver.activate_ready(key),
210            CompletionTarget::Callback(target, callback) => unsafe { callback(target) },
211        }
212    }
213}
214
215impl<'d> CompletionSlot<'d> {
216    pub const fn empty() -> Self {
217        Self {
218            wake: Cell::new(None),
219        }
220    }
221
222    pub fn set(&self, wake: CompletionWaker<'d>) {
223        self.wake.set(Some(wake.0));
224    }
225
226    pub fn clear(&self) {
227        self.wake.set(None);
228    }
229
230    pub fn is_empty(&self) -> bool {
231        self.wake.get().is_none()
232    }
233
234    pub fn matches(&self, wake: &CompletionWaker<'d>) -> bool {
235        self.matches_key(wake.key())
236    }
237
238    pub fn matches_key(&self, key: CompletionKey<'d>) -> bool {
239        self.wake
240            .get()
241            .is_some_and(|current| current.same_target(&key.0))
242    }
243
244    pub fn clear_if(&self, key: CompletionKey<'d>) -> bool {
245        if !self.matches_key(key) {
246            return false;
247        }
248        self.clear();
249        true
250    }
251
252    #[inline]
253    pub fn take(&self) -> Option<CompletionWaker<'d>> {
254        self.wake.take().map(CompletionWaker)
255    }
256
257    pub fn wake(&self) {
258        if let Some(wake) = self.wake.get() {
259            wake.wake();
260        }
261    }
262}
263
264#[derive(Clone, Copy)]
265pub struct ReadyHandle<'d> {
266    arena: &'d Arena,
267    key: ReadyKey<'d>,
268}
269
270impl<'d> ReadyHandle<'d> {
271    pub fn set_target(self, target: Token) {
272        self.arena.set_target(self.key, target);
273    }
274
275    pub fn activate(self) {
276        self.arena.activate(self.key);
277    }
278
279    pub fn key(self) -> ReadyKey<'d> {
280        self.key
281    }
282}
283
284pub struct ReadySlot<'d> {
285    arena: &'d Arena,
286    key: ReadyKey<'d>,
287}
288
289impl<'d> ReadySlot<'d> {
290    fn new(arena: &'d Arena, index: u32) -> Self {
291        arena.live[index as usize].set(true);
292        Self {
293            arena,
294            key: ReadyKey {
295                index,
296                epoch: arena.epochs[index as usize].get(),
297                _arena: PhantomData,
298                _thread: ThreadBound::NEW,
299            },
300        }
301    }
302
303    pub fn set_target(&self, target: Token) {
304        self.arena.set_target(self.key, target);
305    }
306
307    pub fn activate(&self) {
308        self.arena.activate(self.key);
309    }
310
311    pub fn key(&self) -> ReadyKey<'d> {
312        self.key
313    }
314}
315
316impl Drop for ReadySlot<'_> {
317    fn drop(&mut self) {
318        self.arena.release(self.key);
319    }
320}
321
322pub(super) struct Arena {
323    fixed: usize,
324    ready: BatchSet,
325    targets: Box<[Cell<Token>]>,
326    epochs: Box<[Cell<u32>]>,
327    live: Box<[Cell<bool>]>,
328    next_free: Box<[Cell<u32>]>,
329    free: Cell<u32>,
330    free_len: Cell<usize>,
331}
332
333impl Arena {
334    pub(super) fn new(fixed: usize, dynamic: usize) -> Result<Box<Self>> {
335        let capacity = fixed
336            .checked_add(dynamic)
337            .ok_or_else(|| Error::new(ErrorKind::InvalidInput, "dope: ready capacity overflow"))?;
338        if capacity > u32::MAX as usize {
339            return Err(Error::new(
340                ErrorKind::InvalidInput,
341                "dope: ready capacity exceeds u32",
342            ));
343        }
344
345        let dummy = Token::new(ROUTE_FRAMEWORK, SlotIndex::ZERO, Epoch::INITIAL);
346        Ok(Box::new(Self {
347            fixed,
348            ready: BatchSet::with_capacity(capacity),
349            targets: (0..capacity).map(|_| Cell::new(dummy)).collect(),
350            epochs: (0..capacity).map(|_| Cell::new(0)).collect(),
351            live: (0..capacity)
352                .map(|index| Cell::new(index < fixed))
353                .collect(),
354            next_free: (0..capacity)
355                .map(|index| {
356                    let next = if index + 1 < capacity {
357                        index as u32 + 1
358                    } else {
359                        NIL
360                    };
361                    Cell::new(next)
362                })
363                .collect(),
364            free: Cell::new(if fixed < capacity { fixed as u32 } else { NIL }),
365            free_len: Cell::new(capacity - fixed),
366        }))
367    }
368
369    fn valid(&self, key: ReadyKey<'_>) -> bool {
370        let index = key.index as usize;
371        self.live.get(index).is_some_and(Cell::get) && self.epochs[index].get() == key.epoch
372    }
373
374    pub(crate) fn fixed_slot(&self, slot: FdSlot) -> ReadyHandle<'_> {
375        let index = slot.raw();
376        debug_assert!((index as usize) < self.fixed);
377        ReadyHandle {
378            arena: self,
379            key: ReadyKey {
380                index,
381                epoch: 0,
382                _arena: PhantomData,
383                _thread: ThreadBound::NEW,
384            },
385        }
386    }
387
388    pub(crate) fn fd_slot(&self, raw: u32) -> Option<FdSlot> {
389        if (raw as usize) >= self.fixed {
390            return None;
391        }
392        FdSlot::try_from_raw(raw)
393    }
394
395    pub(crate) fn make_slot(&self, target: Token) -> Result<ReadySlot<'_>> {
396        self.make_slot_reserving(target, 0)
397    }
398
399    pub(crate) fn make_slots<I>(&self, targets: I) -> Result<Box<[ReadySlot<'_>]>>
400    where
401        I: IntoIterator<Item = Token>,
402        I::IntoIter: ExactSizeIterator,
403    {
404        let targets = targets.into_iter();
405        let requested = targets.len();
406        let available = self.free_len.get();
407        if requested > available {
408            return Err(Self::capacity_error(requested, available));
409        }
410        targets
411            .enumerate()
412            .map(|(index, target)| self.make_slot_reserving(target, requested - index - 1))
413            .collect()
414    }
415
416    pub(crate) fn make_slot_reserving(
417        &self,
418        target: Token,
419        reserve: usize,
420    ) -> Result<ReadySlot<'_>> {
421        if self.free_len.get() <= reserve {
422            return Err(Self::capacity_error(1, self.free_len.get()));
423        }
424        let index = self.free.get();
425        if index == NIL {
426            return Err(Self::capacity_error(1, 0));
427        }
428        self.free.set(self.next_free[index as usize].get());
429        self.free_len.set(self.free_len.get() - 1);
430        self.targets[index as usize].set(target);
431        Ok(ReadySlot::new(self, index))
432    }
433
434    fn capacity_error(requested: usize, available: usize) -> Error {
435        Error::new(
436            ErrorKind::WouldBlock,
437            format!(
438                "dope: dynamic ready capacity exhausted: requested {requested}, available {available}"
439            ),
440        )
441    }
442
443    fn set_target(&self, key: ReadyKey<'_>, target: Token) {
444        if self.valid(key) {
445            self.targets[key.index as usize].set(target);
446        }
447    }
448
449    pub(crate) fn activate(&self, key: ReadyKey<'_>) {
450        if self.valid(key) {
451            self.ready.insert(key.index as usize);
452        }
453    }
454
455    pub(crate) fn arm_recv_credit(&self, key: ReadyKey<'_>, target: Token) -> bool {
456        use super::token::kind::RECV_CREDIT_HELD;
457
458        if !self.valid(key) {
459            return false;
460        }
461        let current = self.targets[key.index as usize].get();
462        if !current.same_target(target) {
463            return false;
464        }
465        self.targets[key.index as usize].set(target.with_kind(RECV_CREDIT_HELD));
466        true
467    }
468
469    pub(crate) fn release_recv_credit(&self, key: ReadyKey<'_>, target: Token) {
470        use super::token::kind::{RECV_CREDIT_HELD, RECV_CREDIT_RELEASED};
471
472        if !self.valid(key) {
473            return;
474        }
475        let current = self.targets[key.index as usize].get();
476        if current != target.with_kind(RECV_CREDIT_HELD) {
477            return;
478        }
479        self.targets[key.index as usize].set(target.with_kind(RECV_CREDIT_RELEASED));
480        self.ready.insert(key.index as usize);
481    }
482
483    pub(crate) fn take_recv_credit(&self, key: ReadyKey<'_>, target: Token) -> bool {
484        use super::token::kind::RECV_CREDIT_RELEASED;
485
486        if !self.valid(key) {
487            return false;
488        }
489        let current = self.targets[key.index as usize].get();
490        if current != target.with_kind(RECV_CREDIT_RELEASED) {
491            return false;
492        }
493        self.targets[key.index as usize].set(target.with_kind(0));
494        true
495    }
496
497    fn release(&self, key: ReadyKey<'_>) {
498        if !self.valid(key) {
499            return;
500        }
501        let index = key.index as usize;
502        self.live[index].set(false);
503        self.ready.remove(index);
504        let Some(epoch) = key.epoch.checked_add(1) else {
505            self.next_free[index].set(NIL);
506            return;
507        };
508        self.epochs[index].set(epoch);
509        self.next_free[index].set(self.free.get());
510        self.free.set(key.index);
511        self.free_len.set(self.free_len.get() + 1);
512    }
513
514    pub(crate) fn drain(&self, mut activate: impl FnMut(Token)) {
515        let Some(mut ready) = self.ready.drain_batch() else {
516            return;
517        };
518        for index in &mut ready {
519            if self.live[index].get() {
520                activate(self.targets[index].get());
521            }
522        }
523    }
524
525    pub(crate) fn has_ready(&self) -> bool {
526        !self.ready.is_empty()
527    }
528}
529
530#[cfg(test)]
531mod tests {
532    use super::Arena;
533    use crate::driver::token::kind::{RECV_CREDIT_HELD, RECV_CREDIT_RELEASED};
534    use crate::driver::token::{Epoch, SlotIndex, Token};
535    use crate::io::fd::FdSlot;
536
537    fn target(epoch: Epoch) -> Token {
538        Token::new(7, SlotIndex::ZERO, epoch)
539    }
540
541    #[test]
542    fn application_wake_does_not_release_held_receive_credit() {
543        let arena = Arena::new(1, 0).unwrap();
544        let handle = arena.fixed_slot(FdSlot::try_from_raw(0).unwrap());
545        let target = target(Epoch::INITIAL);
546        handle.set_target(target);
547
548        assert!(arena.arm_recv_credit(handle.key(), target));
549        handle.activate();
550        let mut activated = Vec::new();
551        arena.drain(|token| activated.push(token));
552        assert_eq!(activated, [target.with_kind(RECV_CREDIT_HELD)]);
553        assert!(!arena.take_recv_credit(handle.key(), target));
554
555        arena.release_recv_credit(handle.key(), target);
556        activated.clear();
557        arena.drain(|token| activated.push(token));
558        assert_eq!(activated, [target.with_kind(RECV_CREDIT_RELEASED)]);
559        assert!(arena.take_recv_credit(handle.key(), target));
560
561        handle.activate();
562        activated.clear();
563        arena.drain(|token| activated.push(token));
564        assert_eq!(activated, [target]);
565    }
566
567    #[test]
568    fn stale_guard_cannot_release_credit_for_a_reused_connection() {
569        let arena = Arena::new(1, 0).unwrap();
570        let handle = arena.fixed_slot(FdSlot::try_from_raw(0).unwrap());
571        let old = target(Epoch::INITIAL);
572        let new = target(Epoch::INITIAL.next().unwrap());
573        handle.set_target(old);
574        assert!(arena.arm_recv_credit(handle.key(), old));
575
576        handle.set_target(new);
577        arena.release_recv_credit(handle.key(), old);
578        assert!(!arena.has_ready());
579
580        handle.activate();
581        let mut activated = Vec::new();
582        arena.drain(|token| activated.push(token));
583        assert_eq!(activated, [new]);
584    }
585}
dope_core/driver/ready.rs

dope_core/driver/
ready.rs
