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, OwnedTasks, Schedule, SpawnLocation, Task, TaskHarnessScheduleHooks,
7};
8use crate::runtime::{
9 blocking, context, Config, MetricsBatch, SchedulerMetrics, TaskHooks, TaskMeta, WorkerMetrics,
10};
11use crate::sync::notify::Notify;
12use crate::util::atomic_cell::AtomicCell;
13use crate::util::{waker_ref, RngSeedGenerator, Wake, WakerRef};
14
15use std::cell::RefCell;
16use std::collections::VecDeque;
17use std::future::{poll_fn, Future};
18use std::sync::atomic::Ordering::{AcqRel, Acquire, Release};
19use std::task::Poll::{Pending, Ready};
20use std::task::Waker;
21use std::thread::ThreadId;
22use std::time::Duration;
23use std::{fmt, thread};
24
25pub(crate) struct CurrentThread {
27 core: AtomicCell<Core>,
29
30 notify: Notify,
33}
34
35pub(crate) struct Handle {
37 name: Option<String>,
39
40 shared: Shared,
42
43 pub(crate) driver: driver::Handle,
45
46 pub(crate) blocking_spawner: blocking::Spawner,
48
49 pub(crate) seed_generator: RngSeedGenerator,
51
52 pub(crate) task_hooks: TaskHooks,
54
55 pub(crate) local_tid: Option<ThreadId>,
57}
58
59struct Core {
62 tasks: VecDeque<Notified>,
64
65 tick: u32,
67
68 driver: Option<Driver>,
72
73 metrics: MetricsBatch,
75
76 global_queue_interval: u32,
78
79 unhandled_panic: bool,
82}
83
84struct Shared {
86 inject: Inject<Arc<Handle>>,
88
89 owned: OwnedTasks<Arc<Handle>>,
91
92 woken: AtomicBool,
94
95 config: Config,
97
98 scheduler_metrics: SchedulerMetrics,
100
101 worker_metrics: WorkerMetrics,
103}
104
105pub(crate) struct Context {
109 handle: Arc<Handle>,
111
112 core: RefCell<Option<Box<Core>>>,
115
116 pub(crate) defer: Defer,
118}
119
120type Notified = task::Notified<Arc<Handle>>;
121
122const INITIAL_CAPACITY: usize = 64;
124
125const DEFAULT_GLOBAL_QUEUE_INTERVAL: u32 = 31;
129
130impl CurrentThread {
131 pub(crate) fn new(
132 driver: Driver,
133 driver_handle: driver::Handle,
134 blocking_spawner: blocking::Spawner,
135 seed_generator: RngSeedGenerator,
136 config: Config,
137 local_tid: Option<ThreadId>,
138 name: Option<String>,
139 ) -> (CurrentThread, Arc<Handle>) {
140 let worker_metrics = WorkerMetrics::from_config(&config);
141 worker_metrics.set_thread_id(thread::current().id());
142
143 let global_queue_interval = config
145 .global_queue_interval
146 .unwrap_or(DEFAULT_GLOBAL_QUEUE_INTERVAL);
147
148 let handle = Arc::new(Handle {
149 name,
150 task_hooks: TaskHooks {
151 task_spawn_callback: config.before_spawn.clone(),
152 task_terminate_callback: config.after_termination.clone(),
153 #[cfg(tokio_unstable)]
154 before_poll_callback: config.before_poll.clone(),
155 #[cfg(tokio_unstable)]
156 after_poll_callback: config.after_poll.clone(),
157 },
158 shared: Shared {
159 inject: Inject::new(),
160 owned: OwnedTasks::new(1),
161 woken: AtomicBool::new(false),
162 config,
163 scheduler_metrics: SchedulerMetrics::new(),
164 worker_metrics,
165 },
166 driver: driver_handle,
167 blocking_spawner,
168 seed_generator,
169 local_tid,
170 });
171
172 let core = AtomicCell::new(Some(Box::new(Core {
173 tasks: VecDeque::with_capacity(INITIAL_CAPACITY),
174 tick: 0,
175 driver: Some(driver),
176 metrics: MetricsBatch::new(&handle.shared.worker_metrics),
177 global_queue_interval,
178 unhandled_panic: false,
179 })));
180
181 let scheduler = CurrentThread {
182 core,
183 notify: Notify::new(),
184 };
185
186 (scheduler, handle)
187 }
188
189 #[track_caller]
190 pub(crate) fn block_on<F: Future>(&self, handle: &scheduler::Handle, future: F) -> F::Output {
191 pin!(future);
192
193 crate::runtime::context::enter_runtime(handle, false, |blocking| {
194 let handle = handle.as_current_thread();
195
196 loop {
200 if let Some(core) = self.take_core(handle) {
201 handle
202 .shared
203 .worker_metrics
204 .set_thread_id(thread::current().id());
205 return core.block_on(future);
206 } else {
207 let notified = self.notify.notified();
208 pin!(notified);
209
210 if let Some(out) = blocking
211 .block_on(poll_fn(|cx| {
212 if notified.as_mut().poll(cx).is_ready() {
213 return Ready(None);
214 }
215
216 if let Ready(out) = future.as_mut().poll(cx) {
217 return Ready(Some(out));
218 }
219
220 Pending
221 }))
222 .expect("Failed to `Enter::block_on`")
223 {
224 return out;
225 }
226 }
227 }
228 })
229 }
230
231 fn take_core(&self, handle: &Arc<Handle>) -> Option<CoreGuard<'_>> {
232 let core = self.core.take()?;
233
234 Some(CoreGuard {
235 context: scheduler::Context::CurrentThread(Context {
236 handle: handle.clone(),
237 core: RefCell::new(Some(core)),
238 defer: Defer::new(),
239 }),
240 scheduler: self,
241 })
242 }
243
244 pub(crate) fn shutdown(&mut self, handle: &scheduler::Handle) {
245 let handle = handle.as_current_thread();
246
247 let core = match self.take_core(handle) {
251 Some(core) => core,
252 None if std::thread::panicking() => return,
253 None => panic!("Oh no! We never placed the Core back, this is a bug!"),
254 };
255
256 let tls_available = context::with_current(|_| ()).is_ok();
258
259 if tls_available {
260 core.enter(|core, _context| {
261 let core = shutdown2(core, handle);
262 (core, ())
263 });
264 } else {
265 let context = core.context.expect_current_thread();
269 let core = context.core.borrow_mut().take().unwrap();
270
271 let core = shutdown2(core, handle);
272 *context.core.borrow_mut() = Some(core);
273 }
274 }
275}
276
277fn shutdown2(mut core: Box<Core>, handle: &Handle) -> Box<Core> {
278 handle.shared.owned.close_and_shutdown_all(0);
282
283 while let Some(task) = core.next_local_task(handle) {
286 drop(task);
287 }
288
289 handle.shared.inject.close();
291
292 while let Some(task) = handle.shared.inject.pop() {
294 drop(task);
295 }
296
297 assert!(handle.shared.owned.is_empty());
298
299 core.submit_metrics(handle);
301
302 if let Some(driver) = core.driver.as_mut() {
304 driver.shutdown(&handle.driver);
305 }
306
307 core
308}
309
310impl fmt::Debug for CurrentThread {
311 fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
312 fmt.debug_struct("CurrentThread").finish()
313 }
314}
315
316impl Core {
319 fn tick(&mut self) {
321 self.tick = self.tick.wrapping_add(1);
322 }
323
324 fn next_task(&mut self, handle: &Handle) -> Option<Notified> {
325 if self.tick % self.global_queue_interval == 0 {
326 handle
327 .next_remote_task()
328 .or_else(|| self.next_local_task(handle))
329 } else {
330 self.next_local_task(handle)
331 .or_else(|| handle.next_remote_task())
332 }
333 }
334
335 fn next_local_task(&mut self, handle: &Handle) -> Option<Notified> {
336 let ret = self.tasks.pop_front();
337 handle
338 .shared
339 .worker_metrics
340 .set_queue_depth(self.tasks.len());
341 ret
342 }
343
344 fn push_task(&mut self, handle: &Handle, task: Notified) {
345 self.tasks.push_back(task);
346 self.metrics.inc_local_schedule_count();
347 handle
348 .shared
349 .worker_metrics
350 .set_queue_depth(self.tasks.len());
351 }
352
353 fn submit_metrics(&mut self, handle: &Handle) {
354 self.metrics.submit(&handle.shared.worker_metrics, 0);
355 }
356}
357
358#[cfg(feature = "taskdump")]
359fn wake_deferred_tasks_and_free(context: &Context) {
360 let wakers = context.defer.take_deferred();
361 for waker in wakers {
362 waker.wake();
363 }
364}
365
366impl Context {
369 fn run_task<R>(&self, mut core: Box<Core>, f: impl FnOnce() -> R) -> (Box<Core>, R) {
372 core.metrics.start_poll();
373 let mut ret = self.enter(core, || crate::task::coop::budget(f));
374 ret.0.metrics.end_poll();
375 ret
376 }
377
378 fn park(&self, mut core: Box<Core>, handle: &Handle) -> Box<Core> {
381 let mut driver = core.driver.take().expect("driver missing");
382
383 if let Some(f) = &handle.shared.config.before_park {
384 let (c, ()) = self.enter(core, || f());
385 core = c;
386 }
387
388 if !self.has_pending_work(&core) {
389 core.metrics.about_to_park();
391 core.submit_metrics(handle);
392
393 core = self.park_internal(core, handle, &mut driver, None);
394
395 core.metrics.unparked();
396 core.submit_metrics(handle);
397 } else {
398 core.submit_metrics(handle);
404
405 core = self.park_internal(core, handle, &mut driver, Some(Duration::from_millis(0)));
406 }
407
408 if let Some(f) = &handle.shared.config.after_unpark {
409 let (c, ()) = self.enter(core, || f());
410 core = c;
411 }
412
413 core.driver = Some(driver);
414 core
415 }
416
417 fn park_yield(&self, mut core: Box<Core>, handle: &Handle) -> Box<Core> {
419 let mut driver = core.driver.take().expect("driver missing");
420
421 core.submit_metrics(handle);
422
423 core = self.park_internal(core, handle, &mut driver, Some(Duration::from_millis(0)));
424
425 core.driver = Some(driver);
426 core
427 }
428
429 fn has_pending_work(&self, core: &Core) -> bool {
430 !core.tasks.is_empty() || !self.defer.is_empty() || self.handle.shared.woken.load(Acquire)
431 }
432
433 fn park_internal(
434 &self,
435 core: Box<Core>,
436 handle: &Handle,
437 driver: &mut Driver,
438 duration: Option<Duration>,
439 ) -> Box<Core> {
440 let (core, ()) = self.enter(core, || {
441 match duration {
442 Some(dur) => driver.park_timeout(&handle.driver, dur),
443 None => driver.park(&handle.driver),
444 }
445 self.defer.wake();
446 });
447
448 core
449 }
450
451 fn enter<R>(&self, core: Box<Core>, f: impl FnOnce() -> R) -> (Box<Core>, R) {
452 *self.core.borrow_mut() = Some(core);
456
457 let ret = f();
459
460 let core = self.core.borrow_mut().take().expect("core missing");
462 (core, ret)
463 }
464
465 pub(crate) fn defer(&self, waker: &Waker) {
466 self.defer.defer(waker);
467 }
468}
469
470impl Handle {
473 #[track_caller]
475 pub(crate) fn spawn<F>(
476 me: &Arc<Self>,
477 future: F,
478 id: crate::runtime::task::Id,
479 spawned_at: SpawnLocation,
480 ) -> JoinHandle<F::Output>
481 where
482 F: crate::future::Future + Send + 'static,
483 F::Output: Send + 'static,
484 {
485 let (handle, notified) = me.shared.owned.bind(future, me.clone(), id, spawned_at);
486
487 me.task_hooks.spawn(&TaskMeta {
488 id,
489 spawned_at,
490 _phantom: Default::default(),
491 });
492
493 if let Some(notified) = notified {
494 me.schedule(notified);
495 }
496
497 handle
498 }
499
500 #[track_caller]
508 pub(crate) unsafe fn spawn_local<F>(
509 me: &Arc<Self>,
510 future: F,
511 id: crate::runtime::task::Id,
512 spawned_at: SpawnLocation,
513 ) -> JoinHandle<F::Output>
514 where
515 F: crate::future::Future + 'static,
516 F::Output: 'static,
517 {
518 let (handle, notified) = unsafe {
520 me.shared
521 .owned
522 .bind_local(future, me.clone(), id, spawned_at)
523 };
524
525 me.task_hooks.spawn(&TaskMeta {
526 id,
527 spawned_at,
528 _phantom: Default::default(),
529 });
530
531 if let Some(notified) = notified {
532 me.schedule(notified);
533 }
534
535 handle
536 }
537
538 #[cfg(all(
540 tokio_unstable,
541 feature = "taskdump",
542 target_os = "linux",
543 any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
544 ))]
545 pub(crate) fn dump(&self) -> crate::runtime::Dump {
546 use crate::runtime::dump;
547 use task::trace::trace_current_thread;
548
549 let mut traces = vec![];
550
551 context::with_scheduler(|maybe_context| {
553 let context = if let Some(context) = maybe_context {
555 context.expect_current_thread()
556 } else {
557 return;
558 };
559 let mut maybe_core = context.core.borrow_mut();
560 let core = if let Some(core) = maybe_core.as_mut() {
561 core
562 } else {
563 return;
564 };
565 let local = &mut core.tasks;
566
567 if self.shared.inject.is_closed() {
568 return;
569 }
570
571 traces = trace_current_thread(&self.shared.owned, local, &self.shared.inject)
572 .into_iter()
573 .map(|(id, trace)| dump::Task::new(id, trace))
574 .collect();
575
576 drop(maybe_core);
578
579 wake_deferred_tasks_and_free(context);
583 });
584
585 dump::Dump::new(traces)
586 }
587
588 fn next_remote_task(&self) -> Option<Notified> {
589 self.shared.inject.pop()
590 }
591
592 fn waker_ref(me: &Arc<Self>) -> WakerRef<'_> {
593 me.shared.woken.store(true, Release);
596 waker_ref(me)
597 }
598
599 pub(crate) fn reset_woken(&self) -> bool {
601 self.shared.woken.swap(false, AcqRel)
602 }
603
604 pub(crate) fn num_alive_tasks(&self) -> usize {
605 self.shared.owned.num_alive_tasks()
606 }
607
608 pub(crate) fn injection_queue_depth(&self) -> usize {
609 self.shared.inject.len()
610 }
611
612 pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics {
613 assert_eq!(0, worker);
614 &self.shared.worker_metrics
615 }
616}
617
618cfg_unstable_metrics! {
619 impl Handle {
620 pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics {
621 &self.shared.scheduler_metrics
622 }
623
624 pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize {
625 self.worker_metrics(worker).queue_depth()
626 }
627
628 pub(crate) fn num_blocking_threads(&self) -> usize {
629 self.blocking_spawner.num_threads()
630 }
631
632 pub(crate) fn num_idle_blocking_threads(&self) -> usize {
633 self.blocking_spawner.num_idle_threads()
634 }
635
636 pub(crate) fn blocking_queue_depth(&self) -> usize {
637 self.blocking_spawner.queue_depth()
638 }
639
640 cfg_64bit_metrics! {
641 pub(crate) fn spawned_tasks_count(&self) -> u64 {
642 self.shared.owned.spawned_tasks_count()
643 }
644 }
645 }
646}
647
648use std::num::NonZeroU64;
649
650impl Handle {
651 pub(crate) fn owned_id(&self) -> NonZeroU64 {
652 self.shared.owned.id
653 }
654
655 pub(crate) fn name(&self) -> Option<&str> {
656 self.name.as_deref()
657 }
658}
659
660impl fmt::Debug for Handle {
661 fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
662 fmt.debug_struct("current_thread::Handle { ... }").finish()
663 }
664}
665
666impl Schedule for Arc<Handle> {
669 fn release(&self, task: &Task<Self>) -> Option<Task<Self>> {
670 self.shared.owned.remove(task)
671 }
672
673 fn schedule(&self, task: task::Notified<Self>) {
674 use scheduler::Context::CurrentThread;
675
676 context::with_scheduler(|maybe_cx| match maybe_cx {
677 Some(CurrentThread(cx)) if Arc::ptr_eq(self, &cx.handle) => {
678 let mut core = cx.core.borrow_mut();
679
680 if let Some(core) = core.as_mut() {
683 core.push_task(self, task);
684 }
685 }
686 _ => {
687 self.shared.scheduler_metrics.inc_remote_schedule_count();
689
690 self.shared.inject.push(task);
692 self.driver.unpark();
693 }
694 });
695 }
696
697 fn hooks(&self) -> TaskHarnessScheduleHooks {
698 TaskHarnessScheduleHooks {
699 task_terminate_callback: self.task_hooks.task_terminate_callback.clone(),
700 }
701 }
702
703 cfg_unstable! {
704 fn unhandled_panic(&self) {
705 use crate::runtime::UnhandledPanic;
706
707 match self.shared.config.unhandled_panic {
708 UnhandledPanic::Ignore => {
709 }
711 UnhandledPanic::ShutdownRuntime => {
712 use scheduler::Context::CurrentThread;
713
714 context::with_scheduler(|maybe_cx| match maybe_cx {
719 Some(CurrentThread(cx)) if Arc::ptr_eq(self, &cx.handle) => {
720 let mut core = cx.core.borrow_mut();
721
722 if let Some(core) = core.as_mut() {
724 core.unhandled_panic = true;
725 self.shared.owned.close_and_shutdown_all(0);
726 }
727 }
728 _ => unreachable!("runtime core not set in CURRENT thread-local"),
729 })
730 }
731 }
732 }
733 }
734}
735
736impl Wake for Handle {
737 fn wake(arc_self: Arc<Self>) {
738 Wake::wake_by_ref(&arc_self);
739 }
740
741 fn wake_by_ref(arc_self: &Arc<Self>) {
743 let already_woken = arc_self.shared.woken.swap(true, Release);
744
745 if !already_woken {
746 use scheduler::Context::CurrentThread;
747
748 context::with_scheduler(|maybe_cx| match maybe_cx {
751 Some(CurrentThread(cx)) if Arc::ptr_eq(arc_self, &cx.handle) => {}
752 _ => {
753 arc_self.driver.unpark();
754 }
755 });
756 }
757 }
758}
759
760struct CoreGuard<'a> {
765 context: scheduler::Context,
766 scheduler: &'a CurrentThread,
767}
768
769impl CoreGuard<'_> {
770 #[track_caller]
771 fn block_on<F: Future>(self, future: F) -> F::Output {
772 let ret = self.enter(|mut core, context| {
773 let waker = Handle::waker_ref(&context.handle);
774 let mut cx = std::task::Context::from_waker(&waker);
775
776 pin!(future);
777
778 core.metrics.start_processing_scheduled_tasks();
779
780 'outer: loop {
781 let handle = &context.handle;
782
783 if handle.reset_woken() {
784 let (c, res) = context.enter(core, || {
785 crate::task::coop::budget(|| future.as_mut().poll(&mut cx))
786 });
787
788 core = c;
789
790 if let Ready(v) = res {
791 return (core, Some(v));
792 }
793 }
794
795 for _ in 0..handle.shared.config.event_interval {
796 if core.unhandled_panic {
798 return (core, None);
799 }
800
801 core.tick();
802
803 let entry = core.next_task(handle);
804
805 let task = match entry {
806 Some(entry) => entry,
807 None => {
808 core.metrics.end_processing_scheduled_tasks();
809
810 core = if context.has_pending_work(&core) {
811 context.park_yield(core, handle)
812 } else {
813 context.park(core, handle)
814 };
815
816 core.metrics.start_processing_scheduled_tasks();
817
818 continue 'outer;
820 }
821 };
822
823 let task = context.handle.shared.owned.assert_owner(task);
824
825 #[cfg(tokio_unstable)]
826 let task_meta = task.task_meta();
827
828 let (c, ()) = context.run_task(core, || {
829 #[cfg(tokio_unstable)]
830 context.handle.task_hooks.poll_start_callback(&task_meta);
831
832 task.run();
833
834 #[cfg(tokio_unstable)]
835 context.handle.task_hooks.poll_stop_callback(&task_meta);
836 });
837
838 core = c;
839 }
840
841 core.metrics.end_processing_scheduled_tasks();
842
843 core = context.park_yield(core, handle);
846
847 core.metrics.start_processing_scheduled_tasks();
848 }
849 });
850
851 match ret {
852 Some(ret) => ret,
853 None => {
854 panic!("a spawned task panicked and the runtime is configured to shut down on unhandled panic");
856 }
857 }
858 }
859
860 fn enter<F, R>(self, f: F) -> R
863 where
864 F: FnOnce(Box<Core>, &Context) -> (Box<Core>, R),
865 {
866 let context = self.context.expect_current_thread();
867
868 let core = context.core.borrow_mut().take().expect("core missing");
870
871 let (core, ret) = context::set_scheduler(&self.context, || f(core, context));
873
874 *context.core.borrow_mut() = Some(core);
875
876 ret
877 }
878}
879
880impl Drop for CoreGuard<'_> {
881 fn drop(&mut self) {
882 let context = self.context.expect_current_thread();
883
884 if let Some(core) = context.core.borrow_mut().take() {
885 self.scheduler.core.set(core);
888
889 self.scheduler.notify.notify_one();
891 }
892 }
893}