use std::cell::RefCell;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Wake, Waker};
type LocalTask = Pin<Box<dyn Future<Output = ()>>>;
type WakeHook = Arc<dyn Fn() + Send + Sync>;
static WAKE_HOOK: Mutex<Option<WakeHook>> = Mutex::new(None);
thread_local! {
static QUEUE: RefCell<Vec<(LocalTask, Arc<TaskWaker>)>> = const { RefCell::new(Vec::new()) };
}
pub fn set_wake(hook: impl Fn() + Send + Sync + 'static) {
*WAKE_HOOK.lock().unwrap_or_else(|p| p.into_inner()) = Some(Arc::new(hook));
}
pub fn request_wake() {
let hook = WAKE_HOOK.lock().unwrap_or_else(|p| p.into_inner()).clone();
if let Some(hook) = hook {
hook();
}
}
struct TaskWaker {
ready: AtomicBool,
thread: Option<std::thread::Thread>,
}
impl Wake for TaskWaker {
fn wake(self: Arc<Self>) {
self.wake_by_ref();
}
fn wake_by_ref(self: &Arc<Self>) {
self.ready.store(true, Ordering::Release);
if let Some(thread) = &self.thread {
thread.unpark();
}
request_wake();
}
}
pub fn spawn(task: impl Future<Output = ()> + 'static) {
let mut task: LocalTask = Box::pin(task);
let waker = Arc::new(TaskWaker { ready: AtomicBool::new(false), thread: None });
if poll_once(&mut task, &waker).is_pending() {
QUEUE.with(|queue| queue.borrow_mut().push((task, waker)));
}
}
fn poll_once(task: &mut LocalTask, waker: &Arc<TaskWaker>) -> Poll<()> {
let waker_handle = Waker::from(waker.clone());
let mut cx = Context::from_waker(&waker_handle);
task.as_mut().poll(&mut cx)
}
pub fn run_pending() -> usize {
loop {
let tasks = QUEUE.with(|queue| std::mem::take(&mut *queue.borrow_mut()));
let mut progressed = false;
let mut still = Vec::with_capacity(tasks.len());
for (mut task, waker) in tasks {
if waker.ready.swap(false, Ordering::AcqRel) {
progressed = true;
if poll_once(&mut task, &waker).is_ready() {
continue;
}
}
still.push((task, waker));
}
let remaining = QUEUE.with(|queue| {
let mut queue = queue.borrow_mut();
still.append(&mut queue);
*queue = still;
queue.len()
});
if !progressed {
return remaining;
}
}
}
pub fn pending() -> usize {
QUEUE.with(|queue| queue.borrow().len())
}
pub fn block_on<T>(future: impl Future<Output = T>) -> T {
match wait(future, None) {
Some(value) => value,
None => unreachable!("a wait with no deadline only returns with the value"),
}
}
pub fn block_on_for<T>(future: impl Future<Output = T>, limit: std::time::Duration) -> Option<T> {
wait(future, Some(std::time::Instant::now() + limit))
}
fn wait<T>(future: impl Future<Output = T>, deadline: Option<std::time::Instant>) -> Option<T> {
let mut future = std::pin::pin!(future);
let waker = Arc::new(TaskWaker { ready: AtomicBool::new(true), thread: Some(std::thread::current()) });
let waker_handle = Waker::from(waker.clone());
let mut cx = Context::from_waker(&waker_handle);
loop {
if waker.ready.swap(false, Ordering::AcqRel) {
if let Poll::Ready(value) = future.as_mut().poll(&mut cx) {
return Some(value);
}
}
run_pending();
if deadline.is_some_and(|deadline| std::time::Instant::now() >= deadline) {
return None;
}
if !waker.ready.load(Ordering::Acquire) {
std::thread::park_timeout(std::time::Duration::from_millis(5));
}
}
}