tokio/runtime/scheduler/multi_thread/
handle.rs

1use crate::future::Future;
2use crate::loom::sync::Arc;
3use crate::runtime::scheduler::multi_thread::worker;
4use crate::runtime::task::{Notified, Task, TaskHarnessScheduleHooks};
5use crate::runtime::{
6    blocking, driver,
7    task::{self, JoinHandle, SpawnLocation},
8    TaskHooks, TaskMeta,
9};
10use crate::util::RngSeedGenerator;
11
12use std::fmt;
13
14mod metrics;
15
16cfg_taskdump! {
17    mod taskdump;
18}
19
20/// Handle to the multi thread scheduler
21pub(crate) struct Handle {
22    /// Task spawner
23    pub(super) shared: worker::Shared,
24
25    /// Resource driver handles
26    pub(crate) driver: driver::Handle,
27
28    /// Blocking pool spawner
29    pub(crate) blocking_spawner: blocking::Spawner,
30
31    /// Current random number generator seed
32    pub(crate) seed_generator: RngSeedGenerator,
33
34    /// User-supplied hooks to invoke for things
35    pub(crate) task_hooks: TaskHooks,
36}
37
38impl Handle {
39    /// Spawns a future onto the thread pool
40    pub(crate) fn spawn<F>(
41        me: &Arc<Self>,
42        future: F,
43        id: task::Id,
44        spawned_at: SpawnLocation,
45    ) -> JoinHandle<F::Output>
46    where
47        F: crate::future::Future + Send + 'static,
48        F::Output: Send + 'static,
49    {
50        Self::bind_new_task(me, future, id, spawned_at)
51    }
52
53    pub(crate) fn shutdown(&self) {
54        self.close();
55    }
56
57    #[track_caller]
58    pub(super) fn bind_new_task<T>(
59        me: &Arc<Self>,
60        future: T,
61        id: task::Id,
62        spawned_at: SpawnLocation,
63    ) -> JoinHandle<T::Output>
64    where
65        T: Future + Send + 'static,
66        T::Output: Send + 'static,
67    {
68        let (handle, notified) = me.shared.owned.bind(future, me.clone(), id, spawned_at);
69
70        me.task_hooks.spawn(&TaskMeta {
71            id,
72            spawned_at,
73            _phantom: Default::default(),
74        });
75
76        me.schedule_option_task_without_yield(notified);
77
78        handle
79    }
80}
81
82impl task::Schedule for Arc<Handle> {
83    fn release(&self, task: &Task<Self>) -> Option<Task<Self>> {
84        self.shared.owned.remove(task)
85    }
86
87    fn schedule(&self, task: Notified<Self>) {
88        self.schedule_task(task, false);
89    }
90
91    fn hooks(&self) -> TaskHarnessScheduleHooks {
92        TaskHarnessScheduleHooks {
93            task_terminate_callback: self.task_hooks.task_terminate_callback.clone(),
94        }
95    }
96
97    fn yield_now(&self, task: Notified<Self>) {
98        self.schedule_task(task, true);
99    }
100}
101
102cfg_unstable! {
103    use std::num::NonZeroU64;
104
105    impl Handle {
106        pub(crate) fn owned_id(&self) -> NonZeroU64 {
107            self.shared.owned.id
108        }
109    }
110}
111
112impl fmt::Debug for Handle {
113    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
114        fmt.debug_struct("multi_thread::Handle { ... }").finish()
115    }
116}