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, 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};
80
81#[cfg(not(target_arch = "wasm32"))]
85const DEFAULT_INIT_WORKER_SIZE: usize = 2;
86#[cfg(target_arch = "wasm32")]
87const DEFAULT_INIT_WORKER_SIZE: usize = 1;
88
89const DEFAULT_WORKER_THREAD_PREFIX: &str = "Default-Multi-RT";
93
94const DEFAULT_THREAD_STACK_SIZE: usize = 1024 * 1024;
98
99const DEFAULT_WORKER_THREAD_SLEEP_TIME: u64 = 10;
103
104const DEFAULT_RUNTIME_SLEEP_TIME: u64 = 1000;
108
109const DEFAULT_MAX_WEIGHT: u8 = 254;
113
114const DEFAULT_MIN_WEIGHT: u8 = 1;
118
119const MULTI_THREAD_WORKER_ID_MASK: usize = 0xffffffff;
123
124#[derive(Clone, Copy)]
147struct MultiThreadWorkerContext {
148 thread_id: usize,
149 pool: *const (),
150}
151
152impl MultiThreadWorkerContext {
153 const UNBOUND: Self = MultiThreadWorkerContext {
154 thread_id: usize::MAX,
155 pool: std::ptr::null(),
156 };
157
158 #[inline]
163 const fn runtime_id(self) -> usize {
164 self.thread_id >> 32
165 }
166
167 #[inline]
183 fn owner_worker_id<P>(self, pool: &P) -> usize {
184 let expected = pool as *const P as *const ();
185 if self.pool != expected {
186 panic!(
187 "Multi-thread task pool owner mismatch: owner-only worker state requires the worker bound to this pool"
188 );
189 }
190
191 self.thread_id & MULTI_THREAD_WORKER_ID_MASK
192 }
193}
194
195thread_local! {
196 static PI_ASYNC_MULTI_THREAD_WORKER_CONTEXT: Cell<MultiThreadWorkerContext>
201 = Cell::new(MultiThreadWorkerContext::UNBOUND);
202}
203
204#[inline]
210fn current_multi_thread_worker_context() -> MultiThreadWorkerContext {
211 match PI_ASYNC_MULTI_THREAD_WORKER_CONTEXT.try_with(|context| context.get()) {
212 Ok(context) => context,
213 Err(e) => {
214 panic!(
215 "Get multi-thread worker context failed, thread: {:?}, reason: {:?}",
216 thread::current(),
217 e
218 );
219 },
220 }
221}
222
223struct ComputationalTaskQueue<O: Default + 'static> {
227 stack: Worker<Arc<AsyncTask<ComputationalTaskPool<O>, O>>>, queue: SegQueue<Arc<AsyncTask<ComputationalTaskPool<O>, O>>>, thread_waker: Arc<(AtomicBool, Mutex<()>, Condvar)>, }
231
232impl<O: Default + 'static> ComputationalTaskQueue<O> {
233 pub fn new(thread_waker: Arc<(AtomicBool, Mutex<()>, Condvar)>) -> Self {
235 let stack = Worker::new_lifo();
236 let queue = SegQueue::new();
237
238 ComputationalTaskQueue {
239 stack,
240 queue,
241 thread_waker,
242 }
243 }
244
245 pub fn len(&self) -> usize {
247 self.stack.len() + self.queue.len()
248 }
249}
250
251pub struct ComputationalTaskPool<O: Default + 'static> {
286 workers: Vec<ComputationalTaskQueue<O>>, waits: Option<Arc<ArrayQueue<Arc<(AtomicBool, Mutex<()>, Condvar)>>>>, consume_count: Arc<AtomicUsize>, produce_count: Arc<AtomicUsize>, }
291
292unsafe impl<O: Default + 'static> Send for ComputationalTaskPool<O> {}
298unsafe impl<O: Default + 'static> Sync for ComputationalTaskPool<O> {}
301
302impl<O: Default + 'static> Default for ComputationalTaskPool<O> {
303 fn default() -> Self {
304 #[cfg(not(target_arch = "wasm32"))]
305 let core_len = num_cpus::get(); #[cfg(target_arch = "wasm32")]
307 let core_len = 1; ComputationalTaskPool::new(core_len)
309 }
310}
311
312impl<O: Default + 'static> AsyncTaskPool<O> for ComputationalTaskPool<O> {
313 type Pool = ComputationalTaskPool<O>;
314
315 #[inline]
321 fn get_thread_id(&self) -> usize {
322 match PI_ASYNC_THREAD_LOCAL_ID.try_with(move |thread_id| unsafe {
323 *thread_id.get()
325 }) {
326 Err(e) => {
327 panic!(
328 "Get thread id failed, thread: {:?}, reason: {:?}",
329 thread::current(),
330 e
331 );
332 }
333 Ok(id) => id,
334 }
335 }
336
337 #[inline]
338 fn len(&self) -> usize {
339 if let Some(len) = self
340 .produce_count
341 .load(Ordering::Relaxed)
342 .checked_sub(self.consume_count.load(Ordering::Relaxed))
343 {
344 len
345 } else {
346 0
347 }
348 }
349
350 #[inline]
351 fn push(&self, task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
352 let index = self.produce_count.fetch_add(1, Ordering::Relaxed) % self.workers.len();
353 self.workers[index].queue.push(task);
354 Ok(())
355 }
356
357 #[inline]
364 fn push_local(&self, task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
365 let context = current_multi_thread_worker_context();
366 let rt_uid = task.owner();
367 if context.runtime_id() == rt_uid {
368 let worker = &self.workers[context.owner_worker_id(self)];
370 worker.queue.push(task);
371
372 self.produce_count.fetch_add(1, Ordering::Relaxed);
373 Ok(())
374 } else {
375 self.push(task)
377 }
378 }
379
380 #[inline]
387 fn push_priority(&self,
388 priority: usize,
389 task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
390 if priority >= DEFAULT_MAX_HIGH_PRIORITY_BOUNDED {
391 let context = current_multi_thread_worker_context();
393 let rt_uid = task.owner();
394 if context.runtime_id() == rt_uid {
395 let worker = &self.workers[context.owner_worker_id(self)];
396 worker.stack.push(task);
397
398 self.produce_count.fetch_add(1, Ordering::Relaxed);
399 Ok(())
400 } else {
401 self.push(task)
402 }
403 } else if priority >= DEFAULT_HIGH_PRIORITY_BOUNDED {
404 self.push_local(task)
406 } else {
407 self.push(task)
409 }
410 }
411
412 #[inline]
413 fn push_keep(&self, task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
414 self.push_priority(DEFAULT_HIGH_PRIORITY_BOUNDED, task)
415 }
416
417 #[inline]
424 fn try_pop(&self) -> Option<Arc<AsyncTask<Self::Pool, O>>> {
425 let id = current_multi_thread_worker_context().owner_worker_id(self);
426 let worker = &self.workers[id];
427 let task = worker.stack.pop();
428 if task.is_some() {
429 self.consume_count.fetch_add(1, Ordering::Relaxed);
431 return task;
432 }
433
434 let task = worker.queue.pop();
435 if task.is_some() {
436 self.consume_count.fetch_add(1, Ordering::Relaxed);
437 }
438
439 task
440 }
441
442 #[inline]
448 fn try_pop_all(&self) -> IntoIter<Arc<AsyncTask<Self::Pool, O>>> {
449 let mut tasks = Vec::with_capacity(self.len());
450 while let Some(task) = self.try_pop() {
451 tasks.push(task);
452 }
453
454 tasks.into_iter()
455 }
456
457 #[inline]
458 fn get_thread_waker(&self) -> Option<&Arc<(AtomicBool, Mutex<()>, Condvar)>> {
459 None
461 }
462}
463
464impl<O: Default + 'static> AsyncTaskPoolExt<O> for ComputationalTaskPool<O> {
465 #[inline]
466 fn set_waits(&mut self, waits: Arc<ArrayQueue<Arc<(AtomicBool, Mutex<()>, Condvar)>>>) {
467 self.waits = Some(waits);
468 }
469
470 #[inline]
471 fn get_waits(&self) -> Option<&Arc<ArrayQueue<Arc<(AtomicBool, Mutex<()>, Condvar)>>>> {
472 self.waits.as_ref()
473 }
474
475 #[inline]
476 fn worker_len(&self) -> usize {
477 self.workers.len()
478 }
479
480 #[inline]
486 fn clone_thread_waker(&self) -> Option<Arc<(AtomicBool, Mutex<()>, Condvar)>> {
487 let worker = &self.workers[current_multi_thread_worker_context().owner_worker_id(self)];
488 Some(worker.thread_waker.clone())
489 }
490}
491
492impl<O: Default + 'static> ComputationalTaskPool<O> {
493 pub fn new(mut size: usize) -> Self {
510 if size < DEFAULT_INIT_WORKER_SIZE {
511 size = DEFAULT_INIT_WORKER_SIZE;
513 }
514
515 let mut workers = Vec::with_capacity(size);
516 for _ in 0..size {
517 let thread_waker = Arc::new((AtomicBool::new(false), Mutex::new(()), Condvar::new()));
518 let worker = ComputationalTaskQueue::new(thread_waker);
519 workers.push(worker);
520 }
521 let consume_count = Arc::new(AtomicUsize::new(0));
522 let produce_count = Arc::new(AtomicUsize::new(0));
523
524 ComputationalTaskPool {
525 workers,
526 waits: None,
527 consume_count,
528 produce_count,
529 }
530 }
531}
532
533struct StealableTaskQueue<O: Default + 'static> {
537 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)>, }
543
544impl<O: Default + 'static> StealableTaskQueue<O> {
545 pub fn new(
548 init_queue_capacity: usize,
549 thread_waker: Arc<(AtomicBool, Mutex<()>, Condvar)>,
550 ) -> (Self,
551 FIFOStealer<Arc<AsyncTask<StealableTaskPool<O>, O>>>,
552 Stealer<Arc<AsyncTask<StealableTaskPool<O>, O>>>) {
553 let stack = UnsafeCell::new(None);
554 let internal = FIFOWorker::new(init_queue_capacity);
555 let external = Worker::new_fifo();
556 let internal_stealer = internal.stealer();
557 let external_stealer = external.stealer();
558 let selector = UnsafeCell::new(IWRRSelector::new([2, 1]));
559
560 (
561 StealableTaskQueue {
562 stack,
563 internal,
564 external,
565 selector,
566 thread_waker,
567 },
568 internal_stealer,
569 external_stealer
570 )
571 }
572
573 pub const fn stack_capacity(&self) -> usize {
575 1
576 }
577
578 pub fn internal_capacity(&self) -> usize {
580 self.internal.capacity()
581 }
582
583 pub fn remaining_internal_capacity(&self) -> usize {
585 self.internal.spare_capacity()
586 }
587
588 #[inline]
595 pub fn stack_len(&self) -> usize {
596 unsafe {
597 if (&*self.stack.get()).is_some() {
600 1
601 } else {
602 0
603 }
604 }
605 }
606
607 pub fn internal_len(&self) -> usize {
609 self
610 .internal_capacity()
611 .checked_sub(self.remaining_internal_capacity())
612 .unwrap_or(0)
613 }
614
615 pub fn external_len(&self) -> usize {
617 self.external.len()
618 }
619}
620
621pub struct StealableTaskPool<O: Default + 'static> {
667 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)>>>>, }
683
684unsafe impl<O: Default + 'static> Send for StealableTaskPool<O> {}
690unsafe impl<O: Default + 'static> Sync for StealableTaskPool<O> {}
693
694impl<O: Default + 'static> Default for StealableTaskPool<O> {
695 fn default() -> Self {
696 StealableTaskPool::new()
697 }
698}
699
700impl<O: Default + 'static> AsyncTaskPool<O> for StealableTaskPool<O> {
701 type Pool = StealableTaskPool<O>;
702
703 #[inline]
709 fn get_thread_id(&self) -> usize {
710 match PI_ASYNC_THREAD_LOCAL_ID.try_with(move |thread_id| unsafe {
711 *thread_id.get()
713 }) {
714 Err(e) => {
715 panic!(
716 "Get thread id failed, thread: {:?}, reason: {:?}",
717 thread::current(),
718 e
719 );
720 }
721 Ok(id) => id,
722 }
723 }
724
725 #[inline]
726 fn len(&self) -> usize {
727 self.internal_produce
728 .load(Ordering::Relaxed)
729 .checked_sub(self.internal_consume.load(Ordering::Relaxed))
730 .unwrap_or(0)
731 +
732 self.external_produce
733 .load(Ordering::Relaxed)
734 .checked_sub(self.external_consume.load(Ordering::Relaxed))
735 .unwrap_or(0)
736 }
737
738 #[inline]
739 fn push(&self, task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
740 self.public.push(task);
741
742 self
743 .external_produce
744 .fetch_add(1, Ordering::Relaxed);
745 Ok(())
746 }
747
748 #[inline]
755 fn push_local(&self, task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
756 let context = current_multi_thread_worker_context();
757 let rt_uid = task.owner();
758 if context.runtime_id() == rt_uid {
759 let worker = &self.workers[context.owner_worker_id(self)];
761 if worker.remaining_internal_capacity() > 0 {
762 let _ = worker.internal.push(task);
764
765 self
766 .internal_produce
767 .fetch_add(1, Ordering::Relaxed);
768 Ok(())
769 } else {
770 self.push(task)
772 }
773 } else {
774 self.push(task)
776 }
777 }
778
779 #[inline]
786 fn push_priority(&self,
787 priority: usize,
788 task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
789 if priority >= DEFAULT_MAX_HIGH_PRIORITY_BOUNDED {
790 let context = current_multi_thread_worker_context();
792 let rt_uid = task.owner();
793 if context.runtime_id() == rt_uid {
794 let worker = &self.workers[context.owner_worker_id(self)];
796 if worker.stack_len() < 1 {
797 unsafe {
799 *worker.stack.get() = Some(task);
802 }
803 } else if worker.remaining_internal_capacity() > 0 {
804 let _ = worker.internal.push(task);
806 } else {
807 return self.push(task);
809 }
810
811 self
812 .internal_produce
813 .fetch_add(1, Ordering::Relaxed);
814 Ok(())
815 } else {
816 self.push(task)
818 }
819 } else if priority >= DEFAULT_HIGH_PRIORITY_BOUNDED {
820 self.push_local(task)
822 } else {
823 self.push(task)
825 }
826 }
827
828 #[inline]
829 fn push_keep(&self, task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
830 self.push_priority(DEFAULT_HIGH_PRIORITY_BOUNDED, task)
831 }
832
833 #[inline]
841 fn try_pop(&self) -> Option<Arc<AsyncTask<Self::Pool, O>>> {
842 let id = current_multi_thread_worker_context().owner_worker_id(self);
843 let worker = &self.workers[id];
844 let task = unsafe {
845 (&mut *worker
848 .stack
849 .get())
850 .take()
851 };
852 if task.is_some() {
853 return task;
855 }
856
857 try_pop_by_weight(self, worker, id)
859 }
860
861 #[inline]
867 fn try_pop_all(&self) -> IntoIter<Arc<AsyncTask<Self::Pool, O>>> {
868 let mut tasks = Vec::with_capacity(self.len());
869 while let Some(task) = self.try_pop() {
870 tasks.push(task);
871 }
872
873 tasks.into_iter()
874 }
875
876 #[inline]
877 fn get_thread_waker(&self) -> Option<&Arc<(AtomicBool, Mutex<()>, Condvar)>> {
878 None
880 }
881}
882
883const fn get_msb(n: usize) -> usize {
885 usize::BITS as usize - n.leading_zeros() as usize
886}
887
888fn try_pop_by_weight<O: Default + 'static>(pool: &StealableTaskPool<O>,
913 local_worker: &StealableTaskQueue<O>,
914 local_worker_id: usize)
915 -> Option<Arc<AsyncTask<StealableTaskPool<O>, O>>> {
916 unsafe {
917 let duration = pool
921 .clock
922 .recent()
923 .duration_since(pool.last_time.load())
924 .as_millis() as usize;
925 if duration >= pool.interval {
926 let new_external_traffic_statistics = pool
928 .external_produce
929 .load(Ordering::Relaxed);
930 let new_internal_traffic_statistics = pool
931 .internal_produce
932 .load(Ordering::Relaxed);
933
934 let external_delta = if new_external_traffic_statistics == 0 {
936 1
938 } else {
939 new_external_traffic_statistics
941 .checked_sub(pool
942 .external_traffic_statistics
943 .load(Ordering::Relaxed))
944 .unwrap_or(1)
945 };
946 pool
947 .external_traffic_statistics
948 .store(new_external_traffic_statistics, Ordering::Relaxed); let internal_delta = if new_internal_traffic_statistics == 0 {
950 1
952 } else {
953 new_internal_traffic_statistics
955 .checked_sub(pool
956 .internal_traffic_statistics
957 .load(Ordering::Relaxed))
958 .unwrap_or(1)
959 };
960 pool
961 .internal_traffic_statistics
962 .store(new_internal_traffic_statistics, Ordering::Relaxed); let selector = &mut *local_worker.selector.get();
966 if external_delta > internal_delta {
967 let msb = get_msb(internal_delta);
969 let internal_weight
970 = (internal_delta >> msb.checked_sub(2).unwrap_or(0)).max(1);
971 let external_weight
972 = ((external_delta >> msb).min(DEFAULT_MAX_WEIGHT as usize)).max(1);
973
974 selector.change_weight(0, external_weight as u8);
975 selector.change_weight(1, internal_weight as u8);
976 } else if external_delta < internal_delta {
977 let msb = get_msb(external_delta);
979 let external_weight
980 = (external_delta >> msb.checked_sub(2).unwrap_or(0)).max(1);
981 let internal_weight
982 = ((internal_delta >> msb).min(DEFAULT_MAX_WEIGHT as usize)).max(1);
983
984 selector.change_weight(0, external_weight as u8);
985 selector.change_weight(1, internal_weight as u8);
986 } else {
987 selector.change_weight(0, 1);
989 selector.change_weight(1, 1);
990 }
991
992 pool.last_time.store(pool.clock.recent()); }
994
995 match (&mut *local_worker.selector.get()).select() {
997 0 => {
998 let task = try_pop_external(pool, local_worker, local_worker_id);
1000 if task.is_some() {
1001 task
1002 } else {
1003 try_pop_internal(pool, local_worker, local_worker_id)
1005 }
1006 },
1007 _ => {
1008 let task = try_pop_internal(pool, local_worker, local_worker_id);
1010 if task.is_some() {
1011 task
1012 } else {
1013 try_pop_external(pool, local_worker, local_worker_id)
1015 }
1016 },
1017 }
1018 }
1019}
1020
1021#[inline]
1023fn try_pop_internal<O: Default + 'static>(pool: &StealableTaskPool<O>,
1024 local_worker: &StealableTaskQueue<O>,
1025 local_worker_id: usize)
1026 -> Option<Arc<AsyncTask<StealableTaskPool<O>, O>>> {
1027 let task = local_worker
1028 .internal
1029 .pop();
1030 if task.is_some() {
1031 pool
1033 .internal_consume
1034 .fetch_add(1, Ordering::Relaxed);
1035 task
1036 } else {
1037 let mut gen = thread_rng();
1039 let mut worker_stealers: Vec<&FIFOStealer<Arc<AsyncTask<StealableTaskPool<O>, O>>>> = pool
1040 .internal_stealers
1041 .iter()
1042 .enumerate()
1043 .filter_map(|(index, other)| {
1044 if index != local_worker_id {
1045 Some(other)
1046 } else {
1047 None
1049 }
1050 })
1051 .collect();
1052
1053 let remaining_len = local_worker.remaining_internal_capacity();
1054 loop {
1055 if worker_stealers.len() == 0 {
1057 break;
1059 }
1060
1061 let index = gen.gen_range(0..worker_stealers.len());
1062 let worker_stealer = worker_stealers.swap_remove(index);
1063
1064 match worker_stealer.steal_and_pop(&local_worker.internal,
1065 |count| {
1066 let stealable_len = count / 2;
1067 if stealable_len <= remaining_len {
1068 if stealable_len == 0 {
1070 1
1071 } else {
1072 stealable_len
1073 }
1074 } else {
1075 remaining_len
1077 }
1078 }) {
1079 Err(StealError::Empty) => {
1080 continue;
1082 },
1083 Err(StealError::Busy) => {
1084 continue;
1086 },
1087 Ok((task, _)) => {
1088 pool.internal_consume.fetch_add(1, Ordering::Relaxed);
1090 return Some(task);
1091 },
1092 }
1093 }
1094
1095 None
1096 }
1097}
1098
1099#[inline]
1101fn try_pop_external<O: Default + 'static>(pool: &StealableTaskPool<O>,
1102 local_worker: &StealableTaskQueue<O>,
1103 local_worker_id: usize)
1104 -> Option<Arc<AsyncTask<StealableTaskPool<O>, O>>> {
1105 let task = local_worker
1106 .external
1107 .pop();
1108 if task.is_some() {
1109 pool
1111 .external_consume
1112 .fetch_add(1, Ordering::Relaxed);
1113 task
1114 } else {
1115 let task = try_pop_public(pool, local_worker);
1117 if task.is_some() {
1118 pool
1120 .external_consume
1121 .fetch_add(1, Ordering::Relaxed);
1122 task
1123 } else {
1124 let mut gen = thread_rng();
1126 let mut worker_stealers: Vec<&Stealer<Arc<AsyncTask<StealableTaskPool<O>, O>>>> = pool
1127 .external_stealers
1128 .iter()
1129 .enumerate()
1130 .filter_map(|(index, other)| {
1131 if index != local_worker_id {
1132 Some(other)
1133 } else {
1134 None
1136 }
1137 })
1138 .collect();
1139
1140 loop {
1141 if worker_stealers.len() == 0 {
1143 break;
1145 }
1146
1147 let index = gen.gen_range(0..worker_stealers.len());
1148 let worker_stealer = worker_stealers.swap_remove(index);
1149
1150 match worker_stealer.steal_batch_and_pop(&local_worker.external) {
1151 Steal::Success(task) => {
1152 pool.external_consume.fetch_add(1, Ordering::Relaxed);
1154 return Some(task);
1155 },
1156 Steal::Retry => {
1157 continue;
1159 },
1160 Steal::Empty => {
1161 continue;
1163 },
1164 }
1165 }
1166
1167 None
1168 }
1169 }
1170}
1171
1172#[inline]
1174fn try_pop_public<O: Default + 'static>(pool: &StealableTaskPool<O>,
1175 local_worker: &StealableTaskQueue<O>)
1176 -> Option<Arc<AsyncTask<StealableTaskPool<O>, O>>> {
1177 loop {
1178 match pool.public.steal_batch_and_pop(&local_worker.external) {
1179 Steal::Empty => {
1180 return None;
1182 },
1183 Steal::Retry => {
1184 continue;
1186 },
1187 Steal::Success(task) => {
1188 pool.external_consume.fetch_add(1, Ordering::Relaxed);
1190 return Some(task);
1191 },
1192 }
1193 }
1194}
1195
1196impl<O: Default + 'static> AsyncTaskPoolExt<O> for StealableTaskPool<O> {
1197 #[inline]
1198 fn set_waits(&mut self, waits: Arc<ArrayQueue<Arc<(AtomicBool, Mutex<()>, Condvar)>>>) {
1199 self.waits = Some(waits);
1200 }
1201
1202 #[inline]
1203 fn get_waits(&self) -> Option<&Arc<ArrayQueue<Arc<(AtomicBool, Mutex<()>, Condvar)>>>> {
1204 self.waits.as_ref()
1205 }
1206
1207 #[inline]
1208 fn worker_len(&self) -> usize {
1209 self.workers.len()
1210 }
1211
1212 #[inline]
1219 fn clone_thread_waker(&self) -> Option<Arc<(AtomicBool, Mutex<()>, Condvar)>> {
1220 let id = current_multi_thread_worker_context().owner_worker_id(self);
1221 if let Some(worker) = self.workers.get(id) {
1222 return Some(worker.thread_waker.clone());
1223 }
1224
1225 None
1226 }
1227}
1228
1229impl<O: Default + 'static> StealableTaskPool<O> {
1230 pub fn new() -> Self {
1238 #[cfg(not(target_arch = "wasm32"))]
1239 let size = num_cpus::get_physical() * 2; #[cfg(target_arch = "wasm32")]
1241 let size = 1; StealableTaskPool::with(size,
1243 0x8000,
1244 [1, 1],
1245 3000)
1246 }
1247
1248 pub fn with(worker_size: usize,
1270 internal_queue_capacity: usize,
1271 weights: [u8; 2],
1272 interval: usize) -> Self {
1273 if worker_size == 0 {
1274 panic!(
1276 "Create WorkerTaskPool failed, worker size: {}, reason: invalid worker size",
1277 worker_size
1278 );
1279 }
1280 if interval == 0 {
1281 panic!(
1282 "Create WorkerTaskPool failed, interval: {}, reason: invalid interval",
1283 worker_size
1284 );
1285 }
1286
1287 let public = Injector::new();
1288 let mut workers = Vec::with_capacity(worker_size);
1289 let mut internal_stealers = Vec::with_capacity(worker_size);
1290 let mut external_stealers = Vec::with_capacity(worker_size);
1291 for _ in 0..worker_size {
1292 let thread_waker = Arc::new((AtomicBool::new(false), Mutex::new(()), Condvar::new()));
1294 let (worker,
1295 internal_stealer,
1296 external_stealer) =
1297 StealableTaskQueue::new(internal_queue_capacity,
1298 thread_waker);
1299 workers.push(worker);
1300 internal_stealers.push(internal_stealer);
1301 external_stealers.push(external_stealer);
1302 }
1303 let internal_consume = AtomicUsize::new(0);
1304 let internal_produce = AtomicUsize::new(0);
1305 let internal_traffic_statistics = AtomicUsize::new(0);
1306 let external_consume = AtomicUsize::new(0);
1307 let external_produce = AtomicUsize::new(0);
1308 let external_traffic_statistics = AtomicUsize::new(0);
1309 let clock = Clock::new();
1310 let last_time = AtomicCell::new(clock.recent());
1311
1312 StealableTaskPool {
1313 public,
1314 workers,
1315 internal_stealers,
1316 external_stealers,
1317 internal_consume,
1318 internal_produce,
1319 internal_traffic_statistics,
1320 external_consume,
1321 external_produce,
1322 external_traffic_statistics,
1323 weights,
1324 clock,
1325 interval,
1326 last_time,
1327 waits: None,
1328 }
1329 }
1330}
1331
1332pub struct MultiTaskRuntime<
1336 O: Default + 'static = (),
1337 P: AsyncTaskPoolExt<O> + AsyncTaskPool<O> = StealableTaskPool<O>,
1338>(
1339 Arc<(
1340 usize, Arc<P>, Option<
1343 Vec<(
1344 Sender<(usize, AsyncTimingTask<P, O>)>,
1345 Arc<AsyncTaskTimerByNotCancel<P, O>>,
1346 )>,
1347 >, AtomicUsize, Arc<ArrayQueue<Arc<(AtomicBool, Mutex<()>, Condvar)>>>, AtomicUsize, AtomicUsize, )>,
1353);
1354
1355unsafe impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O>> Send
1356 for MultiTaskRuntime<O, P>
1357{
1358}
1359unsafe impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O>> Sync
1360 for MultiTaskRuntime<O, P>
1361{
1362}
1363
1364impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O>> Clone
1365 for MultiTaskRuntime<O, P>
1366{
1367 fn clone(&self) -> Self {
1368 MultiTaskRuntime(self.0.clone())
1369 }
1370}
1371
1372impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>> AsyncRuntime<O>
1373 for MultiTaskRuntime<O, P>
1374{
1375 type Pool = P;
1376
1377 fn shared_pool(&self) -> Arc<Self::Pool> {
1379 (self.0).1.clone()
1380 }
1381
1382 fn get_id(&self) -> usize {
1384 (self.0).0
1385 }
1386
1387 fn wait_len(&self) -> usize {
1389 (self.0)
1390 .5
1391 .load(Ordering::Relaxed)
1392 .checked_sub((self.0).6.load(Ordering::Relaxed))
1393 .unwrap_or(0)
1394 }
1395
1396 fn len(&self) -> usize {
1398 (self.0).1.len()
1399 }
1400
1401 fn alloc<R: 'static>(&self) -> TaskId {
1403 TaskId(UnsafeCell::new((TaskHandle::<R>::default().into_raw() as u128) << 64 | self.get_id() as u128 & 0xffffffffffffffff))
1404 }
1405
1406 fn spawn<F>(&self, future: F) -> Result<TaskId>
1408 where
1409 F: Future<Output = O> + Send + 'static,
1410 {
1411 let task_id = self.alloc::<F::Output>();
1412 if let Err(e) = self.spawn_by_id(task_id.clone(), future) {
1413 return Err(e);
1414 }
1415
1416 Ok(task_id)
1417 }
1418
1419 fn spawn_local<F>(&self, future: F) -> Result<TaskId>
1421 where
1422 F: Future<Output=O> + Send + 'static {
1423 let task_id = self.alloc::<F::Output>();
1424 if let Err(e) = self.spawn_local_by_id(task_id.clone(), future) {
1425 return Err(e);
1426 }
1427
1428 Ok(task_id)
1429 }
1430
1431 fn spawn_priority<F>(&self, priority: usize, future: F) -> Result<TaskId>
1433 where
1434 F: Future<Output=O> + Send + 'static {
1435 let task_id = self.alloc::<F::Output>();
1436 if let Err(e) = self.spawn_priority_by_id(task_id.clone(), priority, future) {
1437 return Err(e);
1438 }
1439
1440 Ok(task_id)
1441 }
1442
1443 fn spawn_yield<F>(&self, future: F) -> Result<TaskId>
1445 where
1446 F: Future<Output=O> + Send + 'static {
1447 let task_id = self.alloc::<F::Output>();
1448 if let Err(e) = self.spawn_yield_by_id(task_id.clone(), future) {
1449 return Err(e);
1450 }
1451
1452 Ok(task_id)
1453 }
1454
1455 fn spawn_timing<F>(&self, future: F, time: usize) -> Result<TaskId>
1457 where
1458 F: Future<Output = O> + Send + 'static,
1459 {
1460 let task_id = self.alloc::<F::Output>();
1461 if let Err(e) = self.spawn_timing_by_id(task_id.clone(), future, time) {
1462 return Err(e);
1463 }
1464
1465 Ok(task_id)
1466 }
1467
1468 fn spawn_by_id<F>(&self, task_id: TaskId, future: F) -> Result<()>
1470 where
1471 F: Future<Output=O> + Send + 'static {
1472 let result = {
1473 (self.0).1.push(Arc::new(AsyncTask::new(
1474 task_id,
1475 (self.0).1.clone(),
1476 DEFAULT_MAX_LOW_PRIORITY_BOUNDED,
1477 Some(future.boxed()),
1478 )))
1479 };
1480
1481 let _ = wake_waiting_worker(&(self.0).4);
1482
1483 result
1484 }
1485
1486 fn spawn_local_by_id<F>(&self, task_id: TaskId, future: F) -> Result<()>
1487 where
1488 F: Future<Output=O> + Send + 'static {
1489 let should_wake = PI_ASYNC_THREAD_LOCAL_ID
1490 .try_with(|thread_id| unsafe { ((*thread_id.get()) >> 32) != self.get_id() })
1491 .unwrap_or(true);
1492 let result = (self.0).1.push_local(Arc::new(AsyncTask::new(
1493 task_id,
1494 (self.0).1.clone(),
1495 DEFAULT_HIGH_PRIORITY_BOUNDED,
1496 Some(future.boxed()),
1497 )));
1498
1499 if should_wake {
1500 let _ = wake_waiting_worker(&(self.0).4);
1501 }
1502
1503 result
1504 }
1505
1506 fn spawn_priority_by_id<F>(&self,
1508 task_id: TaskId,
1509 priority: usize,
1510 future: F) -> Result<()>
1511 where
1512 F: Future<Output=O> + Send + 'static {
1513 let result = {
1514 (self.0).1.push_priority(priority, Arc::new(AsyncTask::new(
1515 task_id,
1516 (self.0).1.clone(),
1517 priority,
1518 Some(future.boxed()),
1519 )))
1520 };
1521
1522 let _ = wake_waiting_worker(&(self.0).4);
1523
1524 result
1525 }
1526
1527 #[inline]
1529 fn spawn_yield_by_id<F>(&self, task_id: TaskId, future: F) -> Result<()>
1530 where
1531 F: Future<Output=O> + Send + 'static {
1532 self.spawn_priority_by_id(task_id,
1533 DEFAULT_HIGH_PRIORITY_BOUNDED,
1534 future)
1535 }
1536
1537 fn spawn_timing_by_id<F>(&self,
1539 task_id: TaskId,
1540 future: F,
1541 time: usize) -> Result<()>
1542 where
1543 F: Future<Output=O> + Send + 'static {
1544 let rt = self.clone();
1545 self.spawn_by_id(task_id, async move {
1546 if let Some(timers) = &(rt.0).2 {
1547 let id = (rt.0).1.get_thread_id() & 0xffffffff;
1549 let (_, timer) = &timers[id];
1550 timer.set_timer(
1551 AsyncTimingTask::WaitRun(Arc::new(AsyncTask::new(
1552 rt.alloc::<F::Output>(),
1553 (rt.0).1.clone(),
1554 DEFAULT_MAX_HIGH_PRIORITY_BOUNDED,
1555 Some(future.boxed()),
1556 ))),
1557 time,
1558 );
1559
1560 (rt.0).5.fetch_add(1, Ordering::Relaxed);
1561 }
1562
1563 Default::default()
1564 })
1565 }
1566
1567 fn pending<Output: 'static>(&self, task_id: &TaskId, waker: Waker) -> Poll<Output> {
1569 task_id.set_waker::<Output>(waker);
1570 Poll::Pending
1571 }
1572
1573 fn wakeup<Output: 'static>(&self, task_id: &TaskId) {
1575 task_id.wakeup::<Output>();
1576 }
1577
1578 fn wait<V: Send + 'static>(&self) -> AsyncWait<V> {
1580 AsyncWait(self.wait_any(2))
1581 }
1582
1583 fn wait_any<V: Send + 'static>(&self, capacity: usize) -> AsyncWaitAny<V> {
1585 let (producor, consumer) = async_bounded(capacity);
1586
1587 AsyncWaitAny {
1588 capacity,
1589 producor,
1590 consumer,
1591 }
1592 }
1593
1594 fn wait_any_callback<V: Send + 'static>(&self, capacity: usize) -> AsyncWaitAnyCallback<V> {
1596 let (producor, consumer) = async_bounded(capacity);
1597
1598 AsyncWaitAnyCallback {
1599 capacity,
1600 producor,
1601 consumer,
1602 }
1603 }
1604
1605 fn map_reduce<V: Send + 'static>(&self, capacity: usize) -> AsyncMapReduce<V> {
1607 let (producor, consumer) = async_bounded(capacity);
1608
1609 AsyncMapReduce {
1610 count: 0,
1611 capacity,
1612 producor,
1613 consumer,
1614 }
1615 }
1616
1617 fn timeout(&self, timeout: usize) -> BoxFuture<'static, ()> {
1619 let rt = self.clone();
1620
1621 if let Some(timers) = &(self.0).2 {
1622 match PI_ASYNC_THREAD_LOCAL_ID.try_with(move |thread_id| {
1624 let thread_id = unsafe { *thread_id.get() };
1626 let index = thread_id & 0xffffffff;
1627 if index > timers.len() {
1628 TimerTaskProducor::Foreign(timers[(self.0).3.load(Ordering::Relaxed) % timers.len()].0.clone())
1630 } else {
1631 TimerTaskProducor::Local(timers[index].1.clone())
1632 }
1633 }) {
1634 Err(_) => {
1635 panic!("Multi thread runtime timeout failed, reason: local thread id not match")
1636 }
1637 Ok(producor) => match producor {
1638 TimerTaskProducor::Local(timer) => {
1639 LocalAsyncWaitTimeout::new(rt, timer, timeout).boxed()
1640 },
1641 TimerTaskProducor::Foreign(producor) => {
1642 AsyncWaitTimeout::new(rt, producor, timeout).boxed()
1643 },
1644 },
1645 }
1646 } else {
1647 async move {
1649 thread::sleep(Duration::from_millis(timeout as u64));
1650 }
1651 .boxed()
1652 }
1653 }
1654
1655 fn yield_now(&self) -> BoxFuture<'static, ()> {
1657 async move {
1658 YieldNow(false).await;
1659 }.boxed()
1660 }
1661
1662 fn pipeline<S, SO, F, FO>(&self, input: S, mut filter: F) -> BoxStream<'static, FO>
1664 where
1665 S: Stream<Item = SO> + Send + 'static,
1666 SO: Send + 'static,
1667 F: FnMut(SO) -> AsyncPipelineResult<FO> + Send + 'static,
1668 FO: Send + 'static,
1669 {
1670 let output = stream! {
1671 for await value in input {
1672 match filter(value) {
1673 AsyncPipelineResult::Disconnect => {
1674 break;
1676 },
1677 AsyncPipelineResult::Filtered(result) => {
1678 yield result;
1679 },
1680 }
1681 }
1682 };
1683
1684 output.boxed()
1685 }
1686
1687 fn close(&self) -> bool {
1689 false
1690 }
1691}
1692
1693impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>> AsyncRuntimeExt<O>
1694 for MultiTaskRuntime<O, P>
1695{
1696 fn spawn_with_context<F, C>(&self, task_id: TaskId, future: F, context: C) -> Result<()>
1697 where
1698 F: Future<Output = O> + Send + 'static,
1699 C: 'static,
1700 {
1701 let task = Arc::new(AsyncTask::with_context(
1702 task_id,
1703 (self.0).1.clone(),
1704 DEFAULT_MAX_LOW_PRIORITY_BOUNDED,
1705 Some(future.boxed()),
1706 context,
1707 ));
1708 let result = (self.0).1.push(task);
1709
1710 let _ = wake_waiting_worker(&(self.0).4);
1711
1712 result
1713 }
1714
1715 fn spawn_timing_with_context<F, C>(
1716 &self,
1717 task_id: TaskId,
1718 future: F,
1719 context: C,
1720 time: usize,
1721 ) -> Result<()>
1722 where
1723 F: Future<Output = O> + Send + 'static,
1724 C: Send + 'static,
1725 {
1726 let rt = self.clone();
1727 self.spawn_by_id(task_id, async move {
1728 if let Some(timers) = &(rt.0).2 {
1729 let id = (rt.0).1.get_thread_id() & 0xffffffff;
1731 let (_, timer) = &timers[id];
1732 timer.set_timer(
1733 AsyncTimingTask::WaitRun(Arc::new(AsyncTask::with_context(
1734 rt.alloc::<F::Output>(),
1735 (rt.0).1.clone(),
1736 DEFAULT_MAX_HIGH_PRIORITY_BOUNDED,
1737 Some(future.boxed()),
1738 context,
1739 ))),
1740 time,
1741 );
1742
1743 (rt.0).5.fetch_add(1, Ordering::Relaxed);
1744 }
1745
1746 Default::default()
1747 })
1748 }
1749
1750 fn block_on<F>(&self, future: F) -> Result<F::Output>
1751 where
1752 F: Future + Send + 'static,
1753 <F as Future>::Output: Default + Send + 'static,
1754 {
1755 if let Some(local_rt) = local_async_runtime::<F::Output>() {
1757 if local_rt.get_id() == self.get_id() {
1759 return Err(Error::new(
1761 ErrorKind::WouldBlock,
1762 format!("Block on failed, reason: would block"),
1763 ));
1764 }
1765 }
1766
1767 let (sender, receiver) = bounded(1);
1768 if let Err(e) = self.spawn(async move {
1769 let r = future.await;
1771 sender.send(r);
1772
1773 Default::default()
1774 }) {
1775 return Err(Error::new(
1776 ErrorKind::Other,
1777 format!("Block on failed, reason: {:?}", e),
1778 ));
1779 }
1780
1781 match receiver.recv() {
1783 Err(e) => Err(Error::new(
1784 ErrorKind::Other,
1785 format!("Block on failed, reason: {:?}", e),
1786 )),
1787 Ok(result) => Ok(result),
1788 }
1789 }
1790}
1791
1792impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>>
1793 MultiTaskRuntime<O, P>
1794{
1795 pub fn idler_len(&self) -> usize {
1797 (self.0).1.idler_len()
1798 }
1799
1800 pub fn worker_len(&self) -> usize {
1802 (self.0).1.worker_len()
1803 }
1804
1805 pub fn buffer_len(&self) -> usize {
1807 (self.0).1.buffer_len()
1808 }
1809
1810 pub fn to_local_runtime(&self) -> LocalAsyncRuntime<O> {
1812 LocalAsyncRuntime {
1813 inner: self.as_raw(),
1814 get_id_func: MultiTaskRuntime::<O, P>::get_id_raw,
1815 spawn_func: MultiTaskRuntime::<O, P>::spawn_raw,
1816 spawn_local_func: MultiTaskRuntime::<O, P>::spawn_local_raw,
1817 spawn_timing_func: MultiTaskRuntime::<O, P>::spawn_timing_raw,
1818 timeout_func: MultiTaskRuntime::<O, P>::timeout_raw,
1819 }
1820 }
1821
1822 #[inline]
1824 pub(crate) fn as_raw(&self) -> *const () {
1825 Arc::into_raw(self.0.clone()) as *const ()
1826 }
1827
1828 #[inline]
1830 pub(crate) fn from_raw(raw: *const ()) -> Self {
1831 let inner = unsafe {
1832 Arc::from_raw(
1833 raw as *const (
1834 usize,
1835 Arc<P>,
1836 Option<
1837 Vec<(
1838 Sender<(usize, AsyncTimingTask<P, O>)>,
1839 Arc<AsyncTaskTimerByNotCancel<P, O>>,
1840 )>,
1841 >,
1842 AtomicUsize,
1843 Arc<ArrayQueue<Arc<(AtomicBool, Mutex<()>, Condvar)>>>,
1844 AtomicUsize,
1845 AtomicUsize,
1846 ),
1847 )
1848 };
1849 MultiTaskRuntime(inner)
1850 }
1851
1852 pub(crate) fn get_id_raw(raw: *const ()) -> usize {
1854 let rt = MultiTaskRuntime::<O, P>::from_raw(raw);
1855 let id = rt.get_id();
1856 Arc::into_raw(rt.0); id
1858 }
1859
1860 pub(crate) fn spawn_raw<F>(raw: *const (), future: F) -> Result<()>
1862 where
1863 F: Future<Output = O> + Send + 'static,
1864 {
1865 let rt = MultiTaskRuntime::<O, P>::from_raw(raw);
1866 let result = rt.spawn_by_id(rt.alloc::<F::Output>(), future);
1867 Arc::into_raw(rt.0); result
1869 }
1870
1871 pub(crate) fn spawn_local_raw<F>(raw: *const (), future: F) -> Result<()>
1873 where
1874 F: Future<Output = O> + Send + 'static,
1875 {
1876 let rt = MultiTaskRuntime::<O, P>::from_raw(raw);
1877 let result = rt.spawn_local_by_id(rt.alloc::<F::Output>(), future);
1878 Arc::into_raw(rt.0); result
1880 }
1881
1882 pub(crate) fn spawn_timing_raw(
1884 raw: *const (),
1885 future: BoxFuture<'static, O>,
1886 timeout: usize,
1887 ) -> Result<()> {
1888 let rt = MultiTaskRuntime::<O, P>::from_raw(raw);
1889 let result = rt.spawn_timing_by_id(rt.alloc::<O>(), future, timeout);
1890 Arc::into_raw(rt.0); result
1892 }
1893
1894 pub(crate) fn timeout_raw(raw: *const (), timeout: usize) -> BoxFuture<'static, ()> {
1896 let rt = MultiTaskRuntime::<O, P>::from_raw(raw);
1897 let boxed = rt.timeout(timeout);
1898 Arc::into_raw(rt.0); boxed
1900 }
1901}
1902
1903pub struct MultiTaskRuntimeBuilder<
1907 O: Default + 'static = (),
1908 P: AsyncTaskPoolExt<O> + AsyncTaskPool<O> = StealableTaskPool<O>,
1909> {
1910 pool: P, prefix: String, init: usize, min: usize, max: usize, stack_size: usize, timeout: u64, interval: Option<usize>, marker: PhantomData<O>,
1919}
1920
1921unsafe impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O>> Send
1922 for MultiTaskRuntimeBuilder<O, P>
1923{
1924}
1925unsafe impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O>> Sync
1926 for MultiTaskRuntimeBuilder<O, P>
1927{
1928}
1929
1930impl<O: Default + 'static> Default for MultiTaskRuntimeBuilder<O> {
1931 fn default() -> Self {
1933 #[cfg(not(target_arch = "wasm32"))]
1934 let core_len = num_cpus::get(); #[cfg(target_arch = "wasm32")]
1936 let core_len = 1; let pool = StealableTaskPool::with(core_len,
1938 65535,
1939 [1, 1],
1940 3000);
1941 MultiTaskRuntimeBuilder::new(pool)
1942 .thread_stack_size(2 * 1024 * 1024)
1943 .set_timer_interval(1)
1944 }
1945}
1946
1947impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>>
1948 MultiTaskRuntimeBuilder<O, P>
1949{
1950 pub fn new(mut pool: P) -> Self {
1952 #[cfg(not(target_arch = "wasm32"))]
1953 let core_len = num_cpus::get(); #[cfg(target_arch = "wasm32")]
1955 let core_len = 1; MultiTaskRuntimeBuilder {
1958 pool,
1959 prefix: DEFAULT_WORKER_THREAD_PREFIX.to_string(),
1960 init: core_len,
1961 min: core_len,
1962 max: core_len,
1963 stack_size: DEFAULT_THREAD_STACK_SIZE,
1964 timeout: DEFAULT_WORKER_THREAD_SLEEP_TIME,
1965 interval: None,
1966 marker: PhantomData,
1967 }
1968 }
1969
1970 pub fn thread_prefix(mut self, prefix: &str) -> Self {
1972 self.prefix = prefix.to_string();
1973 self
1974 }
1975
1976 pub fn thread_stack_size(mut self, stack_size: usize) -> Self {
1978 self.stack_size = stack_size;
1979 self
1980 }
1981
1982 pub fn init_worker_size(mut self, mut init: usize) -> Self {
1984 if init == 0 {
1985 init = DEFAULT_INIT_WORKER_SIZE;
1987 }
1988
1989 self.init = init;
1990 self
1991 }
1992
1993 pub fn set_worker_limit(mut self, mut min: usize, mut max: usize) -> Self {
1995 if self.init > max {
1996 max = self.init;
1998 }
1999
2000 if min == 0 || min > max {
2001 min = max;
2003 }
2004
2005 self.min = min;
2006 self.max = max;
2007 self
2008 }
2009
2010 pub fn set_timeout(mut self, timeout: u64) -> Self {
2012 self.timeout = timeout;
2013 self
2014 }
2015
2016 pub fn set_timer_interval(mut self, interval: usize) -> Self {
2018 self.interval = Some(interval);
2019 self
2020 }
2021
2022 pub fn build(mut self) -> MultiTaskRuntime<O, P> {
2056 let pool_worker_len = self.pool.worker_len();
2057 if pool_worker_len == 0 {
2058 panic!("Build multi thread runtime failed, reason: worker pool is empty");
2059 }
2060 if self.init > pool_worker_len {
2061 self.init = pool_worker_len;
2062 }
2063 if self.max > pool_worker_len {
2064 self.max = pool_worker_len;
2065 }
2066 if self.min > self.max {
2067 self.min = self.max;
2068 }
2069
2070 let interval = self.interval;
2072 let mut timers = if let Some(_) = interval {
2073 Some(Vec::with_capacity(self.max))
2074 } else {
2075 None
2076 };
2077 for _ in 0..self.max {
2078 if let Some(vec) = &mut timers {
2080 let timer = AsyncTaskTimerByNotCancel::new();
2081 let producor = timer.producor.clone();
2082 let timer = Arc::new(timer);
2083 vec.push((producor, timer));
2084 };
2085 }
2086
2087 let rt_uid = alloc_rt_uid();
2089 let waits = Arc::new(ArrayQueue::new(self.max));
2090 let mut pool = self.pool;
2091 pool.set_waits(waits.clone()); let pool = Arc::new(pool);
2093 let runtime = MultiTaskRuntime(Arc::new((
2094 rt_uid,
2095 pool,
2096 timers,
2097 AtomicUsize::new(0),
2098 waits,
2099 AtomicUsize::new(0),
2100 AtomicUsize::new(0),
2101 )));
2102
2103 let mut builders = Vec::with_capacity(self.init);
2105 for index in 0..self.init {
2106 let builder = Builder::new()
2107 .name(self.prefix.clone() + "-" + index.to_string().as_str())
2108 .stack_size(self.stack_size);
2109 builders.push(builder);
2110 }
2111
2112 let min = self.min;
2114 for index in 0..builders.len() {
2115 let builder = builders.remove(0);
2116 let runtime = runtime.clone();
2117 let timeout = self.timeout;
2118 let timer = if let Some(timers) = &(runtime.0).2 {
2119 let (_, timer) = &timers[index];
2120 Some(timer.clone())
2121 } else {
2122 None
2123 };
2124
2125 spawn_worker_thread(builder, index, runtime, min, timeout, interval, timer);
2126 }
2127
2128 runtime
2129 }
2130}
2131
2132fn bind_multi_thread_worker_context<P>(thread_id: usize, pool: &P) {
2151 if let Err(e) = PI_ASYNC_THREAD_LOCAL_ID.try_with(|local_thread_id| unsafe {
2152 *local_thread_id.get() = thread_id;
2155 }) {
2156 panic!(
2157 "Multi thread runtime startup failed, thread id: {:?}, reason: {:?}",
2158 thread_id & MULTI_THREAD_WORKER_ID_MASK,
2159 e
2160 );
2161 }
2162
2163 let context = MultiThreadWorkerContext {
2164 thread_id,
2165 pool: pool as *const P as *const (),
2166 };
2167 if let Err(e) = PI_ASYNC_MULTI_THREAD_WORKER_CONTEXT.try_with(|current| {
2168 current.set(context);
2169 }) {
2170 panic!(
2171 "Bind multi-thread worker pool failed, thread id: {:?}, reason: {:?}",
2172 thread_id & MULTI_THREAD_WORKER_ID_MASK,
2173 e
2174 );
2175 }
2176}
2177
2178fn spawn_worker_thread<
2189 O: Default + 'static,
2190 P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>,
2191>(
2192 builder: Builder,
2193 index: usize,
2194 runtime: MultiTaskRuntime<O, P>,
2195 min: usize,
2196 timeout: u64,
2197 interval: Option<usize>,
2198 timer: Option<Arc<AsyncTaskTimerByNotCancel<P, O>>>,
2199) {
2200 if let Some(timer) = timer {
2201 let rt_uid = runtime.get_id();
2203 let _ = builder.spawn(move || {
2204 let thread_id = rt_uid << 32 | index & MULTI_THREAD_WORKER_ID_MASK;
2206 bind_multi_thread_worker_context(thread_id, (runtime.0).1.as_ref());
2207
2208 let runtime_copy = runtime.clone();
2210 match PI_ASYNC_LOCAL_THREAD_ASYNC_RUNTIME.try_with(move |rt| {
2211 let raw = Arc::into_raw(Arc::new(runtime_copy.to_local_runtime()))
2212 as *mut LocalAsyncRuntime<O> as *mut ();
2213 rt.store(raw, Ordering::Relaxed);
2214 }) {
2215 Err(e) => {
2216 panic!("Bind multi runtime to local thread failed, reason: {:?}", e);
2217 }
2218 Ok(_) => (),
2219 }
2220
2221 timer_work_loop(
2223 runtime,
2224 index,
2225 min,
2226 timeout,
2227 interval.unwrap() as u64,
2228 timer,
2229 );
2230 });
2231 } else {
2232 let rt_uid = runtime.get_id();
2234 let _ = builder.spawn(move || {
2235 let thread_id = rt_uid << 32 | index & MULTI_THREAD_WORKER_ID_MASK;
2237 bind_multi_thread_worker_context(thread_id, (runtime.0).1.as_ref());
2238
2239 let runtime_copy = runtime.clone();
2241 match PI_ASYNC_LOCAL_THREAD_ASYNC_RUNTIME.try_with(move |rt| {
2242 let raw = Arc::into_raw(Arc::new(runtime_copy.to_local_runtime()))
2243 as *mut LocalAsyncRuntime<O> as *mut ();
2244 rt.store(raw, Ordering::Relaxed);
2245 }) {
2246 Err(e) => {
2247 panic!("Bind multi runtime to local thread failed, reason: {:?}", e);
2248 }
2249 Ok(_) => (),
2250 }
2251
2252 work_loop(runtime, index, min, timeout);
2254 });
2255 }
2256}
2257
2258enum WorkerWaitResult<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>> {
2274 TimedOut,
2275 NotSlept,
2276 Task(Arc<AsyncTask<P, O>>),
2277}
2278
2279#[inline]
2346fn worker_wait_for_task<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>>(
2347 runtime: &MultiTaskRuntime<O, P>,
2348 worker_waker: &Arc<(AtomicBool, Mutex<()>, Condvar)>,
2349 sleep_timeout: u64,
2350) -> WorkerWaitResult<O, P> {
2351 let (is_sleep, lock, condvar) = &**worker_waker;
2352
2353 loop {
2354 let _locked = lock.lock();
2355 if is_sleep.load(Ordering::Acquire) {
2356 break;
2357 }
2358
2359 if register_waiting_worker(&(runtime.0).4, worker_waker) {
2360 is_sleep.store(true, Ordering::Release);
2361 break;
2362 }
2363
2364 drop(_locked);
2365 if prune_stale_waiting_workers(&(runtime.0).4) == 0 {
2366 return WorkerWaitResult::NotSlept;
2367 }
2368 }
2369
2370 if let Some(task) = (runtime.0).1.try_pop() {
2371 is_sleep.store(false, Ordering::Release);
2372 return WorkerWaitResult::Task(task);
2373 }
2374
2375 if runtime.len() > 0 {
2376 is_sleep.store(false, Ordering::Release);
2377 return WorkerWaitResult::NotSlept;
2378 }
2379
2380 let mut locked = lock.lock();
2381 if !is_sleep.load(Ordering::Acquire) {
2382 return WorkerWaitResult::NotSlept;
2383 }
2384
2385 let timed_out = condvar
2386 .wait_for(&mut locked, Duration::from_millis(sleep_timeout))
2387 .timed_out();
2388 is_sleep.store(false, Ordering::Release);
2389
2390 if timed_out {
2391 WorkerWaitResult::TimedOut
2392 } else {
2393 WorkerWaitResult::NotSlept
2394 }
2395}
2396
2397fn timer_work_loop<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>>(
2399 runtime: MultiTaskRuntime<O, P>,
2400 index: usize,
2401 min: usize,
2402 sleep_timeout: u64,
2403 timer_interval: u64,
2404 timer: Arc<AsyncTaskTimerByNotCancel<P, O>>,
2405) {
2406 let pool = (runtime.0).1.clone();
2408 let worker_waker = pool.clone_thread_waker().unwrap();
2409
2410 let mut sleep_count = 0; let clock = Clock::new();
2412 loop {
2413 let timer_run_millis = clock.recent(); let mut pop_len = 0;
2416 (runtime.0)
2417 .5
2418 .fetch_add(timer.consume(),
2419 Ordering::Relaxed);
2420 loop {
2421 let current_time = timer.is_require_pop();
2422 if let Some(current_time) = current_time {
2423 loop {
2425 let timed_out = timer.pop(current_time);
2426 if let Some(timing_task) = timed_out {
2427 match timing_task {
2428 AsyncTimingTask::Pended(expired) => {
2429 runtime.wakeup::<O>(&expired);
2431 }
2432 AsyncTimingTask::WaitRun(expired) => {
2433 (runtime.0)
2435 .1
2436 .push_priority(DEFAULT_MAX_HIGH_PRIORITY_BOUNDED,
2437 expired);
2438 if let Some(task) = pool.try_pop() {
2439 sleep_count = 0; run_task(&runtime, task);
2441 }
2442 }
2443 AsyncTimingTask::TimeoutWake(waiter) => {
2444 waiter.fire();
2446 }
2447 }
2448 pop_len += 1;
2449
2450 if let Some(task) = pool.try_pop() {
2451 sleep_count = 0; run_task(&runtime, task);
2454 }
2455 } else {
2456 break;
2458 }
2459 }
2460 } else {
2461 break;
2463 }
2464 }
2465 (runtime.0)
2466 .6
2467 .fetch_add(pop_len,
2468 Ordering::Relaxed);
2469
2470 match pool.try_pop() {
2472 None => {
2473 if runtime.len() > 0 {
2474 continue;
2476 }
2477
2478 let diff_time = clock
2480 .recent()
2481 .duration_since(timer_run_millis)
2482 .as_millis() as u64; let real_timeout = if timer.len() == 0 {
2484 sleep_timeout
2486 } else {
2487 if diff_time >= timer_interval {
2489 continue;
2491 } else {
2492 timer_interval - diff_time
2494 }
2495 };
2496
2497 match worker_wait_for_task(&runtime, &worker_waker, real_timeout) {
2499 WorkerWaitResult::TimedOut => {
2500 sleep_count += 1;
2502 },
2503 WorkerWaitResult::Task(task) => {
2504 sleep_count = 0; run_task(&runtime, task);
2506 },
2507 WorkerWaitResult::NotSlept => (),
2508 }
2509 }
2510 Some(task) => {
2511 sleep_count = 0; run_task(&runtime, task);
2514 }
2515 }
2516 }
2517
2518 (runtime.0).1.close_worker();
2520 warn!(
2521 "Worker of runtime closed, runtime: {}, worker: {}, thread: {:?}",
2522 runtime.get_id(),
2523 index,
2524 thread::current()
2525 );
2526}
2527
2528fn work_loop<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>>(
2530 runtime: MultiTaskRuntime<O, P>,
2531 index: usize,
2532 min: usize,
2533 sleep_timeout: u64,
2534) {
2535 let pool = (runtime.0).1.clone();
2537 let worker_waker = pool.clone_thread_waker().unwrap();
2538
2539 let mut sleep_count = 0; loop {
2541 match pool.try_pop() {
2542 None => {
2543 if runtime.len() > 0 {
2545 continue;
2547 }
2548
2549 match worker_wait_for_task(&runtime, &worker_waker, sleep_timeout) {
2550 WorkerWaitResult::TimedOut => {
2551 sleep_count += 1;
2553 },
2554 WorkerWaitResult::Task(task) => {
2555 sleep_count = 0; run_task(&runtime, task);
2557 },
2558 WorkerWaitResult::NotSlept => (),
2559 }
2560 }
2561 Some(task) => {
2562 sleep_count = 0; run_task(&runtime, task);
2565 }
2566 }
2567 }
2568
2569 (runtime.0).1.close_worker();
2571 warn!(
2572 "Worker of runtime closed, runtime: {}, worker: {}, thread: {:?}",
2573 runtime.get_id(),
2574 index,
2575 thread::current()
2576 );
2577}
2578
2579#[inline]
2581fn run_task<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>>(
2582 runtime: &MultiTaskRuntime<O, P>,
2583 task: Arc<AsyncTask<P, O>>,
2584) {
2585 let waker = waker_ref(&task);
2586 let mut context = Context::from_waker(&*waker);
2587 if let Some(mut future) = task.get_inner() {
2588 if let Poll::Pending = future.as_mut().poll(&mut context) {
2589 task.set_inner(Some(future));
2591 }
2592 } else {
2593 (runtime.0).1.push(task);
2595 }
2596}
2597
2598enum TimerTaskProducor<
2600 O: Default + 'static = (),
2601 P: AsyncTaskPoolExt<O> + AsyncTaskPool<O> = StealableTaskPool<O>,
2602> {
2603 Local(Arc<AsyncTaskTimerByNotCancel<P, O>>), Foreign(Sender<(usize, AsyncTimingTask<P, O>)>), }