1use std::thread;
21use std::sync::Arc;
22use std::vec::IntoIter;
23use std::future::Future;
24use std::cell::UnsafeCell;
25use std::task::{Context, Poll, Waker};
26use std::io::{Error, ErrorKind, Result};
27use std::collections::vec_deque::VecDeque;
28use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
29
30use async_stream::stream;
31use crossbeam_channel::Sender;
32use crossbeam_queue::SegQueue;
33use flume::bounded as async_bounded;
34use futures::{
35 future::{BoxFuture, FutureExt},
36 stream::{BoxStream, Stream, StreamExt},
37 task::waker_ref,
38};
39use parking_lot::{Condvar, Mutex};
40use quanta::Clock;
41
42use wrr::IWRRSelector;
43
44use super::{
45 PI_ASYNC_THREAD_LOCAL_ID, DEFAULT_MAX_HIGH_PRIORITY_BOUNDED, DEFAULT_HIGH_PRIORITY_BOUNDED, DEFAULT_MAX_LOW_PRIORITY_BOUNDED, alloc_rt_uid, AsyncMapReduce, AsyncPipelineResult, AsyncRuntime,
46 AsyncRuntimeExt, AsyncTask, AsyncTaskPollClaim, AsyncTaskPollGuard, AsyncTaskPool, AsyncTaskPoolExt, AsyncTaskTimer, AsyncWait,
47 AsyncWaitAny, AsyncWaitAnyCallback, AsyncWaitTimeout, LocalAsyncRuntime, TaskId, YieldNow,
48 requeue_runtime_task
49};
50use crate::rt::{TaskHandle, AsyncTimingTask};
51
52pub struct SingleTaskPool<O: Default + 'static> {
56 id: usize, public: SegQueue<Arc<AsyncTask<SingleTaskPool<O>, O>>>, internal: UnsafeCell<VecDeque<Arc<AsyncTask<SingleTaskPool<O>, O>>>>, stack: UnsafeCell<Vec<Arc<AsyncTask<SingleTaskPool<O>, O>>>>, selector: UnsafeCell<IWRRSelector<2>>, consume_count: AtomicUsize, produce_count: AtomicUsize, thread_waker: Option<Arc<(AtomicBool, Mutex<()>, Condvar)>>, }
65
66unsafe impl<O: Default + 'static> Send for SingleTaskPool<O> {}
67unsafe impl<O: Default + 'static> Sync for SingleTaskPool<O> {}
68
69impl<O: Default + 'static> Default for SingleTaskPool<O> {
70 fn default() -> Self {
71 SingleTaskPool::new([1, 1])
72 }
73}
74
75impl<O: Default + 'static> AsyncTaskPool<O> for SingleTaskPool<O> {
76 type Pool = SingleTaskPool<O>;
77
78 #[inline]
79 fn get_thread_id(&self) -> usize {
80 let rt_uid = self.id;
81 match PI_ASYNC_THREAD_LOCAL_ID.try_with(move |thread_id| {
82 let current = unsafe { *thread_id.get() };
83 if current == usize::MAX {
84 unsafe {
86 *thread_id.get() = rt_uid << 32;
87 *thread_id.get()
88 }
89 } else {
90 current
91 }
92 }) {
93 Err(e) => {
94 panic!(
96 "Get thread id failed, thread: {:?}, reason: {:?}",
97 thread::current(),
98 e
99 );
100 }
101 Ok(id) => id,
102 }
103 }
104
105 #[inline]
106 fn len(&self) -> usize {
107 if let Some(len) = self
108 .produce_count
109 .load(Ordering::Relaxed)
110 .checked_sub(self.consume_count.load(Ordering::Relaxed))
111 {
112 len
113 } else {
114 0
115 }
116 }
117
118 #[inline]
119 fn push(&self, task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
120 self.public.push(task);
121 self.produce_count.fetch_add(1, Ordering::Relaxed);
122 Ok(())
123 }
124
125 #[inline]
126 fn push_local(&self, task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
127 let id = self.get_thread_id();
128 let rt_uid = task.owner();
129 if (id >> 32) == rt_uid {
130 unsafe {{
132 (&mut *self.internal.get()).push_back(task);
133 }}
134 self.produce_count.fetch_add(1, Ordering::Relaxed);
135 Ok(())
136 } else {
137 self.push(task)
139 }
140 }
141
142 #[inline]
143 fn push_priority(&self,
144 priority: usize,
145 task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
146 if priority >= DEFAULT_MAX_HIGH_PRIORITY_BOUNDED {
147 let id = self.get_thread_id();
149 let rt_uid = task.owner();
150 if (id >> 32) == rt_uid {
151 unsafe {
153 let stack = (&mut *self.stack.get());
154 if stack
155 .capacity()
156 .checked_sub(stack.len())
157 .unwrap_or(0) >= 0 {
158 (&mut *self.stack.get()).push(task);
160 } else {
161 (&mut *self.internal.get()).push_back(task);
163 }
164 }
165
166 self.produce_count.fetch_add(1, Ordering::Relaxed);
167 Ok(())
168 } else {
169 self.push(task)
171 }
172 } else if priority >= DEFAULT_HIGH_PRIORITY_BOUNDED {
173 self.push_local(task)
175 } else {
176 self.push(task)
178 }
179 }
180
181 #[inline]
182 fn push_keep(&self, task: Arc<AsyncTask<Self::Pool, O>>) -> Result<()> {
183 self.push_priority(DEFAULT_HIGH_PRIORITY_BOUNDED, task)
184 }
185
186 #[inline]
187 fn try_pop(&self) -> Option<Arc<AsyncTask<Self::Pool, O>>> {
188 let task = unsafe { (&mut *self
189 .stack
190 .get())
191 .pop()
192 };
193 if task.is_some() {
194 self.consume_count.fetch_add(1, Ordering::Relaxed);
196 return task;
197 }
198
199 let task = try_pop_by_weight(self);
201 if task.is_some() {
202 self
203 .consume_count
204 .fetch_add(1, Ordering::Relaxed);
205 }
206 task
207 }
208
209 #[inline]
210 fn try_pop_all(&self) -> IntoIter<Arc<AsyncTask<Self::Pool, O>>> {
211 let mut all = Vec::with_capacity(self.len());
212
213 let internal = unsafe { (&mut *self.internal.get()) };
214 for _ in 0..internal.len() {
215 if let Some(task) = internal.pop_front() {
216 all.push(task);
217 }
218 }
219
220 let public_len = self.public.len();
221 for _ in 0..public_len {
222 if let Some(task) = self.public.pop() {
223 all.push(task);
224 }
225 }
226
227 all.into_iter()
228 }
229
230 #[inline]
231 fn get_thread_waker(&self) -> Option<&Arc<(AtomicBool, Mutex<()>, Condvar)>> {
232 self.thread_waker.as_ref()
233 }
234}
235
236fn try_pop_by_weight<O: Default + 'static>(pool: &SingleTaskPool<O>)
238 -> Option<Arc<AsyncTask<SingleTaskPool<O>, O>>> {
239 unsafe {
240 match (&mut *pool.selector.get()).select() {
242 0 => {
243 let task = try_pop_external(pool);
245 if task.is_some() {
246 task
247 } else {
248 try_pop_internal(pool)
250 }
251 },
252 _ => {
253 let task = try_pop_internal(pool);
255 if task.is_some() {
256 task
257 } else {
258 try_pop_external(pool)
260 }
261 },
262 }
263 }
264}
265
266#[inline]
268fn try_pop_internal<O: Default + 'static>(pool: &SingleTaskPool<O>)
269 -> Option<Arc<AsyncTask<SingleTaskPool<O>, O>>> {
270 unsafe { (&mut *pool.internal.get()).pop_front() }
271}
272
273#[inline]
275fn try_pop_external<O: Default + 'static>(pool: &SingleTaskPool<O>)
276 -> Option<Arc<AsyncTask<SingleTaskPool<O>, O>>> {
277 pool.public.pop()
278}
279
280impl<O: Default + 'static> AsyncTaskPoolExt<O> for SingleTaskPool<O> {
281 fn set_thread_waker(&mut self, thread_waker: Arc<(AtomicBool, Mutex<()>, Condvar)>) {
282 self.thread_waker = Some(thread_waker);
283 }
284}
285
286impl<O: Default + 'static> SingleTaskPool<O> {
287 pub fn new(weights: [u8; 2]) -> Self {
289 let id = alloc_rt_uid();
290 let public = SegQueue::new();
291 let internal = UnsafeCell::new(VecDeque::new());
292 let stack = UnsafeCell::new(Vec::with_capacity(1));
293 let selector = UnsafeCell::new(IWRRSelector::new(weights));
294 let consume_count = AtomicUsize::new(0);
295 let produce_count = AtomicUsize::new(0);
296
297 SingleTaskPool {
298 id,
299 public,
300 internal,
301 stack,
302 selector,
303 consume_count,
304 produce_count,
305 thread_waker: Some(Arc::new((
306 AtomicBool::new(false),
307 Mutex::new(()),
308 Condvar::new(),
309 ))),
310 }
311 }
312}
313
314pub struct SingleTaskRuntime<
318 O: Default + 'static = (),
319 P: AsyncTaskPoolExt<O> + AsyncTaskPool<O> = SingleTaskPool<O>,
320>(
321 Arc<(
322 usize, Arc<P>, Sender<(usize, AsyncTimingTask<P, O>)>, AsyncTaskTimer<P, O>, AtomicUsize, AtomicUsize, )>,
329);
330
331unsafe impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O>> Send
332 for SingleTaskRuntime<O, P>
333{
334}
335unsafe impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O>> Sync
336 for SingleTaskRuntime<O, P>
337{
338}
339
340impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O>> Clone
341 for SingleTaskRuntime<O, P>
342{
343 fn clone(&self) -> Self {
344 SingleTaskRuntime(self.0.clone())
345 }
346}
347
348impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>> AsyncRuntime<O>
349 for SingleTaskRuntime<O, P>
350{
351 type Pool = P;
352
353 fn shared_pool(&self) -> Arc<Self::Pool> {
355 (self.0).1.clone()
356 }
357
358 fn get_id(&self) -> usize {
360 (self.0).0
361 }
362
363 fn wait_len(&self) -> usize {
365 (self.0)
366 .4
367 .load(Ordering::Relaxed)
368 .checked_sub((self.0).5.load(Ordering::Relaxed))
369 .unwrap_or(0)
370 }
371
372 fn len(&self) -> usize {
374 (self.0).1.len()
375 }
376
377 fn alloc<R: 'static>(&self) -> TaskId {
379 TaskId(UnsafeCell::new((TaskHandle::<R>::default().into_raw() as u128) << 64 | self.get_id() as u128 & 0xffffffffffffffff))
380 }
381
382 fn spawn<F>(&self, future: F) -> Result<TaskId>
384 where
385 F: Future<Output = O> + Send + 'static,
386 {
387 let task_id = self.alloc::<F::Output>();
388 if let Err(e) = self.spawn_by_id(task_id.clone(), future) {
389 return Err(e);
390 }
391
392 Ok(task_id)
393 }
394
395 fn spawn_local<F>(&self, future: F) -> Result<TaskId>
397 where
398 F: Future<Output = O> + Send + 'static {
399 let task_id = self.alloc::<F::Output>();
400 if let Err(e) = self.spawn_local_by_id(task_id.clone(), future) {
401 return Err(e);
402 }
403
404 Ok(task_id)
405 }
406
407 fn spawn_priority<F>(&self, priority: usize, future: F) -> Result<TaskId>
409 where
410 F: Future<Output = O> + Send + 'static {
411 let task_id = self.alloc::<F::Output>();
412 if let Err(e) = self.spawn_priority_by_id(task_id.clone(), priority, future) {
413 return Err(e);
414 }
415
416 Ok(task_id)
417 }
418
419 fn spawn_yield<F>(&self, future: F) -> Result<TaskId>
421 where
422 F: Future<Output = O> + Send + 'static {
423 let task_id = self.alloc::<F::Output>();
424 if let Err(e) = self.spawn_yield_by_id(task_id.clone(), future) {
425 return Err(e);
426 }
427
428 Ok(task_id)
429 }
430
431 fn spawn_timing<F>(&self, future: F, time: usize) -> Result<TaskId>
433 where
434 F: Future<Output = O> + Send + 'static,
435 {
436 let task_id = self.alloc::<F::Output>();
437 if let Err(e) = self.spawn_timing_by_id(task_id.clone(), future, time) {
438 return Err(e);
439 }
440
441 Ok(task_id)
442 }
443
444 fn spawn_by_id<F>(&self, task_id: TaskId, future: F) -> Result<()>
446 where
447 F: Future<Output = O> + Send + 'static {
448 if let Err(e) = (self.0).1.push(Arc::new(AsyncTask::new(
449 task_id,
450 (self.0).1.clone(),
451 DEFAULT_MAX_LOW_PRIORITY_BOUNDED,
452 Some(future.boxed()),
453 ))) {
454 return Err(Error::new(ErrorKind::Other, e));
455 }
456
457 Ok(())
458 }
459
460 fn spawn_local_by_id<F>(&self, task_id: TaskId, future: F) -> Result<()>
462 where
463 F: Future<Output = O> + Send + 'static {
464 (self.0).1.push_local(Arc::new(AsyncTask::new(
465 task_id,
466 (self.0).1.clone(),
467 DEFAULT_HIGH_PRIORITY_BOUNDED,
468 Some(future.boxed()))))
469 }
470
471 fn spawn_priority_by_id<F>(&self,
473 task_id: TaskId,
474 priority: usize,
475 future: F) -> Result<()>
476 where
477 F: Future<Output = O> + Send + 'static {
478 (self.0).1.push_priority(priority, Arc::new(AsyncTask::new(
479 task_id,
480 (self.0).1.clone(),
481 priority,
482 Some(future.boxed()))))
483 }
484
485 #[inline]
487 fn spawn_yield_by_id<F>(&self, task_id: TaskId, future: F) -> Result<()>
488 where
489 F: Future<Output = O> + Send + 'static {
490 self.spawn_priority_by_id(task_id,
491 DEFAULT_HIGH_PRIORITY_BOUNDED,
492 future)
493 }
494
495 fn spawn_timing_by_id<F>(&self,
497 task_id: TaskId,
498 future: F,
499 time: usize) -> Result<()>
500 where
501 F: Future<Output = O> + Send + 'static {
502 let rt = self.clone();
503 self.spawn_by_id(task_id, async move {
504 (rt.0).3.set_timer(
505 AsyncTimingTask::WaitRun(Arc::new(AsyncTask::new(
506 rt.alloc::<F::Output>(),
507 (rt.0).1.clone(),
508 DEFAULT_HIGH_PRIORITY_BOUNDED,
509 Some(future.boxed()),
510 ))),
511 time,
512 );
513
514 (rt.0).4.fetch_add(1, Ordering::Relaxed);
515 Default::default()
516 })
517 }
518
519 fn pending<Output: 'static>(&self, task_id: &TaskId, waker: Waker) -> Poll<Output> {
521 task_id.set_waker::<Output>(waker);
522 Poll::Pending
523 }
524
525 fn wakeup<Output: 'static>(&self, task_id: &TaskId) {
527 task_id.wakeup::<Output>();
528 }
529
530 fn wait<V: Send + 'static>(&self) -> AsyncWait<V> {
532 AsyncWait(self.wait_any(2))
533 }
534
535 fn wait_any<V: Send + 'static>(&self, capacity: usize) -> AsyncWaitAny<V> {
537 let (producor, consumer) = async_bounded(capacity);
538
539 AsyncWaitAny {
540 capacity,
541 producor,
542 consumer,
543 }
544 }
545
546 fn wait_any_callback<V: Send + 'static>(&self, capacity: usize) -> AsyncWaitAnyCallback<V> {
548 let (producor, consumer) = async_bounded(capacity);
549
550 AsyncWaitAnyCallback {
551 capacity,
552 producor,
553 consumer,
554 }
555 }
556
557 fn map_reduce<V: Send + 'static>(&self, capacity: usize) -> AsyncMapReduce<V> {
559 let (producor, consumer) = async_bounded(capacity);
560
561 AsyncMapReduce {
562 count: 0,
563 capacity,
564 producor,
565 consumer,
566 }
567 }
568
569 fn timeout(&self, timeout: usize) -> BoxFuture<'static, ()> {
571 let rt = self.clone();
572 let producor = (self.0).2.clone();
573
574 AsyncWaitTimeout::new(rt, producor, timeout).boxed()
575 }
576
577 fn yield_now(&self) -> BoxFuture<'static, ()> {
579 async move {
580 YieldNow(false).await;
581 }.boxed()
582 }
583
584 fn pipeline<S, SO, F, FO>(&self, input: S, mut filter: F) -> BoxStream<'static, FO>
586 where
587 S: Stream<Item = SO> + Send + 'static,
588 SO: Send + 'static,
589 F: FnMut(SO) -> AsyncPipelineResult<FO> + Send + 'static,
590 FO: Send + 'static,
591 {
592 let output = stream! {
593 for await value in input {
594 match filter(value) {
595 AsyncPipelineResult::Disconnect => {
596 break;
598 },
599 AsyncPipelineResult::Filtered(result) => {
600 yield result;
601 },
602 }
603 }
604 };
605
606 output.boxed()
607 }
608
609 fn close(&self) -> bool {
611 false
612 }
613}
614
615impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>> AsyncRuntimeExt<O>
616 for SingleTaskRuntime<O, P>
617{
618 fn spawn_with_context<F, C>(&self, task_id: TaskId, future: F, context: C) -> Result<()>
619 where
620 F: Future<Output = O> + Send + 'static,
621 C: 'static,
622 {
623 if let Err(e) = (self.0).1.push(Arc::new(AsyncTask::with_context(
624 task_id,
625 (self.0).1.clone(),
626 DEFAULT_MAX_LOW_PRIORITY_BOUNDED,
627 Some(future.boxed()),
628 context,
629 ))) {
630 return Err(Error::new(ErrorKind::Other, e));
631 }
632
633 Ok(())
634 }
635
636 fn spawn_timing_with_context<F, C>(
637 &self,
638 task_id: TaskId,
639 future: F,
640 context: C,
641 time: usize,
642 ) -> Result<()>
643 where
644 F: Future<Output = O> + Send + 'static,
645 C: Send + 'static,
646 {
647 let rt = self.clone();
648 self.spawn_by_id(task_id, async move {
649 (rt.0).3.set_timer(
650 AsyncTimingTask::WaitRun(Arc::new(AsyncTask::with_context(
651 rt.alloc::<F::Output>(),
652 (rt.0).1.clone(),
653 DEFAULT_MAX_HIGH_PRIORITY_BOUNDED,
654 Some(future.boxed()),
655 context,
656 ))),
657 time,
658 );
659
660 (rt.0).4.fetch_add(1, Ordering::Relaxed);
661 Default::default()
662 })
663 }
664
665 fn block_on<F>(&self, future: F) -> Result<F::Output>
666 where
667 F: Future + Send + 'static,
668 <F as Future>::Output: Default + Send + 'static,
669 {
670 let runner = SingleTaskRunner {
671 is_running: AtomicBool::new(true),
672 runtime: self.clone(),
673 clock: Clock::new(),
674 };
675 let mut result: Option<<F as Future>::Output> = None;
676 let result_raw = (&mut result) as *mut Option<<F as Future>::Output> as usize;
677
678 self.spawn(async move {
679 let r = future.await;
681 unsafe {
682 *(result_raw as *mut Option<<F as Future>::Output>) = Some(r);
683 }
684
685 Default::default()
686 });
687
688 loop {
689 while runner.run()? > 0 {}
691
692 if let Some(result) = result.take() {
694 return Ok(result);
696 }
697 }
698 }
699}
700
701impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>>
702 SingleTaskRuntime<O, P>
703{
704 pub fn to_local_runtime(&self) -> LocalAsyncRuntime<O> {
706 LocalAsyncRuntime {
707 inner: self.as_raw(),
708 get_id_func: SingleTaskRuntime::<O, P>::get_id_raw,
709 spawn_func: SingleTaskRuntime::<O, P>::spawn_raw,
710 spawn_local_func: SingleTaskRuntime::<O, P>::spawn_local_raw,
711 spawn_timing_func: SingleTaskRuntime::<O, P>::spawn_timing_raw,
712 timeout_func: SingleTaskRuntime::<O, P>::timeout_raw,
713 }
714 }
715
716 #[inline]
718 pub(crate) fn as_raw(&self) -> *const () {
719 Arc::into_raw(self.0.clone()) as *const ()
720 }
721
722 #[inline]
724 pub(crate) fn from_raw(raw: *const ()) -> Self {
725 let inner = unsafe {
726 Arc::from_raw(
727 raw as *const (
728 usize,
729 Arc<P>,
730 Sender<(usize, AsyncTimingTask<P, O>)>,
731 AsyncTaskTimer<P, O>,
732 AtomicUsize,
733 AtomicUsize,
734 ),
735 )
736 };
737 SingleTaskRuntime(inner)
738 }
739
740 pub(crate) fn get_id_raw(raw: *const ()) -> usize {
742 let rt = SingleTaskRuntime::<O, P>::from_raw(raw);
743 let id = rt.get_id();
744 Arc::into_raw(rt.0); id
746 }
747
748 pub(crate) fn spawn_raw(raw: *const (), future: BoxFuture<'static, O>) -> Result<()> {
750 let rt = SingleTaskRuntime::<O, P>::from_raw(raw);
751 let result = rt.spawn_by_id(rt.alloc::<O>(), future);
752 Arc::into_raw(rt.0); result
754 }
755
756 pub(crate) fn spawn_local_raw(raw: *const (), future: BoxFuture<'static, O>) -> Result<()> {
758 let rt = SingleTaskRuntime::<O, P>::from_raw(raw);
759 let result = rt.spawn_local_by_id(rt.alloc::<O>(), future);
760 Arc::into_raw(rt.0); result
762 }
763
764 pub(crate) fn spawn_timing_raw(
766 raw: *const (),
767 future: BoxFuture<'static, O>,
768 timeout: usize,
769 ) -> Result<()> {
770 let rt = SingleTaskRuntime::<O, P>::from_raw(raw);
771 let result = rt.spawn_timing_by_id(rt.alloc::<O>(), future, timeout);
772 Arc::into_raw(rt.0); result
774 }
775
776 pub(crate) fn timeout_raw(raw: *const (), timeout: usize) -> BoxFuture<'static, ()> {
778 let rt = SingleTaskRuntime::<O, P>::from_raw(raw);
779 let boxed = rt.timeout(timeout);
780 Arc::into_raw(rt.0); boxed
782 }
783}
784
785pub struct SingleTaskRunner<
789 O: Default + 'static,
790 P: AsyncTaskPoolExt<O> + AsyncTaskPool<O> = SingleTaskPool<O>,
791> {
792 is_running: AtomicBool, runtime: SingleTaskRuntime<O, P>, clock: Clock, }
796
797unsafe impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O>> Send
798 for SingleTaskRunner<O, P>
799{
800}
801unsafe impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O>> Sync
802 for SingleTaskRunner<O, P>
803{
804}
805
806impl<O: Default + 'static> Default for SingleTaskRunner<O> {
807 fn default() -> Self {
808 SingleTaskRunner::new(SingleTaskPool::default())
809 }
810}
811
812impl<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>>
813 SingleTaskRunner<O, P>
814{
815 pub fn new(pool: P) -> Self {
817 let rt_uid = pool.get_thread_id() >> 32;
818 let pool = Arc::new(pool);
819
820 let timer = AsyncTaskTimer::new();
822 let producor = timer.producor.clone();
823 let timer_producor_count = AtomicUsize::new(0);
824 let timer_consume_count = AtomicUsize::new(0);
825
826 let runtime = SingleTaskRuntime(Arc::new((rt_uid,
828 pool,
829 producor,
830 timer,
831 timer_producor_count,
832 timer_consume_count)));
833
834 SingleTaskRunner {
835 is_running: AtomicBool::new(false),
836 runtime,
837 clock: Clock::new(),
838 }
839 }
840
841 pub fn get_thread_waker(&self) -> Option<Arc<(AtomicBool, Mutex<()>, Condvar)>> {
843 (self.runtime.0).1.get_thread_waker().cloned()
844 }
845
846 pub fn startup(&self) -> Option<SingleTaskRuntime<O, P>> {
848 if cfg!(target_arch = "aarch64") {
849 match self
850 .is_running
851 .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
852 {
853 Ok(false) => {
854 Some(self.runtime.clone())
856 }
857 _ => {
858 None
860 }
861 }
862 } else {
863 match self.is_running.compare_exchange(
864 false,
865 true,
866 Ordering::SeqCst,
867 Ordering::SeqCst,
868 ) {
869 Ok(false) => {
870 Some(self.runtime.clone())
872 }
873 _ => {
874 None
876 }
877 }
878 }
879 }
880
881 pub fn run_once(&self) -> Result<usize> {
883 if !self.is_running.load(Ordering::Relaxed) {
884 return Err(Error::new(
886 ErrorKind::Other,
887 "Single thread runtime not running",
888 ));
889 }
890
891 let mut pop_len = 0;
893 (self.runtime.0)
894 .4
895 .fetch_add((self.runtime.0).3.consume(),
896 Ordering::Relaxed);
897 loop {
898 let current_time = (self.runtime.0).3.is_require_pop();
899 if let Some(current_time) = current_time {
900 let timed_out = (self.runtime.0).3.pop(current_time);
902 if let Some((handle, timing_task)) = timed_out {
903 match timing_task {
904 AsyncTimingTask::Pended(expired) => {
905 self.runtime.wakeup::<O>(&expired);
907 if let Some(task) = (self.runtime.0).1.try_pop() {
908 run_task(task);
909 }
910 }
911 AsyncTimingTask::WaitRun(expired) => {
912 (self.runtime.0).1.push_priority(handle, expired);
914 if let Some(task) = (self.runtime.0).1.try_pop() {
915 run_task(task);
916 }
917 }
918 AsyncTimingTask::TimeoutWake(waiter) => {
919 waiter.fire();
921 if let Some(task) = (self.runtime.0).1.try_pop() {
922 run_task(task);
923 }
924 }
925 }
926 pop_len += 1;
927 }
928 } else {
929 break;
931 }
932 }
933 (self.runtime.0)
934 .5
935 .fetch_add(pop_len,
936 Ordering::Relaxed);
937
938 match (self.runtime.0).1.try_pop() {
940 None => {
941 return Ok(0);
943 }
944 Some(task) => {
945 run_task(task);
946 }
947 }
948
949 Ok((self.runtime.0).1.len())
950 }
951
952 pub fn run(&self) -> Result<usize> {
954 if !self.is_running.load(Ordering::Relaxed) {
955 return Err(Error::new(
957 ErrorKind::Other,
958 "Single thread runtime not running",
959 ));
960 }
961
962 loop {
963 let mut pop_len = 0;
965 let mut start_run_millis = self.clock.recent(); (self.runtime.0)
967 .4
968 .fetch_add((self.runtime.0).3.consume(),
969 Ordering::Relaxed);
970 loop {
971 let current_time = (self.runtime.0).3.is_require_pop();
972 if let Some(current_time) = current_time {
973 let timed_out = (self.runtime.0).3.pop(current_time);
975 if let Some((handle, timing_task)) = timed_out {
976 match timing_task {
977 AsyncTimingTask::Pended(expired) => {
978 self.runtime.wakeup::<O>(&expired);
980 if let Some(task) = (self.runtime.0).1.try_pop() {
981 run_task(task);
982 }
983 }
984 AsyncTimingTask::WaitRun(expired) => {
985 (self.runtime.0).1.push_priority(handle, expired);
987 if let Some(task) = (self.runtime.0).1.try_pop() {
988 run_task(task);
989 }
990 }
991 AsyncTimingTask::TimeoutWake(waiter) => {
992 waiter.fire();
994 if let Some(task) = (self.runtime.0).1.try_pop() {
995 run_task(task);
996 }
997 }
998 }
999 pop_len += 1;
1000 }
1001 } else {
1002 break;
1004 }
1005 }
1006 (self.runtime.0)
1007 .5
1008 .fetch_add(pop_len,
1009 Ordering::Relaxed);
1010
1011 while self
1013 .clock
1014 .recent()
1015 .duration_since(start_run_millis)
1016 .as_millis() < 1 {
1017 match (self.runtime.0).1.try_pop() {
1018 None => {
1019 return Ok((self.runtime.0).1.len());
1021 }
1022 Some(task) => {
1023 run_task(task);
1024 }
1025 }
1026 }
1027 }
1028 }
1029
1030 pub fn into_local(self) -> SingleTaskRuntime<O, P> {
1032 self.runtime
1033 }
1034}
1035
1036#[inline]
1050fn run_task<O: Default + 'static, P: AsyncTaskPoolExt<O> + AsyncTaskPool<O, Pool = P>>(
1051 task: Arc<AsyncTask<P, O>>,
1052) {
1053 match task.try_begin_runtime_poll() {
1054 AsyncTaskPollClaim::Discard => return,
1055 AsyncTaskPollClaim::Legacy => {
1056 let waker = waker_ref(&task);
1057 let mut context = Context::from_waker(&*waker);
1058 if let Some(mut future) = task.get_inner() {
1059 if let Poll::Pending = future.as_mut().poll(&mut context) {
1060 task.set_inner(Some(future));
1061 }
1062 }
1063 return;
1064 },
1065 AsyncTaskPollClaim::Managed => (),
1066 }
1067
1068 let guard = AsyncTaskPollGuard::new(&task);
1071 let waker = waker_ref(&task);
1072 let mut context = Context::from_waker(&*waker);
1073 let mut future = match task.take_inner_for_runtime_poll() {
1074 Some(future) => future,
1075 None => {
1076 guard.finish_ready();
1077 return;
1078 },
1079 };
1080
1081 match future.as_mut().poll(&mut context) {
1082 Poll::Pending => {
1083 task.restore_inner_after_runtime_poll(future);
1084 if guard.finish_pending() {
1085 requeue_runtime_task(task.get_pool(), &task);
1086 }
1087 },
1088 Poll::Ready(_) => guard.finish_ready(),
1089 }
1090}