sequoia_async/

runtime.rs

#![forbid(unsafe_code)]

use std::any::Any;
use std::cell::{Cell, RefCell};
use std::collections::HashMap;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::rc::Rc;
use std::sync::mpsc::{Receiver, Sender, TryRecvError, channel};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use futures_core::Future;
use sequoia_reactive::Signal;

use crate::emitter::{Emitter, StreamKind, StreamMsgAny};
use crate::task::{TaskHandle, TaskId};
use crate::wake::{WakeFn, noop_wake};

enum Completion {
    Done(u64, Box<dyn Any + Send>),
    Panicked(u64),
}

type Apply = Box<dyn FnOnce(Box<dyn Any + Send>)>;
type StreamApply = std::rc::Rc<dyn Fn(Box<dyn Any + Send>)>;
type Job = Box<dyn FnOnce() + Send + 'static>;

const THREAD_NAME: &str = "sequoia-async";

#[derive(Clone)]
pub struct AsyncRuntime {
    inner: Rc<Inner>,
}

struct Inner {
    pool: Arc<Pool>,
    done: Receiver<Completion>,
    applies: RefCell<HashMap<u64, Apply>>,
    streams: RefCell<HashMap<u64, StreamApply>>,
    stream_tx: Sender<StreamMsgAny>,
    stream_rx: Receiver<StreamMsgAny>,
    pending: Signal<u64>,
    next_id: Cell<u64>,
    wake: WakeFn,
}

impl AsyncRuntime {
    pub fn new(threads: usize) -> Self {
        Self::with_wake(threads, noop_wake())
    }

    pub fn with_default_threads() -> Self {
        Self::new(default_threads())
    }

    pub fn with_wake(threads: usize, wake: WakeFn) -> Self {
        let (tx, rx) = channel::<Completion>();
        let (stx, srx) = channel::<StreamMsgAny>();
        Self {
            inner: Rc::new(Inner {
                pool: Pool::new(threads, tx),
                done: rx,
                applies: RefCell::new(HashMap::new()),
                streams: RefCell::new(HashMap::new()),
                stream_tx: stx,
                stream_rx: srx,
                pending: Signal::new(0),
                next_id: Cell::new(1),
                wake,
            }),
        }
    }

    fn reserve(&self) -> u64 {
        let id = self.inner.next_id.get();
        self.inner.next_id.set(id + 1);
        self.inner.pending.update(|n| *n += 1);
        id
    }

    pub fn spawn_future<T, F, A>(&self, future: F, apply: A) -> TaskHandle
    where
        T: Send + 'static,
        F: Future<Output = T> + Send + 'static,
        A: FnOnce(T) + 'static,
    {
        let id = self.reserve();
        self.inner.applies.borrow_mut().insert(
            id,
            Box::new(move |any: Box<dyn Any + Send>| {
                if let Ok(value) = any.downcast::<T>() {
                    apply(*value);
                }
            }),
        );
        let tx = self.inner.pool.results_tx();
        let wake = self.inner.wake.clone();
        self.inner.pool.submit(Box::new(move || {
            let msg = match catch_unwind(AssertUnwindSafe(|| pollster::block_on(future))) {
                Ok(value) => Completion::Done(id, Box::new(value)),
                Err(_) => Completion::Panicked(id),
            };
            let _ = tx.send(msg);
            wake();
        }));
        let rt = self.clone();
        TaskHandle::new(TaskId::new(id), move |tid| rt.cancel_raw(tid.get()))
    }

    pub fn spawn_task<F>(&self, future: F) -> TaskHandle
    where
        F: Future<Output = ()> + Send + 'static,
    {
        self.spawn_future(future, |()| {})
    }

    pub fn spawn_blocking<T, W, A>(&self, work: W, apply: A) -> TaskHandle
    where
        T: Send + 'static,
        W: FnOnce() -> T + Send + 'static,
        A: FnOnce(T) + 'static,
    {
        let id = self.reserve();
        self.inner.applies.borrow_mut().insert(
            id,
            Box::new(move |any: Box<dyn Any + Send>| {
                if let Ok(value) = any.downcast::<T>() {
                    apply(*value);
                }
            }),
        );
        let tx = self.inner.pool.results_tx();
        let wake = self.inner.wake.clone();
        self.inner.pool.submit(Box::new(move || {
            let msg = match catch_unwind(AssertUnwindSafe(work)) {
                Ok(value) => Completion::Done(id, Box::new(value)),
                Err(_) => Completion::Panicked(id),
            };
            let _ = tx.send(msg);
            wake();
        }));
        let rt = self.clone();
        TaskHandle::new(TaskId::new(id), move |tid| rt.cancel_raw(tid.get()))
    }

    fn cancel_raw(&self, id: u64) -> bool {
        let mut hit = false;
        if self.inner.applies.borrow_mut().remove(&id).is_some() {
            hit = true;
        }
        if self.inner.streams.borrow_mut().remove(&id).is_some() {
            hit = true;
        }
        if hit {
            self.inner.pending.update(|n| *n = n.saturating_sub(1));
        }
        hit
    }

