johnandersen777.bsky.social/asyncai

Switch repository

Loading…

use std::future::Future;

use tokio::task::{JoinError, JoinSet};

/// Dynamic task multiplexer.
///
/// Mirrors the Python `concurrently()` helper: callers insert futures while
/// consuming completed ones, and the set is re-read after each yield.
pub struct WorkSet<C, T> {
    tasks: JoinSet<(C, T)>,
}

impl<C, T> Default for WorkSet<C, T>
where
    C: Send + 'static,
    T: Send + 'static,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<C, T> WorkSet<C, T>
where
    C: Send + 'static,
    T: Send + 'static,
{
    pub fn new() -> Self {
        Self {
            tasks: JoinSet::new(),
        }
    }

    pub fn spawn<F>(&mut self, context: C, future: F)
    where
        F: Future<Output = T> + Send + 'static,
    {
        self.tasks.spawn(async move { (context, future.await) });
    }

    /// Wait for the next task to complete. New tasks may be inserted between
    /// calls. Returns `None` once the set is empty.
    pub async fn next(&mut self) -> Option<Result<(C, T), JoinError>> {
        self.tasks.join_next().await
    }

    pub fn is_empty(&self) -> bool {
        self.tasks.is_empty()
    }

    pub fn len(&self) -> usize {
        self.tasks.len()
    }

    pub fn abort_all(&mut self) {
        self.tasks.abort_all();
    }
}

#[cfg(test)]
mod tests {
    use std::future::pending;
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::sync::Arc;
    use std::time::Duration;

    use super::WorkSet;

    #[tokio::test]
    async fn yields_first_completed_task() {
        let mut work = WorkSet::new();
        work.spawn("slow", async {
            tokio::time::sleep(Duration::from_millis(80)).await;
            1
        });
        work.spawn("fast", async {
            tokio::time::sleep(Duration::from_millis(5)).await;
            2
        });

        let (context, value) = work.next().await.unwrap().unwrap();
        assert_eq!((context, value), ("fast", 2));
        let (context, value) = work.next().await.unwrap().unwrap();
        assert_eq!((context, value), ("slow", 1));
        assert!(work.next().await.is_none());
    }

    #[tokio::test]
    async fn tasks_can_be_added_while_consuming() {
        let mut work = WorkSet::new();
        work.spawn("first", async { 1 });

        let (context, value) = work.next().await.unwrap().unwrap();
        assert_eq!((context, value), ("first", 1));

        work.spawn("second", async { 2 });
        let (context, value) = work.next().await.unwrap().unwrap();
        assert_eq!((context, value), ("second", 2));
        assert!(work.next().await.is_none());
    }

    #[tokio::test]
    async fn panics_surface_as_join_errors() {
        let mut work = WorkSet::<(), ()>::new();
        work.spawn((), async { panic!("backend exploded") });

        let error = work.next().await.unwrap().unwrap_err();
        assert!(error.is_panic());
    }

    #[tokio::test]
    async fn dropping_work_set_aborts_tasks() {
        struct DropFlag(Arc<AtomicBool>);
        impl Drop for DropFlag {
            fn drop(&mut self) {
                self.0.store(true, Ordering::SeqCst);
            }
        }

        let flag = Arc::new(AtomicBool::new(false));
        {
            let mut work = WorkSet::<(), ()>::new();
            let flag_clone = Arc::clone(&flag);
            work.spawn((), async move {
                let _guard = DropFlag(flag_clone);
                pending::<()>().await;
            });
            tokio::time::sleep(Duration::from_millis(10)).await;
        }

        for _ in 0..20 {
            if flag.load(Ordering::SeqCst) {
                break;
            }
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
        assert!(flag.load(Ordering::SeqCst), "task was not cancelled");
    }
}