tokio/runtime/blocking/
schedule.rs

1#[cfg(feature = "test-util")]
2use crate::runtime::scheduler;
3use crate::runtime::task::{self, Task, TaskHarnessScheduleHooks};
4use crate::runtime::Handle;
5
6/// `task::Schedule` implementation that does nothing (except some bookkeeping
7/// in test-util builds). This is unique to the blocking scheduler as tasks
8/// scheduled are not really futures but blocking operations.
9///
10/// We avoid storing the task by forgetting it in `bind` and re-materializing it
11/// in `release`.
12pub(crate) struct BlockingSchedule {
13    #[cfg(feature = "test-util")]
14    handle: Handle,
15    hooks: TaskHarnessScheduleHooks,
16}
17
18impl BlockingSchedule {
19    #[cfg_attr(not(feature = "test-util"), allow(unused_variables))]
20    pub(crate) fn new(handle: &Handle) -> Self {
21        #[cfg(feature = "test-util")]
22        {
23            match &handle.inner {
24                scheduler::Handle::CurrentThread(handle) => {
25                    handle.driver.clock.inhibit_auto_advance();
26                }
27                #[cfg(feature = "rt-multi-thread")]
28                scheduler::Handle::MultiThread(_) => {}
29                #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))]
30                scheduler::Handle::MultiThreadAlt(_) => {}
31            }
32        }
33        BlockingSchedule {
34            #[cfg(feature = "test-util")]
35            handle: handle.clone(),
36            hooks: TaskHarnessScheduleHooks {
37                task_terminate_callback: handle.inner.hooks().task_terminate_callback.clone(),
38            },
39        }
40    }
41}
42
43impl task::Schedule for BlockingSchedule {
44    fn release(&self, _task: &Task<Self>) -> Option<Task<Self>> {
45        #[cfg(feature = "test-util")]
46        {
47            match &self.handle.inner {
48                scheduler::Handle::CurrentThread(handle) => {
49                    handle.driver.clock.allow_auto_advance();
50                    handle.driver.unpark();
51                }
52                #[cfg(feature = "rt-multi-thread")]
53                scheduler::Handle::MultiThread(_) => {}
54                #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))]
55                scheduler::Handle::MultiThreadAlt(_) => {}
56            }
57        }
58        None
59    }
60
61    fn schedule(&self, _task: task::Notified<Self>) {
62        unreachable!();
63    }
64
65    fn hooks(&self) -> TaskHarnessScheduleHooks {
66        TaskHarnessScheduleHooks {
67            task_terminate_callback: self.hooks.task_terminate_callback.clone(),
68        }
69    }
70}