1use std::sync::Arc;
44use std::vec::IntoIter;
45use std::time::Duration;
46use std::future::Future;
47use std::cell::{Cell, UnsafeCell};
48use std::marker::PhantomData;
49use std::io::{Error, ErrorKind, Result};
50use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
51use std::task::{Context, Poll, Waker};
52use std::thread::{self, Builder};
53
54use async_stream::stream;
55use crossbeam_channel::{bounded, Sender};
56use crossbeam_deque::{Injector, Steal, Stealer, Worker};
57use crossbeam_queue::{ArrayQueue, SegQueue};
58use crossbeam_utils::atomic::AtomicCell;
59use st3::{StealError,
60 fifo::{Worker as FIFOWorker, Stealer as FIFOStealer}};
61use flume::bounded as async_bounded;
62use futures::{
63 future::{BoxFuture, FutureExt},
64 stream::{BoxStream, Stream, StreamExt},
65 task::waker_ref,
66 TryFuture,
67};
68use parking_lot::{Condvar, Mutex};
69use rand::{Rng, thread_rng};
70use num_cpus;
71use wrr::IWRRSelector;
72use quanta::{Clock, Instant as QInstant};
73use log::warn;
74
75use super::{
76 PI_ASYNC_LOCAL_THREAD_ASYNC_RUNTIME, PI_ASYNC_THREAD_LOCAL_ID, DEFAULT_MAX_HIGH_PRIORITY_BOUNDED, DEFAULT_HIGH_PRIORITY_BOUNDED, DEFAULT_MAX_LOW_PRIORITY_BOUNDED, alloc_rt_uid, local_async_runtime, AsyncMapReduce, AsyncPipelineResult, AsyncRuntime,
77 AsyncRuntimeExt, AsyncTask, AsyncTaskPollClaim, AsyncTaskPollGuard, AsyncTaskPool, AsyncTaskPoolExt, AsyncTaskTimerByNotCancel, AsyncTimingTask,
78 AsyncWait, AsyncWaitAny, AsyncWaitAnyCallback, AsyncWaitTimeout, LocalAsyncWaitTimeout, LocalAsyncRuntime, TaskId, TaskHandle, YieldNow, prune_stale_waiting_workers, register_waiting_worker, wake_waiting_worker,
79 requeue_runtime_task
80};
81
82#[cfg(not(target_arch = "wasm32"))]
86const DEFAULT_INIT_WORKER_SIZE: usize = 2;
87#[cfg(target_arch = "wasm32")]
88const DEFAULT_INIT_WORKER_SIZE: usize = 1;
89
90const DEFAULT_WORKER_THREAD_PREFIX: &str = "Default-Multi-RT";
94
95const DEFAULT_THREAD_STACK_SIZE: usize = 1024 * 1024;
99
100const DEFAULT_WORKER_THREAD_SLEEP_TIME: u64 = 10;
104
105const DEFAULT_RUNTIME_SLEEP_TIME: u64 = 1000;
109
110const DEFAULT_MAX_WEIGHT: u8 = 254;
114
115const DEFAULT_MIN_WEIGHT: u8 = 1;
119
120const MULTI_THREAD_WORKER_ID_MASK: usize = 0xffffffff;
124
125#[derive(Clone, Copy)]
148struct MultiThreadWorkerContext {
149 thread_id: usize,
150 pool: *const (),
151}
152
153impl MultiThreadWorkerContext {
154 const UNBOUND: Self = MultiThreadWorkerContext {
155 thread_id: usize::MAX,
156 pool: std::ptr::null(),
157 };
158
159 #[inline]
164 const fn runtime_id(self) -> usize {
165 self.thread_id >> 32
166 }
167
168 #[inline]
184 fn owner_worker_id<P>(self, pool: &P) -> usize {
185 let expected = pool as *const P as *const ();
186 if self.pool != expected {
187 panic!(
188 "Multi-thread task pool owner mismatch: owner-only worker state requires the worker bound to this pool"
189 );
190 }
191
192 self.thread_id & MULTI_THREAD_WORKER_ID_MASK
193 }
194}
195
196thread_local! {
197 static PI_ASYNC_MULTI_THREAD_WORKER_CONTEXT: Cell<MultiThreadWorkerContext>
202 = Cell::new(MultiThreadWorkerContext::UNBOUND);
203}
204
205#[inline]
211fn current_multi_thread_worker_context() -> MultiThreadWorkerContext {
212 match PI_ASYNC_MULTI_THREAD_WORKER_CONTEXT.try_with(|context| context.get()) {
213 Ok(context) => context,
214 Err(e) => {
215 panic!(
216 "Get multi-thread worker context failed, thread: {:?}, reason: {:?}",
217 thread::current(),
218 e
219 );
220 },
221 }
222}
223
224struct ComputationalTaskQueue<O: Default + 'static> {
228 stack: Worker<Arc<AsyncTask<ComputationalTaskPool<O>, O>>>, queue: SegQueue<Arc<AsyncTask<ComputationalTaskPool<O>, O>>>, thread_waker: Arc<(AtomicBool, Mutex<()>, Condvar)>, }
232
233impl<O: Default + 'static> ComputationalTaskQueue<O> {
234 pub fn new(thread_waker: Arc<(AtomicBool, Mutex<()>, Condvar)>) -> Self {
236 let stack = Worker::new_lifo();
237 let queue = SegQueue::new();
238
239 ComputationalTaskQueue {
240 stack,
241 queue,
242 thread_waker,
243 }
244 }
245
246 pub fn len(&self) -> usize {
248 self.stack.len() + self.queue.len()
249 }
250}
251
252pub struct ComputationalTaskPool<O: Default + 'static> {
287 workers: Vec<ComputationalTaskQueue<O>>, waits: Option<Arc<ArrayQueue<Arc<(AtomicBool, Mutex<()>, Condvar)>>>>, consume_count: Arc<AtomicUsize>, produce_count: Arc<AtomicUsize>, }
292
293unsafe impl<O: Default + 'static> Send for ComputationalTaskPool<O> {}
299unsafe impl<O: Default + 'static> Sync for ComputationalTaskPool<O> {}
302
303impl<O: Default + 'static> Default for ComputationalTaskPool<O> {
304 fn default() -> Self {
305 #[cfg(not(target_arch = "wasm32"))]
306 let core_len = num_cpus::get(); #[cfg(target_arch = "wasm32")]
308 let core_len = 1; ComputationalTaskPool::new(core_len)
310 }
311}
312
313impl<O: Default + 'static> AsyncTaskPool<O> for ComputationalTaskPool<O> {
314 type Pool = ComputationalTaskPool<O>;
315
316 #[inline]
322 fn get_thread_id(&self) -> usize {
323 match PI_ASYNC_THREAD_LOCAL_ID.try_with(move |thread_id| unsafe {
324 *thread_id.get()
326 }) {
327 Err(e) => {
328 panic!(
329 "Get thread id failed, thread: {:?}, reason: {:?}",
330 thread::current(),
331 e
332 );
333 }
334 Ok(id) => id,
335 }
336 }
337
338 #[inline]
339 fn len(&self) -> usize {
340 if let Some(len) = self
341 .produce_count
342 .load(Ordering::Relaxed)
343 .checked_sub(self.consume_count.load(Ordering::Relaxed))
344 {
345 len
346 } else {
347 0
348 }
349 }
350
351 #[inline]
352 fn push(&self, task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
353 let index = self.produce_count.fetch_add(1, Ordering::Relaxed) % self.workers.len();
354 self.workers[index].queue.push(task);
355 Ok(())
356 }
357
358 #[inline]
365 fn push_local(&self, task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
366 let context = current_multi_thread_worker_context();
367 let rt_uid = task.owner();
368 if context.runtime_id() == rt_uid {
369 let worker = &self.workers[context.owner_worker_id(self)];
371 worker.queue.push(task);
372
373 self.produce_count.fetch_add(1, Ordering::Relaxed);
374 Ok(())
375 } else {
376 self.push(task)
378 }
379 }
380
381 #[inline]
388 fn push_priority(&self,
389 priority: usize,
390 task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
391 if priority >= DEFAULT_MAX_HIGH_PRIORITY_BOUNDED {
392 let context = current_multi_thread_worker_context();
394 let rt_uid = task.owner();
395 if context.runtime_id() == rt_uid {
396 let worker = &self.workers[context.owner_worker_id(self)];
397 worker.stack.push(task);
398
399 self.produce_count.fetch_add(1, Ordering::Relaxed);
400 Ok(())
401 } else {
402 self.push(task)
403 }
404 } else if priority >= DEFAULT_HIGH_PRIORITY_BOUNDED {
405 self.push_local(task)
407 } else {
408 self.push(task)
410 }
411 }
412
413 #[inline]
414 fn push_keep(&self, task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
415 self.push_priority(DEFAULT_HIGH_PRIORITY_BOUNDED, task)
416 }
417
418 #[inline]
425 fn try_pop(&self) -> Option<Arc<AsyncTask<Self::Pool, O>>> {
426 let id = current_multi_thread_worker_context().owner_worker_id(self);
427 let worker = &self.workers[id];
428 let task = worker.stack.pop();
429 if task.is_some() {
430 self.consume_count.fetch_add(1, Ordering::Relaxed);
432 return task;
433 }
434
435 let task = worker.queue.pop();
436 if task.is_some() {
437 self.consume_count.fetch_add(1, Ordering::Relaxed);
438 }
439
440 task
441 }
442
443 #[inline]
449 fn try_pop_all(&self) -> IntoIter<Arc<AsyncTask<Self::Pool, O>>> {
450 let mut tasks = Vec::with_capacity(self.len());
451 while let Some(task) = self.try_pop() {
452 tasks.push(task);
453 }
454
455 tasks.into_iter()
456 }
457
458 #[inline]
459 fn get_thread_waker(&self) -> Option<&Arc<(AtomicBool, Mutex<()>, Condvar)>> {
460 None
462 }
463}
464
465impl<O: Default + 'static> AsyncTaskPoolExt<O> for ComputationalTaskPool<O> {
466 #[inline]
467 fn set_waits(&mut self, waits: Arc<ArrayQueue<Arc<(AtomicBool, Mutex<()>, Condvar)>>>) {
468 self.waits = Some(waits);
469 }
470
471 #[inline]
472 fn get_waits(&self) -> Option<&Arc<ArrayQueue<Arc<(AtomicBool, Mutex<()>, Condvar)>>>> {
473 self.waits.as_ref()
474 }
475
476 #[inline]
477 fn worker_len(&self) -> usize {
478 self.workers.len()
479 }
480
481 #[inline]
487 fn clone_thread_waker(&self) -> Option<Arc<(AtomicBool, Mutex<()>, Condvar)>> {
488 let worker = &self.workers[current_multi_thread_worker_context().owner_worker_id(self)];
489 Some(worker.thread_waker.clone())
490 }
491}
492
493impl<O: Default + 'static> ComputationalTaskPool<O> {
494 pub fn new(mut size: usize) -> Self {
511 if size < DEFAULT_INIT_WORKER_SIZE {
512 size = DEFAULT_INIT_WORKER_SIZE;
514 }
515
516 let mut workers = Vec::with_capacity(size);
517 for _ in 0..size {
518 let thread_waker = Arc::new((AtomicBool::new(false), Mutex::new(()), Condvar::new()));
519 let worker = ComputationalTaskQueue::new(thread_waker);
520 workers.push(worker);
521 }
522 let consume_count = Arc::new(AtomicUsize::new(0));
523 let produce_count = Arc::new(AtomicUsize::new(0));
524
525 ComputationalTaskPool {
526 workers,
527 waits: None,
528 consume_count,
529 produce_count,
530 }
531 }
532}
533
534struct StealableTaskQueue<O: Default + 'static> {
538 stack: UnsafeCell<Option<Arc<AsyncTask<StealableTaskPool<O>, O>>>>, internal: FIFOWorker<Arc<AsyncTask<StealableTaskPool<O>, O>>>, external: Worker<Arc<AsyncTask<StealableTaskPool<O>, O>>>, selector: UnsafeCell<IWRRSelector<2>>, thread_waker: Arc<(AtomicBool, Mutex<()>, Condvar)>, }
544
545impl<O: Default + 'static> StealableTaskQueue<O> {
546 pub fn new(
549 init_queue_capacity: usize,
550 thread_waker: Arc<(AtomicBool, Mutex<()>, Condvar)>,
551 ) -> (Self,
552 FIFOStealer<Arc<AsyncTask<StealableTaskPool<O>, O>>>,
553 Stealer<Arc<AsyncTask<StealableTaskPool<O>, O>>>) {
554 let stack = UnsafeCell::new(None);
555 let internal = FIFOWorker::new(init_queue_capacity);
556 let external = Worker::new_fifo();
557 let internal_stealer = internal.stealer();
558 let external_stealer = external.stealer();
559 let selector = UnsafeCell::new(IWRRSelector::new([2, 1]));
560
561 (
562 StealableTaskQueue {
563 stack,
564 internal,
565 external,
566 selector,
567 thread_waker,
568 },
569 internal_stealer,
570 external_stealer
571 )
572 }
573
574 pub const fn stack_capacity(&self) -> usize {
576 1
577 }
578
579 pub fn internal_capacity(&self) -> usize {
581 self.internal.capacity()
582 }
583
584 pub fn remaining_internal_capacity(&self) -> usize {
586 self.internal.spare_capacity()
587 }
588
589 #[inline]
596 pub fn stack_len(&self) -> usize {
597 unsafe {
598 if (&*self.stack.get()).is_some() {
601 1
602 } else {
603 0
604 }
605 }
606 }
607
608 pub fn internal_len(&self) -> usize {
610 self
611 .internal_capacity()
612 .checked_sub(self.remaining_internal_capacity())
613 .unwrap_or(0)
614 }
615
616 pub fn external_len(&self) -> usize {
618 self.external.len()
619 }
620}
621
622pub struct StealableTaskPool<O: Default + 'static> {
668 public: Injector<Arc<AsyncTask<StealableTaskPool<O>, O>>>, workers: Vec<StealableTaskQueue<O>>, internal_stealers: Vec<FIFOStealer<Arc<AsyncTask<StealableTaskPool<O>, O>>>>, external_stealers: Vec<Stealer<Arc<AsyncTask<StealableTaskPool<O>, O>>>>, internal_consume: AtomicUsize, internal_produce: AtomicUsize, internal_traffic_statistics: AtomicUsize, external_consume: AtomicUsize, external_produce: AtomicUsize, external_traffic_statistics: AtomicUsize, weights: [u8; 2], clock: Clock, interval: usize, last_time: AtomicCell<QInstant>, waits: Option<Arc<ArrayQueue<Arc<(AtomicBool, Mutex<()>, Condvar)>>>>, }
684
685unsafe impl<O: Default + 'static> Send for StealableTaskPool<O> {}
691unsafe impl<O: Default + 'static> Sync for StealableTaskPool<O> {}
694
695impl<O: Default + 'static> Default for StealableTaskPool<O> {
696 fn default() -> Self {
697 StealableTaskPool::new()
698 }
699}
700
701impl<O: Default + 'static> AsyncTaskPool<O> for StealableTaskPool<O> {
702 type Pool = StealableTaskPool<O>;
703
704 #[inline]
710 fn get_thread_id(&self) -> usize {
711 match PI_ASYNC_THREAD_LOCAL_ID.try_with(move |thread_id| unsafe {
712 *thread_id.get()
714 }) {
715 Err(e) => {
716 panic!(
717 "Get thread id failed, thread: {:?}, reason: {:?}",
718 thread::current(),
719 e
720 );
721 }
722 Ok(id) => id,
723 }
724 }
725
726 #[inline]
727 fn len(&self) -> usize {
728 self.internal_produce
729 .load(Ordering::Relaxed)
730 .checked_sub(self.internal_consume.load(Ordering::Relaxed))
731 .unwrap_or(0)
732 +
733 self.external_produce
734 .load(Ordering::Relaxed)
735 .checked_sub(self.external_consume.load(Ordering::Relaxed))
736 .unwrap_or(0)
737 }
738
739 #[inline]
740 fn push(&self, task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
741 self.public.push(task);
742
743 self
744 .external_produce
745 .fetch_add(1, Ordering::Relaxed);
746 Ok(())
747 }
748
749 #[inline]
756 fn push_local(&self, task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
757 let context = current_multi_thread_worker_context();
758 let rt_uid = task.owner();
759 if context.runtime_id() == rt_uid {
760 let worker = &self.workers[context.owner_worker_id(self)];
762 if worker.remaining_internal_capacity() > 0 {
763 let _ = worker.internal.push(task);
765
766 self
767 .internal_produce
768 .fetch_add(1, Ordering::Relaxed);
769 Ok(())
770 } else {
771 self.push(task)
773 }
774 } else {
775 self.push(task)
777 }
778 }
779
780 #[inline]
787 fn push_priority(&self,
788 priority: usize,
789 task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
790 if priority >= DEFAULT_MAX_HIGH_PRIORITY_BOUNDED {
791 let context = current_multi_thread_worker_context();
793 let rt_uid = task.owner();
794 if context.runtime_id() == rt_uid {
795 let worker = &self.workers[context.owner_worker_id(self)];
797 if worker.stack_len() < 1 {
798 unsafe {
800 *worker.stack.get() = Some(task);
803 }
804 } else if worker.remaining_internal_capacity() > 0 {
805 let _ = worker.internal.push(task);
807 } else {
808 return self.push(task);
810 }
811
812 self
813 .internal_produce
814 .fetch_add(1, Ordering::Relaxed);
815 Ok(())
816 } else {
817 self.push(task)
819 }
820 } else if priority >= DEFAULT_HIGH_PRIORITY_BOUNDED {
821 self.push_local(task)
823 } else {
824 self.push(task)
826 }
827 }
828
829 #[inline]
830 fn push_keep(&self, task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
831 self.push_priority(DEFAULT_HIGH_PRIORITY_BOUNDED, task)
832 }
833
834 #[inline]
842 fn try_pop(&self) -> Option<Arc<AsyncTask<Self::Pool, O>>> {
843 let id = current_multi_thread_worker_context().owner_worker_id(self);
844 let worker = &self.workers[id];
845 let task = unsafe {
846 (&mut *worker
849 .stack
850 .get())
851 .take()
852 };
853 if task.is_some() {
854 return task;
856 }
857
858 try_pop_by_weight(self, worker, id)
860 }
861
862 #[inline]
868 fn try_pop_all(&self) -> IntoIter<Arc<AsyncTask<Self::Pool, O>>> {
869 let mut tasks = Vec::with_capacity(self.len());
870 while let Some(task) = self.try_pop() {
871 tasks.push(task);
872 }
873
874 tasks.into_iter()
875 }
876
877 #[inline]
878 fn get_thread_waker(&self) -> Option<&Arc<(AtomicBool, Mutex<()>, Condvar)>> {
879 None
881 }
882}
883
884const fn get_msb(n: usize) -> usize {
886 usize::BITS as usize - n.leading_zeros() as usize
887}
888
889fn try_pop_by_weight<O: Default + 'static>(pool: &StealableTaskPool<O>,
914 local_worker: &StealableTaskQueue<O>,
915 local_worker_id: usize)
916 -> Option<Arc<AsyncTask<StealableTaskPool<O>, O>>> {
917 unsafe {
918 let duration = pool
922 .clock
923 .recent()
924 .duration_since(pool.last_time.load())
925 .as_millis() as usize;
926 if duration >= pool.interval {
927 let new_external_traffic_statistics = pool
929 .external_produce
930 .load(Ordering::Relaxed);
931 let new_internal_traffic_statistics = pool
932 .internal_produce
933 .load(Ordering::Relaxed);
934
935 let external_delta = if new_external_traffic_statistics == 0 {
937 1
939 } else {
940 new_external_traffic_statistics
942 .checked_sub(pool
943 .external_traffic_statistics
944 .load(Ordering::Relaxed))
945 .unwrap_or(1)
946 };
947 pool
948 .external_traffic_statistics
949 .store(new_external_traffic_statistics, Ordering::Relaxed); let internal_delta = if new_internal_traffic_statistics == 0 {
951 1
953 } else {
954 new_internal_traffic_statistics
956 .checked_sub(pool
957 .internal_traffic_statistics
958 .load(Ordering::Relaxed))
959 .unwrap_or(1)
960 };
961 pool
962 .internal_traffic_statistics
963 .store(new_internal_traffic_statistics, Ordering::Relaxed); let selector = &mut *local_worker.selector.get();
967 if external_delta > internal_delta {
968 let msb = get_msb(internal_delta);
970 let internal_weight
971 = (internal_delta >> msb.checked_sub(2).unwrap_or(0)).max(1);
972 let external_weight
973 = ((external_delta >> msb).min(DEFAULT_MAX_WEIGHT as usize)).max(1);
974
975 selector.change_weight(0, external_weight as u8);
976 selector.change_weight(1, internal_weight as u8);
977 } else if external_delta < internal_delta {
978 let msb = get_msb(external_delta);
980 let external_weight
981 = (external_delta >> msb.checked_sub(2).unwrap_or(0)).max(1);
982 let internal_weight
983 = ((internal_delta >> msb).min(DEFAULT_MAX_WEIGHT as usize)).max(1);
984
985 selector.change_weight(0, external_weight as u8);
986 selector.change_weight(1, internal_weight as u8);
987 } else {
988 selector.change_weight(0, 1);
990 selector.change_weight(1, 1);
991 }
992
993 pool.last_time.store(pool.clock.recent()); }
995
996 match (&mut *local_worker.selector.get()).select() {
998 0 => {
999 let task = try_pop_external(pool, local_worker, local_worker_id);
1001 if task.is_some() {
1002 task
1003 } else {
1004 try_pop_internal(pool, local_worker, local_worker_id)
1006 }
1007 },
1008 _ => {
1009 let task = try_pop_internal(pool, local_worker, local_worker_id);
1011 if task.is_some() {
1012 task
1013 } else {
1014 try_pop_external(pool, local_worker, local_worker_id)
1016 }
1017 },
1018 }
1019 }
1020}
1021
1022#[inline]
1024fn try_pop_internal<O: Default + 'static>(pool: &StealableTaskPool<O>,
1025 local_worker: &StealableTaskQueue<O>,
1026 local_worker_id: usize)
1027 -> Option<Arc<AsyncTask<StealableTaskPool<O>, O>>> {
1028 let task = local_worker
1029 .internal
1030 .pop();
1031 if task.is_some() {
1032 pool
1034 .internal_consume
1035 .fetch_add(1, Ordering::Relaxed);
1036 task
1037 } else {
1038 let mut gen = thread_rng();
1040 let mut worker_stealers: Vec<&FIFOStealer<Arc<AsyncTask<StealableTaskPool<O>, O>>>> = pool
1041 .internal_stealers
1042 .iter()
1043 .enumerate()
1044 .filter_map(|(index, other)| {
1045 if index != local_worker_id {
1046 Some(other)
1047 } else {
1048 None
1050 }
1051 })
1052 .collect();
1053
1054 let remaining_len = local_worker.remaining_internal_capacity();
1055 loop {
1056 if worker_stealers.len() == 0 {
1058 break;
1060 }
1061
1062 let index = gen.gen_range(0..worker_stealers.len());
1063 let worker_stealer = worker_stealers.swap_remove(index);
1064
1065 match worker_stealer.steal_and_pop(&local_worker.internal,
1066 |count| {
1067 let stealable_len = count / 2;
1068 if stealable_len <= remaining_len {
1069 if stealable_len == 0 {
1071 1
1072 } else {
1073 stealable_len
1074 }
1075 } else {
1076 remaining_len
1078 }
1079 }) {
1080 Err(StealError::Empty) => {
1081 continue;
1083 },
1084 Err(StealError::Busy) => {
1085 continue;
1087 },
1088 Ok((task, _)) => {
1089 pool.internal_consume.fetch_add(1, Ordering::Relaxed);
1091 return Some(task);
1092 },
1093 }
1094 }
1095
1096 None
1097 }
1098}
1099
1100#[inline]
1102fn try_pop_external<O: Default + 'static>(pool: &StealableTaskPool<O>,
1103 local_worker: &StealableTaskQueue<O>,
1104 local_worker_id: usize)
1105 -> Option<Arc<AsyncTask<StealableTaskPool<O>, O>>> {
1106 let task = local_worker
1107 .external
1108 .pop();
1109 if task.is_some() {
1110 pool
1112 .external_consume
1113 .fetch_add(1, Ordering::Relaxed);
1114 task
1115 } else {
1116 let task = try_pop_public(pool, local_worker);
1118 if task.is_some() {
1119 pool
1121 .external_consume
1122 .fetch_add(1, Ordering::Relaxed);
1123 task
1124 } else {
1125 let mut gen = thread_rng();
1127 let mut worker_stealers: Vec<&Stealer<Arc<AsyncTask<StealableTaskPool<O>, O>>>> = pool
1128 .external_stealers
1129 .iter()
1130 .enumerate()
1131 .filter_map(|(index, other)| {
1132 if index != local_worker_id {
1133 Some(other)
1134 } else {
1135 None
1137 }
1138 })
1139 .collect();
1140
1141 loop {
1142 if worker_stealers.len() == 0 {
1144 break;
1146 }
1147
1148 let index = gen.gen_range(0..worker_stealers.len());
1149 let worker_stealer = worker_stealers.swap_remove(index);
1150
1151 match worker_stealer.steal_batch_and_pop(&local_worker.external) {
1152 Steal::Success(task) => {
1153 pool.external_consume.fetch_add(1, Ordering::Relaxed);
1155 return Some(task);
1156 },
1157 Steal::Retry => {
1158 continue;
1160 },
1161 Steal::Empty => {
1162 continue;
1164 },
1165 }
1166 }
1167
1168 None
1169 }
1170 }
1171}
1172
1173#[inline]
1175fn try_pop_public<O: Default + 'static>(pool: &StealableTaskPool<O>,
1176 local_worker: &StealableTaskQueue<O>)
1177 -> Option<Arc<AsyncTask<StealableTaskPool<O>, O>>> {
1178 loop {
1179 match pool.public.steal_batch_and_pop(&local_worker.external) {
1180 Steal::Empty => {
1181 return None;
1183 },
1184 Steal::Retry => {
1185 continue;
1187 },
1188 Steal::Success(task) => {
1189 pool.external_consume.fetch_add(1, Ordering::Relaxed);
1191 return Some(task);
1192 },
1193 }
1194 }
1195}
1196
1197impl<O: Default + 'static> AsyncTaskPoolExt<O> for StealableTaskPool<O> {
1198 #[inline]
1199 fn set_waits(&mut self, waits: Arc<ArrayQueue<Arc<(AtomicBool, Mutex<()>, Condvar)>>>) {
1200 self.waits = Some(waits);
1201 }
1202
1203 #[inline]
1204 fn get_waits(&self) -> Option<&Arc<ArrayQueue<Arc<(AtomicBool, Mutex<()>, Condvar)>>>> {
1205 self.waits.as_ref()
1206 }
1207
1208 #[inline]
1209 fn worker_len(&self) -> usize {
1210 self.workers.len()
1211 }
1212
1213 #[inline]
1220 fn clone_thread_waker(&self) -> Option<Arc<(AtomicBool, Mutex<()>, Condvar)>> {
1221 let id = current_multi_thread_worker_context().owner_worker_id(self);
1222 if let Some(worker) = self.workers.get(id) {
1223 return Some(worker.thread_waker.clone());
1224 }
1225
1226 None
1227 }
1228}
1229
1230impl<O: Default + 'static> StealableTaskPool<O> {
1231 pub fn new() -> Self {
1239 #[cfg(not(target_arch = "wasm32"))]
1240 let size = num_cpus::get_physical() * 2; #[cfg(target_arch = "wasm32")]
1242 let size = 1; StealableTaskPool::with(size,
1244 0x8000,
1245 [1, 1],
1246 3000)
1247 }
1248
1249 pub fn with(worker_size: usize,
1271 internal_queue_capacity: usize,
1272 weights: [u8; 2],
1273 interval: usize) -> Self {
1274 if worker_size == 0 {
1275 panic!(
1277 "Create WorkerTaskPool failed, worker size: {}, reason: invalid worker size",
1278 worker_size
1279 );
1280 }
1281 if interval == 0 {
1282 panic!(
1283 "Create WorkerTaskPool failed, interval: {}, reason: invalid interval",
1284 worker_size
1285 );
1286 }
1287
1288 let public = Injector::new();
1289 let mut workers = Vec::with_capacity(worker_size);
1290 let mut internal_stealers = Vec::with_capacity(worker_size);
1291 let mut external_stealers = Vec::with_capacity(worker_size);
1292 for _ in 0..worker_size {
1293 let thread_waker = Arc::new((AtomicBool::new(false), Mutex::new(()), Condvar::new()));
1295 let (worker,
1296 internal_stealer,
1297 external_stealer) =
1298 StealableTaskQueue::new(internal_queue_capacity,
1299 thread_waker);
1300 workers.push(worker);
1301 internal_stealers.push(internal_stealer);
1302 external_stealers.push(external_stealer);
1303 }
1304 let internal_consume = AtomicUsize::new(0);
1305 let internal_produce = AtomicUsize::new(0);
1306 let internal_traffic_statistics = AtomicUsize::new(0);
1307 let external_consume = AtomicUsize::new(0);
1308 let external_produce = AtomicUsize::new(0);
1309 let external_traffic_statistics = AtomicUsize::new(0);
1310 let clock = Clock::new();
1311 let last_time = AtomicCell::new(clock.recent());
1312
1313 StealableTaskPool {
1314 public,
1315 workers,
1316 internal_stealers,
1317 external_stealers,
1318 internal_consume,
1319 internal_produce,
1320 internal_traffic_statistics,
1321 external_consume,
1322 external_produce,
1323 external_traffic_statistics,
1324 weights,
1325 clock,
1326 interval,
1327 last_time,
1328 waits: None,
1329 }
1330 }
1331}
1332
1333pub struct MultiTaskRuntime<
1337 O: Default + 'static = (),
1338 P: AsyncTaskPoolExt<O> + AsyncTaskPool<O> = StealableTaskPool<O>,
1339>(
1340 Arc<(
1341 usize, Arc<P>, Option<
1344 Vec<(
1345 Sender<(usize, AsyncTimingTask<P, O>)>,
1346 Arc<AsyncTaskTimerByNotCancel<P, O>>,
1347 )>,
1348 >, AtomicUsize, Arc<ArrayQueue<Arc<(AtomicBool, Mutex<()>, Condvar)>>>, AtomicUsize, AtomicUsize, )>,
1354);
1355
1356unsafe impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O>> Send
1357 for MultiTaskRuntime<O, P>
1358{
1359}
1360unsafe impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O>> Sync
1361 for MultiTaskRuntime<O, P>
1362{
1363}
1364
1365impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O>> Clone
1366 for MultiTaskRuntime<O, P>
1367{
1368 fn clone(&self) -> Self {
1369 MultiTaskRuntime(self.0.clone())
1370 }
1371}
1372
1373impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>> AsyncRuntime<O>
1374 for MultiTaskRuntime<O, P>
1375{
1376 type Pool = P;
1377
1378 fn shared_pool(&self) -> Arc<Self::Pool> {
1380 (self.0).1.clone()
1381 }
1382
1383 fn get_id(&self) -> usize {
1385 (self.0).0
1386 }
1387
1388 fn wait_len(&self) -> usize {
1390 (self.0)
1391 .5
1392 .load(Ordering::Relaxed)
1393 .checked_sub((self.0).6.load(Ordering::Relaxed))
1394 .unwrap_or(0)
1395 }
1396
1397 fn len(&self) -> usize {
1399 (self.0).1.len()
1400 }
1401
1402 fn alloc<R: 'static>(&self) -> TaskId {
1404 TaskId(UnsafeCell::new((TaskHandle::<R>::default().into_raw() as u128) << 64 | self.get_id() as u128 & 0xffffffffffffffff))
1405 }
1406
1407 fn spawn<F>(&self, future: F) -> Result<TaskId>
1409 where
1410 F: Future<Output = O> + Send + 'static,
1411 {
1412 let task_id = self.alloc::<F::Output>();
1413 if let Err(e) = self.spawn_by_id(task_id.clone(), future) {
1414 return Err(e);
1415 }
1416
1417 Ok(task_id)
1418 }
1419
1420 fn spawn_local<F>(&self, future: F) -> Result<TaskId>
1422 where
1423 F: Future<Output=O> + Send + 'static {
1424 let task_id = self.alloc::<F::Output>();
1425 if let Err(e) = self.spawn_local_by_id(task_id.clone(), future) {
1426 return Err(e);
1427 }
1428
1429 Ok(task_id)
1430 }
1431
1432 fn spawn_priority<F>(&self, priority: usize, future: F) -> Result<TaskId>
1434 where
1435 F: Future<Output=O> + Send + 'static {
1436 let task_id = self.alloc::<F::Output>();
1437 if let Err(e) = self.spawn_priority_by_id(task_id.clone(), priority, future) {
1438 return Err(e);
1439 }
1440
1441 Ok(task_id)
1442 }
1443
1444 fn spawn_yield<F>(&self, future: F) -> Result<TaskId>
1446 where
1447 F: Future<Output=O> + Send + 'static {
1448 let task_id = self.alloc::<F::Output>();
1449 if let Err(e) = self.spawn_yield_by_id(task_id.clone(), future) {
1450 return Err(e);
1451 }
1452
1453 Ok(task_id)
1454 }
1455
1456 fn spawn_timing<F>(&self, future: F, time: usize) -> Result<TaskId>
1458 where
1459 F: Future<Output = O> + Send + 'static,
1460 {
1461 let task_id = self.alloc::<F::Output>();
1462 if let Err(e) = self.spawn_timing_by_id(task_id.clone(), future, time) {
1463 return Err(e);
1464 }
1465
1466 Ok(task_id)
1467 }
1468
1469 fn spawn_by_id<F>(&self, task_id: TaskId, future: F) -> Result<()>
1471 where
1472 F: Future<Output=O> + Send + 'static {
1473 let result = {
1474 (self.0).1.push(Arc::new(AsyncTask::new(
1475 task_id,
1476 (self.0).1.clone(),
1477 DEFAULT_MAX_LOW_PRIORITY_BOUNDED,
1478 Some(future.boxed()),
1479 )))
1480 };
1481
1482 let _ = wake_waiting_worker(&(self.0).4);
1483
1484 result
1485 }
1486
1487 fn spawn_local_by_id<F>(&self, task_id: TaskId, future: F) -> Result<()>
1488 where
1489 F: Future<Output=O> + Send + 'static {
1490 let should_wake = PI_ASYNC_THREAD_LOCAL_ID
1491 .try_with(|thread_id| unsafe { ((*thread_id.get()) >> 32) != self.get_id() })
1492 .unwrap_or(true);
1493 let result = (self.0).1.push_local(Arc::new(AsyncTask::new(
1494 task_id,
1495 (self.0).1.clone(),
1496 DEFAULT_HIGH_PRIORITY_BOUNDED,
1497 Some(future.boxed()),
1498 )));
1499
1500 if should_wake {
1501 let _ = wake_waiting_worker(&(self.0).4);
1502 }
1503
1504 result
1505 }
1506
1507 fn spawn_priority_by_id<F>(&self,
1509 task_id: TaskId,
1510 priority: usize,
1511 future: F) -> Result<()>
1512 where
1513 F: Future<Output=O> + Send + 'static {
1514 let result = {
1515 (self.0).1.push_priority(priority, Arc::new(AsyncTask::new(
1516 task_id,
1517 (self.0).1.clone(),
1518 priority,
1519 Some(future.boxed()),
1520 )))
1521 };
1522
1523 let _ = wake_waiting_worker(&(self.0).4);
1524
1525 result
1526 }
1527
1528 #[inline]
1530 fn spawn_yield_by_id<F>(&self, task_id: TaskId, future: F) -> Result<()>
1531 where
1532 F: Future<Output=O> + Send + 'static {
1533 self.spawn_priority_by_id(task_id,
1534 DEFAULT_HIGH_PRIORITY_BOUNDED,
1535 future)
1536 }
1537
1538 fn spawn_timing_by_id<F>(&self,
1540 task_id: TaskId,
1541 future: F,
1542 time: usize) -> Result<()>
1543 where
1544 F: Future<Output=O> + Send + 'static {
1545 let rt = self.clone();
1546 self.spawn_by_id(task_id, async move {
1547 if let Some(timers) = &(rt.0).2 {
1548 let id = (rt.0).1.get_thread_id() & 0xffffffff;
1550 let (_, timer) = &timers[id];
1551 timer.set_timer(
1552 AsyncTimingTask::WaitRun(Arc::new(AsyncTask::new(
1553 rt.alloc::<F::Output>(),
1554 (rt.0).1.clone(),
1555 DEFAULT_MAX_HIGH_PRIORITY_BOUNDED,
1556 Some(future.boxed()),
1557 ))),
1558 time,
1559 );
1560
1561 (rt.0).5.fetch_add(1, Ordering::Relaxed);
1562 }
1563
1564 Default::default()
1565 })
1566 }
1567
1568 fn pending<Output: 'static>(&self, task_id: &TaskId, waker: Waker) -> Poll<Output> {
1570 task_id.set_waker::<Output>(waker);
1571 Poll::Pending
1572 }
1573
1574 fn wakeup<Output: 'static>(&self, task_id: &TaskId) {
1576 task_id.wakeup::<Output>();
1577 }
1578
1579 fn wait<V: Send + 'static>(&self) -> AsyncWait<V> {
1581 AsyncWait(self.wait_any(2))
1582 }
1583
1584 fn wait_any<V: Send + 'static>(&self, capacity: usize) -> AsyncWaitAny<V> {
1586 let (producor, consumer) = async_bounded(capacity);
1587
1588 AsyncWaitAny {
1589 capacity,
1590 producor,
1591 consumer,
1592 }
1593 }
1594
1595 fn wait_any_callback<V: Send + 'static>(&self, capacity: usize) -> AsyncWaitAnyCallback<V> {
1597 let (producor, consumer) = async_bounded(capacity);
1598
1599 AsyncWaitAnyCallback {
1600 capacity,
1601 producor,
1602 consumer,
1603 }
1604 }
1605
1606 fn map_reduce<V: Send + 'static>(&self, capacity: usize) -> AsyncMapReduce<V> {
1608 let (producor, consumer) = async_bounded(capacity);
1609
1610 AsyncMapReduce {
1611 count: 0,
1612 capacity,
1613 producor,
1614 consumer,
1615 }
1616 }
1617
1618 fn timeout(&self, timeout: usize) -> BoxFuture<'static, ()> {
1620 let rt = self.clone();
1621
1622 if let Some(timers) = &(self.0).2 {
1623 match PI_ASYNC_THREAD_LOCAL_ID.try_with(move |thread_id| {
1625 let thread_id = unsafe { *thread_id.get() };
1627 let index = thread_id & 0xffffffff;
1628 if index > timers.len() {
1629 TimerTaskProducor::Foreign(timers[(self.0).3.load(Ordering::Relaxed) % timers.len()].0.clone())
1631 } else {
1632 TimerTaskProducor::Local(timers[index].1.clone())
1633 }
1634 }) {
1635 Err(_) => {
1636 panic!("Multi thread runtime timeout failed, reason: local thread id not match")
1637 }
1638 Ok(producor) => match producor {
1639 TimerTaskProducor::Local(timer) => {
1640 LocalAsyncWaitTimeout::new(rt, timer, timeout).boxed()
1641 },
1642 TimerTaskProducor::Foreign(producor) => {
1643 AsyncWaitTimeout::new(rt, producor, timeout).boxed()
1644 },
1645 },
1646 }
1647 } else {
1648 async move {
1650 thread::sleep(Duration::from_millis(timeout as u64));
1651 }
1652 .boxed()
1653 }
1654 }
1655
1656 fn yield_now(&self) -> BoxFuture<'static, ()> {
1658 async move {
1659 YieldNow(false).await;
1660 }.boxed()
1661 }
1662
1663 fn pipeline<S, SO, F, FO>(&self, input: S, mut filter: F) -> BoxStream<'static, FO>
1665 where
1666 S: Stream<Item = SO> + Send + 'static,
1667 SO: Send + 'static,
1668 F: FnMut(SO) -> AsyncPipelineResult<FO> + Send + 'static,
1669 FO: Send + 'static,
1670 {
1671 let output = stream! {
1672 for await value in input {
1673 match filter(value) {
1674 AsyncPipelineResult::Disconnect => {
1675 break;
1677 },
1678 AsyncPipelineResult::Filtered(result) => {
1679 yield result;
1680 },
1681 }
1682 }
1683 };
1684
1685 output.boxed()
1686 }
1687
1688 fn close(&self) -> bool {
1690 false
1691 }
1692}
1693
1694impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>> AsyncRuntimeExt<O>
1695 for MultiTaskRuntime<O, P>
1696{
1697 fn spawn_with_context<F, C>(&self, task_id: TaskId, future: F, context: C) -> Result<()>
1698 where
1699 F: Future<Output = O> + Send + 'static,
1700 C: 'static,
1701 {
1702 let task = Arc::new(AsyncTask::with_context(
1703 task_id,
1704 (self.0).1.clone(),
1705 DEFAULT_MAX_LOW_PRIORITY_BOUNDED,
1706 Some(future.boxed()),
1707 context,
1708 ));
1709 let result = (self.0).1.push(task);
1710
1711 let _ = wake_waiting_worker(&(self.0).4);
1712
1713 result
1714 }
1715
1716 fn spawn_timing_with_context<F, C>(
1717 &self,
1718 task_id: TaskId,
1719 future: F,
1720 context: C,
1721 time: usize,
1722 ) -> Result<()>
1723 where
1724 F: Future<Output = O> + Send + 'static,
1725 C: Send + 'static,
1726 {
1727 let rt = self.clone();
1728 self.spawn_by_id(task_id, async move {
1729 if let Some(timers) = &(rt.0).2 {
1730 let id = (rt.0).1.get_thread_id() & 0xffffffff;
1732 let (_, timer) = &timers[id];
1733 timer.set_timer(
1734 AsyncTimingTask::WaitRun(Arc::new(AsyncTask::with_context(
1735 rt.alloc::<F::Output>(),
1736 (rt.0).1.clone(),
1737 DEFAULT_MAX_HIGH_PRIORITY_BOUNDED,
1738 Some(future.boxed()),
1739 context,
1740 ))),
1741 time,
1742 );
1743
1744 (rt.0).5.fetch_add(1, Ordering::Relaxed);
1745 }
1746
1747 Default::default()
1748 })
1749 }
1750
1751 fn block_on<F>(&self, future: F) -> Result<F::Output>
1752 where
1753 F: Future + Send + 'static,
1754 <F as Future>::Output: Default + Send + 'static,
1755 {
1756 if let Some(local_rt) = local_async_runtime::<F::Output>() {
1758 if local_rt.get_id() == self.get_id() {
1760 return Err(Error::new(
1762 ErrorKind::WouldBlock,
1763 format!("Block on failed, reason: would block"),
1764 ));
1765 }
1766 }
1767
1768 let (sender, receiver) = bounded(1);
1769 if let Err(e) = self.spawn(async move {
1770 let r = future.await;
1772 sender.send(r);
1773
1774 Default::default()
1775 }) {
1776 return Err(Error::new(
1777 ErrorKind::Other,
1778 format!("Block on failed, reason: {:?}", e),
1779 ));
1780 }
1781
1782 match receiver.recv() {
1784 Err(e) => Err(Error::new(
1785 ErrorKind::Other,
1786 format!("Block on failed, reason: {:?}", e),
1787 )),
1788 Ok(result) => Ok(result),
1789 }
1790 }
1791}
1792
1793impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>>
1794 MultiTaskRuntime<O, P>
1795{
1796 pub fn idler_len(&self) -> usize {
1798 (self.0).1.idler_len()
1799 }
1800
1801 pub fn worker_len(&self) -> usize {
1803 (self.0).1.worker_len()
1804 }
1805
1806 pub fn buffer_len(&self) -> usize {
1808 (self.0).1.buffer_len()
1809 }
1810
1811 pub fn to_local_runtime(&self) -> LocalAsyncRuntime<O> {
1813 LocalAsyncRuntime {
1814 inner: self.as_raw(),
1815 get_id_func: MultiTaskRuntime::<O, P>::get_id_raw,
1816 spawn_func: MultiTaskRuntime::<O, P>::spawn_raw,
1817 spawn_local_func: MultiTaskRuntime::<O, P>::spawn_local_raw,
1818 spawn_timing_func: MultiTaskRuntime::<O, P>::spawn_timing_raw,
1819 timeout_func: MultiTaskRuntime::<O, P>::timeout_raw,
1820 }
1821 }
1822
1823 #[inline]
1825 pub(crate) fn as_raw(&self) -> *const () {
1826 Arc::into_raw(self.0.clone()) as *const ()
1827 }
1828
1829 #[inline]
1831 pub(crate) fn from_raw(raw: *const ()) -> Self {
1832 let inner = unsafe {
1833 Arc::from_raw(
1834 raw as *const (
1835 usize,
1836 Arc<P>,
1837 Option<
1838 Vec<(
1839 Sender<(usize, AsyncTimingTask<P, O>)>,
1840 Arc<AsyncTaskTimerByNotCancel<P, O>>,
1841 )>,
1842 >,
1843 AtomicUsize,
1844 Arc<ArrayQueue<Arc<(AtomicBool, Mutex<()>, Condvar)>>>,
1845 AtomicUsize,
1846 AtomicUsize,
1847 ),
1848 )
1849 };
1850 MultiTaskRuntime(inner)
1851 }
1852
1853 pub(crate) fn get_id_raw(raw: *const ()) -> usize {
1855 let rt = MultiTaskRuntime::<O, P>::from_raw(raw);
1856 let id = rt.get_id();
1857 Arc::into_raw(rt.0); id
1859 }
1860
1861 pub(crate) fn spawn_raw<F>(raw: *const (), future: F) -> Result<()>
1863 where
1864 F: Future<Output = O> + Send + 'static,
1865 {
1866 let rt = MultiTaskRuntime::<O, P>::from_raw(raw);
1867 let result = rt.spawn_by_id(rt.alloc::<F::Output>(), future);
1868 Arc::into_raw(rt.0); result
1870 }
1871
1872 pub(crate) fn spawn_local_raw<F>(raw: *const (), future: F) -> Result<()>
1874 where
1875 F: Future<Output = O> + Send + 'static,
1876 {
1877 let rt = MultiTaskRuntime::<O, P>::from_raw(raw);
1878 let result = rt.spawn_local_by_id(rt.alloc::<F::Output>(), future);
1879 Arc::into_raw(rt.0); result
1881 }
1882
1883 pub(crate) fn spawn_timing_raw(
1885 raw: *const (),
1886 future: BoxFuture<'static, O>,
1887 timeout: usize,
1888 ) -> Result<()> {
1889 let rt = MultiTaskRuntime::<O, P>::from_raw(raw);
1890 let result = rt.spawn_timing_by_id(rt.alloc::<O>(), future, timeout);
1891 Arc::into_raw(rt.0); result
1893 }
1894
1895 pub(crate) fn timeout_raw(raw: *const (), timeout: usize) -> BoxFuture<'static, ()> {
1897 let rt = MultiTaskRuntime::<O, P>::from_raw(raw);
1898 let boxed = rt.timeout(timeout);
1899 Arc::into_raw(rt.0); boxed
1901 }
1902}
1903
1904pub struct MultiTaskRuntimeBuilder<
1908 O: Default + 'static = (),
1909 P: AsyncTaskPoolExt<O> + AsyncTaskPool<O> = StealableTaskPool<O>,
1910> {
1911 pool: P, prefix: String, init: usize, min: usize, max: usize, stack_size: usize, timeout: u64, interval: Option<usize>, marker: PhantomData<O>,
1920}
1921
1922unsafe impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O>> Send
1923 for MultiTaskRuntimeBuilder<O, P>
1924{
1925}
1926unsafe impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O>> Sync
1927 for MultiTaskRuntimeBuilder<O, P>
1928{
1929}
1930
1931impl<O: Default + 'static> Default for MultiTaskRuntimeBuilder<O> {
1932 fn default() -> Self {
1934 #[cfg(not(target_arch = "wasm32"))]
1935 let core_len = num_cpus::get(); #[cfg(target_arch = "wasm32")]
1937 let core_len = 1; let pool = StealableTaskPool::with(core_len,
1939 65535,
1940 [1, 1],
1941 3000);
1942 MultiTaskRuntimeBuilder::new(pool)
1943 .thread_stack_size(2 * 1024 * 1024)
1944 .set_timer_interval(1)
1945 }
1946}
1947
1948impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>>
1949 MultiTaskRuntimeBuilder<O, P>
1950{
1951 pub fn new(mut pool: P) -> Self {
1953 #[cfg(not(target_arch = "wasm32"))]
1954 let core_len = num_cpus::get(); #[cfg(target_arch = "wasm32")]
1956 let core_len = 1; MultiTaskRuntimeBuilder {
1959 pool,
1960 prefix: DEFAULT_WORKER_THREAD_PREFIX.to_string(),
1961 init: core_len,
1962 min: core_len,
1963 max: core_len,
1964 stack_size: DEFAULT_THREAD_STACK_SIZE,
1965 timeout: DEFAULT_WORKER_THREAD_SLEEP_TIME,
1966 interval: None,
1967 marker: PhantomData,
1968 }
1969 }
1970
1971 pub fn thread_prefix(mut self, prefix: &str) -> Self {
1973 self.prefix = prefix.to_string();
1974 self
1975 }
1976
1977 pub fn thread_stack_size(mut self, stack_size: usize) -> Self {
1979 self.stack_size = stack_size;
1980 self
1981 }
1982
1983 pub fn init_worker_size(mut self, mut init: usize) -> Self {
1985 if init == 0 {
1986 init = DEFAULT_INIT_WORKER_SIZE;
1988 }
1989
1990 self.init = init;
1991 self
1992 }
1993
1994 pub fn set_worker_limit(mut self, mut min: usize, mut max: usize) -> Self {
1996 if self.init > max {
1997 max = self.init;
1999 }
2000
2001 if min == 0 || min > max {
2002 min = max;
2004 }
2005
2006 self.min = min;
2007 self.max = max;
2008 self
2009 }
2010
2011 pub fn set_timeout(mut self, timeout: u64) -> Self {
2013 self.timeout = timeout;
2014 self
2015 }
2016
2017 pub fn set_timer_interval(mut self, interval: usize) -> Self {
2019 self.interval = Some(interval);
2020 self
2021 }
2022
2023 pub fn build(mut self) -> MultiTaskRuntime<O, P> {
2057 let pool_worker_len = self.pool.worker_len();
2058 if pool_worker_len == 0 {
2059 panic!("Build multi thread runtime failed, reason: worker pool is empty");
2060 }
2061 if self.init > pool_worker_len {
2062 self.init = pool_worker_len;
2063 }
2064 if self.max > pool_worker_len {
2065 self.max = pool_worker_len;
2066 }
2067 if self.min > self.max {
2068 self.min = self.max;
2069 }
2070
2071 let interval = self.interval;
2073 let mut timers = if let Some(_) = interval {
2074 Some(Vec::with_capacity(self.max))
2075 } else {
2076 None
2077 };
2078 for _ in 0..self.max {
2079 if let Some(vec) = &mut timers {
2081 let timer = AsyncTaskTimerByNotCancel::new();
2082 let producor = timer.producor.clone();
2083 let timer = Arc::new(timer);
2084 vec.push((producor, timer));
2085 };
2086 }
2087
2088 let rt_uid = alloc_rt_uid();
2090 let waits = Arc::new(ArrayQueue::new(self.max));
2091 let mut pool = self.pool;
2092 pool.set_waits(waits.clone()); let pool = Arc::new(pool);
2094 let runtime = MultiTaskRuntime(Arc::new((
2095 rt_uid,
2096 pool,
2097 timers,
2098 AtomicUsize::new(0),
2099 waits,
2100 AtomicUsize::new(0),
2101 AtomicUsize::new(0),
2102 )));
2103
2104 let mut builders = Vec::with_capacity(self.init);
2106 for index in 0..self.init {
2107 let builder = Builder::new()
2108 .name(self.prefix.clone() + "-" + index.to_string().as_str())
2109 .stack_size(self.stack_size);
2110 builders.push(builder);
2111 }
2112
2113 let min = self.min;
2115 for index in 0..builders.len() {
2116 let builder = builders.remove(0);
2117 let runtime = runtime.clone();
2118 let timeout = self.timeout;
2119 let timer = if let Some(timers) = &(runtime.0).2 {
2120 let (_, timer) = &timers[index];
2121 Some(timer.clone())
2122 } else {
2123 None
2124 };
2125
2126 spawn_worker_thread(builder, index, runtime, min, timeout, interval, timer);
2127 }
2128
2129 runtime
2130 }
2131}
2132
2133fn bind_multi_thread_worker_context<P>(thread_id: usize, pool: &P) {
2152 if let Err(e) = PI_ASYNC_THREAD_LOCAL_ID.try_with(|local_thread_id| unsafe {
2153 *local_thread_id.get() = thread_id;
2156 }) {
2157 panic!(
2158 "Multi thread runtime startup failed, thread id: {:?}, reason: {:?}",
2159 thread_id & MULTI_THREAD_WORKER_ID_MASK,
2160 e
2161 );
2162 }
2163
2164 let context = MultiThreadWorkerContext {
2165 thread_id,
2166 pool: pool as *const P as *const (),
2167 };
2168 if let Err(e) = PI_ASYNC_MULTI_THREAD_WORKER_CONTEXT.try_with(|current| {
2169 current.set(context);
2170 }) {
2171 panic!(
2172 "Bind multi-thread worker pool failed, thread id: {:?}, reason: {:?}",
2173 thread_id & MULTI_THREAD_WORKER_ID_MASK,
2174 e
2175 );
2176 }
2177}
2178
2179fn spawn_worker_thread<
2190 O: Default + 'static,
2191 P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>,
2192>(
2193 builder: Builder,
2194 index: usize,
2195 runtime: MultiTaskRuntime<O, P>,
2196 min: usize,
2197 timeout: u64,
2198 interval: Option<usize>,
2199 timer: Option<Arc<AsyncTaskTimerByNotCancel<P, O>>>,
2200) {
2201 if let Some(timer) = timer {
2202 let rt_uid = runtime.get_id();
2204 let _ = builder.spawn(move || {
2205 let thread_id = rt_uid << 32 | index & MULTI_THREAD_WORKER_ID_MASK;
2207 bind_multi_thread_worker_context(thread_id, (runtime.0).1.as_ref());
2208
2209 let runtime_copy = runtime.clone();
2211 match PI_ASYNC_LOCAL_THREAD_ASYNC_RUNTIME.try_with(move |rt| {
2212 let raw = Arc::into_raw(Arc::new(runtime_copy.to_local_runtime()))
2213 as *mut LocalAsyncRuntime<O> as *mut ();
2214 rt.store(raw, Ordering::Relaxed);
2215 }) {
2216 Err(e) => {
2217 panic!("Bind multi runtime to local thread failed, reason: {:?}", e);
2218 }
2219 Ok(_) => (),
2220 }
2221
2222 timer_work_loop(
2224 runtime,
2225 index,
2226 min,
2227 timeout,
2228 interval.unwrap() as u64,
2229 timer,
2230 );
2231 });
2232 } else {
2233 let rt_uid = runtime.get_id();
2235 let _ = builder.spawn(move || {
2236 let thread_id = rt_uid << 32 | index & MULTI_THREAD_WORKER_ID_MASK;
2238 bind_multi_thread_worker_context(thread_id, (runtime.0).1.as_ref());
2239
2240 let runtime_copy = runtime.clone();
2242 match PI_ASYNC_LOCAL_THREAD_ASYNC_RUNTIME.try_with(move |rt| {
2243 let raw = Arc::into_raw(Arc::new(runtime_copy.to_local_runtime()))
2244 as *mut LocalAsyncRuntime<O> as *mut ();
2245 rt.store(raw, Ordering::Relaxed);
2246 }) {
2247 Err(e) => {
2248 panic!("Bind multi runtime to local thread failed, reason: {:?}", e);
2249 }
2250 Ok(_) => (),
2251 }
2252
2253 work_loop(runtime, index, min, timeout);
2255 });
2256 }
2257}
2258
2259enum WorkerWaitResult<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>> {
2275 TimedOut,
2276 NotSlept,
2277 Task(Arc<AsyncTask<P, O>>),
2278}
2279
2280#[inline]
2347fn worker_wait_for_task<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>>(
2348 runtime: &MultiTaskRuntime<O, P>,
2349 worker_waker: &Arc<(AtomicBool, Mutex<()>, Condvar)>,
2350 sleep_timeout: u64,
2351) -> WorkerWaitResult<O, P> {
2352 let (is_sleep, lock, condvar) = &**worker_waker;
2353
2354 loop {
2355 let _locked = lock.lock();
2356 if is_sleep.load(Ordering::Acquire) {
2357 break;
2358 }
2359
2360 if register_waiting_worker(&(runtime.0).4, worker_waker) {
2361 is_sleep.store(true, Ordering::Release);
2362 break;
2363 }
2364
2365 drop(_locked);
2366 if prune_stale_waiting_workers(&(runtime.0).4) == 0 {
2367 return WorkerWaitResult::NotSlept;
2368 }
2369 }
2370
2371 if let Some(task) = (runtime.0).1.try_pop() {
2372 is_sleep.store(false, Ordering::Release);
2373 return WorkerWaitResult::Task(task);
2374 }
2375
2376 if runtime.len() > 0 {
2377 is_sleep.store(false, Ordering::Release);
2378 return WorkerWaitResult::NotSlept;
2379 }
2380
2381 let mut locked = lock.lock();
2382 if !is_sleep.load(Ordering::Acquire) {
2383 return WorkerWaitResult::NotSlept;
2384 }
2385
2386 let timed_out = condvar
2387 .wait_for(&mut locked, Duration::from_millis(sleep_timeout))
2388 .timed_out();
2389 is_sleep.store(false, Ordering::Release);
2390
2391 if timed_out {
2392 WorkerWaitResult::TimedOut
2393 } else {
2394 WorkerWaitResult::NotSlept
2395 }
2396}
2397
2398fn timer_work_loop<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>>(
2400 runtime: MultiTaskRuntime<O, P>,
2401 index: usize,
2402 min: usize,
2403 sleep_timeout: u64,
2404 timer_interval: u64,
2405 timer: Arc<AsyncTaskTimerByNotCancel<P, O>>,
2406) {
2407 let pool = (runtime.0).1.clone();
2409 let worker_waker = pool.clone_thread_waker().unwrap();
2410
2411 let mut sleep_count = 0; let clock = Clock::new();
2413 loop {
2414 let timer_run_millis = clock.recent(); let mut pop_len = 0;
2417 (runtime.0)
2418 .5
2419 .fetch_add(timer.consume(),
2420 Ordering::Relaxed);
2421 loop {
2422 let current_time = timer.is_require_pop();
2423 if let Some(current_time) = current_time {
2424 loop {
2426 let timed_out = timer.pop(current_time);
2427 if let Some(timing_task) = timed_out {
2428 match timing_task {
2429 AsyncTimingTask::Pended(expired) => {
2430 runtime.wakeup::<O>(&expired);
2432 }
2433 AsyncTimingTask::WaitRun(expired) => {
2434 (runtime.0)
2436 .1
2437 .push_priority(DEFAULT_MAX_HIGH_PRIORITY_BOUNDED,
2438 expired);
2439 if let Some(task) = pool.try_pop() {
2440 sleep_count = 0; run_task(&runtime, task);
2442 }
2443 }
2444 AsyncTimingTask::TimeoutWake(waiter) => {
2445 waiter.fire();
2447 }
2448 }
2449 pop_len += 1;
2450
2451 if let Some(task) = pool.try_pop() {
2452 sleep_count = 0; run_task(&runtime, task);
2455 }
2456 } else {
2457 break;
2459 }
2460 }
2461 } else {
2462 break;
2464 }
2465 }
2466 (runtime.0)
2467 .6
2468 .fetch_add(pop_len,
2469 Ordering::Relaxed);
2470
2471 match pool.try_pop() {
2473 None => {
2474 if runtime.len() > 0 {
2475 continue;
2477 }
2478
2479 let diff_time = clock
2481 .recent()
2482 .duration_since(timer_run_millis)
2483 .as_millis() as u64; let real_timeout = if timer.len() == 0 {
2485 sleep_timeout
2487 } else {
2488 if diff_time >= timer_interval {
2490 continue;
2492 } else {
2493 timer_interval - diff_time
2495 }
2496 };
2497
2498 match worker_wait_for_task(&runtime, &worker_waker, real_timeout) {
2500 WorkerWaitResult::TimedOut => {
2501 sleep_count += 1;
2503 },
2504 WorkerWaitResult::Task(task) => {
2505 sleep_count = 0; run_task(&runtime, task);
2507 },
2508 WorkerWaitResult::NotSlept => (),
2509 }
2510 }
2511 Some(task) => {
2512 sleep_count = 0; run_task(&runtime, task);
2515 }
2516 }
2517 }
2518
2519 (runtime.0).1.close_worker();
2521 warn!(
2522 "Worker of runtime closed, runtime: {}, worker: {}, thread: {:?}",
2523 runtime.get_id(),
2524 index,
2525 thread::current()
2526 );
2527}
2528
2529fn work_loop<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>>(
2531 runtime: MultiTaskRuntime<O, P>,
2532 index: usize,
2533 min: usize,
2534 sleep_timeout: u64,
2535) {
2536 let pool = (runtime.0).1.clone();
2538 let worker_waker = pool.clone_thread_waker().unwrap();
2539
2540 let mut sleep_count = 0; loop {
2542 match pool.try_pop() {
2543 None => {
2544 if runtime.len() > 0 {
2546 continue;
2548 }
2549
2550 match worker_wait_for_task(&runtime, &worker_waker, sleep_timeout) {
2551 WorkerWaitResult::TimedOut => {
2552 sleep_count += 1;
2554 },
2555 WorkerWaitResult::Task(task) => {
2556 sleep_count = 0; run_task(&runtime, task);
2558 },
2559 WorkerWaitResult::NotSlept => (),
2560 }
2561 }
2562 Some(task) => {
2563 sleep_count = 0; run_task(&runtime, task);
2566 }
2567 }
2568 }
2569
2570 (runtime.0).1.close_worker();
2572 warn!(
2573 "Worker of runtime closed, runtime: {}, worker: {}, thread: {:?}",
2574 runtime.get_id(),
2575 index,
2576 thread::current()
2577 );
2578}
2579
2580#[inline]
2597fn run_task<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>>(
2598 runtime: &MultiTaskRuntime<O, P>,
2599 task: Arc<AsyncTask<P, O>>,
2600) {
2601 match task.try_begin_runtime_poll() {
2602 AsyncTaskPollClaim::Discard => return,
2603 AsyncTaskPollClaim::Legacy => {
2604 let waker = waker_ref(&task);
2605 let mut context = Context::from_waker(&*waker);
2606 if let Some(mut future) = task.get_inner() {
2607 if let Poll::Pending = future.as_mut().poll(&mut context) {
2608 task.set_inner(Some(future));
2609 }
2610 } else {
2611 (runtime.0).1.push(task);
2613 }
2614 return;
2615 },
2616 AsyncTaskPollClaim::Managed => (),
2617 }
2618
2619 let guard = AsyncTaskPollGuard::new(&task);
2622 let waker = waker_ref(&task);
2623 let mut context = Context::from_waker(&*waker);
2624 let mut future = match task.take_inner_for_runtime_poll() {
2625 Some(future) => future,
2626 None => {
2627 guard.finish_ready();
2630 return;
2631 },
2632 };
2633
2634 match future.as_mut().poll(&mut context) {
2635 Poll::Pending => {
2636 task.restore_inner_after_runtime_poll(future);
2637 if guard.finish_pending() {
2638 requeue_runtime_task((runtime.0).1.as_ref(), &task);
2639 }
2640 },
2641 Poll::Ready(_) => guard.finish_ready(),
2642 }
2643}
2644
2645enum TimerTaskProducor<
2647 O: Default + 'static = (),
2648 P: AsyncTaskPoolExt<O> + AsyncTaskPool<O> = StealableTaskPool<O>,
2649> {
2650 Local(Arc<AsyncTaskTimerByNotCancel<P, O>>), Foreign(Sender<(usize, AsyncTimingTask<P, O>)>), }