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

dope_core/driver/
ready.rs