    pub fn cancel(&self, id: TaskId) -> bool {
        self.cancel_raw(id.get())
    }

    pub fn spawn_stream<T, F, Make, C>(&self, make_future: Make, on_item: C) -> TaskHandle
    where
        T: Send + 'static,
        F: Future<Output = ()> + Send + 'static,
        Make: FnOnce(Emitter<T>) -> F + Send + 'static,
        C: Fn(T) + 'static,
    {
        let id = self.reserve();
        let apply: StreamApply = std::rc::Rc::new(move |any: Box<dyn Any + Send>| {
            if let Ok(value) = any.downcast::<T>() {
                on_item(*value);
            }
        });
        self.inner.streams.borrow_mut().insert(id, apply);
        let stx = self.inner.stream_tx.clone();
        let wake = self.inner.wake.clone();
        let emitter = Emitter::new(id, stx.clone(), wake.clone());
        let future = make_future(emitter);
        self.inner.pool.submit(Box::new(move || {
            let outcome = catch_unwind(AssertUnwindSafe(|| pollster::block_on(future)));
            let kind = match outcome {
                Ok(()) => StreamKind::Finished,
                Err(_) => StreamKind::Panicked,
            };
            let _ = stx.send(StreamMsgAny { id, kind });
            wake();
        }));
        let rt = self.clone();
        TaskHandle::new(TaskId::new(id), move |tid| rt.cancel_raw(tid.get()))
    }

    pub fn spawn_to_signal<T, F>(&self, signal: &Signal<T>, future: F) -> TaskHandle
    where
        T: Clone + PartialEq + Send + 'static,
        F: Future<Output = T> + Send + 'static,
    {
        let signal = signal.clone();
        self.spawn_future(future, move |v| signal.set(v))
    }

    pub fn pending(&self) -> Signal<u64> {
        self.inner.pending.clone()
    }

    pub fn is_busy(&self) -> bool {
        self.inner.pending.get_untracked() > 0
    }

    pub fn threads(&self) -> usize {
        self.inner.pool.threads()
    }

    pub fn drain(&self) -> usize {
        let mut handled = 0;
        loop {
            let msg = match self.inner.done.try_recv() {
                Ok(msg) => msg,
                Err(TryRecvError::Empty | TryRecvError::Disconnected) => break,
            };
            let id = match &msg {
                Completion::Done(id, _) | Completion::Panicked(id) => *id,
            };
            let apply = self.inner.applies.borrow_mut().remove(&id);
            match (msg, apply) {
                (Completion::Done(_, value), Some(apply)) => {
                    apply(value);
                    handled += 1;
                    self.inner.pending.update(|n| *n = n.saturating_sub(1));
                }
                (Completion::Panicked(..), Some(_)) => {
                    log::error!("sequoia-async: task #{id} panicked — result dropped");
                    handled += 1;
                    self.inner.pending.update(|n| *n = n.saturating_sub(1));
                }
                (Completion::Done(..), None) | (Completion::Panicked(..), None) => {}
            }
        }
        handled += self.drain_streams();
        handled
    }

    fn drain_streams(&self) -> usize {
        let mut handled = 0;
        loop {
            let msg = match self.inner.stream_rx.try_recv() {
                Ok(msg) => msg,
                Err(TryRecvError::Empty | TryRecvError::Disconnected) => break,
            };
            match msg.kind {
                StreamKind::Item(value) => {
                    if let Some(apply) = self.inner.streams.borrow().get(&msg.id).cloned() {
                        apply(value);
                        handled += 1;
                    }
                }
                StreamKind::Finished | StreamKind::Panicked => {
                    let was = self.inner.streams.borrow_mut().remove(&msg.id).is_some();
                    if was {
                        if matches!(msg.kind, StreamKind::Panicked) {
                            log::error!("sequoia-async: stream #{} panicked — closed", msg.id);
                        }
                        self.inner.pending.update(|n| *n = n.saturating_sub(1));
                        handled += 1;
                    }
                }
            }
        }
        handled
    }

    pub fn drain_until(&self, timeout: Duration) -> usize {
        crate::testing::drain_until(self, timeout)
    }
}

struct Pool {
    jobs: Mutex<Option<Sender<Job>>>,
    results: Mutex<Sender<Completion>>,
    threads: usize,
}

impl Pool {
    fn new(threads: usize, results_tx: Sender<Completion>) -> Arc<Self> {
        let threads = threads.max(1);
        let (jobs_tx, jobs_rx) = channel::<Job>();
        let jobs_rx = Arc::new(Mutex::new(jobs_rx));
        let mut started = 0;
        for index in 0..threads {
            let rx = Arc::clone(&jobs_rx);
            let spawned = std::thread::Builder::new()
                .name(format!("{THREAD_NAME}-{index}"))
                .spawn(move || {
                    loop {
                        let job = rx.lock().unwrap_or_else(|e| e.into_inner()).recv();
                        match job {
                            Ok(job) => job(),
                            Err(_) => break,
                        }
                    }
                });
            if spawned.is_ok() {
                started += 1;
            } else {
                log::error!("sequoia-async: failed to start worker thread");
            }
        }
        Arc::new(Self {
            jobs: Mutex::new(Some(jobs_tx)),
            results: Mutex::new(results_tx),
            threads: started.max(1),
        })
    }

