rayon_cancel/

lib.rs

//! Rayon does not natively support interrupting long-running computations.
//! This crate provides an iterator adapter for rayon that can be interrupted during computation.
//!
//! The provided [`CancelAdapter`] can be used on any rayon iterator and can be used to interrupt processing new items at any given point.
//! The adapter provides a handle to cancel the computation and a handle to access the number of processed items.
//!
//! By design, the adapter cannot interrupt processing individual items.
//! Once the computation is cancelled, the adapter will stop producing or consuming new items.
//! Which items are processed before the computation stops is non-deterministic and depends on the way rayon distributes the work.
//!
//! Using this adapter may be less efficient than using the underlying iterator directly as the
//! number of items produced by the iterator cannot be known in advance.
//!
//! If you only need access to the number of processed items, you may want to have a look at the [rayon-progress](https://crates.io/crates/rayon-progress) crate.
//!
//! # Example
//! ```
//! use rayon::prelude::*;
//! use rayon_cancel::CancelAdapter;
//! let adapter = CancelAdapter::new(0..100000);
//! let canceller = adapter.canceller();
//! let progress = adapter.counter();
//! std::thread::spawn(move || {
//!     while progress.get() < 1000 {
//!        std::thread::sleep(std::time::Duration::from_millis(2));
//!     }
//!     canceller.cancel();
//! });
//! let count = adapter.counter();
//! // some expensive computation
//! let processed: Vec<_> = adapter.filter(|_| true).map(|i| {
//!     std::thread::sleep(std::time::Duration::from_millis(20));
//!     i
//! }).collect();
//! assert!(count.get() > 1000);
//! assert!(count.get() < 100000);
//! // `processed` contains `count` items, but which ones is non-deterministic
//! assert_eq!(processed.len(), count.get());
//! ```
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use rayon::iter::IntoParallelIterator;
use rayon::iter::plumbing::{Consumer, Folder, UnindexedConsumer};
use rayon::prelude::ParallelIterator;

/// Handle that allows cancelling a running iterator computation.
#[derive(Debug, Clone)]
pub struct CancelHandle {
    cancelled: Arc<AtomicBool>,
}

impl CancelHandle {
    /// Cancel the running computation.
    pub fn cancel(&self) {
        self.cancelled.store(true, Ordering::Relaxed);
    }

    /// Returns whether the computation is cancelled.
    pub fn is_cancelled(&self) -> bool {
        self.cancelled.load(Ordering::Relaxed)
    }
}

/// Access the number of items processed by the iterator.
///
/// Note that this provides only an upper bound on the number as this only counts the number of items produced, not the number consumed by following operations.
#[derive(Debug, Clone)]
pub struct CountHandle {
    count: Arc<AtomicUsize>,
}

impl CountHandle {
    pub fn get(&self) -> usize {
        self.count.load(Ordering::Relaxed)
    }
}

/// Iterator adapter that can be interrupted during computation.
///
/// The adapter allows provides the number of items processed by the iterator which can be used to estimate the progress.
///
#[derive(Debug, Clone)]
pub struct CancelAdapter<I> {
    inner: I,
    cancel: Arc<AtomicBool>,
    count: Arc<AtomicUsize>,
}

impl<I: ParallelIterator> CancelAdapter<I> {
    /// Wrap the given iterator with a cancel adapter.
    ///
    /// `inner` can be any type that implements [`IntoParallelIterator`].
    pub fn new<It: IntoParallelIterator<Iter=I>>(inner: It) -> Self {
        Self {
            inner: inner.into_par_iter(),
            cancel: Arc::new(AtomicBool::new(false)),
            count: Arc::new(AtomicUsize::new(0)),
        }
    }

    /// Handle to cancel the computation.
    pub fn canceller(&self) -> CancelHandle {
        CancelHandle {
            cancelled: self.cancel.clone(),
        }
    }

    /// Handle to the number of processed items.
    pub fn counter(&self) -> CountHandle {
        CountHandle {
            count: self.count.clone(),
        }
    }
}

