gpu_async_reactor/core/wait.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! The completion port and the driver that runs a [`WaitPlan`] over it.
4//!
5//! One wait, given the plan and the sleep the point's estimate asks for:
6//! 1. [`Park::Block`]: the port's blocking wait, nothing else.
7//! 2. Otherwise one status query; done → `Immediate` (an exact residual).
8//! 3. [`Park::FlagWait`]: arm the device word behind the work.
9//! 4. With a sleep: await the [`Timer`] to `start + sleep`, then query; done → `Late`
10//! (censored).
11//! 5. Park: [`Park::QueryPoll`] → the port polls to done; [`Park::FlagWait`] → bounded parks
12//! on the word with a status query every `backstop_ns`, then the port polls to done (the
13//! query is the authority). The outcome is an exact residual.
14//!
15//! Invariants: never returns `Ok` before the port's query or blocking wait reported done; a
16//! port error ends the wait with that error; no sleep is longer than the estimate allows.
17
18use crate::core::estimator::Outcome;
19use crate::core::plan::{Park, WaitPlan};
20use crate::core::time::{Clock, DurNs, MonoNs, Timer};
21
22/// One wait on device work, implemented per backend (the core names no vendor). `query` answers
23/// at once; the others take time and are awaited: a host that runs them on the waiting thread
24/// returns a ready future, an executor's pend. Errors are the port's own type and pass through
25/// [`wait`] unchanged.
26#[expect(
27 async_fn_in_trait,
28 reason = "the waiter is single-threaded by design; no Send bound is wanted"
29)]
30pub trait CompletionPort {
31 /// The backend's error.
32 type Error;
33
34 /// The status: `Ok(true)` done, `Ok(false)` pending; `Err` = the work failed.
35 ///
36 /// # Errors
37 /// The device reports a failure.
38 fn query(&mut self) -> Result<bool, Self::Error>;
39
40 /// [`Park::Block`]: the vendor's own blocking wait.
41 ///
42 /// # Errors
43 /// The device reports a failure.
44 async fn block(&mut self) -> Result<(), Self::Error>;
45
46 /// [`Park::QueryPoll`] and the tail of [`Park::FlagWait`]: query until done.
47 ///
48 /// # Errors
49 /// The device reports a failure.
50 async fn poll(&mut self) -> Result<(), Self::Error>;
51
52 /// [`Park::FlagWait`]: enqueue the completion word's write behind the pending work.
53 ///
54 /// # Errors
55 /// The enqueue fails.
56 fn arm(&mut self) -> Result<(), Self::Error>;
57
58 /// [`Park::FlagWait`]: park on the armed word for at most `bound`; `Ok(true)` once the word
59 /// has been written. May return `Ok(false)` early (a spurious wake).
60 ///
61 /// # Errors
62 /// The park itself fails (never the device status: that is `query`'s).
63 async fn park(&mut self, bound: DurNs) -> Result<bool, Self::Error>;
64}
65
66/// How one wait ended.
67#[derive(Clone, Copy, Debug, PartialEq, Eq)]
68pub enum WaitKind {
69 /// [`Park::Block`]: the port's blocking wait.
70 Blocked,
71 /// The first query found the work done.
72 Immediate,
73 /// The query right after the sleep found the work done (the sleep overshot).
74 Late,
75 /// The park saw the work finish.
76 Parked,
77}
78
79/// What one wait did, for the estimator (`outcome`) and the host's counters.
80#[derive(Clone, Copy, Debug, PartialEq, Eq)]
81pub struct WaitReport {
82 /// How it ended.
83 pub kind: WaitKind,
84 /// The sample for the point's estimate (`None` for [`WaitKind::Blocked`]).
85 pub outcome: Option<Outcome>,
86 /// Start to return.
87 pub wait_ns: u64,
88 /// Time asleep on the timer (0 when the wait did not sleep).
89 pub slept_ns: u64,
90 /// Parks on the device word (flag waits only).
91 pub parks: u64,
92 /// Status queries a flag wait's backstop issued.
93 pub backstop_queries: u64,
94}
95
96/// Run `plan` over `port`. `sleep_ns` is the point's estimate (`PointTable::sleep_ns`)
97/// and is used only when `plan.predict` is set.
98///
99/// # Errors
100/// Any port error, unchanged.
101pub async fn wait<P, C, T>(
102 plan: &WaitPlan,
103 sleep_ns: Option<u64>,
104 port: &mut P,
105 clock: &C,
106 timer: &T,
107) -> Result<WaitReport, P::Error>
108where
109 P: CompletionPort,
110 C: Clock,
111 T: Timer,
112{
113 let t0 = clock.monotonic_ns();
114 let since = |c: &C| c.monotonic_ns().since(t0).0;
115 let mut report = WaitReport {
116 kind: WaitKind::Blocked,
117 outcome: None,
118 wait_ns: 0,
119 slept_ns: 0,
120 parks: 0,
121 backstop_queries: 0,
122 };
123 if plan.park == Park::Block {
124 port.block().await?;
125 report.wait_ns = since(clock);
126 return Ok(report);
127 }
128 if port.query()? {
129 report.wait_ns = since(clock);
130 report.kind = WaitKind::Immediate;
131 report.outcome = Some(Outcome::Exact {
132 residual_ns: report.wait_ns,
133 });
134 return Ok(report);
135 }
136 if matches!(plan.park, Park::FlagWait { .. }) {
137 port.arm()?;
138 }
139 if let Some(s) = plan.predict.and(sleep_ns) {
140 timer.sleep_until(MonoNs(t0.0.saturating_add(s))).await;
141 report.slept_ns = since(clock);
142 if port.query()? {
143 report.wait_ns = since(clock);
144 report.kind = WaitKind::Late;
145 report.outcome = Some(Outcome::Late {
146 woke_ns: report.slept_ns,
147 });
148 return Ok(report);
149 }
150 }
151 if let Park::FlagWait {
152 park_ns,
153 backstop_ns,
154 } = plan.park
155 {
156 let mut last_query = clock.monotonic_ns();
157 loop {
158 report.parks += 1;
159 if port.park(DurNs(park_ns)).await? {
160 break;
161 }
162 let now = clock.monotonic_ns();
163 if now.since(last_query).0 >= backstop_ns {
164 report.backstop_queries += 1;
165 last_query = now;
166 if port.query()? {
167 break;
168 }
169 }
170 }
171 }
172 port.poll().await?;
173 report.wait_ns = since(clock);
174 report.kind = WaitKind::Parked;
175 report.outcome = Some(Outcome::Exact {
176 residual_ns: report.wait_ns,
177 });
178 Ok(report)
179}
180
181#[cfg(test)]
182#[path = "wait_tests.rs"]
183mod tests;