Skip to main content

dope_core/driver/schedule/timer/
registration.rs

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