Skip to main content

dope_core/driver/schedule/ready/wake/
queue.rs

1use std::{mem, num};
2
3use o3::depot::queue::resident::{
4    self, fifo,
5    fifo::{Index, Scope},
6};
7
8use crate::{
9    driver,
10    driver::schedule::{
11        app,
12        ready::{
13            self, wake,
14            wake::{
15                Drain, Once, Waker,
16                fanout::storage::Step,
17                resident::{drain, one},
18            },
19        },
20    },
21};
22
23/// Bounded FIFO wakes with a constructor-proved lane extent.
24pub struct Queue<'d, L = num::NonZero<usize>> {
25    driver: driver::Reference<'d>,
26    inner: fifo::Fifo<ready::Retained<'d>, L>,
27}
28
29/// Linear ownership of one resident queue slot.
30#[repr(transparent)]
31pub struct Registration<'target, 'd: 'target, L = num::NonZero<usize>> {
32    pub(super) inner: fifo::Registration<'target, ready::Retained<'d>, L>,
33}
34
35#[derive(Clone, Copy, Debug, Eq, PartialEq)]
36#[must_use = "a rejected wake registration must be handled"]
37pub enum Error {
38    Lane(resident::Bounds<num::NonZero<u32>>),
39    Active { position: wake::Position },
40    Draining { position: wake::Position },
41    Full { position: wake::Position },
42}
43
44impl From<fifo::Reject> for Error {
45    fn from(error: fifo::Reject) -> Self {
46        match error {
47            fifo::Reject::Active { index } => Self::Active {
48                position: wake::Position::new(index),
49            },
50            fifo::Reject::Draining { index } => Self::Draining {
51                position: wake::Position::new(index),
52            },
53            fifo::Reject::Full { index } => Self::Full {
54                position: wake::Position::new(index),
55            },
56        }
57    }
58}
59
60impl<'d, L> Queue<'d, L> {
61    pub fn try_new(
62        driver: driver::Reference<'d>,
63        capacity: usize,
64        lanes: L,
65    ) -> Result<Self, resident::Error>
66    where
67        L: resident::Extent,
68    {
69        Ok(Self {
70            driver,
71            inner: fifo::Fifo::try_new(capacity, lanes)?,
72        })
73    }
74}
75
76impl<'d> Queue<'d> {
77    /// Creates one queue and borrows its source under a fresh invariant scope.
78    pub fn try_with<S, R>(
79        driver: driver::Reference<'d>,
80        capacity: usize,
81        lanes: num::NonZero<usize>,
82        source: S,
83        run: impl for<'scope> FnOnce(&'scope S, Queue<'d, Scope<'scope>>) -> R,
84    ) -> Result<R, resident::Error> {
85        fifo::Fifo::try_with(capacity, lanes, source, |source, inner| {
86            run(source, Queue { driver, inner })
87        })
88    }
89
90    pub fn try_register<'target>(
91        &'target self,
92        position: usize,
93        slot: &mut Registration<'target, 'd>,
94        wake: &Waker<'d>,
95    ) -> Result<(), Error> {
96        let inner = self.inner.try_lane(position).map_err(Error::Lane)?;
97        inner
98            .try_keep(&mut slot.inner, *wake.key())
99            .map_err(Error::from)
100    }
101
102    pub fn try_wake_one(
103        &self,
104        position: usize,
105    ) -> Result<Once, resident::Bounds<num::NonZero<u32>>> {
106        let lane = self.inner.try_lane(position)?;
107        Ok(one(self.driver, lane.take_one()))
108    }
109}
110
111impl<'scope, 'd> Queue<'d, Scope<'scope>> {
112    pub fn try_index(
113        &self,
114        position: usize,
115    ) -> Result<Index<'scope>, resident::Bounds<num::NonZero<u32>>> {
116        self.inner.try_index(position)
117    }
118
119    pub fn try_register<'target>(
120        &'target self,
121        index: Index<'scope>,
122        slot: &mut Registration<'target, 'd, Scope<'scope>>,
123        wake: &Waker<'d>,
124    ) -> Result<(), fifo::Reject> {
125        self.inner
126            .lane(index)
127            .try_keep(&mut slot.inner, *wake.key())
128    }
129
130    pub fn wake_one(&self, index: Index<'scope>) -> Once {
131        one(self.driver, self.inner.lane(index).take_one())
132    }
133}
134
135impl<'d, const N: usize> Queue<'d, resident::Fixed<N>> {
136    pub fn try_register<'target, const I: usize>(
137        &'target self,
138        slot: &mut Registration<'target, 'd, resident::Fixed<N>>,
139        wake: &Waker<'d>,
140    ) -> Result<(), fifo::Reject> {
141        self.inner
142            .lane::<I>()
143            .try_keep(&mut slot.inner, *wake.key())
144    }
145
146    pub fn wake_one<const I: usize>(&self) -> Once {
147        one(self.driver, self.inner.lane::<I>().take_one())
148    }
149}
150
151impl<'d, L> Queue<'d, L> {
152    /// Cancels the registration, forwarding a taken wake within its original lane.
153    /// Active drains retain their work; successful admission uses `clear`.
154    pub fn cancel(&self, slot: &mut Registration<'_, 'd, L>) -> Once {
155        let Some(target) = slot.inner.cancel() else {
156            return Once::Empty;
157        };
158        target.wake(self.driver);
159        Once::Woke
160    }
161
162    /// Registers one wake in the constructor-proved first lane.
163    pub fn try_first<'target>(
164        &'target self,
165        slot: &mut Registration<'target, 'd, L>,
166        wake: &Waker<'d>,
167    ) -> Result<(), fifo::Reject> {
168        self.inner.first().try_keep(&mut slot.inner, *wake.key())
169    }
170
171    /// Wakes one registration from lane zero.
172    pub fn wake_first(&self) -> Once {
173        one(self.driver, self.inner.first().take_one())
174    }
175
176    pub fn wake_next(&self) -> Once {
177        let Some(target) = self.inner.take_next() else {
178            return Once::Empty;
179        };
180        target.wake(self.driver);
181        Once::Woke
182    }
183
184    pub fn has_waiters(&self) -> bool {
185        self.inner.has_waiters()
186    }
187
188    pub fn begin_drain(&self) {
189        self.inner.begin_drain();
190    }
191
192    pub fn is_draining(&self) -> bool {
193        self.inner.is_draining()
194    }
195
196    pub fn drain_one(&self, work: app::Work<'_, 'd>) -> Drain {
197        drain::one(
198            self.driver,
199            work,
200            || self.inner.is_draining() && self.inner.has_waiters(),
201            || match self.inner.drain_next() {
202                resident::Drain::Done => None,
203                resident::Drain::Item(target) => Some(Step::Item(target)),
204            },
205        )
206    }
207}
208
209impl<L> Default for Registration<'_, '_, L> {
210    fn default() -> Self {
211        Self::new()
212    }
213}
214
215impl<L> Registration<'_, '_, L> {
216    pub const fn new() -> Self {
217        Self {
218            inner: fifo::Registration::new(),
219        }
220    }
221
222    pub fn is_active(&self) -> bool {
223        self.inner.is_active()
224    }
225
226    pub fn clear(&mut self) {
227        *self = Self::new();
228    }
229}
230
231const _: () = {
232    assert!(mem::size_of::<Registration<'static, 'static>>() == 2 * mem::size_of::<usize>());
233    assert!(mem::align_of::<Registration<'static, 'static>>() == mem::align_of::<usize>());
234    assert!(
235        size_of::<Registration<'static, 'static, resident::Fixed<1>>>() == 2 * size_of::<usize>()
236    );
237    assert!(size_of::<Queue<'static, resident::Fixed<1>>>() == size_of::<Queue<'static>>());
238    assert!(size_of::<Queue<'static, Scope<'static>>>() == size_of::<Queue<'static>>());
239    assert!(size_of::<Index<'static>>() == size_of::<u32>());
240    assert!(size_of::<Registration<'static, 'static, Scope<'static>>>() == 2 * size_of::<usize>());
241};