use std::future::Future;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
use std::sync::{Arc, Mutex, Weak};
use std::task::{Context, Poll, Wake, Waker};
use super::callback_executor::{Callback, CallbackHandle, CallbackPriority};
pub const DEFAULT_SPAWN_PRIORITY: CallbackPriority = CallbackPriority::Medium;
const IDLE: u8 = 0;
const SCHEDULED: u8 = 1;
const RUNNING: u8 = 2;
const RUNNING_NOTIFIED: u8 = 3;
const DONE: u8 = 4;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct JoinError {
kind: JoinErrorKind,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum JoinErrorKind {
Cancelled,
Panicked,
}
impl JoinError {
fn cancelled() -> Self {
JoinError {
kind: JoinErrorKind::Cancelled,
}
}
fn panicked() -> Self {
JoinError {
kind: JoinErrorKind::Panicked,
}
}
pub fn is_cancelled(&self) -> bool {
matches!(self.kind, JoinErrorKind::Cancelled)
}
}
impl std::fmt::Display for JoinError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self.kind {
JoinErrorKind::Cancelled => f.write_str("task was cancelled"),
JoinErrorKind::Panicked => f.write_str("task panicked"),
}
}
}
impl std::error::Error for JoinError {}
trait Schedulable: Send + Sync {
fn schedule(self: Arc<Self>);
}
struct Control {
abort: AtomicBool,
finished: AtomicBool,
task: Mutex<Option<Weak<dyn Schedulable>>>,
}
impl Control {
fn new() -> Arc<Self> {
Arc::new(Control {
abort: AtomicBool::new(false),
finished: AtomicBool::new(false),
task: Mutex::new(None),
})
}
}
fn request_abort(control: &Control) {
control.abort.store(true, Ordering::Release);
let task = control.task.lock().unwrap().clone();
if let Some(task) = task.and_then(|w| w.upgrade()) {
task.schedule();
}
}
struct Slot<T> {
result: Option<Result<T, JoinError>>,
finalized: bool,
join_waker: Option<Waker>,
}
struct Shared<T> {
control: Arc<Control>,
slot: Mutex<Slot<T>>,
}
impl<T> Shared<T> {
fn new() -> Arc<Self> {
Arc::new(Shared {
control: Control::new(),
slot: Mutex::new(Slot {
result: None,
finalized: false,
join_waker: None,
}),
})
}
fn finalize(&self, result: Result<T, JoinError>) {
let waker = {
let mut slot = self.slot.lock().unwrap();
if slot.finalized {
return;
}
slot.finalized = true;
slot.result = Some(result);
let waker = slot.join_waker.take();
self.control.finished.store(true, Ordering::Release);
waker
};
if let Some(w) = waker {
w.wake();
}
}
}
struct Task<T> {
state: AtomicU8,
future: Mutex<Option<Pin<Box<dyn Future<Output = T> + Send>>>>,
shared: Arc<Shared<T>>,
callbacks: CallbackHandle,
priority: CallbackPriority,
#[cfg(test)]
enqueues: std::sync::atomic::AtomicUsize,
}
impl<T: Send + 'static> Task<T> {
fn enqueue(self: &Arc<Self>) {
#[cfg(test)]
self.enqueues.fetch_add(1, Ordering::Relaxed);
let mut entry = Entry {
task: Some(Arc::clone(self)),
};
let cb: Callback = Box::new(move || entry.run());
if self.callbacks.request(self.priority, cb).is_err() {
self.state.store(DONE, Ordering::Release);
tracing::error!(
target: "epics_base_rs::runtime::future_exec",
"spawn_future: callback ring full; task dropped, handle resolves cancelled"
);
}
}
fn run(self: Arc<Self>) {
if self
.state
.compare_exchange(SCHEDULED, RUNNING, Ordering::AcqRel, Ordering::Acquire)
.is_err()
{
return;
}
let Some(mut fut) = self.future.lock().unwrap().take() else {
self.state.store(DONE, Ordering::Release);
return;
};
if self.shared.control.abort.load(Ordering::Acquire) {
drop(fut);
self.state.store(DONE, Ordering::Release);
self.shared.finalize(Err(JoinError::cancelled()));
return;
}
let waker = Waker::from(Arc::clone(&self));
let polled = catch_unwind(AssertUnwindSafe(|| {
let mut cx = Context::from_waker(&waker);
fut.as_mut().poll(&mut cx)
}));
match polled {
Ok(Poll::Ready(value)) => {
drop(fut);
self.state.store(DONE, Ordering::Release);
self.shared.finalize(Ok(value));
}
Err(_panic) => {
drop(fut);
self.state.store(DONE, Ordering::Release);
self.shared.finalize(Err(JoinError::panicked()));
}
Ok(Poll::Pending) => {
*self.future.lock().unwrap() = Some(fut);
if self
.state
.compare_exchange(RUNNING, IDLE, Ordering::AcqRel, Ordering::Acquire)
.is_err()
{
self.state.store(SCHEDULED, Ordering::Release);
self.enqueue();
}
}
}
}
}
impl<T: Send + 'static> Schedulable for Task<T> {
fn schedule(self: Arc<Self>) {
loop {
match self.state.load(Ordering::Acquire) {
IDLE => {
if self
.state
.compare_exchange(IDLE, SCHEDULED, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
{
self.enqueue();
return;
}
}
RUNNING => {
if self
.state
.compare_exchange(
RUNNING,
RUNNING_NOTIFIED,
Ordering::AcqRel,
Ordering::Acquire,
)
.is_ok()
{
return;
}
}
_ => return,
}
}
}
}
impl<T: Send + 'static> Wake for Task<T> {
fn wake(self: Arc<Self>) {
Schedulable::schedule(self);
}
fn wake_by_ref(self: &Arc<Self>) {
Schedulable::schedule(Arc::clone(self));
}
}
impl<T> Drop for Task<T> {
fn drop(&mut self) {
self.shared.finalize(Err(JoinError::cancelled()));
}
}
struct Entry<T> {
task: Option<Arc<Task<T>>>,
}
impl<T: Send + 'static> Entry<T> {
fn run(&mut self) {
if let Some(task) = self.task.take() {
task.run();
}
}
}
impl<T> Drop for Entry<T> {
fn drop(&mut self) {
if let Some(task) = self.task.take() {
task.state.store(DONE, Ordering::Release);
let _ = task.future.lock().unwrap().take();
task.shared.finalize(Err(JoinError::cancelled()));
}
}
}
pub struct JoinFuture<T> {
shared: Arc<Shared<T>>,
}
impl<T> JoinFuture<T> {
pub fn abort(&self) {
request_abort(&self.shared.control);
}
pub fn is_finished(&self) -> bool {
self.shared.control.finished.load(Ordering::Acquire)
}
pub fn abort_handle(&self) -> AbortHandle {
AbortHandle {
control: Arc::clone(&self.shared.control),
}
}
}
impl<T> Future for JoinFuture<T> {
type Output = Result<T, JoinError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut slot = self.shared.slot.lock().unwrap();
match slot.result.take() {
Some(result) => Poll::Ready(result),
None => {
slot.join_waker = Some(cx.waker().clone());
Poll::Pending
}
}
}
}
#[derive(Clone)]
pub struct AbortHandle {
control: Arc<Control>,
}
impl std::fmt::Debug for AbortHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AbortHandle")
.field("aborted", &self.control.abort.load(Ordering::Relaxed))
.field("finished", &self.control.finished.load(Ordering::Acquire))
.finish()
}
}
impl AbortHandle {
pub fn abort(&self) {
request_abort(&self.control);
}
pub fn is_finished(&self) -> bool {
self.control.finished.load(Ordering::Acquire)
}
}
pub fn spawn_future<F>(
callbacks: &CallbackHandle,
priority: CallbackPriority,
fut: F,
) -> JoinFuture<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
spawn_task(callbacks, priority, fut).0
}
fn spawn_task<F>(
callbacks: &CallbackHandle,
priority: CallbackPriority,
fut: F,
) -> (JoinFuture<F::Output>, Arc<Task<F::Output>>)
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
let shared = Shared::new();
let task = Arc::new(Task {
state: AtomicU8::new(SCHEDULED),
future: Mutex::new(Some(Box::pin(fut))),
shared: Arc::clone(&shared),
callbacks: callbacks.clone(),
priority,
#[cfg(test)]
enqueues: std::sync::atomic::AtomicUsize::new(0),
});
*shared.control.task.lock().unwrap() = Some(Arc::downgrade(&task) as Weak<dyn Schedulable>);
task.enqueue();
(JoinFuture { shared }, task)
}
pub fn spawn_blocking_on<F, R>(
callbacks: &CallbackHandle,
priority: CallbackPriority,
f: F,
) -> JoinFuture<R>
where
F: FnOnce() -> R + Send + 'static,
R: Send + 'static,
{
let shared = Shared::new();
let task_shared = Arc::clone(&shared);
let mut guard = FinalizeOnDrop {
shared: Some(Arc::clone(&shared)),
};
let callback: Callback = Box::new(move || {
guard.defuse();
if task_shared.control.abort.load(Ordering::Acquire) {
task_shared.finalize(Err(JoinError::cancelled()));
return;
}
let outcome = catch_unwind(AssertUnwindSafe(f));
let result = match outcome {
Ok(value) => Ok(value),
Err(_panic) => Err(JoinError::panicked()),
};
task_shared.finalize(result);
});
if callbacks.request(priority, callback).is_err() {
tracing::error!(
target: "epics_base_rs::runtime::future_exec",
"spawn_blocking_on: callback ring full; closure dropped, handle resolves cancelled"
);
}
JoinFuture { shared }
}
struct FinalizeOnDrop<R> {
shared: Option<Arc<Shared<R>>>,
}
impl<R> FinalizeOnDrop<R> {
fn defuse(&mut self) {
self.shared = None;
}
}
impl<R> Drop for FinalizeOnDrop<R> {
fn drop(&mut self) {
if let Some(shared) = self.shared.take() {
shared.finalize(Err(JoinError::cancelled()));
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::runtime::background::callback_executor::{
CallbackPool, DEFAULT_QUEUE_SIZE, DEFAULT_THREADS_PER_PRIORITY,
};
use crate::runtime::background::delayed_timer::DelayedTimer;
use crate::runtime::background::timer_sleep::sleep;
use crate::runtime::task::park_on_interruptible as drive;
use std::sync::mpsc;
use std::time::Duration;
const T: Duration = Duration::from_secs(5);
fn join<T>(jf: JoinFuture<T>) -> Result<T, JoinError> {
drive(jf, || false).expect("uncancelled join returned None")
}
fn jf_reborrow<T>(jf: &JoinFuture<T>) -> JoinFuture<T> {
JoinFuture {
shared: Arc::clone(&jf.shared),
}
}
struct YieldN(usize);
impl Future for YieldN {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
if self.0 == 0 {
return Poll::Ready(());
}
self.0 -= 1;
cx.waker().wake_by_ref();
Poll::Pending
}
}
#[test]
fn future_runs_to_completion() {
let pool = CallbackPool::new();
let jf = spawn_future(&pool.handle(), DEFAULT_SPAWN_PRIORITY, async { 42u32 });
assert_eq!(join(jf).unwrap(), 42);
}
#[test]
fn future_awaiting_cross_thread_primitive_completes() {
let pool = CallbackPool::new();
let (tx, rx) = tokio::sync::oneshot::channel::<u32>();
let jf = spawn_future(&pool.handle(), DEFAULT_SPAWN_PRIORITY, async move {
rx.await.unwrap()
});
std::thread::sleep(Duration::from_millis(20));
tx.send(7).unwrap();
assert_eq!(join(jf).unwrap(), 7);
}
#[test]
fn panic_in_task_does_not_kill_the_worker() {
let pool = CallbackPool::new();
let jf = spawn_future(&pool.handle(), CallbackPriority::Medium, async {
panic!("boom");
});
let err = join(jf).unwrap_err();
assert!(!err.is_cancelled(), "a panic is not a cancellation");
let jf2 = spawn_future(&pool.handle(), CallbackPriority::Medium, async { 99u32 });
assert_eq!(join(jf2).unwrap(), 99);
}
#[test]
fn panic_after_a_yield_is_isolated_too() {
let pool = CallbackPool::new();
let jf = spawn_future(&pool.handle(), CallbackPriority::Medium, async {
YieldN(3).await;
panic!("late boom");
});
assert!(!join::<()>(jf).unwrap_err().is_cancelled());
let jf2 = spawn_future(&pool.handle(), CallbackPriority::Medium, async { 7u32 });
assert_eq!(join(jf2).unwrap(), 7);
}
#[test]
fn synchronously_ready_future_is_enqueued_exactly_once() {
let pool = CallbackPool::new();
let (jf, task) = spawn_task(&pool.handle(), CallbackPriority::Medium, async { 1u8 });
assert_eq!(join(jf).unwrap(), 1);
assert_eq!(
task.enqueues.load(Ordering::Relaxed),
1,
"a future that completes on its first poll must not round-trip the ring"
);
}
#[test]
fn each_suspension_costs_exactly_one_re_enqueue() {
let pool = CallbackPool::new();
let (jf, task) = spawn_task(&pool.handle(), CallbackPriority::Medium, YieldN(3));
join(jf).unwrap();
assert_eq!(
task.enqueues.load(Ordering::Relaxed),
4,
"expected 1 spawn + 3 wakes"
);
}
#[test]
fn a_yielding_task_releases_the_worker_to_a_queued_task() {
assert_eq!(DEFAULT_THREADS_PER_PRIORITY, 1);
let pool = CallbackPool::new();
let (tx, rx) = mpsc::channel::<&'static str>();
let tx_slow = tx.clone();
let slow = spawn_future(&pool.handle(), CallbackPriority::Medium, async move {
YieldN(20).await;
tx_slow.send("slow").unwrap();
});
let quick = spawn_future(&pool.handle(), CallbackPriority::Medium, async move {
tx.send("quick").unwrap();
});
assert_eq!(
rx.recv_timeout(T).unwrap(),
"quick",
"a suspended task must not hold the band's only worker"
);
assert_eq!(
rx.recv_timeout(T).unwrap(),
"slow",
"a repeatedly re-enqueued task must still make progress"
);
join(quick).unwrap();
join(slow).unwrap();
}
#[test]
fn many_self_waking_tasks_all_finish_on_one_worker() {
assert_eq!(DEFAULT_THREADS_PER_PRIORITY, 1);
let pool = CallbackPool::new();
let handles: Vec<_> = (0..8)
.map(|i| {
spawn_future(&pool.handle(), CallbackPriority::Medium, async move {
YieldN(25).await;
i
})
})
.collect();
for (i, jf) in handles.into_iter().enumerate() {
assert_eq!(join(jf).unwrap(), i);
}
}
#[test]
fn spawned_future_awaiting_sleep_does_not_self_deadlock() {
let pool = CallbackPool::new();
let timer = DelayedTimer::new(pool.handle());
let th = timer.handle();
let (done_tx, done_rx) = mpsc::channel::<()>();
let jf = spawn_future(&pool.handle(), CallbackPriority::Medium, async move {
sleep(&th, Duration::from_millis(50)).await;
done_tx.send(()).unwrap();
});
done_rx.recv_timeout(T).expect(
"spawned future awaiting sleep deadlocked (wake starved behind its own worker)",
);
assert!(join(jf).is_ok());
}
#[test]
fn a_sleep_wake_needs_no_pool_worker() {
assert_eq!(DEFAULT_THREADS_PER_PRIORITY, 1);
let pool = CallbackPool::new();
let timer = DelayedTimer::new(pool.handle());
let (pinned_tx, pinned_rx) = mpsc::channel::<()>();
let (release_tx, release_rx) = mpsc::channel::<()>();
pool.request(
CallbackPriority::Medium,
Box::new(move || {
pinned_tx.send(()).unwrap();
release_rx.recv().unwrap();
}),
)
.unwrap();
pinned_rx.recv_timeout(T).unwrap();
let (woke_tx, woke_rx) = mpsc::channel::<()>();
let th = timer.handle();
let sleeper = std::thread::spawn(move || {
drive(sleep(&th, Duration::from_millis(30)), || false).unwrap();
woke_tx.send(()).unwrap();
});
woke_rx
.recv_timeout(T)
.expect("sleep wake did not arrive while the band's worker was pinned");
release_tx.send(()).unwrap();
sleeper.join().unwrap();
}
#[test]
fn three_sleeping_tails_share_one_worker() {
assert_eq!(DEFAULT_THREADS_PER_PRIORITY, 1);
let pool = CallbackPool::with_config(DEFAULT_QUEUE_SIZE, DEFAULT_THREADS_PER_PRIORITY);
let timer = DelayedTimer::new(pool.handle());
#[derive(Debug, PartialEq, Eq)]
enum Ev {
Started,
Done(u8),
}
let (tx, rx) = mpsc::channel::<Ev>();
let handles: Vec<_> = (0..3u8)
.map(|i| {
let th = timer.handle();
let tx = tx.clone();
spawn_future(&pool.handle(), CallbackPriority::Medium, async move {
tx.send(Ev::Started).unwrap();
sleep(&th, Duration::from_millis(120)).await;
tx.send(Ev::Done(i)).unwrap();
i
})
})
.collect();
drop(tx);
for n in 0..3 {
assert_eq!(
rx.recv_timeout(T).unwrap(),
Ev::Started,
"task {n} had not started before the first task finished — the \
worker was held across a suspension"
);
}
let mut done: Vec<u8> = (0..3)
.map(|_| match rx.recv_timeout(T).unwrap() {
Ev::Done(i) => i,
Ev::Started => panic!("only three tasks exist"),
})
.collect();
done.sort_unstable();
assert_eq!(done, vec![0, 1, 2], "all three tails must complete");
for (i, jf) in handles.into_iter().enumerate() {
assert_eq!(join(jf).unwrap() as usize, i);
}
}
#[test]
fn abort_while_idle_cancels_cleanly() {
let pool = CallbackPool::new();
let (_tx, rx) = tokio::sync::oneshot::channel::<()>();
let (ran_tx, ran_rx) = mpsc::channel();
let jf = spawn_future(&pool.handle(), CallbackPriority::Medium, async move {
ran_tx.send(()).unwrap();
let _ = rx.await;
});
ran_rx.recv_timeout(T).unwrap();
std::thread::sleep(Duration::from_millis(20));
jf.abort();
let err = join(jf_reborrow(&jf)).unwrap_err();
assert!(err.is_cancelled(), "aborted task must report cancelled");
assert!(jf.is_finished());
}
#[test]
fn abort_before_the_first_poll_cancels_without_running() {
assert_eq!(DEFAULT_THREADS_PER_PRIORITY, 1);
let pool = CallbackPool::new();
let (pinned_tx, pinned_rx) = mpsc::channel::<()>();
let (release_tx, release_rx) = mpsc::channel::<()>();
pool.request(
CallbackPriority::Medium,
Box::new(move || {
pinned_tx.send(()).unwrap();
release_rx.recv().unwrap();
}),
)
.unwrap();
pinned_rx.recv_timeout(T).unwrap();
let ran = Arc::new(AtomicBool::new(false));
let flag = Arc::clone(&ran);
let jf = spawn_future(&pool.handle(), CallbackPriority::Medium, async move {
flag.store(true, Ordering::SeqCst);
});
jf.abort();
release_tx.send(()).unwrap();
assert!(join(jf_reborrow(&jf)).unwrap_err().is_cancelled());
assert!(
!ran.load(Ordering::SeqCst),
"an abort observed before the first poll must not run the future"
);
}
#[test]
fn abort_during_a_poll_is_observed_at_the_next_one() {
let pool = CallbackPool::new();
let (in_poll_tx, in_poll_rx) = mpsc::channel::<()>();
let (go_tx, go_rx) = mpsc::channel::<()>();
let polls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let counter = Arc::clone(&polls);
let jf = spawn_future(&pool.handle(), CallbackPriority::Medium, async move {
counter.fetch_add(1, Ordering::SeqCst);
in_poll_tx.send(()).unwrap();
go_rx.recv().unwrap();
YieldN(1).await; counter.fetch_add(100, Ordering::SeqCst);
});
in_poll_rx.recv_timeout(T).unwrap();
jf.abort(); go_tx.send(()).unwrap();
assert!(join(jf_reborrow(&jf)).unwrap_err().is_cancelled());
assert_eq!(
polls.load(Ordering::SeqCst),
1,
"the code after the suspension point must not have run"
);
}
#[test]
fn abort_handle_cancels_the_task() {
let pool = CallbackPool::new();
let (_tx, rx) = tokio::sync::oneshot::channel::<()>();
let jf = spawn_future(&pool.handle(), CallbackPriority::Medium, async move {
let _ = rx.await;
});
let ah = jf.abort_handle();
std::thread::sleep(Duration::from_millis(20));
assert!(!ah.is_finished());
ah.abort();
assert!(join(jf).unwrap_err().is_cancelled());
assert!(ah.is_finished());
}
#[test]
fn abort_after_completion_does_not_rewrite_the_outcome() {
let pool = CallbackPool::new();
let jf = spawn_future(&pool.handle(), CallbackPriority::Medium, async { 5u32 });
while !jf.is_finished() {
std::thread::sleep(Duration::from_millis(2));
}
jf.abort();
assert_eq!(join(jf).unwrap(), 5);
}
#[test]
fn full_ring_resolves_the_handle_as_cancelled() {
let pool = CallbackPool::with_config(1, 1);
let (pinned_tx, pinned_rx) = mpsc::channel::<()>();
let (release_tx, release_rx) = mpsc::channel::<()>();
pool.request(
CallbackPriority::Medium,
Box::new(move || {
pinned_tx.send(()).unwrap();
release_rx.recv().unwrap();
}),
)
.unwrap();
pinned_rx.recv_timeout(T).unwrap();
pool.request(CallbackPriority::Medium, Box::new(|| {}))
.unwrap();
let jf = spawn_future(&pool.handle(), CallbackPriority::Medium, async { 1u32 });
assert!(
join(jf).unwrap_err().is_cancelled(),
"a rejected spawn must resolve its handle, not strand it"
);
release_tx.send(()).unwrap();
}
#[test]
fn spawn_blocking_returns_value_and_isolates_panic() {
let pool = CallbackPool::new();
let jf = spawn_blocking_on(&pool.handle(), CallbackPriority::Medium, || 123u32);
assert_eq!(join(jf).unwrap(), 123);
let jf = spawn_blocking_on(&pool.handle(), CallbackPriority::Medium, || {
panic!("blocking boom")
});
assert!(!join::<()>(jf).unwrap_err().is_cancelled());
let jf = spawn_blocking_on(&pool.handle(), CallbackPriority::Medium, || 5u32);
assert_eq!(join(jf).unwrap(), 5);
}
}