Skip to main content

tokio/runtime/scheduler/current_thread/
mod.rs

1use crate::loom::sync::atomic::AtomicBool;
2use crate::loom::sync::Arc;
3use crate::runtime::driver::{self, Driver};
4use crate::runtime::scheduler::{self, Defer, Inject};
5use crate::runtime::task::{
6    self, JoinHandle, LocalNotified, OwnedTasks, Schedule, SpawnLocation, Task,
7    TaskHarnessScheduleHooks,
8};
9use crate::runtime::{
10    blocking, context, Config, MetricsBatch, SchedulerMetrics, TaskHooks, TaskMeta, WorkerMetrics,
11};
12use crate::sync::notify::Notify;
13use crate::util::atomic_cell::AtomicCell;
14use crate::util::{waker_ref, RngSeedGenerator, Wake, WakerRef};
15
16use std::cell::RefCell;
17use std::collections::VecDeque;
18use std::future::{poll_fn, Future};
19use std::pin::Pin;
20use std::sync::atomic::Ordering::{AcqRel, Acquire, Release};
21use std::task::Poll::{Pending, Ready};
22use std::task::Waker;
23use std::thread::ThreadId;
24use std::time::Duration;
25use std::time::Instant;
26use std::{fmt, thread};
27
28/// Executes tasks on the current thread
29pub(crate) struct CurrentThread {
30    /// Core scheduler data is acquired by a thread entering `block_on`.
31    core: AtomicCell<Core>,
32
33    /// Notifier for waking up other threads to steal the
34    /// driver.
35    notify: Notify,
36}
37
38/// Handle to the current thread scheduler
39pub(crate) struct Handle {
40    /// The name of the runtime
41    name: Option<String>,
42
43    /// Scheduler state shared across threads
44    shared: Shared,
45
46    /// Resource driver handles
47    pub(crate) driver: driver::Handle,
48
49    /// Blocking pool spawner
50    pub(crate) blocking_spawner: blocking::Spawner,
51
52    /// Current random number generator seed
53    pub(crate) seed_generator: RngSeedGenerator,
54
55    /// User-supplied hooks to invoke for things
56    pub(crate) task_hooks: TaskHooks,
57
58    /// If this is a `LocalRuntime`, flags the owning thread ID.
59    pub(crate) local_tid: Option<ThreadId>,
60}
61
62/// Data required for executing the scheduler. The struct is passed around to
63/// a function that will perform the scheduling work and acts as a capability token.
64struct Core {
65    /// Scheduler run queue
66    tasks: VecDeque<Notified>,
67
68    /// Current tick
69    tick: u32,
70
71    /// Runtime driver
72    ///
73    /// The driver is removed before starting to park the thread
74    driver: Option<Driver>,
75
76    /// Metrics batch
77    metrics: MetricsBatch,
78
79    /// How often to check the global queue
80    global_queue_interval: u32,
81
82    /// True if a task panicked without being handled and the runtime is
83    /// configured to shutdown on unhandled panic.
84    unhandled_panic: bool,
85}
86
87/// Scheduler state shared between threads.
88struct Shared {
89    /// Remote run queue
90    inject: Inject<Arc<Handle>>,
91
92    /// Collection of all active tasks spawned onto this executor.
93    owned: OwnedTasks<Arc<Handle>>,
94
95    /// Indicates whether the blocked on thread was woken.
96    woken: AtomicBool,
97
98    /// Scheduler configuration options
99    config: Config,
100
101    /// Keeps track of various runtime metrics.
102    scheduler_metrics: SchedulerMetrics,
103
104    /// This scheduler only has one worker.
105    worker_metrics: WorkerMetrics,
106
107    /// Startup time of this scheduler.
108    ///
109    /// This instant is used as the basis of task `scheduled_at` measurements.
110    schedule_latency_start: Option<Instant>,
111}
112
113/// Thread-local context.
114///
115/// pub(crate) to store in `runtime::context`.
116pub(crate) struct Context {
117    /// Scheduler handle
118    handle: Arc<Handle>,
119
120    /// Scheduler core, enabling the holder of `Context` to execute the
121    /// scheduler.
122    core: RefCell<Option<Box<Core>>>,
123
124    /// Deferred tasks, usually ones that called `task::yield_now()`.
125    pub(crate) defer: Defer,
126}
127
128type Notified = task::Notified<Arc<Handle>>;
129
130/// Initial queue capacity.
131const INITIAL_CAPACITY: usize = 64;
132
133/// Used if none is specified. This is a temporary constant and will be removed
134/// as we unify tuning logic between the multi-thread and current-thread
135/// schedulers.
136const DEFAULT_GLOBAL_QUEUE_INTERVAL: u32 = 31;
137
138impl CurrentThread {
139    pub(crate) fn new(
140        driver: Driver,
141        driver_handle: driver::Handle,
142        blocking_spawner: blocking::Spawner,
143        seed_generator: RngSeedGenerator,
144        config: Config,
145        local_tid: Option<ThreadId>,
146        name: Option<String>,
147    ) -> (CurrentThread, Arc<Handle>) {
148        let worker_metrics = WorkerMetrics::from_config(&config);
149        worker_metrics.set_thread_id(thread::current().id());
150
151        // Get the configured global queue interval, or use the default.
152        let global_queue_interval = config
153            .global_queue_interval
154            .unwrap_or(DEFAULT_GLOBAL_QUEUE_INTERVAL);
155
156        let schedule_latency_start = config.track_task_schedule_latency.then(Instant::now);
157
158        let handle = Arc::new(Handle {
159            name,
160            task_hooks: TaskHooks {
161                task_spawn_callback: config.before_spawn.clone(),
162                task_terminate_callback: config.after_termination.clone(),
163                #[cfg(tokio_unstable)]
164                before_poll_callback: config.before_poll.clone(),
165                #[cfg(tokio_unstable)]
166                after_poll_callback: config.after_poll.clone(),
167            },
168            shared: Shared {
169                inject: Inject::new(),
170                owned: OwnedTasks::new(1),
171                woken: AtomicBool::new(false),
172                config,
173                scheduler_metrics: SchedulerMetrics::new(),
174                worker_metrics,
175                schedule_latency_start,
176            },
177            driver: driver_handle,
178            blocking_spawner,
179            seed_generator,
180            local_tid,
181        });
182
183        let core = AtomicCell::new(Some(Box::new(Core {
184            tasks: VecDeque::with_capacity(INITIAL_CAPACITY),
185            tick: 0,
186            driver: Some(driver),
187            metrics: MetricsBatch::new(&handle.shared.worker_metrics),
188            global_queue_interval,
189            unhandled_panic: false,
190        })));
191
192        let scheduler = CurrentThread {
193            core,
194            notify: Notify::new(),
195        };
196
197        (scheduler, handle)
198    }
199
200    #[track_caller]
201    pub(crate) fn block_on<F: Future>(&self, handle: &scheduler::Handle, future: F) -> F::Output {
202        pin!(future);
203
204        crate::runtime::context::enter_runtime(handle, false, |blocking| {
205            let handle = handle.as_current_thread();
206
207            // Attempt to steal the scheduler core and block_on the future if we can
208            // there, otherwise, lets select on a notification that the core is
209            // available or the future is complete.
210            loop {
211                if let Some(core) = self.take_core(handle) {
212                    handle
213                        .shared
214                        .worker_metrics
215                        .set_thread_id(thread::current().id());
216                    return core.block_on(future);
217                } else {
218                    let notified = self.notify.notified();
219                    pin!(notified);
220
221                    if let Some(out) = blocking
222                        .block_on(poll_fn(|cx| {
223                            if notified.as_mut().poll(cx).is_ready() {
224                                return Ready(None);
225                            }
226
227                            if let Ready(out) = future.as_mut().poll(cx) {
228                                return Ready(Some(out));
229                            }
230
231                            Pending
232                        }))
233                        .expect("Failed to `Enter::block_on`")
234                    {
235                        return out;
236                    }
237                }
238            }
239        })
240    }
241
242    fn take_core(&self, handle: &Arc<Handle>) -> Option<CoreGuard<'_>> {
243        let core = self.core.take()?;
244
245        Some(CoreGuard {
246            context: scheduler::Context::CurrentThread(Context {
247                handle: handle.clone(),
248                core: RefCell::new(Some(core)),
249                defer: Defer::new(),
250            }),
251            scheduler: self,
252        })
253    }
254
255    pub(crate) fn shutdown(&mut self, handle: &scheduler::Handle) {
256        let handle = handle.as_current_thread();
257
258        // Avoid a double panic if we are currently panicking and
259        // the lock may be poisoned.
260
261        let core = match self.take_core(handle) {
262            Some(core) => core,
263            None if std::thread::panicking() => return,
264            None => panic!("Oh no! We never placed the Core back, this is a bug!"),
265        };
266
267        // Check that the thread-local is not being destroyed
268        let tls_available = context::with_current(|_| ()).is_ok();
269
270        if tls_available {
271            core.enter(|core, _context| {
272                let core = shutdown2(core, handle);
273                (core, ())
274            });
275        } else {
276            // Shutdown without setting the context. `tokio::spawn` calls will
277            // fail, but those will fail either way because the thread-local is
278            // not available anymore.
279            let context = core.context.expect_current_thread();
280            let core = context.core.borrow_mut().take().unwrap();
281
282            let core = shutdown2(core, handle);
283            *context.core.borrow_mut() = Some(core);
284        }
285    }
286}
287
288fn shutdown2(mut core: Box<Core>, handle: &Handle) -> Box<Core> {
289    // Drain the OwnedTasks collection. This call also closes the
290    // collection, ensuring that no tasks are ever pushed after this
291    // call returns.
292    handle.shared.owned.close_and_shutdown_all(0);
293
294    // Drain local queue
295    // We already shut down every task, so we just need to drop the task.
296    while let Some(task) = core.next_local_task(handle) {
297        drop(task);
298    }
299
300    // Close the injection queue
301    handle.shared.inject.close();
302
303    // Drain remote queue
304    while let Some(task) = handle.shared.inject.pop() {
305        drop(task);
306    }
307
308    assert!(handle.shared.owned.is_empty());
309
310    // Submit metrics
311    core.submit_metrics(handle);
312
313    // Shutdown the resource drivers
314    if let Some(driver) = core.driver.as_mut() {
315        driver.shutdown(&handle.driver);
316    }
317
318    core
319}
320
321impl fmt::Debug for CurrentThread {
322    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
323        fmt.debug_struct("CurrentThread").finish()
324    }
325}
326
327// ===== impl Core =====
328
329impl Core {
330    /// Get and increment the current tick
331    fn tick(&mut self) {
332        self.tick = self.tick.wrapping_add(1);
333    }
334
335    fn next_task(&mut self, handle: &Handle) -> Option<Notified> {
336        if self.tick % self.global_queue_interval == 0 {
337            handle
338                .next_remote_task()
339                .or_else(|| self.next_local_task(handle))
340        } else {
341            self.next_local_task(handle)
342                .or_else(|| handle.next_remote_task())
343        }
344    }
345
346    fn next_local_task(&mut self, handle: &Handle) -> Option<Notified> {
347        let ret = self.tasks.pop_front();
348        handle
349            .shared
350            .worker_metrics
351            .set_queue_depth(self.tasks.len());
352        ret
353    }
354
355    fn push_task(&mut self, handle: &Handle, task: Notified) {
356        self.tasks.push_back(task);
357        self.metrics.inc_local_schedule_count();
358        handle
359            .shared
360            .worker_metrics
361            .set_queue_depth(self.tasks.len());
362    }
363
364    fn submit_metrics(&mut self, handle: &Handle) {
365        self.metrics.submit(&handle.shared.worker_metrics, 0);
366    }
367}
368
369#[cfg(all(tokio_unstable, feature = "taskdump"))]
370fn wake_deferred_tasks_and_free(context: &Context) {
371    let wakers = context.defer.take_deferred();
372    for waker in wakers {
373        waker.wake();
374    }
375}
376
377// ===== impl Context =====
378
379impl Context {
380    /// Execute the closure with the given scheduler core stored in the
381    /// thread-local context.
382    fn run_task(&self, task: LocalNotified<Arc<Handle>>, mut core: Box<Core>) -> Box<Core> {
383        let schedule_latency_context = task
384            .get_scheduled_at()
385            .prepare(self.handle.shared.schedule_latency_start);
386        let _task_schedule_latency = core.metrics.start_poll(schedule_latency_context);
387
388        #[cfg(tokio_unstable)]
389        let task_meta = task.task_meta(_task_schedule_latency);
390
391        let (mut c, ()) = self.enter(core, || {
392            crate::task::coop::budget(|| {
393                #[cfg(tokio_unstable)]
394                self.handle.task_hooks.poll_start_callback(&task_meta);
395
396                task.run();
397
398                #[cfg(tokio_unstable)]
399                self.handle.task_hooks.poll_stop_callback(&task_meta);
400            })
401        });
402        c.metrics.end_poll();
403        c
404    }
405
406    /// Blocks the current thread until an event is received by the driver,
407    /// including I/O events, timer events, ...
408    fn park(&self, mut core: Box<Core>, handle: &Handle) -> Box<Core> {
409        let mut driver = core.driver.take().expect("driver missing");
410
411        if let Some(f) = &handle.shared.config.before_park {
412            let (c, ()) = self.enter(core, || f());
413            core = c;
414        }
415
416        if !self.has_pending_work(&core) {
417            // Park until the thread is signaled
418            core.metrics.about_to_park();
419            core.submit_metrics(handle);
420
421            core = self.park_internal(core, handle, &mut driver, None);
422
423            core.metrics.unparked();
424            core.submit_metrics(handle);
425        } else {
426            // `before_park` scheduled work (e.g. an `on_thread_park` hook that woke the
427            // `block_on` future), so we don't block. We must still poll the driver once
428            // without blocking, or timer and I/O events would stall under a runtime driven
429            // by repeated short `block_on` calls. See
430            // <https://github.com/tokio-rs/tokio/issues/8212>.
431            core.submit_metrics(handle);
432
433            core = self.park_internal(core, handle, &mut driver, Some(Duration::from_millis(0)));
434        }
435
436        if let Some(f) = &handle.shared.config.after_unpark {
437            let (c, ()) = self.enter(core, || f());
438            core = c;
439        }
440
441        core.driver = Some(driver);
442        core
443    }
444
445    /// Checks the driver for new events without blocking the thread.
446    fn park_yield(&self, mut core: Box<Core>, handle: &Handle) -> Box<Core> {
447        let mut driver = core.driver.take().expect("driver missing");
448
449        core.submit_metrics(handle);
450
451        core = self.park_internal(core, handle, &mut driver, Some(Duration::from_millis(0)));
452
453        core.driver = Some(driver);
454        core
455    }
456
457    fn has_pending_work(&self, core: &Core) -> bool {
458        !core.tasks.is_empty() || !self.defer.is_empty() || self.handle.shared.woken.load(Acquire)
459    }
460
461    fn park_internal(
462        &self,
463        core: Box<Core>,
464        handle: &Handle,
465        driver: &mut Driver,
466        duration: Option<Duration>,
467    ) -> Box<Core> {
468        let (core, ()) = self.enter(core, || {
469            match duration {
470                Some(dur) => driver.park_timeout(&handle.driver, dur),
471                None => driver.park(&handle.driver),
472            }
473            self.defer.wake();
474        });
475
476        core
477    }
478
479    fn enter<R>(&self, core: Box<Core>, f: impl FnOnce() -> R) -> (Box<Core>, R) {
480        // Store the scheduler core in the thread-local context
481        //
482        // A drop-guard is employed at a higher level.
483        *self.core.borrow_mut() = Some(core);
484
485        // Execute the closure while tracking the execution budget
486        let ret = f();
487
488        // Take the scheduler core back
489        let core = self.core.borrow_mut().take().expect("core missing");
490        (core, ret)
491    }
492
493    pub(crate) fn defer(&self, waker: &Waker) {
494        self.defer.defer(waker);
495    }
496}
497
498// ===== impl Handle =====
499
500impl Handle {
501    /// Spawns a future onto the `CurrentThread` scheduler
502    #[track_caller]
503    pub(crate) fn spawn<F>(
504        me: &Arc<Self>,
505        future: F,
506        id: crate::runtime::task::Id,
507        spawned_at: SpawnLocation,
508    ) -> JoinHandle<F::Output>
509    where
510        F: crate::future::Future + Send + 'static,
511        F::Output: Send + 'static,
512    {
513        let (handle, notified) = me.shared.owned.bind(future, me.clone(), id, spawned_at);
514
515        me.task_hooks.spawn(&TaskMeta {
516            id,
517            spawned_at,
518            #[cfg(all(tokio_unstable, feature = "schedule-latency"))]
519            schedule_latency: None,
520            _phantom: Default::default(),
521        });
522
523        if let Some(notified) = notified {
524            me.schedule(notified);
525        }
526
527        handle
528    }
529
530    /// Spawn a task which isn't safe to send across thread boundaries onto the runtime.
531    ///
532    /// # Safety
533    ///
534    /// This should only be used when this is a `LocalRuntime` or in another case where the runtime
535    /// provably cannot be driven from or moved to different threads from the one on which the task
536    /// is spawned.
537    #[track_caller]
538    pub(crate) unsafe fn spawn_local<F>(
539        me: &Arc<Self>,
540        future: F,
541        id: crate::runtime::task::Id,
542        spawned_at: SpawnLocation,
543    ) -> JoinHandle<F::Output>
544    where
545        F: crate::future::Future + 'static,
546        F::Output: 'static,
547    {
548        // Safety: the caller guarantees that this is only called on a `LocalRuntime`.
549        let (handle, notified) = unsafe {
550            me.shared
551                .owned
552                .bind_local(future, me.clone(), id, spawned_at)
553        };
554
555        me.task_hooks.spawn(&TaskMeta {
556            id,
557            spawned_at,
558            #[cfg(all(tokio_unstable, feature = "schedule-latency"))]
559            schedule_latency: None,
560            _phantom: Default::default(),
561        });
562
563        if let Some(notified) = notified {
564            me.schedule(notified);
565        }
566
567        handle
568    }
569
570    /// Capture a snapshot of this runtime's state.
571    #[cfg(all(
572        tokio_unstable,
573        feature = "taskdump",
574        target_os = "linux",
575        any(
576            target_arch = "aarch64",
577            target_arch = "x86",
578            target_arch = "x86_64",
579            target_arch = "s390x"
580        )
581    ))]
582    pub(crate) fn dump(self: &Arc<Self>) -> crate::runtime::Dump {
583        use crate::runtime::dump;
584        use scheduler::Context::CurrentThread;
585        use task::trace::trace_current_thread;
586
587        let mut traces = vec![];
588
589        // todo: how to make this work outside of a runtime context?
590        context::with_scheduler(|maybe_context| {
591            // drain the local queue
592            let context = match maybe_context {
593                Some(CurrentThread(context)) if Arc::ptr_eq(self, &context.handle) => context,
594                _ => return,
595            };
596            let mut maybe_core = context.core.borrow_mut();
597            let core = if let Some(core) = maybe_core.as_mut() {
598                core
599            } else {
600                return;
601            };
602            let local = &mut core.tasks;
603
604            if self.shared.inject.is_closed() {
605                return;
606            }
607
608            traces = trace_current_thread(&self.shared.owned, local, &self.shared.inject)
609                .into_iter()
610                .map(|(id, trace)| dump::Task::new(id, trace))
611                .collect();
612
613            // Avoid double borrow panic
614            drop(maybe_core);
615
616            // Taking a taskdump could wakes every task, but we probably don't want
617            // the `yield_now` vector to be that large under normal circumstances.
618            // Therefore, we free its allocation.
619            wake_deferred_tasks_and_free(context);
620        });
621
622        dump::Dump::new(traces)
623    }
624
625    fn next_remote_task(&self) -> Option<Notified> {
626        self.shared.inject.pop()
627    }
628
629    fn waker_ref(me: &Arc<Self>) -> WakerRef<'_> {
630        // Set woken to true when enter block_on, ensure outer future
631        // be polled for the first time when enter loop
632        me.shared.woken.store(true, Release);
633        waker_ref(me)
634    }
635
636    // reset woken to false and return original value
637    pub(crate) fn reset_woken(&self) -> bool {
638        self.shared.woken.swap(false, AcqRel)
639    }
640
641    pub(crate) fn num_alive_tasks(&self) -> usize {
642        self.shared.owned.num_alive_tasks()
643    }
644
645    pub(crate) fn injection_queue_depth(&self) -> usize {
646        self.shared.inject.len()
647    }
648
649    pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics {
650        assert_eq!(0, worker);
651        &self.shared.worker_metrics
652    }
653}
654
655cfg_unstable_metrics! {
656    impl Handle {
657        pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics {
658            &self.shared.scheduler_metrics
659        }
660
661        pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize {
662            self.worker_metrics(worker).queue_depth()
663        }
664
665        pub(crate) fn num_blocking_threads(&self) -> usize {
666            self.blocking_spawner.num_blocking_threads()
667        }
668
669        pub(crate) fn num_idle_blocking_threads(&self) -> usize {
670            self.blocking_spawner.num_idle_threads()
671        }
672
673        pub(crate) fn blocking_queue_depth(&self) -> usize {
674            self.blocking_spawner.queue_depth()
675        }
676
677        cfg_64bit_metrics! {
678            pub(crate) fn spawned_tasks_count(&self) -> u64 {
679                self.shared.owned.spawned_tasks_count()
680            }
681        }
682    }
683}
684
685use crate::runtime::metrics::ScheduleLatencyInstant;
686use std::num::NonZeroU64;
687
688impl Handle {
689    pub(crate) fn owned_id(&self) -> NonZeroU64 {
690        self.shared.owned.id
691    }
692
693    pub(crate) fn name(&self) -> Option<&str> {
694        self.name.as_deref()
695    }
696}
697
698impl fmt::Debug for Handle {
699    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
700        fmt.debug_struct("current_thread::Handle { ... }").finish()
701    }
702}
703
704// ===== impl Shared =====
705
706impl Schedule for Arc<Handle> {
707    fn release(&self, task: &Task<Self>) -> Option<Task<Self>> {
708        self.shared.owned.remove(task)
709    }
710
711    fn schedule(&self, task: task::Notified<Self>) {
712        use scheduler::Context::CurrentThread;
713
714        if self.shared.schedule_latency_start.is_some() {
715            task.set_scheduled_at(ScheduleLatencyInstant::new(
716                self.shared.schedule_latency_start,
717            ));
718        }
719
720        context::with_scheduler(|maybe_cx| match maybe_cx {
721            Some(CurrentThread(cx)) if Arc::ptr_eq(self, &cx.handle) => {
722                let mut core = cx.core.borrow_mut();
723
724                // If `None`, the runtime is shutting down, so there is no need
725                // to schedule the task.
726                if let Some(core) = core.as_mut() {
727                    core.push_task(self, task);
728                }
729                self.driver.unpark_self();
730            }
731            _ => {
732                // Track that a task was scheduled from **outside** of the runtime.
733                self.shared.scheduler_metrics.inc_remote_schedule_count();
734
735                // Schedule the task
736                self.shared.inject.push(task);
737                self.driver.unpark();
738            }
739        });
740    }
741
742    fn hooks(&self) -> TaskHarnessScheduleHooks {
743        TaskHarnessScheduleHooks {
744            task_terminate_callback: self.task_hooks.task_terminate_callback.clone(),
745        }
746    }
747
748    cfg_unstable! {
749        fn unhandled_panic(&self) {
750            use crate::runtime::UnhandledPanic;
751
752            match self.shared.config.unhandled_panic {
753                UnhandledPanic::Ignore => {
754                    // Do nothing
755                }
756                UnhandledPanic::ShutdownRuntime => {
757                    use scheduler::Context::CurrentThread;
758
759                    // This hook is only called from within the runtime, so
760                    // `context::with_scheduler` should match with `&self`, i.e.
761                    // there is no opportunity for a nested scheduler to be
762                    // called.
763                    context::with_scheduler(|maybe_cx| match maybe_cx {
764                        Some(CurrentThread(cx)) if Arc::ptr_eq(self, &cx.handle) => {
765                            let mut core = cx.core.borrow_mut();
766
767                            // If `None`, the runtime is shutting down, so there is no need to signal shutdown
768                            if let Some(core) = core.as_mut() {
769                                core.unhandled_panic = true;
770                                self.shared.owned.close_and_shutdown_all(0);
771                            }
772                        }
773                        _ => unreachable!("runtime core not set in CURRENT thread-local"),
774                    })
775                }
776            }
777        }
778    }
779}
780
781impl Wake for Handle {
782    fn wake(arc_self: Arc<Self>) {
783        Wake::wake_by_ref(&arc_self);
784    }
785
786    /// Wake by reference
787    fn wake_by_ref(arc_self: &Arc<Self>) {
788        let already_woken = arc_self.shared.woken.swap(true, Release);
789
790        if !already_woken {
791            use scheduler::Context::CurrentThread;
792
793            // If we are already running on the runtime, then it's not required to wake up the
794            // runtime.
795            context::with_scheduler(|maybe_cx| match maybe_cx {
796                Some(CurrentThread(cx)) if Arc::ptr_eq(arc_self, &cx.handle) => {
797                    arc_self.driver.unpark_self();
798                }
799                _ => {
800                    arc_self.driver.unpark();
801                }
802            });
803        }
804    }
805}
806
807// ===== CoreGuard =====
808
809/// Used to ensure we always place the `Core` value back into its slot in
810/// `CurrentThread`, even if the future panics.
811struct CoreGuard<'a> {
812    context: scheduler::Context,
813    scheduler: &'a CurrentThread,
814}
815
816impl CoreGuard<'_> {
817    #[track_caller]
818    fn block_on<F: Future>(self, mut future: Pin<&mut F>) -> F::Output {
819        let ret = self.enter(|mut core, context| {
820            let waker = Handle::waker_ref(&context.handle);
821            let mut cx = std::task::Context::from_waker(&waker);
822
823            core.metrics.start_processing_scheduled_tasks();
824
825            'outer: loop {
826                let handle = &context.handle;
827
828                if handle.reset_woken() {
829                    let (c, res) = context.enter(core, || {
830                        crate::task::coop::budget(|| future.as_mut().poll(&mut cx))
831                    });
832
833                    core = c;
834
835                    if let Ready(v) = res {
836                        return (core, Some(v));
837                    }
838                }
839
840                for _ in 0..handle.shared.config.event_interval {
841                    // Make sure we didn't hit an unhandled_panic
842                    if core.unhandled_panic {
843                        return (core, None);
844                    }
845
846                    core.tick();
847
848                    let entry = core.next_task(handle);
849
850                    let task = match entry {
851                        Some(entry) => entry,
852                        None => {
853                            core.metrics.end_processing_scheduled_tasks();
854
855                            core = if context.has_pending_work(&core) {
856                                context.park_yield(core, handle)
857                            } else {
858                                context.park(core, handle)
859                            };
860
861                            core.metrics.start_processing_scheduled_tasks();
862
863                            // Try polling the `block_on` future next
864                            continue 'outer;
865                        }
866                    };
867
868                    let task = context.handle.shared.owned.assert_owner(task);
869
870                    let c = context.run_task(task, core);
871
872                    core = c;
873                }
874
875                core.metrics.end_processing_scheduled_tasks();
876
877                // Yield to the driver, this drives the timer and pulls any
878                // pending I/O events.
879                core = context.park_yield(core, handle);
880
881                core.metrics.start_processing_scheduled_tasks();
882            }
883        });
884
885        match ret {
886            Some(ret) => ret,
887            None => {
888                // `block_on` panicked.
889                panic!("a spawned task panicked and the runtime is configured to shut down on unhandled panic");
890            }
891        }
892    }
893
894    /// Enters the scheduler context. This sets the queue and other necessary
895    /// scheduler state in the thread-local.
896    fn enter<F, R>(self, f: F) -> R
897    where
898        F: FnOnce(Box<Core>, &Context) -> (Box<Core>, R),
899    {
900        let context = self.context.expect_current_thread();
901
902        // Remove `core` from `context` to pass into the closure.
903        let core = context.core.borrow_mut().take().expect("core missing");
904
905        // Call the closure and place `core` back
906        let (core, ret) = context::set_scheduler(&self.context, || f(core, context));
907
908        *context.core.borrow_mut() = Some(core);
909
910        ret
911    }
912}
913
914impl Drop for CoreGuard<'_> {
915    fn drop(&mut self) {
916        let context = self.context.expect_current_thread();
917
918        if let Some(core) = context.core.borrow_mut().take() {
919            // Replace old scheduler back into the state to allow
920            // other threads to pick it up and drive it.
921            self.scheduler.core.set(core);
922
923            // Wake up other possible threads that could steal the driver.
924            self.scheduler.notify.notify_one();
925        }
926    }
927}