    fn submit(&self, job: Job) {
        let guard = self.jobs.lock().unwrap_or_else(|e| e.into_inner());
        if let Some(tx) = guard.as_ref()
            && tx.send(job).is_err()
        {
            log::error!("sequoia-async: failed to queue task");
        }
    }

    fn results_tx(&self) -> Sender<Completion> {
        self.results
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .clone()
    }

    fn threads(&self) -> usize {
        self.threads
    }
}

impl Drop for Pool {
    fn drop(&mut self) {
        self.jobs.lock().unwrap_or_else(|e| e.into_inner()).take();
    }
}

pub(crate) fn default_threads() -> usize {
    std::thread::available_parallelism()
        .map(|n| n.get().saturating_sub(1).max(1))
        .unwrap_or(1)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn future_runs_off_thread_and_apply_on_caller() {
        use std::sync::{Arc, Mutex};
        let rt = AsyncRuntime::new(1);
        let ui = std::thread::current().id();
        let worker = Arc::new(Mutex::new(None));
        let out = Arc::new(Mutex::new(None));
        let w = Arc::clone(&worker);
        let o = Arc::clone(&out);
        rt.spawn_future(
            async move {
                w.lock().unwrap().replace(std::thread::current().id());
                21 * 2
            },
            move |v: i32| {
                o.lock().unwrap().replace((v, std::thread::current().id()));
            },
        );
        assert!(rt.is_busy());
        assert_eq!(rt.drain_until(Duration::from_secs(5)), 1);
        assert_eq!(*out.lock().unwrap(), Some((42, ui)));
        assert_ne!(worker.lock().unwrap().unwrap(), ui);
        assert!(!rt.is_busy());
    }

    #[test]
    fn cancel_drops_result() {
        use std::sync::Arc;
        use std::sync::atomic::{AtomicBool, Ordering};
        let rt = AsyncRuntime::new(1);
        let flag = Arc::new(AtomicBool::new(false));
        let f = Arc::clone(&flag);
        let h = rt.spawn_future(
            async move {
                crate::time::sleep(Duration::from_millis(30)).await;
            },
            move |()| f.store(true, Ordering::SeqCst),
        );
        assert!(h.cancel());
        std::thread::sleep(Duration::from_millis(120));
        assert_eq!(rt.drain(), 0);
        assert!(!flag.load(Ordering::SeqCst));
        assert_eq!(rt.pending().get_untracked(), 0);
    }

    #[test]
    fn threads_never_below_one() {
        let rt = AsyncRuntime::new(0);
        assert_eq!(rt.threads(), 1);
        assert_eq!(rt.drain(), 0);
    }

    #[test]
    fn stream_delivers_multiple_items_then_finishes() {
        use std::sync::{Arc, Mutex};
        let rt = AsyncRuntime::new(2);
        let items = Arc::new(Mutex::new(Vec::new()));
        let probe = Arc::clone(&items);
        let _h = rt.spawn_stream(
            |em: Emitter<i32>| async move {
                for v in [1, 2, 3] {
                    em.emit(v).unwrap();
                }
            },
            move |v: i32| probe.lock().unwrap().push(v),
        );
        assert!(rt.is_busy());
        let ok = crate::testing::run_until(&rt, Duration::from_secs(5), || {
            items.lock().unwrap().len() == 3 && !rt.is_busy()
        });
        assert!(ok, "stream items + finish drained");
        assert_eq!(*items.lock().unwrap(), vec![1, 2, 3]);
        assert_eq!(rt.pending().get_untracked(), 0);
    }

    #[test]
    fn stream_cancel_stops_delivery() {
        use std::sync::{Arc, Mutex};
        let rt = AsyncRuntime::new(1);
        let items = Arc::new(Mutex::new(Vec::new()));
        let probe = Arc::clone(&items);
        let h = rt.spawn_stream(
            |em: Emitter<i32>| async move {
                crate::time::sleep(Duration::from_millis(50)).await;
                let _ = em.emit(99);
            },
            move |v: i32| probe.lock().unwrap().push(v),
        );
        assert!(h.cancel());
        std::thread::sleep(Duration::from_millis(150));
        rt.drain();
        assert!(items.lock().unwrap().is_empty());
        assert_eq!(rt.pending().get_untracked(), 0);
    }

    #[test]
    fn spawn_to_signal_writes_value() {
        let rt = AsyncRuntime::new(1);
        let s = Signal::new(0u32);
        rt.spawn_to_signal(&s, async { 11u32 });
        assert!(crate::testing::run_until(
            &rt,
            Duration::from_secs(5),
            || { s.get_untracked() == 11 }
        ));
    }
}