tokio_threadpool/notifier.rs
1use pool::Pool;
2use task::Task;
3
4use std::mem;
5use std::ops;
6use std::sync::Arc;
7
8use futures::executor::Notify;
9
10/// Implements the future `Notify` API.
11///
12/// This is how external events are able to signal the task, informing it to try
13/// to poll the future again.
14#[derive(Debug)]
15pub(crate) struct Notifier {
16    pub pool: Arc<Pool>,
17}
18
19/// A guard that ensures that the inner value gets forgotten.
20#[derive(Debug)]
21struct Forget<T>(Option<T>);
22
23impl Notify for Notifier {
24    fn notify(&self, id: usize) {
25        trace!("Notifier::notify; id=0x{:x}", id);
26
27        unsafe {
28            let ptr = id as *const Task;
29
30            // We did not actually take ownership of the `Arc` in this function
31            // so we must ensure that the Arc is forgotten.
32            let task = Forget::new(Arc::from_raw(ptr));
33
34            // TODO: Unify this with Task::notify
35            if task.schedule() {
36                // TODO: Check if the pool is still running
37                //
38                // Bump the ref count
39                let task = task.clone();
40
41                let _ = self.pool.submit(task, &self.pool);
42            }
43        }
44    }
45
46    fn clone_id(&self, id: usize) -> usize {
47        let ptr = id as *const Task;
48
49        // This function doesn't actually get a strong ref to the task here.
50        // However, the only method we have to convert a raw pointer -> &Arc<T>
51        // is to call `Arc::from_raw` which returns a strong ref. So, to
52        // maintain the invariants, `t1` has to be forgotten. This prevents the
53        // ref count from being decremented.
54        let t1 = Forget::new(unsafe { Arc::from_raw(ptr) });
55
56        // The clone is forgotten so that the fn exits without decrementing the ref
57        // count. The caller of `clone_id` ensures that `drop_id` is called when
58        // the ref count needs to be decremented.
59        let _ = Forget::new(t1.clone());
60
61        id
62    }
63
64    fn drop_id(&self, id: usize) {
65        unsafe {
66            let ptr = id as *const Task;
67            let _ = Arc::from_raw(ptr);
68        }
69    }
70}
71
72// ===== impl Forget =====
73
74impl<T> Forget<T> {
75    fn new(t: T) -> Self {
76        Forget(Some(t))
77    }
78}
79
80impl<T> ops::Deref for Forget<T> {
81    type Target = T;
82
83    fn deref(&self) -> &T {
84        self.0.as_ref().unwrap()
85    }
86}
87
88impl<T> Drop for Forget<T> {
89    fn drop(&mut self) {
90        mem::forget(self.0.take());
91    }
92}