dope_core/driver/timer/
registration.rs

1use std::cell::Cell;
2use std::pin::Pin;
3use std::task::Poll;
4use std::time::Instant;
5
6use pin_project::{pin_project, pinned_drop};
7
8use super::{Ticket, Timer, Waiter};
9use crate::driver::ready::CompletionWaker;
10
11/// An owning, starvation-safe timer registration.
12///
13/// The registration keeps the timer ticket linear and falls back to the
14/// driver's intrusive deadline queue when all fixed timer slots are busy.
15/// Consequently timer capacity is a performance bound, never a correctness
16/// bound.
17#[pin_project(PinnedDrop, !Unpin)]
18pub struct Registration<'timer, 'd> {
19    deadline: Cell<Option<Instant>>,
20    ticket: Cell<Option<Ticket>>,
21    #[pin]
22    waiter: Waiter<'timer, 'd>,
23}
24
25impl<'timer, 'd> Registration<'timer, 'd> {
26    pub fn new(timer: &'timer Timer<'d>) -> Self {
27        Self {
28            deadline: Cell::new(None),
29            ticket: Cell::new(None),
30            waiter: timer.waiter(),
31        }
32    }
33
34    pub fn with_deadline(timer: &'timer Timer<'d>, deadline: Instant) -> Self {
35        Self {
36            deadline: Cell::new(Some(deadline)),
37            ticket: Cell::new(None),
38            waiter: timer.waiter(),
39        }
40    }
41
42    pub fn is_armed(self: Pin<&Self>) -> bool {
43        self.project_ref().deadline.get().is_some()
44    }
45
46    pub fn arm(self: Pin<&Self>, deadline: Instant, wake: CompletionWaker<'d>) {
47        let this = self.project_ref();
48        let timer = this.waiter.as_ref().timer();
49        if let Some(ticket) = this.ticket.take() {
50            timer.cancel(ticket);
51        }
52        this.waiter.as_ref().unregister();
53        this.deadline.set(Some(deadline));
54        self.register(wake);
55    }
56
57    pub fn poll(self: Pin<&Self>, now: Instant, wake: CompletionWaker<'d>) -> Poll<()> {
58        let this = self.project_ref();
59        let Some(deadline) = this.deadline.get() else {
60            return Poll::Pending;
61        };
62        let timer = this.waiter.as_ref().timer();
63        if now >= deadline
64            || this
65                .ticket
66                .get()
67                .is_some_and(|ticket| timer.is_fired(ticket))
68        {
69            if let Some(ticket) = this.ticket.take() {
70                timer.cancel(ticket);
71            }
72            this.waiter.as_ref().unregister();
73            this.deadline.set(None);
74            return Poll::Ready(());
75        }
76        if let Some(ticket) = this.ticket.get() {
77            timer.replace_waker(ticket, wake);
78            this.waiter.as_ref().unregister();
79        } else {
80            self.register(wake);
81        }
82        Poll::Pending
83    }
84
85    pub fn cancel(self: Pin<&Self>) -> bool {
86        let this = self.project_ref();
87        let was_armed = this.deadline.take().is_some();
88        if let Some(ticket) = this.ticket.take() {
89            this.waiter.as_ref().timer().cancel(ticket);
90        }
91        this.waiter.as_ref().unregister();
92        was_armed
93    }
94
95    fn register(self: Pin<&Self>, wake: CompletionWaker<'d>) {
96        let this = self.project_ref();
97        let Some(deadline) = this.deadline.get() else {
98            return;
99        };
100        let waiter = this.waiter.as_ref();
101        match waiter.timer().try_arm(deadline, wake) {
102            Ok(armed) => {
103                this.ticket.set(Some(armed));
104                waiter.unregister();
105            }
106            Err(wake) => waiter.register(deadline, wake),
107        }
108    }
109}
110
111#[pinned_drop]
112impl PinnedDrop for Registration<'_, '_> {
113    fn drop(self: Pin<&mut Self>) {
114        self.as_ref().cancel();
115    }
116}
dope_core/driver/timer/registration.rs

dope_core/driver/timer/
registration.rs
