1 | use std::cell::RefCell; |
2 | use std::task::Waker; |
3 | |
4 | pub(crate) struct Defer { |
5 | deferred: RefCell<Vec<Waker>>, |
6 | } |
7 | |
8 | impl Defer { |
9 | pub(crate) fn new() -> Defer { |
10 | Defer { |
11 | deferred: RefCell::default(), |
12 | } |
13 | } |
14 | |
15 | pub(crate) fn defer(&self, waker: &Waker) { |
16 | let mut deferred = self.deferred.borrow_mut(); |
17 | |
18 | // If the same task adds itself a bunch of times, then only add it once. |
19 | if let Some(last) = deferred.last() { |
20 | if last.will_wake(waker) { |
21 | return; |
22 | } |
23 | } |
24 | |
25 | deferred.push(waker.clone()); |
26 | } |
27 | |
28 | pub(crate) fn is_empty(&self) -> bool { |
29 | self.deferred.borrow().is_empty() |
30 | } |
31 | |
32 | pub(crate) fn wake(&self) { |
33 | while let Some(waker) = self.deferred.borrow_mut().pop() { |
34 | waker.wake(); |
35 | } |
36 | } |
37 | |
38 | #[cfg (tokio_taskdump)] |
39 | pub(crate) fn take_deferred(&self) -> Vec<Waker> { |
40 | let mut deferred = self.deferred.borrow_mut(); |
41 | std::mem::take(&mut *deferred) |
42 | } |
43 | } |
44 | |