tokio/runtime/scheduler/
defer.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
use std::cell::RefCell;
use std::task::Waker;

pub(crate) struct Defer {
    deferred: RefCell<Vec<Waker>>,
}

impl Defer {
    pub(crate) fn new() -> Defer {
        Defer {
            deferred: RefCell::default(),
        }
    }

    pub(crate) fn defer(&self, waker: &Waker) {
        let mut deferred = self.deferred.borrow_mut();

        // If the same task adds itself a bunch of times, then only add it once.
        if let Some(last) = deferred.last() {
            if last.will_wake(waker) {
                return;
            }
        }

        deferred.push(waker.clone());
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.deferred.borrow().is_empty()
    }

    pub(crate) fn wake(&self) {
        while let Some(waker) = self.deferred.borrow_mut().pop() {
            waker.wake();
        }
    }

    #[cfg(tokio_taskdump)]
    pub(crate) fn take_deferred(&self) -> Vec<Waker> {
        let mut deferred = self.deferred.borrow_mut();
        std::mem::take(&mut *deferred)
    }
}