impl<I> ParallelIterator for CancelAdapter<I>
where
    I: ParallelIterator,
{
    type Item = I::Item;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
    where
        C: UnindexedConsumer<Self::Item>,
    {
        if self.cancel.load(Ordering::Relaxed) {
            return consumer.split_off_left().into_folder().complete();
        }
        self.inner.drive_unindexed(CancelConsumer {
            inner: consumer,
            cancel: self.cancel,
            count: self.count,
        })
    }
}

#[derive(Debug, Clone)]
struct CancelConsumer<C> {
    inner: C,
    cancel: Arc<AtomicBool>,
    count: Arc<AtomicUsize>,
}

#[derive(Debug, Clone)]
struct CancelFolder<F> {
    inner: F,
    cancel: Arc<AtomicBool>,
    count: Arc<AtomicUsize>,
}

impl<Item, C> Consumer<Item> for CancelConsumer<C>
where
    C: Consumer<Item>,
{
    type Folder = CancelFolder<C::Folder>;
    type Reducer = C::Reducer;
    type Result = C::Result;

    fn split_at(self, index: usize) -> (Self, Self, Self::Reducer) {
        let Self { inner, cancel, count } = self;
        let (left, right, reducer) = inner.split_at(index);
        (
            CancelConsumer {
                inner: left,
                cancel: cancel.clone(),
                count: count.clone(),
            },
            CancelConsumer {
                inner: right,
                cancel,
                count,
            },
            reducer,
        )
    }

    fn into_folder(self) -> Self::Folder {
        CancelFolder {
            inner: self.inner.into_folder(),
            cancel: self.cancel,
            count: self.count,
        }
    }

    fn full(&self) -> bool {
        self.cancel.load(Ordering::Relaxed) || self.inner.full()
    }
}

impl<F: Folder<I>, I> Folder<I> for CancelFolder<F> {
    type Result = F::Result;

    fn consume(self, item: I) -> Self {
        if self.cancel.load(Ordering::Relaxed) {
            return self;
        }
        self.count.fetch_add(1, Ordering::Relaxed);
        Self {
            inner: self.inner.consume(item),
            cancel: self.cancel,
            count: self.count,
        }
    }

    fn complete(self) -> Self::Result {
        self.inner.complete()
    }

    fn full(&self) -> bool {
        self.cancel.load(Ordering::Relaxed) || self.inner.full()
    }
}

impl<Item, C> UnindexedConsumer<Item> for CancelConsumer<C>
where
    C: UnindexedConsumer<Item>,
{
    fn split_off_left(&self) -> Self {
        CancelConsumer {
            inner: self.inner.split_off_left(),
            cancel: self.cancel.clone(),
            count: self.count.clone(),
        }
    }

    fn to_reducer(&self) -> Self::Reducer {
        self.inner.to_reducer()
    }
}

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

    #[test]
    fn test_cancel() {
        let adapter = CancelAdapter::new(0..10000);
        let canceller = adapter.canceller();
        rayon::spawn(move || {
            std::thread::sleep(std::time::Duration::from_millis(100));
            canceller.cancel();
        });
        let count = adapter.counter();
        // simulate expensive computation
        let total: usize = adapter.filter(|_| true).map(|i| {
            std::thread::sleep(std::time::Duration::from_millis(20));
            i
        }).count();
        assert!(count.get() < 10000);
        assert_eq!(total, count.get());
        println!("total: {}, count: {}", total, count.get());
    }

    #[test]
    fn test_no_cancel() {
        let adapter = CancelAdapter::new(0..10000);
        let count = adapter.counter();
        // very expensive calculation
        let total: usize = adapter.filter(|_| true).count();
        assert_eq!(total, 10000);
        assert_eq!(total, count.get());
        println!("total: {}, count: {}", total, count.get());
    }

    #[test]
    fn test_blub() {
        let iter = CancelAdapter::new(0..10000);
        let canceller = iter.canceller();
        let count = iter.counter();
        rayon::spawn(move || {
            std::thread::sleep(std::time::Duration::from_millis(510));
            canceller.cancel();
        });
        let items: Vec<usize> = iter.filter(|_| true).map(|i| {
            std::thread::sleep(std::time::Duration::from_millis(20));
            i
        }).collect();
        assert!(count.get() < 10000);
        assert_eq!(items.len(), count.get());
    }
}