Skip to main content

axpoll/
lib.rs

1//! Typed readiness capabilities independent of a queue or scheduler implementation.
2
3#![no_std]
4#![deny(missing_docs)]
5
6extern crate alloc;
7
8use alloc::{boxed::Box, sync::Arc, vec::Vec};
9use core::{marker::PhantomData, task::Waker};
10
11use bitflags::bitflags;
12use linux_raw_sys::general::*;
13
14bitflags! {
15    /// I/O readiness events.
16    #[derive(Debug, Clone, Copy, Eq, PartialEq)]
17    pub struct IoEvents: u32 {
18        /// Available for read.
19        const IN     = POLLIN;
20        /// Urgent data for read.
21        const PRI    = POLLPRI;
22        /// Available for write.
23        const OUT    = POLLOUT;
24        /// Error condition.
25        const ERR    = POLLERR;
26        /// Hang up.
27        const HUP    = POLLHUP;
28        /// Invalid request.
29        const NVAL   = POLLNVAL;
30        /// Equivalent to [`IN`](Self::IN).
31        const RDNORM = POLLRDNORM;
32        /// Priority band data can be read.
33        const RDBAND = POLLRDBAND;
34        /// Equivalent to [`OUT`](Self::OUT).
35        const WRNORM = POLLWRNORM;
36        /// Priority data can be written.
37        const WRBAND = POLLWRBAND;
38        /// Message.
39        const MSG    = POLLMSG;
40        /// Remove.
41        const REMOVE = POLLREMOVE;
42        /// Stream socket peer closed connection, or shut down writing half.
43        const RDHUP  = POLLRDHUP;
44        /// Events reported even when callers did not request them.
45        const ALWAYS_POLL = Self::ERR.bits() | Self::HUP.bits();
46    }
47}
48
49/// Marker for a readiness observer that does not consume an event.
50pub enum SharedObserver {}
51
52/// Marker for a waiter that competes to consume one readiness event.
53pub enum ExclusiveConsumer {}
54
55/// Selection mode attached to one readiness registration.
56#[derive(Clone, Copy, Debug, Eq, PartialEq)]
57pub enum RegistrationMode {
58    /// Every matching observer is notified.
59    Shared,
60    /// One matching consumer is notified by an ordinary wake transaction.
61    Exclusive,
62}
63
64/// An owned source registration whose drop cancels the exact registered entry.
65pub trait PollRegistration: Send {
66    /// Returns whether the source selected this registration for notification.
67    ///
68    /// Once true, the result remains true until the registration is dropped.
69    /// Sources publish this state before invoking the registered waker.
70    fn was_notified(&self) -> bool;
71}
72
73/// A readiness source capable of creating owned registrations.
74pub trait PollSource: Send + Sync {
75    /// Registers `waker` and returns the exact cancellation lease.
76    ///
77    /// # Safety
78    ///
79    /// Registration is task/deferred-context only. The caller must not invoke
80    /// it from hard IRQ, NMI, or a trap callback. A producer must publish its
81    /// readiness state before waking this registration, and the consumer's
82    /// next readiness check must observe that publication through the same
83    /// lock or a matching Release/Acquire synchronization pair.
84    unsafe fn register(
85        &self,
86        waker: &Waker,
87        interests: IoEvents,
88        mode: RegistrationMode,
89    ) -> Option<Box<dyn PollRegistration>>;
90}
91
92impl<T: PollSource + ?Sized> PollSource for Arc<T> {
93    unsafe fn register(
94        &self,
95        waker: &Waker,
96        interests: IoEvents,
97        mode: RegistrationMode,
98    ) -> Option<Box<dyn PollRegistration>> {
99        unsafe { self.as_ref().register(waker, interests, mode) }
100    }
101}
102
103struct OwnedRegistration {
104    lease: Box<dyn PollRegistration>,
105    mode: RegistrationMode,
106}
107
108/// Owns every readiness registration made by one polling attempt.
109///
110/// Dropping or resetting a registrar cancels every still-live source lease.
111#[must_use = "dropping the registrar immediately cancels its poll registrations"]
112pub struct PollRegistrar<M> {
113    waker: Waker,
114    registrations: Vec<OwnedRegistration>,
115    mode: PhantomData<fn() -> M>,
116}
117
118impl<M> PollRegistrar<M> {
119    /// Creates an empty registrar for `waker`.
120    pub fn new(waker: &Waker) -> Self {
121        Self {
122            waker: waker.clone(),
123            registrations: Vec::new(),
124            mode: PhantomData,
125        }
126    }
127
128    /// Cancels the previous polling attempt and starts one with `waker`.
129    pub fn reset(&mut self, waker: &Waker) {
130        self.registrations.clear();
131        self.waker.clone_from(waker);
132    }
133
134    /// Cancels every registration owned by this registrar.
135    pub fn clear(&mut self) {
136        self.registrations.clear();
137    }
138
139    /// Returns whether this registrar currently owns no registration.
140    pub fn is_empty(&self) -> bool {
141        self.registrations.is_empty()
142    }
143
144    unsafe fn register_mode(
145        &mut self,
146        source: &dyn PollSource,
147        interests: IoEvents,
148        mode: RegistrationMode,
149    ) {
150        if interests.is_empty() {
151            return;
152        }
153        if let Some(lease) = unsafe { source.register(&self.waker, interests, mode) } {
154            self.registrations.push(OwnedRegistration { lease, mode });
155        }
156    }
157}
158
159impl PollRegistrar<SharedObserver> {
160    /// Registers this observer in `source` for `interests`.
161    ///
162    /// # Safety
163    ///
164    /// Registration is task/deferred-context only.
165    pub unsafe fn register(&mut self, source: &dyn PollSource, interests: IoEvents) {
166        unsafe { self.register_mode(source, interests, RegistrationMode::Shared) };
167    }
168}
169
170impl PollRegistrar<ExclusiveConsumer> {
171    /// Returns whether an exclusive source selected this polling attempt.
172    ///
173    /// Consumptive sources use this to transfer still-available readiness to
174    /// the next exclusive waiter without turning ordinary wakeups into a
175    /// broadcast.
176    pub fn was_exclusively_notified(&self) -> bool {
177        self.registrations.iter().any(|registration| {
178            registration.mode == RegistrationMode::Exclusive && registration.lease.was_notified()
179        })
180    }
181
182    /// Registers this consumer as an exclusive waiter.
183    ///
184    /// # Safety
185    ///
186    /// Registration is task/deferred-context only.
187    pub unsafe fn register_exclusive(&mut self, source: &dyn PollSource, interests: IoEvents) {
188        unsafe { self.register_mode(source, interests, RegistrationMode::Exclusive) };
189    }
190
191    /// Registers this consumer as a shared observer at a composite boundary.
192    ///
193    /// # Safety
194    ///
195    /// Registration is task/deferred-context only.
196    pub unsafe fn register_shared(&mut self, source: &dyn PollSource, interests: IoEvents) {
197        unsafe { self.register_mode(source, interests, RegistrationMode::Shared) };
198    }
199}
200
201/// Capability for adding shared registrations to an owned attempt.
202pub trait SharedRegistrationSink {
203    /// Returns the waker owned by this registration attempt.
204    fn waker(&self) -> &Waker;
205
206    /// Adds a shared registration owned by this sink.
207    ///
208    /// # Safety
209    ///
210    /// Registration is task/deferred-context only.
211    unsafe fn register_shared(&mut self, source: &dyn PollSource, interests: IoEvents);
212}
213
214/// Capability for adding exclusive registrations to an owned attempt.
215pub trait ExclusiveRegistrationSink {
216    /// Returns the waker owned by this registration attempt.
217    fn waker(&self) -> &Waker;
218
219    /// Adds an exclusive registration owned by this sink.
220    ///
221    /// # Safety
222    ///
223    /// Registration is task/deferred-context only.
224    unsafe fn register_exclusive(&mut self, source: &dyn PollSource, interests: IoEvents);
225
226    /// Borrows this sink's shared-registration capability.
227    fn as_shared(&mut self) -> &mut dyn SharedRegistrationSink;
228}
229
230impl SharedRegistrationSink for PollRegistrar<SharedObserver> {
231    fn waker(&self) -> &Waker {
232        &self.waker
233    }
234
235    unsafe fn register_shared(&mut self, source: &dyn PollSource, interests: IoEvents) {
236        unsafe { self.register(source, interests) };
237    }
238}
239
240impl SharedRegistrationSink for PollRegistrar<ExclusiveConsumer> {
241    fn waker(&self) -> &Waker {
242        &self.waker
243    }
244
245    unsafe fn register_shared(&mut self, source: &dyn PollSource, interests: IoEvents) {
246        unsafe { self.register_shared(source, interests) };
247    }
248}
249
250impl ExclusiveRegistrationSink for PollRegistrar<ExclusiveConsumer> {
251    fn waker(&self) -> &Waker {
252        &self.waker
253    }
254
255    unsafe fn register_exclusive(&mut self, source: &dyn PollSource, interests: IoEvents) {
256        unsafe { self.register_exclusive(source, interests) };
257    }
258
259    fn as_shared(&mut self) -> &mut dyn SharedRegistrationSink {
260        self
261    }
262}
263
264/// A value that reports I/O readiness and publishes owned registrations.
265pub trait Pollable {
266    /// Polls for I/O events.
267    fn poll(&self) -> IoEvents;
268
269    /// Registers a shared readiness observer.
270    ///
271    /// # Safety
272    ///
273    /// Registration is task/deferred-context only.
274    unsafe fn register_shared(&self, sink: &mut dyn SharedRegistrationSink, events: IoEvents);
275
276    /// Registers a consumer that may sleep until readiness changes.
277    ///
278    /// The default preserves shared-observer semantics. Consumptive sources
279    /// override it and use the exclusive capability.
280    ///
281    /// # Safety
282    ///
283    /// Registration is task/deferred-context only.
284    unsafe fn register_exclusive(
285        &self,
286        sink: &mut dyn ExclusiveRegistrationSink,
287        events: IoEvents,
288    ) {
289        unsafe { self.register_shared(sink.as_shared(), events) };
290    }
291}