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
27pub(crate) struct CurrentThread {
29 core: AtomicCell<Core>,
31
32 notify: Notify,
35}
36
37pub(crate) struct Handle {
39 name: Option<String>,
41
42 shared: Shared,
44
45 pub(crate) driver: driver::Handle,
47
48 pub(crate) blocking_spawner: blocking::Spawner,
50
51 pub(crate) seed_generator: RngSeedGenerator,
53
54 pub(crate) task_hooks: TaskHooks,
56
57 pub(crate) local_tid: Option<ThreadId>,
59}
60
61struct Core {
64 tasks: VecDeque<Notified>,
66
67 tick: u32,
69
70 driver: Option<Driver>,
74
75 metrics: MetricsBatch,
77
78 global_queue_interval: u32,
80
81 unhandled_panic: bool,
84}
85
86struct Shared {
88 inject: Inject<Arc<Handle>>,
90
91 owned: OwnedTasks<Arc<Handle>>,
93
94 woken: AtomicBool,
96
97 config: Config,
99
100 scheduler_metrics: SchedulerMetrics,
102
103 worker_metrics: WorkerMetrics,
105
106 started_at: Option<Instant>,
110}
111
112pub(crate) struct Context {
116 handle: Arc<Handle>,
118
119 core: RefCell<Option<Box<Core>>>,
122
123 pub(crate) defer: Defer,
125}
126
127type Notified = task::Notified<Arc<Handle>>;
128
129const INITIAL_CAPACITY: usize = 64;
131
132const 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 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 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 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 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 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 handle.shared.owned.close_and_shutdown_all(0);
295
296 while let Some(task) = core.next_local_task(handle) {
299 drop(task);
300 }
301
302 handle.shared.inject.close();
304
305 while let Some(task) = handle.shared.inject.pop() {
307 drop(task);
308 }
309
310 assert!(handle.shared.owned.is_empty());
311
312 core.submit_metrics(handle);
314
315 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
329impl Core {
332 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
379impl Context {
382 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 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 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 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 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 *self.core.borrow_mut() = Some(core);
486
487 let ret = f();
489
490 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
500impl Handle {
503 #[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 #[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 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 #[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 context::with_scheduler(|maybe_context| {
588 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 drop(maybe_core);
613
614 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 me.shared.woken.store(true, Release);
631 waker_ref(me)
632 }
633
634 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
702impl 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 let Some(core) = core.as_mut() {
728 core.push_task(self, task);
729 }
730 }
731 _ => {
732 self.shared.scheduler_metrics.inc_remote_schedule_count();
734
735 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 }
756 UnhandledPanic::ShutdownRuntime => {
757 use scheduler::Context::CurrentThread;
758
759 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 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 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 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
805struct 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 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 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 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 panic!("a spawned task panicked and the runtime is configured to shut down on unhandled panic");
890 }
891 }
892 }
893
894 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 let core = context.core.borrow_mut().take().expect("core missing");
904
905 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 self.scheduler.core.set(core);
922
923 self.scheduler.notify.notify_one();
925 }
926 }
927}