1use alloc::boxed::Box;
45use alloc::collections::VecDeque;
46use alloc::rc::Rc;
47use alloc::sync::Arc;
48use alloc::task::Wake;
49use alloc::vec::Vec;
50use core::cell::RefCell;
51use core::future::Future;
52use core::pin::Pin;
53use core::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering::SeqCst};
54use core::task::{Context, Poll, Waker};
55
56use rt::sync::SpinLock;
57
58pub use crate::Cancelled;
59
60pub const BLOCKING_MAX: u32 = 8;
62pub const PROGRESS_WAKE_MS: u64 = 30;
64const BLOCKING_STACK: usize = 512 * 1024;
65
66#[cfg(any(windows, target_os = "linux"))]
67fn tid() -> u64 {
68 crate::plat_tid()
69}
70
71#[cfg(not(any(windows, target_os = "linux")))]
73fn tid() -> u64 {
74 0
75}
76
77static VCLOCK: AtomicU64 = AtomicU64::new(u64::MAX); pub fn now_us() -> u64 {
83 match VCLOCK.load(SeqCst) {
84 u64::MAX => crate::time_us(),
85 v => v,
86 }
87}
88
89pub fn clock_set(us: Option<u64>) {
92 VCLOCK.store(us.map_or(u64::MAX, |v| v.min(u64::MAX - 1)), SeqCst);
93}
94
95type Boxed = Pin<Box<dyn Future<Output = ()>>>;
98
99struct Shared {
100 owner: u64,
101 park: Option<Arc<crate::Sem>>,
103 parked: bool,
104 local: SpinLock<Vec<(u32, u32)>>,
105 remote: SpinLock<Vec<(u32, u32)>>,
106 notified: AtomicBool,
107}
108
109impl Shared {
110 fn push(&self, e: (u32, u32)) {
111 let here = tid();
112 if here == self.owner && here != 0 {
113 self.local.lock().push(e);
114 } else if here == 0 {
115 self.local.lock().push(e);
116 self.signal();
117 } else {
118 self.remote.lock().push(e);
119 self.signal();
120 }
121 }
122
123 fn signal(&self) {
124 if !self.parked {
125 crate::wake();
126 } else if !self.notified.swap(true, SeqCst) {
127 if let Some(s) = &self.park {
128 s.post(1);
129 }
130 }
131 }
132}
133
134struct TaskWaker {
135 sh: Arc<Shared>,
136 slot: u32,
137 ver: u32,
138 queued: AtomicBool,
139}
140
141impl Wake for TaskWaker {
142 fn wake(self: Arc<Self>) {
143 self.wake_by_ref()
144 }
145 fn wake_by_ref(self: &Arc<Self>) {
146 if !self.queued.swap(true, SeqCst) {
147 self.sh.push((self.slot, self.ver));
148 }
149 }
150}
151
152struct Running {
153 fut: Boxed,
154 waker: Waker,
155 tw: Arc<TaskWaker>,
156}
157
158struct Slot {
159 ver: u32,
160 task: Option<Running>,
161}
162
163pub struct Executor {
165 sh: Arc<Shared>,
166 slots: Vec<Slot>,
167 free: Vec<u32>,
168 queue: Rc<RefCell<Vec<Boxed>>>,
169 live: usize,
170}
171
172impl Executor {
173 fn make(park: bool) -> Executor {
174 let sem = if park { crate::Sem::new(0).map(Arc::new) } else { None };
175 let sh = Shared {
176 owner: tid(),
177 parked: park,
178 park: sem,
179 local: SpinLock::new(Vec::new()),
180 remote: SpinLock::new(Vec::new()),
181 notified: AtomicBool::new(false),
182 };
183 Executor { sh: Arc::new(sh), slots: Vec::new(), free: Vec::new(), queue: Rc::new(RefCell::new(Vec::new())), live: 0 }
184 }
185
186 pub fn ui() -> Executor {
188 Executor::make(false)
189 }
190
191 pub fn new() -> Executor {
193 Executor::make(true)
194 }
195
196 pub fn spawner(&self) -> Spawner {
198 Spawner { q: self.queue.clone() }
199 }
200
201 pub fn spawn<F: Future + 'static>(&self, fut: F) -> Task<F::Output> {
203 self.spawner().spawn(fut)
204 }
205
206 fn insert(&mut self, fut: Boxed) {
208 let slot = match self.free.pop() {
209 Some(s) => s,
210 None => {
211 self.slots.push(Slot { ver: 0, task: None });
212 (self.slots.len() - 1) as u32
213 }
214 };
215 let ver = self.slots[slot as usize].ver;
216 let tw = Arc::new(TaskWaker { sh: self.sh.clone(), slot, ver, queued: AtomicBool::new(true) });
217 let waker = Waker::from(tw.clone());
218 self.slots[slot as usize].task = Some(Running { fut, waker, tw });
219 self.live += 1;
220 self.sh.local.lock().push((slot, ver));
221 }
222
223 pub fn accept_remote(&mut self) -> bool {
225 let mut r = core::mem::take(&mut *self.sh.remote.lock());
226 if r.is_empty() {
227 return false;
228 }
229 self.sh.local.lock().append(&mut r);
230 true
231 }
232
233 pub fn run(&mut self) -> usize {
236 let spawned = core::mem::take(&mut *self.queue.borrow_mut());
237 for f in spawned {
238 self.insert(f);
239 }
240 fire_timers(now_us());
241 run_inline(1);
242 let mut ready = core::mem::take(&mut *self.sh.local.lock());
243 if self.sh.parked {
244 ready.append(&mut core::mem::take(&mut *self.sh.remote.lock()));
245 }
246 let mut n = 0;
247 for (slot, ver) in ready {
248 let Some(s) = self.slots.get_mut(slot as usize) else { continue };
249 if s.ver != ver {
250 continue; }
252 let Some(r) = s.task.as_mut() else { continue };
253 r.tw.queued.store(false, SeqCst); n += 1;
255 let Running { fut, waker, .. } = r;
256 if fut.as_mut().poll(&mut Context::from_waker(waker)).is_ready() {
257 let done = s.task.take();
258 s.ver = s.ver.wrapping_add(1);
259 self.free.push(slot);
260 self.live -= 1;
261 drop(done); }
263 }
264 n
265 }
266
267 pub fn has_ready(&self) -> bool {
269 !self.queue.borrow().is_empty()
270 || !self.sh.local.lock().is_empty()
271 || (self.sh.parked && !self.sh.remote.lock().is_empty())
272 || inline_pending()
273 }
274
275 pub fn busy(&self) -> bool {
278 self.has_ready() || !self.sh.remote.lock().is_empty() || blocking_pending() > 0
279 }
280
281 pub fn live(&self) -> usize {
283 self.live + self.queue.borrow().len()
284 }
285
286 pub fn block_on<F: Future>(&mut self, fut: F) -> F::Output {
289 let mut fut = core::pin::pin!(fut);
290 let main = Arc::new(MainWaker { sh: self.sh.clone(), woken: AtomicBool::new(true) });
291 let waker = Waker::from(main.clone());
292 loop {
293 self.sh.notified.store(false, SeqCst);
294 if main.woken.swap(false, SeqCst) {
295 if let Poll::Ready(v) = fut.as_mut().poll(&mut Context::from_waker(&waker)) {
296 return v;
297 }
298 }
299 self.run();
300 if main.woken.load(SeqCst) || self.has_ready() {
301 continue;
302 }
303 self.park();
304 }
305 }
306
307 fn park(&self) {
309 let ms = match next_timer() {
310 None => -1,
311 Some(at) => {
312 let now = now_us();
313 if at <= now {
314 return;
315 }
316 ((at - now).saturating_add(999) / 1000).min(i32::MAX as u64) as i32
317 }
318 };
319 match &self.sh.park {
320 Some(s) => {
321 s.wait(ms);
322 }
323 None => crate::thread_yield(),
324 }
325 }
326}
327
328impl Default for Executor {
329 fn default() -> Self {
330 Executor::new()
331 }
332}
333
334struct MainWaker {
335 sh: Arc<Shared>,
336 woken: AtomicBool,
337}
338
339impl Wake for MainWaker {
340 fn wake(self: Arc<Self>) {
341 self.wake_by_ref()
342 }
343 fn wake_by_ref(self: &Arc<Self>) {
344 self.woken.store(true, SeqCst);
345 if tid() != self.sh.owner || tid() == 0 {
346 self.sh.signal();
347 }
348 }
349}
350
351pub fn block_on<F: Future>(fut: F) -> F::Output {
353 Executor::new().block_on(fut)
354}
355
356#[derive(Clone)]
358pub struct Spawner {
359 q: Rc<RefCell<Vec<Boxed>>>,
360}
361
362impl Spawner {
363 pub fn spawn<F: Future + 'static>(&self, fut: F) -> Task<F::Output> {
365 let st = Arc::new(TaskState { flags: AtomicU32::new(0), out: SpinLock::new(None), waiter: SpinLock::new(None), me: SpinLock::new(None) });
366 let w = Wrap { fut: Some(Box::pin(fut)), st: st.clone() };
367 self.q.borrow_mut().push(Box::pin(w));
368 Task { st: Some(st) }
369 }
370}
371
372const T_DONE: u32 = 1;
375const T_CANCEL: u32 = 2;
376const T_DROPPED: u32 = 4; struct TaskState<T> {
379 flags: AtomicU32,
380 out: SpinLock<Option<T>>,
381 waiter: SpinLock<Option<Waker>>,
383 me: SpinLock<Option<Waker>>,
385}
386
387impl<T> TaskState<T> {
388 fn finish(&self, v: Option<T>) {
389 let bits = if v.is_some() { T_DONE } else { T_DONE | T_DROPPED };
390 *self.out.lock() = v;
391 self.flags.fetch_or(bits, SeqCst);
392 *self.me.lock() = None;
393 let w = self.waiter.lock().take();
394 if let Some(w) = w {
395 w.wake();
396 }
397 }
398
399 fn cancel(&self) {
400 if self.flags.load(SeqCst) & (T_DONE | T_CANCEL) != 0 || self.flags.fetch_or(T_CANCEL, SeqCst) & (T_DONE | T_CANCEL) != 0 {
401 return;
402 }
403 let me = self.me.lock().clone();
404 if let Some(w) = me {
405 w.wake();
406 }
407 }
408}
409
410struct Wrap<F: Future> {
412 fut: Option<Pin<Box<F>>>,
413 st: Arc<TaskState<F::Output>>,
414}
415
416impl<F: Future> Future for Wrap<F> {
417 type Output = ();
418 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
419 let this = &mut *self;
420 if this.st.flags.load(SeqCst) & T_CANCEL != 0 {
421 this.fut = None;
422 this.st.finish(None);
423 return Poll::Ready(());
424 }
425 let Some(f) = this.fut.as_mut() else { return Poll::Ready(()) };
426 {
427 let mut me = this.st.me.lock();
428 if !me.as_ref().is_some_and(|w| w.will_wake(cx.waker())) {
429 *me = Some(cx.waker().clone());
430 }
431 }
432 match f.as_mut().poll(cx) {
433 Poll::Ready(v) => {
434 this.fut = None;
435 this.st.finish(Some(v));
436 Poll::Ready(())
437 }
438 Poll::Pending => Poll::Pending,
439 }
440 }
441}
442
443impl<F: Future> Drop for Wrap<F> {
444 fn drop(&mut self) {
445 if self.fut.take().is_some() {
446 self.st.finish(None); }
448 }
449}
450
451pub struct Task<T> {
453 st: Option<Arc<TaskState<T>>>,
454}
455
456impl<T> Task<T> {
457 pub fn is_done(&self) -> bool {
460 self.st.as_ref().is_none_or(|s| s.flags.load(SeqCst) & T_DONE != 0)
461 }
462
463 pub fn is_cancelled(&self) -> bool {
465 self.st.as_ref().is_some_and(|s| s.flags.load(SeqCst) & T_DROPPED != 0)
466 }
467
468 pub fn take(&self) -> Option<T> {
470 let s = self.st.as_ref()?;
471 if s.flags.load(SeqCst) & T_DONE == 0 {
472 return None;
473 }
474 s.out.lock().take()
475 }
476
477 pub fn cancel(&self) {
480 if let Some(s) = &self.st {
481 s.cancel();
482 }
483 }
484
485 pub fn detach(mut self) {
487 self.st = None;
488 }
489}
490
491impl<T> Drop for Task<T> {
492 fn drop(&mut self) {
493 if let Some(s) = &self.st {
494 s.cancel();
495 }
496 }
497}
498
499impl<T> Future for Task<T> {
500 type Output = Result<T, Cancelled>;
501 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
502 let Some(s) = self.st.as_ref() else { return Poll::Ready(Err(Cancelled)) };
503 if s.flags.load(SeqCst) & T_DONE == 0 {
504 *s.waiter.lock() = Some(cx.waker().clone());
505 if s.flags.load(SeqCst) & T_DONE == 0 {
506 return Poll::Pending;
507 }
508 }
509 Poll::Ready(s.out.lock().take().ok_or(Cancelled))
510 }
511}
512
513struct Timer {
516 at: u64,
517 id: u64,
518 owner: u64,
519 waker: Waker,
520}
521
522static TIMERS: SpinLock<Vec<Timer>> = SpinLock::new(Vec::new());
523static TIMER_ID: AtomicU64 = AtomicU64::new(1);
524
525pub fn sleep(ms: u32) -> Sleep {
527 sleep_until(now_us().saturating_add(ms as u64 * 1000))
528}
529
530pub fn sleep_until(at_us: u64) -> Sleep {
532 Sleep { at: at_us, id: 0 }
533}
534
535pub struct Sleep {
537 at: u64,
538 id: u64,
539}
540
541impl Sleep {
542 pub fn deadline(&self) -> u64 {
544 self.at
545 }
546}
547
548impl Future for Sleep {
549 type Output = ();
550 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
551 if now_us() >= self.at {
552 if self.id != 0 {
553 let id = core::mem::replace(&mut self.id, 0);
554 TIMERS.lock().retain(|t| t.id != id);
555 }
556 return Poll::Ready(());
557 }
558 let mut g = TIMERS.lock();
559 let found = if self.id == 0 { None } else { g.iter_mut().find(|t| t.id == self.id) };
560 match found {
561 Some(t) => {
562 if !t.waker.will_wake(cx.waker()) {
563 t.waker = cx.waker().clone();
564 }
565 t.owner = tid();
566 }
567 None => {
568 if self.id == 0 {
569 self.id = TIMER_ID.fetch_add(1, SeqCst);
570 }
571 g.push(Timer { at: self.at, id: self.id, owner: tid(), waker: cx.waker().clone() });
572 }
573 }
574 Poll::Pending
575 }
576}
577
578impl Drop for Sleep {
579 fn drop(&mut self) {
580 if self.id != 0 {
581 TIMERS.lock().retain(|t| t.id != self.id);
582 }
583 }
584}
585
586pub fn fire_timers(now: u64) -> usize {
588 let me = tid();
589 let mut due = Vec::new();
590 {
591 let mut g = TIMERS.lock();
592 let mut i = 0;
593 while i < g.len() {
594 if g[i].owner == me && g[i].at <= now {
595 let t = g.remove(i);
596 due.push((t.at, t.id, t.waker));
597 } else {
598 i += 1;
599 }
600 }
601 }
602 due.sort_by_key(|d| (d.0, d.1)); let n = due.len();
604 for (_, _, w) in due {
605 w.wake();
606 }
607 n
608}
609
610pub fn next_timer() -> Option<u64> {
612 let me = tid();
613 TIMERS.lock().iter().filter(|t| t.owner == me).map(|t| t.at).min()
614}
615
616pub fn yield_now() -> YieldNow {
618 YieldNow(false)
619}
620
621pub struct YieldNow(bool);
623
624impl Future for YieldNow {
625 type Output = ();
626 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
627 if self.0 {
628 return Poll::Ready(());
629 }
630 self.0 = true;
631 cx.waker().wake_by_ref();
632 Poll::Pending
633 }
634}
635
636type BJob = Box<dyn FnOnce() + Send>;
639
640struct BPool {
641 jobs: VecDeque<BJob>,
642 sem: Option<Arc<crate::Sem>>,
643 threads: u32,
644 idle: u32,
645 no_threads: bool,
646}
647
648static BPOOL: SpinLock<BPool> = SpinLock::new(BPool { jobs: VecDeque::new(), sem: None, threads: 0, idle: 0, no_threads: false });
649static INFLIGHT: AtomicU32 = AtomicU32::new(0);
650
651pub fn blocking_pending() -> u32 {
653 INFLIGHT.load(SeqCst)
654}
655
656fn submit(job: BJob) {
657 let mut g = BPOOL.lock();
658 g.jobs.push_back(job);
659 if g.no_threads {
660 return; }
662 if g.sem.is_none() {
663 g.sem = crate::Sem::new(0).map(Arc::new);
664 }
665 let Some(sem) = g.sem.clone() else {
666 g.no_threads = true;
667 return;
668 };
669 let start = g.jobs.len() as u32 > g.idle && g.threads < BLOCKING_MAX;
670 if start {
671 g.threads += 1;
672 }
673 drop(g);
674 if start {
675 let s = sem.clone();
676 if crate::thread_spawn(BLOCKING_STACK, Box::new(move || blocking_worker(s))).is_none() {
678 let mut g = BPOOL.lock();
679 g.threads -= 1;
680 if g.threads == 0 {
681 g.no_threads = true;
682 }
683 }
684 }
685 sem.post(1);
686}
687
688fn blocking_worker(sem: Arc<crate::Sem>) {
689 loop {
690 let job = {
691 let mut g = BPOOL.lock();
692 let j = g.jobs.pop_front();
693 if j.is_none() {
694 g.idle += 1;
695 }
696 j
697 };
698 match job {
699 Some(j) => j(),
700 None => {
701 sem.wait(-1);
702 BPOOL.lock().idle -= 1;
703 }
704 }
705 }
706}
707
708fn inline_pending() -> bool {
709 let g = BPOOL.lock();
710 g.no_threads && !g.jobs.is_empty()
711}
712
713fn run_inline(n: usize) {
715 for _ in 0..n {
716 let job = {
717 let mut g = BPOOL.lock();
718 if !g.no_threads {
719 return;
720 }
721 g.jobs.pop_front()
722 };
723 match job {
724 Some(j) => j(),
725 None => return,
726 }
727 }
728}
729
730struct Cell<T> {
731 val: SpinLock<Option<T>>,
732 waker: SpinLock<Option<Waker>>,
733 done: AtomicBool,
734}
735
736pub struct Blocking<T> {
738 cell: Arc<Cell<T>>,
739 tok: CancelToken,
740 taken: bool,
741}
742
743pub fn spawn_blocking<T: Send + 'static>(f: impl FnOnce() -> T + Send + 'static) -> Blocking<T> {
745 spawn_cancellable(move |_| f())
746}
747
748pub fn spawn_cancellable<T: Send + 'static>(f: impl FnOnce(&CancelToken) -> T + Send + 'static) -> Blocking<T> {
751 let cell = Arc::new(Cell { val: SpinLock::new(None), waker: SpinLock::new(None), done: AtomicBool::new(false) });
752 let tok = CancelToken::new();
753 let (c, t) = (cell.clone(), tok.clone());
754 let trusted = crate::policy::trusted_here();
755 let scopes = crate::policy::scopes_here(); INFLIGHT.fetch_add(1, SeqCst);
757 submit(Box::new(move || {
758 if !t.is_cancelled() {
759 let run = || if trusted { crate::policy::trusted_scope(|| f(&t)) } else { f(&t) };
760 let v = crate::policy::in_scopes(&scopes, run);
761 *c.val.lock() = Some(v);
762 } else {
763 drop(f); }
765 c.done.store(true, SeqCst);
766 let w = c.waker.lock().take();
767 if let Some(w) = w {
768 w.wake();
769 }
770 INFLIGHT.fetch_sub(1, SeqCst); }));
772 Blocking { cell, tok, taken: false }
773}
774
775impl<T> Blocking<T> {
776 pub fn is_done(&self) -> bool {
778 self.cell.done.load(SeqCst)
779 }
780}
781
782impl<T> Future for Blocking<T> {
783 type Output = T;
784 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
785 if !self.cell.done.load(SeqCst) {
786 *self.cell.waker.lock() = Some(cx.waker().clone());
787 if !self.cell.done.load(SeqCst) {
788 return Poll::Pending;
789 }
790 }
791 let v = self.cell.val.lock().take();
792 match v {
793 Some(v) => {
794 self.taken = true;
795 Poll::Ready(v)
796 }
797 None => Poll::Pending, }
799 }
800}
801
802impl<T> Drop for Blocking<T> {
803 fn drop(&mut self) {
804 if !self.taken {
805 self.tok.cancel();
806 }
807 }
808}
809
810struct CancelInner {
813 flag: AtomicBool,
814 waiters: SpinLock<Vec<(u64, Waker)>>,
815}
816
817#[derive(Clone)]
819pub struct CancelToken(Arc<CancelInner>);
820
821static WAIT_ID: AtomicU64 = AtomicU64::new(1);
822
823impl CancelToken {
824 pub fn new() -> CancelToken {
826 CancelToken(Arc::new(CancelInner { flag: AtomicBool::new(false), waiters: SpinLock::new(Vec::new()) }))
827 }
828
829 pub fn cancel(&self) {
831 if self.0.flag.swap(true, SeqCst) {
832 return;
833 }
834 let w = core::mem::take(&mut *self.0.waiters.lock());
835 for (_, w) in w {
836 w.wake();
837 }
838 }
839
840 pub fn is_cancelled(&self) -> bool {
842 self.0.flag.load(SeqCst)
843 }
844
845 pub fn check(&self) -> Result<(), Cancelled> {
847 if self.is_cancelled() { Err(Cancelled) } else { Ok(()) }
848 }
849
850 pub fn cancelled(&self) -> WaitCancelled {
852 WaitCancelled { tok: self.clone(), id: 0 }
853 }
854
855 pub fn run<F: Future>(&self, fut: F) -> Guarded<F> {
857 Guarded { fut: Some(Box::pin(fut)), wait: self.cancelled() }
858 }
859}
860
861impl Default for CancelToken {
862 fn default() -> Self {
863 CancelToken::new()
864 }
865}
866
867pub struct WaitCancelled {
869 tok: CancelToken,
870 id: u64,
871}
872
873impl Future for WaitCancelled {
874 type Output = ();
875 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
876 if self.tok.is_cancelled() {
877 return Poll::Ready(());
878 }
879 if self.id == 0 {
880 self.id = WAIT_ID.fetch_add(1, SeqCst);
881 }
882 {
883 let mut g = self.tok.0.waiters.lock();
884 match g.iter_mut().find(|(i, _)| *i == self.id) {
885 Some((_, w)) => {
886 if !w.will_wake(cx.waker()) {
887 *w = cx.waker().clone();
888 }
889 }
890 None => g.push((self.id, cx.waker().clone())),
891 }
892 }
893 if self.tok.is_cancelled() { Poll::Ready(()) } else { Poll::Pending }
894 }
895}
896
897impl Drop for WaitCancelled {
898 fn drop(&mut self) {
899 if self.id != 0 {
900 let id = self.id;
901 self.tok.0.waiters.lock().retain(|(i, _)| *i != id);
902 }
903 }
904}
905
906pub struct Guarded<F: Future> {
908 fut: Option<Pin<Box<F>>>,
909 wait: WaitCancelled,
910}
911
912impl<F: Future> Unpin for Guarded<F> {}
913
914impl<F: Future> Future for Guarded<F> {
915 type Output = Result<F::Output, Cancelled>;
916 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
917 let this = &mut *self;
918 if this.wait.tok.is_cancelled() {
919 this.fut = None;
920 return Poll::Ready(Err(Cancelled));
921 }
922 let Some(f) = this.fut.as_mut() else { return Poll::Ready(Err(Cancelled)) };
923 if let Poll::Ready(v) = f.as_mut().poll(cx) {
924 this.fut = None;
925 return Poll::Ready(Ok(v));
926 }
927 if Pin::new(&mut this.wait).poll(cx).is_ready() {
928 this.fut = None;
929 return Poll::Ready(Err(Cancelled));
930 }
931 Poll::Pending
932 }
933}
934
935struct ProgInner {
938 done: AtomicU64,
939 total: AtomicU64,
940 seq: AtomicU64,
941 last_wake: AtomicU64,
942 waiters: SpinLock<Vec<Waker>>,
943}
944
945#[derive(Clone)]
947pub struct Progress(Arc<ProgInner>);
948
949impl Progress {
950 pub fn new() -> Progress {
952 Progress(Arc::new(ProgInner {
953 done: AtomicU64::new(0),
954 total: AtomicU64::new(0),
955 seq: AtomicU64::new(0),
956 last_wake: AtomicU64::new(0),
957 waiters: SpinLock::new(Vec::new()),
958 }))
959 }
960
961 pub fn set_total(&self, n: u64) {
963 self.0.total.store(n, SeqCst);
964 self.report();
965 }
966
967 pub fn set(&self, done: u64) {
969 self.0.done.store(done, SeqCst);
970 self.report();
971 }
972
973 pub fn add(&self, n: u64) {
975 self.0.done.fetch_add(n, SeqCst);
976 self.report();
977 }
978
979 pub fn done(&self) -> u64 {
981 self.0.done.load(SeqCst)
982 }
983
984 pub fn total(&self) -> u64 {
986 self.0.total.load(SeqCst)
987 }
988
989 pub fn permille(&self) -> u32 {
991 let (d, t) = (self.done(), self.total());
992 if t == 0 { 0 } else { (d.min(t) as u128 * 1000 / t as u128) as u32 }
993 }
994
995 pub fn reports(&self) -> u64 {
997 self.0.seq.load(SeqCst)
998 }
999
1000 pub fn changed(&self) -> Changed {
1002 Changed { p: self.clone(), seen: self.reports() }
1003 }
1004
1005 fn report(&self) {
1007 self.0.seq.fetch_add(1, SeqCst);
1008 let w = core::mem::take(&mut *self.0.waiters.lock());
1009 for w in w {
1010 w.wake();
1011 }
1012 let now = crate::time_us();
1013 let last = self.0.last_wake.load(SeqCst);
1014 let end = self.total() != 0 && self.done() >= self.total();
1015 if end || now.wrapping_sub(last) >= PROGRESS_WAKE_MS * 1000 {
1016 self.0.last_wake.store(now, SeqCst);
1017 crate::wake();
1018 }
1019 }
1020}
1021
1022impl Default for Progress {
1023 fn default() -> Self {
1024 Progress::new()
1025 }
1026}
1027
1028pub struct Changed {
1030 p: Progress,
1031 seen: u64,
1032}
1033
1034impl Future for Changed {
1035 type Output = ();
1036 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
1037 if self.p.reports() != self.seen {
1038 return Poll::Ready(());
1039 }
1040 {
1041 let mut g = self.p.0.waiters.lock();
1042 if !g.iter().any(|w| w.will_wake(cx.waker())) {
1043 g.push(cx.waker().clone());
1044 }
1045 }
1046 if self.p.reports() != self.seen { Poll::Ready(()) } else { Poll::Pending }
1047 }
1048}
1049
1050enum MaybeDone<F: Future> {
1053 Run(Pin<Box<F>>),
1054 Done(Option<F::Output>),
1055}
1056
1057impl<F: Future> MaybeDone<F> {
1058 fn new(f: F) -> Self {
1059 MaybeDone::Run(Box::pin(f))
1060 }
1061 fn poll(&mut self, cx: &mut Context<'_>) -> bool {
1063 if let MaybeDone::Run(f) = self {
1064 match f.as_mut().poll(cx) {
1065 Poll::Ready(v) => *self = MaybeDone::Done(Some(v)),
1066 Poll::Pending => return false,
1067 }
1068 }
1069 true
1070 }
1071 fn take(&mut self) -> F::Output {
1072 match self {
1073 MaybeDone::Done(v) => v.take().expect("join: polled after completion"),
1074 MaybeDone::Run(_) => unreachable!(),
1075 }
1076 }
1077}
1078
1079pub struct Join<A: Future, B: Future>(MaybeDone<A>, MaybeDone<B>);
1081impl<A: Future, B: Future> Unpin for Join<A, B> {}
1082
1083pub fn join<A: Future, B: Future>(a: A, b: B) -> Join<A, B> {
1085 Join(MaybeDone::new(a), MaybeDone::new(b))
1086}
1087
1088impl<A: Future, B: Future> Future for Join<A, B> {
1089 type Output = (A::Output, B::Output);
1090 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
1091 let this = &mut *self;
1092 let a = this.0.poll(cx);
1093 let b = this.1.poll(cx);
1094 if a && b { Poll::Ready((this.0.take(), this.1.take())) } else { Poll::Pending }
1095 }
1096}
1097
1098pub struct Join3<A: Future, B: Future, C: Future>(MaybeDone<A>, MaybeDone<B>, MaybeDone<C>);
1100impl<A: Future, B: Future, C: Future> Unpin for Join3<A, B, C> {}
1101
1102pub fn join3<A: Future, B: Future, C: Future>(a: A, b: B, c: C) -> Join3<A, B, C> {
1104 Join3(MaybeDone::new(a), MaybeDone::new(b), MaybeDone::new(c))
1105}
1106
1107impl<A: Future, B: Future, C: Future> Future for Join3<A, B, C> {
1108 type Output = (A::Output, B::Output, C::Output);
1109 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
1110 let this = &mut *self;
1111 let a = this.0.poll(cx);
1112 let b = this.1.poll(cx);
1113 let c = this.2.poll(cx);
1114 if a && b && c { Poll::Ready((this.0.take(), this.1.take(), this.2.take())) } else { Poll::Pending }
1115 }
1116}
1117
1118pub struct JoinAll<F: Future>(Vec<MaybeDone<F>>);
1120impl<F: Future> Unpin for JoinAll<F> {}
1121
1122pub fn join_all<F: Future>(v: impl IntoIterator<Item = F>) -> JoinAll<F> {
1124 JoinAll(v.into_iter().map(MaybeDone::new).collect())
1125}
1126
1127impl<F: Future> Future for JoinAll<F> {
1128 type Output = Vec<F::Output>;
1129 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
1130 let mut all = true;
1131 for f in self.0.iter_mut() {
1132 all &= f.poll(cx);
1133 }
1134 if all { Poll::Ready(self.0.iter_mut().map(|f| f.take()).collect()) } else { Poll::Pending }
1135 }
1136}
1137
1138#[derive(Debug, PartialEq, Eq)]
1140pub enum Either<A, B> {
1141 Left(A),
1142 Right(B),
1143}
1144
1145pub struct Select<A: Future, B: Future>(Pin<Box<A>>, Pin<Box<B>>);
1147
1148pub fn select<A: Future, B: Future>(a: A, b: B) -> Select<A, B> {
1150 Select(Box::pin(a), Box::pin(b))
1151}
1152
1153impl<A: Future, B: Future> Future for Select<A, B> {
1154 type Output = Either<A::Output, B::Output>;
1155 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
1156 if let Poll::Ready(v) = self.0.as_mut().poll(cx) {
1157 return Poll::Ready(Either::Left(v));
1158 }
1159 if let Poll::Ready(v) = self.1.as_mut().poll(cx) {
1160 return Poll::Ready(Either::Right(v));
1161 }
1162 Poll::Pending
1163 }
1164}