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            }
30        }
31        BlockingSchedule {
32            #[cfg(feature = "test-util")]
33            handle: handle.clone(),
34            hooks: TaskHarnessScheduleHooks {
35                task_terminate_callback: handle.inner.hooks().task_terminate_callback.clone(),
36            },
37        }
38    }
39}
40
41impl task::Schedule for BlockingSchedule {
42    fn release(&self, _task: &Task<Self>) -> Option<Task<Self>> {
43        #[cfg(feature = "test-util")]
44        {
45            match &self.handle.inner {
46                scheduler::Handle::CurrentThread(handle) => {
47                    handle.driver.clock.allow_auto_advance();
48                    handle.driver.unpark();
49                }
50                #[cfg(feature = "rt-multi-thread")]
51                scheduler::Handle::MultiThread(_) => {}
52            }
53        }
54        None
55    }
56
57    fn schedule(&self, _task: task::Notified<Self>) {
58        unreachable!();
59    }
60
61    fn hooks(&self) -> TaskHarnessScheduleHooks {
62        TaskHarnessScheduleHooks {
63            task_terminate_callback: self.hooks.task_terminate_callback.clone(),
64        }
65    }
66}