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::sync::atomic::Ordering::{AcqRel, Acquire, Release};
20use std::task::Poll::{Pending, Ready};
21use std::task::Waker;
22use std::thread::ThreadId;
23use std::time::Duration;
24use std::time::Instant;
25use std::{fmt, thread};
26
27/// Executes tasks on the current thread
28pub(crate) struct CurrentThread {
29    /// Core scheduler data is acquired by a thread entering `block_on`.
30    core: AtomicCell<Core>,
31
32    /// Notifier for waking up other threads to steal the
33    /// driver.
34    notify: Notify,
35}
36
37/// Handle to the current thread scheduler
38pub(crate) struct Handle {
39    /// The name of the runtime
40    name: Option<String>,
41
42    /// Scheduler state shared across threads
43    shared: Shared,
44
45    /// Resource driver handles
46    pub(crate) driver: driver::Handle,
47
48    /// Blocking pool spawner
49    pub(crate) blocking_spawner: blocking::Spawner,
50
51    /// Current random number generator seed
52    pub(crate) seed_generator: RngSeedGenerator,
53
54    /// User-supplied hooks to invoke for things
55    pub(crate) task_hooks: TaskHooks,
56
57    /// If this is a `LocalRuntime`, flags the owning thread ID.
58    pub(crate) local_tid: Option<ThreadId>,
59}
60
61/// Data required for executing the scheduler. The struct is passed around to
62/// a function that will perform the scheduling work and acts as a capability token.
63struct Core {
64    /// Scheduler run queue
65    tasks: VecDeque<Notified>,
66
67    /// Current tick
68    tick: u32,
69
70    /// Runtime driver
71    ///
72    /// The driver is removed before starting to park the thread
73    driver: Option<Driver>,
74
75    /// Metrics batch
76    metrics: MetricsBatch,
77
78    /// How often to check the global queue
79    global_queue_interval: u32,
80
81    /// True if a task panicked without being handled and the runtime is
82    /// configured to shutdown on unhandled panic.
83    unhandled_panic: bool,
84}
85
86/// Scheduler state shared between threads.
87struct Shared {
88    /// Remote run queue
89    inject: Inject<Arc<Handle>>,
90
91    /// Collection of all active tasks spawned onto this executor.
92    owned: OwnedTasks<Arc<Handle>>,
93
94    /// Indicates whether the blocked on thread was woken.
95    woken: AtomicBool,
96
97    /// Scheduler configuration options
98    config: Config,
99
100    /// Keeps track of various runtime metrics.
101    scheduler_metrics: SchedulerMetrics,
102
103    /// This scheduler only has one worker.
104    worker_metrics: WorkerMetrics,
105
106    /// Startup time of this scheduler.
107    ///
108    /// This instant is used as the basis of task `scheduled_at` measurements.
109    started_at: Option<Instant>,
110}
111
112/// Thread-local context.
113///
114/// pub(crate) to store in `runtime::context`.
115pub(crate) struct Context {
116    /// Scheduler handle
117    handle: Arc<Handle>,
118
119    /// Scheduler core, enabling the holder of `Context` to execute the
120    /// scheduler.
121    core: RefCell<Option<Box<Core>>>,
122
123    /// Deferred tasks, usually ones that called `task::yield_now()`.
124    pub(crate) defer: Defer,
125}
126
127type Notified = task::Notified<Arc<Handle>>;
128
129/// Initial queue capacity.
130const INITIAL_CAPACITY: usize = 64;
131
132/// Used if none is specified. This is a temporary constant and will be removed
133/// as we unify tuning logic between the multi-thread and current-thread
134/// schedulers.
135const DEFAULT_GLOBAL_QUEUE_INTERVAL: u32 = 31;
136
137impl CurrentThread {
138    pub(crate) fn new(
139        driver: Driver,
140        driver_handle: driver::Handle,
141        blocking_spawner: blocking::Spawner,
142        seed_generator: RngSeedGenerator,
143        config: Config,
144        local_tid: Option<ThreadId>,
145        name: Option<String>,
146    ) -> (CurrentThread, Arc<Handle>) {
147        let worker_metrics = WorkerMetrics::from_config(&config);
148        worker_metrics.set_thread_id(thread::current().id());
149
150        // Get the configured global queue interval, or use the default.
151        let global_queue_interval = config
152            .global_queue_interval
153            .unwrap_or(DEFAULT_GLOBAL_QUEUE_INTERVAL);
154
155        let started_at = config
156            .metrics_schedule_latency_histogram
157            .as_ref()
158            .map(|_| Instant::now());
159
160        let handle = Arc::new(Handle {
161            name,
162            task_hooks: TaskHooks {
163                task_spawn_callback: config.before_spawn.clone(),
164                task_terminate_callback: config.after_termination.clone(),
165                #[cfg(tokio_unstable)]
166                before_poll_callback: config.before_poll.clone(),
167                #[cfg(tokio_unstable)]
168                after_poll_callback: config.after_poll.clone(),
169            },
170            shared: Shared {
171                inject: Inject::new(),
172                owned: OwnedTasks::new(1),
173                woken: AtomicBool::new(false),
174                config,
175                scheduler_metrics: SchedulerMetrics::new(),
176                worker_metrics,
177                started_at,
178            },
179            driver: driver_handle,
180            blocking_spawner,
181            seed_generator,
182            local_tid,
183        });
184
185        let core = AtomicCell::new(Some(Box::new(Core {
186            tasks: VecDeque::with_capacity(INITIAL_CAPACITY),
187            tick: 0,
188            driver: Some(driver),
189            metrics: MetricsBatch::new(&handle.shared.worker_metrics),
190            global_queue_interval,
191            unhandled_panic: false,
192        })));
193
194        let scheduler = CurrentThread {
195            core,
196            notify: Notify::new(),
197        };
198
199        (scheduler, handle)
200    }
201
202    #[track_caller]
203    pub(crate) fn block_on<F: Future>(&self, handle: &scheduler::Handle, future: F) -> F::Output {
204        pin!(future);
205
206        crate::runtime::context::enter_runtime(handle, false, |blocking| {
207            let handle = handle.as_current_thread();
208
209            // Attempt to steal the scheduler core and block_on the future if we can
210            // there, otherwise, lets select on a notification that the core is
211            // available or the future is complete.
212            loop {
213                if let Some(core) = self.take_core(handle) {
214                    handle
215                        .shared
216                        .worker_metrics
217                        .set_thread_id(thread::current().id());
218                    return core.block_on(future);
219                } else {
220                    let notified = self.notify.notified();
221                    pin!(notified);
222
223                    if let Some(out) = blocking
224                        .block_on(poll_fn(|cx| {
225                            if notified.as_mut().poll(cx).is_ready() {
226                                return Ready(None);
227                            }
228
229                            if let Ready(out) = future.as_mut().poll(cx) {
230                                return Ready(Some(out));
231                            }
232
233                            Pending
234                        }))
235                        .expect("Failed to `Enter::block_on`")
236                    {
237                        return out;
238                    }
239                }
240            }
241        })
242    }
243
244    fn take_core(&self, handle: &Arc<Handle>) -> Option<CoreGuard<'_>> {
245        let core = self.core.take()?;
246
247        Some(CoreGuard {
248            context: scheduler::Context::CurrentThread(Context {
249                handle: handle.clone(),
250                core: RefCell::new(Some(core)),
251                defer: Defer::new(),
252            }),
253            scheduler: self,
254        })
255    }
256
257    pub(crate) fn shutdown(&mut self, handle: &scheduler::Handle) {
258        let handle = handle.as_current_thread();
259
260        // Avoid a double panic if we are currently panicking and
261        // the lock may be poisoned.
262
263        let core = match self.take_core(handle) {
264            Some(core) => core,
265            None if std::thread::panicking() => return,
266            None => panic!("Oh no! We never placed the Core back, this is a bug!"),
267        };
268
269        // Check that the thread-local is not being destroyed
270        let tls_available = context::with_current(|_| ()).is_ok();
271
272        if tls_available {
273            core.enter(|core, _context| {
274                let core = shutdown2(core, handle);
275                (core, ())
276            });
277        } else {
278            // Shutdown without setting the context. `tokio::spawn` calls will
279            // fail, but those will fail either way because the thread-local is
280            // not available anymore.
281            let context = core.context.expect_current_thread();
282            let core = context.core.borrow_mut().take().unwrap();
283
284            let core = shutdown2(core, handle);
285            *context.core.borrow_mut() = Some(core);
286        }
287    }
288}
289
290fn shutdown2(mut core: Box<Core>, handle: &Handle) -> Box<Core> {
291    // Drain the OwnedTasks collection. This call also closes the
292    // collection, ensuring that no tasks are ever pushed after this
293    // call returns.
294    handle.shared.owned.close_and_shutdown_all(0);
295
296    // Drain local queue
297    // We already shut down every task, so we just need to drop the task.
298    while let Some(task) = core.next_local_task(handle) {
299        drop(task);
300    }
301
302    // Close the injection queue
303    handle.shared.inject.close();
304
305    // Drain remote queue
306    while let Some(task) = handle.shared.inject.pop() {
307        drop(task);
308    }
309
310    assert!(handle.shared.owned.is_empty());
311
312    // Submit metrics
313    core.submit_metrics(handle);
314
315    // Shutdown the resource drivers
316    if let Some(driver) = core.driver.as_mut() {
317        driver.shutdown(&handle.driver);
318    }
319
320    core
321}
322
323impl fmt::Debug for CurrentThread {
324    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
325        fmt.debug_struct("CurrentThread").finish()
326    }
327}
328
329// ===== impl Core =====
330
331impl Core {
332    /// Get and increment the current tick
333    fn tick(&mut self) {
334        self.tick = self.tick.wrapping_add(1);
335    }
336
337    fn next_task(&mut self, handle: &Handle) -> Option<Notified> {
338        if self.tick % self.global_queue_interval == 0 {
339            handle
340                .next_remote_task()
341                .or_else(|| self.next_local_task(handle))
342        } else {
343            self.next_local_task(handle)
344                .or_else(|| handle.next_remote_task())
345        }
346    }
347
348    fn next_local_task(&mut self, handle: &Handle) -> Option<Notified> {
349        let ret = self.tasks.pop_front();
350        handle
351            .shared
352            .worker_metrics
353            .set_queue_depth(self.tasks.len());
354        ret
355    }
356
357    fn push_task(&mut self, handle: &Handle, task: Notified) {
358        self.tasks.push_back(task);
359        self.metrics.inc_local_schedule_count();
360        handle
361            .shared
362            .worker_metrics
363            .set_queue_depth(self.tasks.len());
364    }
365
366    fn submit_metrics(&mut self, handle: &Handle) {
367        self.metrics.submit(&handle.shared.worker_metrics, 0);
368    }
369}
370
371#[cfg(feature = "taskdump")]
372fn wake_deferred_tasks_and_free(context: &Context) {
373    let wakers = context.defer.take_deferred();
374    for waker in wakers {
375        waker.wake();
376    }
377}
378
379// ===== impl Context =====
380
381impl Context {
382    /// Execute the closure with the given scheduler core stored in the
383    /// thread-local context.
384    fn run_task(&self, task: LocalNotified<Arc<Handle>>, mut core: Box<Core>) -> Box<Core> {
385        #[cfg(tokio_unstable)]
386        let task_meta = task.task_meta();
387
388        core.metrics.start_poll(
389            task.get_scheduled_at()
390                .prepare(self.handle.shared.started_at),
391        );
392
393        let (mut c, ()) = self.enter(core, || {
394            crate::task::coop::budget(|| {
395                #[cfg(tokio_unstable)]
396                self.handle.task_hooks.poll_start_callback(&task_meta);
397
398                task.run();
399
400                #[cfg(tokio_unstable)]
401                self.handle.task_hooks.poll_stop_callback(&task_meta);
402            })
403        });
404        c.metrics.end_poll();
405        c
406    }
407
408    /// Blocks the current thread until an event is received by the driver,
409    /// including I/O events, timer events, ...
410    fn park(&self, mut core: Box<Core>, handle: &Handle) -> Box<Core> {
411        let mut driver = core.driver.take().expect("driver missing");
412
413        if let Some(f) = &handle.shared.config.before_park {
414            let (c, ()) = self.enter(core, || f());
415            core = c;
416        }
417
418        if !self.has_pending_work(&core) {
419            // Park until the thread is signaled
420            core.metrics.about_to_park();
421            core.submit_metrics(handle);
422
423            core = self.park_internal(core, handle, &mut driver, None);
424
425            core.metrics.unparked();
426            core.submit_metrics(handle);
427        } else {
428            // `before_park` scheduled work (e.g. an `on_thread_park` hook that woke the
429            // `block_on` future), so we don't block. We must still poll the driver once
430            // without blocking, or timer and I/O events would stall under a runtime driven
431            // by repeated short `block_on` calls. See
432            // <https://github.com/tokio-rs/tokio/issues/8212>.
433            core.submit_metrics(handle);
434
435            core = self.park_internal(core, handle, &mut driver, Some(Duration::from_millis(0)));
436        }
437
438        if let Some(f) = &handle.shared.config.after_unpark {
439            let (c, ()) = self.enter(core, || f());
440            core = c;
441        }
442
443        core.driver = Some(driver);
444        core
445    }
446
447    /// Checks the driver for new events without blocking the thread.
448    fn park_yield(&self, mut core: Box<Core>, handle: &Handle) -> Box<Core> {
449        let mut driver = core.driver.take().expect("driver missing");
450
451        core.submit_metrics(handle);
452
453        core = self.park_internal(core, handle, &mut driver, Some(Duration::from_millis(0)));
454
455        core.driver = Some(driver);
456        core
457    }
458
459    fn has_pending_work(&self, core: &Core) -> bool {
460        !core.tasks.is_empty() || !self.defer.is_empty() || self.handle.shared.woken.load(Acquire)
461    }
462
463    fn park_internal(
464        &self,
465        core: Box<Core>,
466        handle: &Handle,
467        driver: &mut Driver,
468        duration: Option<Duration>,
469    ) -> Box<Core> {
470        let (core, ()) = self.enter(core, || {
471            match duration {
472                Some(dur) => driver.park_timeout(&handle.driver, dur),
473                None => driver.park(&handle.driver),
474            }
475            self.defer.wake();
476        });
477
478        core
479    }
480
481    fn enter<R>(&self, core: Box<Core>, f: impl FnOnce() -> R) -> (Box<Core>, R) {
482        // Store the scheduler core in the thread-local context
483        //
484        // A drop-guard is employed at a higher level.
485        *self.core.borrow_mut() = Some(core);
486
487        // Execute the closure while tracking the execution budget
488        let ret = f();
489
490        // Take the scheduler core back
491        let core = self.core.borrow_mut().take().expect("core missing");
492        (core, ret)
493    }
494
495    pub(crate) fn defer(&self, waker: &Waker) {
496        self.defer.defer(waker);
497    }
498}
499
500// ===== impl Handle =====
501
502impl Handle {
503    /// Spawns a future onto the `CurrentThread` scheduler
504    #[track_caller]
505    pub(crate) fn spawn<F>(
506        me: &Arc<Self>,
507        future: F,
508        id: crate::runtime::task::Id,
509        spawned_at: SpawnLocation,
510    ) -> JoinHandle<F::Output>
511    where
512        F: crate::future::Future + Send + 'static,
513        F::Output: Send + 'static,
514    {
515        let (handle, notified) = me.shared.owned.bind(future, me.clone(), id, spawned_at);
516
517        me.task_hooks.spawn(&TaskMeta {
518            id,
519            spawned_at,
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            _phantom: Default::default(),
559        });
560
561        if let Some(notified) = notified {
562            me.schedule(notified);
563        }
564
565        handle
566    }
567
568    /// Capture a snapshot of this runtime's state.
569    #[cfg(all(
570        tokio_unstable,
571        feature = "taskdump",
572        target_os = "linux",
573        any(
574            target_arch = "aarch64",
575            target_arch = "x86",
576            target_arch = "x86_64",
577            target_arch = "s390x"
578        )
579    ))]
580    pub(crate) fn dump(&self) -> crate::runtime::Dump {
581        use crate::runtime::dump;
582        use task::trace::trace_current_thread;
583
584        let mut traces = vec![];
585
586        // todo: how to make this work outside of a runtime context?
587        context::with_scheduler(|maybe_context| {
588            // drain the local queue
589            let context = if let Some(context) = maybe_context {
590                context.expect_current_thread()
591            } else {
592                return;
593            };
594            let mut maybe_core = context.core.borrow_mut();
595            let core = if let Some(core) = maybe_core.as_mut() {
596                core
597            } else {
598                return;
599            };
600            let local = &mut core.tasks;
601
602            if self.shared.inject.is_closed() {
603                return;
604            }
605
606            traces = trace_current_thread(&self.shared.owned, local, &self.shared.inject)
607                .into_iter()
608                .map(|(id, trace)| dump::Task::new(id, trace))
609                .collect();
610
611            // Avoid double borrow panic
612            drop(maybe_core);
613
614            // Taking a taskdump could wakes every task, but we probably don't want
615            // the `yield_now` vector to be that large under normal circumstances.
616            // Therefore, we free its allocation.
617            wake_deferred_tasks_and_free(context);
618        });
619
620        dump::Dump::new(traces)
621    }
622
623    fn next_remote_task(&self) -> Option<Notified> {
624        self.shared.inject.pop()
625    }
626
627    fn waker_ref(me: &Arc<Self>) -> WakerRef<'_> {
628        // Set woken to true when enter block_on, ensure outer future
629        // be polled for the first time when enter loop
630        me.shared.woken.store(true, Release);
631        waker_ref(me)
632    }
633
634    // reset woken to false and return original value
635    pub(crate) fn reset_woken(&self) -> bool {
636        self.shared.woken.swap(false, AcqRel)
637    }
638
639    pub(crate) fn num_alive_tasks(&self) -> usize {
640        self.shared.owned.num_alive_tasks()
641    }
642
643    pub(crate) fn injection_queue_depth(&self) -> usize {
644        self.shared.inject.len()
645    }
646
647    pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics {
648        assert_eq!(0, worker);
649        &self.shared.worker_metrics
650    }
651}
652
653cfg_unstable_metrics! {
654    impl Handle {
655        pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics {
656            &self.shared.scheduler_metrics
657        }
658
659        pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize {
660            self.worker_metrics(worker).queue_depth()
661        }
662
663        pub(crate) fn num_blocking_threads(&self) -> usize {
664            self.blocking_spawner.num_threads()
665        }
666
667        pub(crate) fn num_idle_blocking_threads(&self) -> usize {
668            self.blocking_spawner.num_idle_threads()
669        }
670
671        pub(crate) fn blocking_queue_depth(&self) -> usize {
672            self.blocking_spawner.queue_depth()
673        }
674
675        cfg_64bit_metrics! {
676            pub(crate) fn spawned_tasks_count(&self) -> u64 {
677                self.shared.owned.spawned_tasks_count()
678            }
679        }
680    }
681}
682
683use crate::runtime::metrics::ScheduleLatencyInstant;
684use std::num::NonZeroU64;
685
686impl Handle {
687    pub(crate) fn owned_id(&self) -> NonZeroU64 {
688        self.shared.owned.id
689    }
690
691    pub(crate) fn name(&self) -> Option<&str> {
692        self.name.as_deref()
693    }
694}
695
696impl fmt::Debug for Handle {
697    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
698        fmt.debug_struct("current_thread::Handle { ... }").finish()
699    }
700}
701
702// ===== impl Shared =====
703
704impl Schedule for Arc<Handle> {
705    fn release(&self, task: &Task<Self>) -> Option<Task<Self>> {
706        self.shared.owned.remove(task)
707    }
708
709    fn schedule(&self, task: task::Notified<Self>) {
710        use scheduler::Context::CurrentThread;
711
712        if self
713            .shared
714            .config
715            .metrics_schedule_latency_histogram
716            .is_some()
717        {
718            task.set_scheduled_at(ScheduleLatencyInstant::new(self.shared.started_at));
719        }
720
721        context::with_scheduler(|maybe_cx| match maybe_cx {
722            Some(CurrentThread(cx)) if Arc::ptr_eq(self, &cx.handle) => {
723                let mut core = cx.core.borrow_mut();
724
725                // If `None`, the runtime is shutting down, so there is no need
726                // to schedule the task.
727                if let Some(core) = core.as_mut() {
728                    core.push_task(self, task);
729                }
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                _ => {
798                    arc_self.driver.unpark();
799                }
800            });
801        }
802    }
803}
804
805// ===== CoreGuard =====
806
807/// Used to ensure we always place the `Core` value back into its slot in
808/// `CurrentThread`, even if the future panics.
809struct CoreGuard<'a> {
810    context: scheduler::Context,
811    scheduler: &'a CurrentThread,
812}
813
814impl CoreGuard<'_> {
815    #[track_caller]
816    fn block_on<F: Future>(self, future: F) -> F::Output {
817        let ret = self.enter(|mut core, context| {
818            let waker = Handle::waker_ref(&context.handle);
819            let mut cx = std::task::Context::from_waker(&waker);
820
821            pin!(future);
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}