agentos_kernel/
poll.rs

1use crate::socket_table::SocketId;
2use std::ops::{BitOr, BitOrAssign};
3use std::sync::{Arc, Condvar, Mutex, MutexGuard};
4use std::time::Duration;
5use web_time::Instant;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
8pub struct PollEvents(u16);
9
10impl PollEvents {
11    pub const fn empty() -> Self {
12        Self(0)
13    }
14
15    pub const fn from_bits(bits: u16) -> Self {
16        Self(bits)
17    }
18
19    pub const fn bits(self) -> u16 {
20        self.0
21    }
22
23    pub const fn is_empty(self) -> bool {
24        self.0 == 0
25    }
26
27    pub const fn contains(self, other: Self) -> bool {
28        self.0 & other.0 == other.0
29    }
30
31    pub const fn intersects(self, other: Self) -> bool {
32        self.0 & other.0 != 0
33    }
34}
35
36impl BitOr for PollEvents {
37    type Output = Self;
38
39    fn bitor(self, rhs: Self) -> Self::Output {
40        Self(self.0 | rhs.0)
41    }
42}
43
44impl BitOrAssign for PollEvents {
45    fn bitor_assign(&mut self, rhs: Self) {
46        self.0 |= rhs.0;
47    }
48}
49
50pub const POLLIN: PollEvents = PollEvents(0x0001);
51pub const POLLOUT: PollEvents = PollEvents(0x0004);
52pub const POLLERR: PollEvents = PollEvents(0x0008);
53pub const POLLHUP: PollEvents = PollEvents(0x0010);
54pub const POLLNVAL: PollEvents = PollEvents(0x0020);
55
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
57pub struct PollFd {
58    pub fd: u32,
59    pub events: PollEvents,
60    pub revents: PollEvents,
61}
62
63impl PollFd {
64    pub const fn new(fd: u32, events: PollEvents) -> Self {
65        Self {
66            fd,
67            events,
68            revents: PollEvents::empty(),
69        }
70    }
71}
72
73#[derive(Debug, Clone, PartialEq, Eq)]
74pub struct PollResult {
75    pub ready_count: usize,
76    pub fds: Vec<PollFd>,
77}
78
79#[derive(Debug, Clone, Copy, PartialEq, Eq)]
80pub enum PollTarget {
81    Fd(u32),
82    Socket(SocketId),
83}
84
85#[derive(Debug, Clone, Copy, PartialEq, Eq)]
86pub struct PollTargetEntry {
87    pub target: PollTarget,
88    pub events: PollEvents,
89    pub revents: PollEvents,
90}
91
92impl PollTargetEntry {
93    pub const fn new(target: PollTarget, events: PollEvents) -> Self {
94        Self {
95            target,
96            events,
97            revents: PollEvents::empty(),
98        }
99    }
100
101    pub const fn fd(fd: u32, events: PollEvents) -> Self {
102        Self::new(PollTarget::Fd(fd), events)
103    }
104
105    pub const fn socket(socket_id: SocketId, events: PollEvents) -> Self {
106        Self::new(PollTarget::Socket(socket_id), events)
107    }
108}
109
110#[derive(Debug, Clone, PartialEq, Eq)]
111pub struct PollTargetResult {
112    pub ready_count: usize,
113    pub targets: Vec<PollTargetEntry>,
114}
115
116/// Cloneable, Send handle onto the kernel's poll notifier so a caller can wait
117/// for "some poll-visible state changed" OFF the thread that owns the kernel
118/// (deferred readiness servicing), then re-check readiness with a zero-timeout
119/// poll on the owning thread. Take `snapshot()` BEFORE the readiness check it
120/// guards so a change landing between the check and the wait wakes immediately
121/// instead of being lost.
122#[derive(Debug, Clone)]
123pub struct PollWaitHandle {
124    notifier: PollNotifier,
125}
126
127impl PollWaitHandle {
128    pub(crate) fn new(notifier: PollNotifier) -> Self {
129        Self { notifier }
130    }
131
132    /// Snapshot the current change generation.
133    pub fn snapshot(&self) -> u64 {
134        self.notifier.snapshot()
135    }
136
137    /// Block until the generation moves past `observed` or `timeout` elapses
138    /// (`None` = wait forever). Returns true when a change was observed.
139    pub fn wait_for_change(&self, observed: u64, timeout: Option<Duration>) -> bool {
140        self.notifier.wait_for_change(observed, timeout)
141    }
142}
143
144#[derive(Debug, Clone, Default)]
145pub(crate) struct PollNotifier {
146    inner: Arc<PollNotifierInner>,
147}
148
149#[derive(Debug, Default)]
150struct PollNotifierInner {
151    generation: Mutex<u64>,
152    waiters: Condvar,
153}
154
155impl PollNotifier {
156    pub(crate) fn notify(&self) {
157        let mut generation = lock_or_recover(&self.inner.generation);
158        *generation = generation.wrapping_add(1);
159        self.inner.waiters.notify_all();
160    }
161
162    pub(crate) fn snapshot(&self) -> u64 {
163        *lock_or_recover(&self.inner.generation)
164    }
165
166    pub(crate) fn wait_for_change(&self, observed: u64, timeout: Option<Duration>) -> bool {
167        let mut generation = lock_or_recover(&self.inner.generation);
168        if *generation != observed {
169            return true;
170        }
171
172        let Some(timeout) = timeout else {
173            while *generation == observed {
174                generation = wait_or_recover(&self.inner.waiters, generation);
175            }
176            return true;
177        };
178
179        if timeout.is_zero() {
180            return *generation != observed;
181        }
182
183        let deadline = Instant::now() + timeout;
184        loop {
185            let now = Instant::now();
186            if now >= deadline {
187                return *generation != observed;
188            }
189
190            let remaining = deadline.saturating_duration_since(now);
191            let (next_generation, wait_result) =
192                wait_timeout_or_recover(&self.inner.waiters, generation, remaining);
193            generation = next_generation;
194            if *generation != observed {
195                return true;
196            }
197            if wait_result.timed_out() {
198                return false;
199            }
200        }
201    }
202}
203
204fn lock_or_recover<'a, T>(mutex: &'a Mutex<T>) -> MutexGuard<'a, T> {
205    match mutex.lock() {
206        Ok(guard) => guard,
207        Err(poisoned) => poisoned.into_inner(),
208    }
209}
210
211fn wait_or_recover<'a, T>(condvar: &Condvar, guard: MutexGuard<'a, T>) -> MutexGuard<'a, T> {
212    match condvar.wait(guard) {
213        Ok(guard) => guard,
214        Err(poisoned) => poisoned.into_inner(),
215    }
216}
217
218fn wait_timeout_or_recover<'a, T>(
219    condvar: &Condvar,
220    guard: MutexGuard<'a, T>,
221    timeout: Duration,
222) -> (MutexGuard<'a, T>, std::sync::WaitTimeoutResult) {
223    match condvar.wait_timeout(guard, timeout) {
224        Ok(result) => result,
225        Err(poisoned) => poisoned.into_inner(),
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::PollNotifier;
232    use std::sync::mpsc;
233    use std::thread;
234    use std::time::Duration;
235
236    #[test]
237    fn infinite_wait_returns_after_notification_without_waiter_storage() {
238        let notifier = PollNotifier::default();
239        let observed = notifier.snapshot();
240        let waiter = notifier.clone();
241        let (started_tx, started_rx) = mpsc::channel();
242        let (done_tx, done_rx) = mpsc::channel();
243
244        let handle = thread::spawn(move || {
245            started_tx.send(()).expect("signal waiter start");
246            let changed = waiter.wait_for_change(observed, None);
247            done_tx.send(changed).expect("signal waiter result");
248        });
249
250        started_rx.recv().expect("waiter should start");
251        assert!(
252            done_rx.recv_timeout(Duration::from_millis(25)).is_err(),
253            "waiter should stay blocked before notification"
254        );
255
256        notifier.notify();
257        assert!(done_rx
258            .recv_timeout(Duration::from_secs(1))
259            .expect("waiter should wake after notification"));
260        handle.join().expect("waiter thread should finish");
261    }
262
263    #[test]
264    fn saturated_generation_still_notifies_waiters() {
265        let notifier = PollNotifier::default();
266        {
267            let mut generation = super::lock_or_recover(&notifier.inner.generation);
268            *generation = u64::MAX;
269        }
270
271        let observed = notifier.snapshot();
272        let waiter = notifier.clone();
273        let (started_tx, started_rx) = mpsc::channel();
274        let (done_tx, done_rx) = mpsc::channel();
275
276        let handle = thread::spawn(move || {
277            started_tx.send(()).expect("signal waiter start");
278            let changed = waiter.wait_for_change(observed, Some(Duration::from_secs(1)));
279            done_tx.send(changed).expect("signal waiter result");
280        });
281
282        started_rx.recv().expect("waiter should start");
283        notifier.notify();
284
285        assert!(
286            done_rx
287                .recv_timeout(Duration::from_secs(2))
288                .expect("waiter should return after saturated notify"),
289            "saturated notify should still wake the waiter"
290        );
291        handle.join().expect("waiter thread should finish");
292    }
293}
agentos_kernel/poll.rs

agentos_kernel/
poll.rs
