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
28pub(crate) struct CurrentThread {
30 core: AtomicCell<Core>,
32
33 notify: Notify,
36}
37
38pub(crate) struct Handle {
40 name: Option<String>,
42
43 shared: Shared,
45
46 pub(crate) driver: driver::Handle,
48
49 pub(crate) blocking_spawner: blocking::Spawner,
51
52 pub(crate) seed_generator: RngSeedGenerator,
54
55 pub(crate) task_hooks: TaskHooks,
57
58 pub(crate) local_tid: Option<ThreadId>,
60}
61
62struct Core {
65 tasks: VecDeque<Notified>,
67
68 tick: u32,
70
71 driver: Option<Driver>,
75
76 metrics: MetricsBatch,
78
79 global_queue_interval: u32,
81
82 unhandled_panic: bool,
85}
86
87struct Shared {
89 inject: Inject<Arc<Handle>>,
91
92 owned: OwnedTasks<Arc<Handle>>,
94
95 woken: AtomicBool,
97
98 config: Config,
100
101 scheduler_metrics: SchedulerMetrics,
103
104 worker_metrics: WorkerMetrics,
106
107 schedule_latency_start: Option<Instant>,
111}
112
113pub(crate) struct Context {
117 handle: Arc<Handle>,
119
120 core: RefCell<Option<Box<Core>>>,
123
124 pub(crate) defer: Defer,
126}
127
128type Notified = task::Notified<Arc<Handle>>;
129
130const INITIAL_CAPACITY: usize = 64;
132
133const 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 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 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 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 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 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 handle.shared.owned.close_and_shutdown_all(0);
293
294 while let Some(task) = core.next_local_task(handle) {
297 drop(task);
298 }
299
300 handle.shared.inject.close();
302
303 while let Some(task) = handle.shared.inject.pop() {
305 drop(task);
306 }
307
308 assert!(handle.shared.owned.is_empty());
309
310 core.submit_metrics(handle);
312
313 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
327impl Core {
330 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
377impl Context {
380 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 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 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 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 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 *self.core.borrow_mut() = Some(core);
484
485 let ret = f();
487
488 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
498impl Handle {
501 #[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 #[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 #[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 #[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 context::with_scheduler(|maybe_context| {
591 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 drop(maybe_core);
615
616 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 me.shared.woken.store(true, Release);
633 waker_ref(me)
634 }
635
636 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
704impl 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 let Some(core) = core.as_mut() {
727 core.push_task(self, task);
728 }
729 self.driver.unpark_self();
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 arc_self.driver.unpark_self();
798 }
799 _ => {
800 arc_self.driver.unpark();
801 }
802 });
803 }
804 }
805}
806
807struct 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 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}