use std::cell::Cell;
use std::future::Future;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
use std::thread;
use std::time::{Duration, Instant};
use async_lock::OnceCell;
use futures_lite::pin;
use waker_fn::waker_fn;
use crate::reactor::Reactor;
static BLOCK_ON_COUNT: AtomicUsize = AtomicUsize::new(0);
fn unparker() -> &'static parking::Unparker {
static UNPARKER: OnceCell<parking::Unparker> = OnceCell::new();
UNPARKER.get_or_init_blocking(|| {
let (parker, unparker) = parking::pair();
thread::Builder::new()
.name("async-io".to_string())
.spawn(move || main_loop(parker))
.expect("cannot spawn async-io thread");
unparker
})
}
pub(crate) fn init() {
let _ = unparker();
}
fn main_loop(parker: parking::Parker) {
let mut last_tick = 0;
let mut sleeps = 0u64;
loop {
let tick = Reactor::get().ticker();
if last_tick == tick {
let reactor_lock = if sleeps >= 10 {
Some(Reactor::get().lock())
} else {
Reactor::get().try_lock()
};
if let Some(mut reactor_lock) = reactor_lock {
log::trace!("main_loop: waiting on I/O");
reactor_lock.react(None).ok();
last_tick = Reactor::get().ticker();
sleeps = 0;
}
} else {
last_tick = tick;
}
if BLOCK_ON_COUNT.load(Ordering::SeqCst) > 0 {
let delay_us = [50, 75, 100, 250, 500, 750, 1000, 2500, 5000]
.get(sleeps as usize)
.unwrap_or(&10_000);
log::trace!("main_loop: sleeping for {} us", delay_us);
if parker.park_timeout(Duration::from_micros(*delay_us)) {
log::trace!("main_loop: notified");
last_tick = Reactor::get().ticker();
sleeps = 0;
} else {
sleeps += 1;
}
}
}
}
pub fn block_on<T>(future: impl Future<Output = T>) -> T {
log::trace!("block_on()");
BLOCK_ON_COUNT.fetch_add(1, Ordering::SeqCst);
let _guard = CallOnDrop(|| {
BLOCK_ON_COUNT.fetch_sub(1, Ordering::SeqCst);
unparker().unpark();
});
let (p, u) = parking::pair();
let io_blocked = Arc::new(AtomicBool::new(false));
thread_local! {
static IO_POLLING: Cell<bool> = Cell::new(false);
}
let waker = waker_fn({
let io_blocked = io_blocked.clone();
move || {
if u.unpark() {
if !IO_POLLING.with(Cell::get) && io_blocked.load(Ordering::SeqCst) {
Reactor::get().notify();
}
}
}
});
let cx = &mut Context::from_waker(&waker);
pin!(future);
loop {
if let Poll::Ready(t) = future.as_mut().poll(cx) {
log::trace!("block_on: completed");
return t;
}
if p.park_timeout(Duration::from_secs(0)) {
log::trace!("block_on: notified");
if let Some(mut reactor_lock) = Reactor::get().try_lock() {
IO_POLLING.with(|io| io.set(true));
let _guard = CallOnDrop(|| {
IO_POLLING.with(|io| io.set(false));
});
reactor_lock.react(Some(Duration::from_secs(0))).ok();
}
continue;
}
if let Some(mut reactor_lock) = Reactor::get().try_lock() {
let start = Instant::now();
loop {
IO_POLLING.with(|io| io.set(true));
io_blocked.store(true, Ordering::SeqCst);
let _guard = CallOnDrop(|| {
IO_POLLING.with(|io| io.set(false));
io_blocked.store(false, Ordering::SeqCst);
});
if p.park_timeout(Duration::from_secs(0)) {
log::trace!("block_on: notified");
break;
}
log::trace!("block_on: waiting on I/O");
reactor_lock.react(None).ok();
if p.park_timeout(Duration::from_secs(0)) {
log::trace!("block_on: notified");
break;
}
if start.elapsed() > Duration::from_micros(500) {
log::trace!("block_on: stops hogging the reactor");
drop(reactor_lock);
unparker().unpark();
p.park();
break;
}
}
} else {
log::trace!("block_on: sleep until notification");
p.park();
}
}
}
struct CallOnDrop<F: Fn()>(F);
impl<F: Fn()> Drop for CallOnDrop<F> {
fn drop(&mut self) {
(self.0)();
}
}