use std::{
cell::Cell,
future::Future,
pin::Pin,
ptr,
task::{Context, Poll, RawWaker, RawWakerVTable, Waker},
thread,
time::Duration,
};
use crate::Error;
const NOOP_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
|_| NOOP_RAW_WAKER,
|_| {},
|_| {},
|_| {},
);
const NOOP_RAW_WAKER: RawWaker = RawWaker::new(ptr::null(), &NOOP_WAKER_VTABLE);
#[derive(Default)]
pub(crate) struct YieldOnce(bool);
impl Future for YieldOnce {
type Output = ();
fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
let flag = &mut std::pin::Pin::into_inner(self).0;
if !*flag {
*flag = true;
Poll::Pending
} else {
Poll::Ready(())
}
}
}
pub(crate) fn block_on<Fut1, Fut2>(
mut fut1: Fut1,
mut fut2: Fut2,
has_made_progress: &Cell<bool>,
) -> Result<(), Error>
where
Fut1: Future<Output = Result<(), Error>>,
Fut2: Future<Output = Result<(), Error>>,
{
let mut fut1 = Some(unsafe { Pin::new_unchecked(&mut fut1) });
let mut fut2 = Some(unsafe { Pin::new_unchecked(&mut fut2) });
let waker = unsafe { Waker::from_raw(NOOP_RAW_WAKER) };
let mut context = Context::from_waker(&waker);
let mut backoff_cnt = 0;
loop {
has_made_progress.set(false);
if let Some(fut) = fut2.as_mut() {
if let Poll::Ready(res) = fut.as_mut().poll(&mut context) {
fut2 = None;
res?;
}
}
if let Some(fut) = fut1.as_mut() {
if let Poll::Ready(res) = fut.as_mut().poll(&mut context) {
fut1 = None;
res?;
}
}
if fut1.is_none() && fut2.is_none() {
return Ok(());
}
if !has_made_progress.get() {
if backoff_cnt > 3 {
let duration = Duration::from_millis(100 * (backoff_cnt - 3).min(10));
thread::sleep(duration);
} else {
thread::yield_now();
}
}
backoff_cnt = if has_made_progress.get() {
0
} else {
backoff_cnt + 1
};
}
}