Merge rayon-rs#644

644: Add ParallelIterator::panic_fuse() r=nikomatsakis a=cuviper

This wraps an iterator with a fuse in case of panics, to halt all
threads as soon as possible. While we always propagate panics anyway,
that doesn't usually synchronize across threads, so iterators may end up
processing many more items before the panic is realized. With
`panic_fuse()`, a shared `AtomicBool` is checked before processing each
item, and set for any panic unwinding in its path.

Co-authored-by: Josh Stone <[email protected]>

Author: bors[bot]

src/compile_fail/must_use.rs

@@ -49,6 +49,7 @@ must_use! {
     map                 /** v.par_iter().map(|x| x); */
     map_with            /** v.par_iter().map_with(0, |_, x| x); */
     map_init            /** v.par_iter().map_init(|| 0, |_, x| x); */
+    panic_fuse          /** v.par_iter().panic_fuse(); */
     rev                 /** v.par_iter().rev(); */
     skip                /** v.par_iter().skip(1); */
     take                /** v.par_iter().take(1); */

src/iter/collect/consumer.rs

@@ -7,6 +7,12 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 pub struct CollectConsumer<'c, T: Send + 'c> {
     /// Tracks how many items we successfully wrote. Used to guarantee
     /// safety in the face of panics or buggy parallel iterators.
+    ///
+    /// In theory we could just produce this as a `CollectConsumer::Result`,
+    /// folding local counts and reducing by addition, but that requires a
+    /// certain amount of trust that the producer driving this will behave
+    /// itself. Since this count is important to the safety of marking the
+    /// memory initialized (`Vec::set_len`), we choose to keep it internal.
     writes: &'c AtomicUsize,
 
     /// A slice covering the target memory, not yet initialized!
@@ -85,7 +91,8 @@ impl<'c, T: Send + 'c> Folder<T> for CollectFolder<'c, T> {
     }
 
     fn complete(self) {
-        assert!(self.target.len() == 0, "too few values pushed to consumer");
+        // NB: We don't explicitly check that the local writes were complete,
+        // but `Collect::complete()` will assert the global write count.
 
         // track total values written
         self.global_writes

src/iter/collect/test.rs

@@ -7,6 +7,7 @@
 
 use super::Collect;
 use iter::plumbing::*;
+use rayon_core::join;
 
 /// Promises to produce 2 items, but then produces 3.  Does not do any
 /// splits at all.
@@ -25,15 +26,18 @@ fn produce_too_many_items() {
 /// Produces fewer items than promised. Does not do any
 /// splits at all.
 #[test]
-#[should_panic(expected = "too few values")]
+#[should_panic(expected = "expected 5 total writes, but got 2")]
 fn produce_fewer_items() {
     let mut v = vec![];
     let mut collect = Collect::new(&mut v, 5);
-    let consumer = collect.as_consumer();
-    let mut folder = consumer.into_folder();
-    folder = folder.consume(22);
-    folder = folder.consume(23);
-    folder.complete();
+    {
+        let consumer = collect.as_consumer();
+        let mut folder = consumer.into_folder();
+        folder = folder.consume(22);
+        folder = folder.consume(23);
+        folder.complete();
+    }
+    collect.complete();
 }
 
 // Complete is not called by the consumer.Hence,the collection vector is not fully initialized.
@@ -129,7 +133,7 @@ fn right_produces_too_many_items() {
 // The left consumer produces fewer items while the right
 // consumer produces correct number.
 #[test]
-#[should_panic(expected = "too few values")]
+#[should_panic(expected = "expected 4 total writes, but got 3")]
 fn left_produces_fewer_items() {
     let mut v = vec![];
     let mut collect = Collect::new(&mut v, 4);
@@ -149,7 +153,7 @@ fn left_produces_fewer_items() {
 // The right consumer produces fewer items while the left
 // consumer produces correct number.
 #[test]
-#[should_panic(expected = "too few values")]
+#[should_panic(expected = "expected 4 total writes, but got 3")]
 fn right_produces_fewer_items() {
     let mut v = vec![];
     let mut collect = Collect::new(&mut v, 4);
@@ -165,3 +169,55 @@ fn right_produces_fewer_items() {
     }
     collect.complete();
 }
+
+// The left consumer panics and the right stops short, like `panic_fuse()`.
+// We should get the left panic without ever reaching `Collect::complete()`.
+#[test]
+#[should_panic(expected = "left consumer panic")]
+fn left_panics() {
+    let mut v = vec![];
+    let mut collect = Collect::new(&mut v, 4);
+    {
+        let consumer = collect.as_consumer();
+        let (left_consumer, right_consumer, _) = consumer.split_at(2);
+        join(
+            || {
+                let mut left_folder = left_consumer.into_folder();
+                left_folder = left_folder.consume(0);
+                panic!("left consumer panic");
+            },
+            || {
+                let mut right_folder = right_consumer.into_folder();
+                right_folder = right_folder.consume(2);
+                right_folder.complete() // early return
+            },
+        );
+    }
+    collect.complete();
+}
+
+// The right consumer panics and the left stops short, like `panic_fuse()`.
+// We should get the right panic without ever reaching `Collect::complete()`.
+#[test]
+#[should_panic(expected = "right consumer panic")]
+fn right_panics() {
+    let mut v = vec![];
+    let mut collect = Collect::new(&mut v, 4);
+    {
+        let consumer = collect.as_consumer();
+        let (left_consumer, right_consumer, _) = consumer.split_at(2);
+        join(
+            || {
+                let mut left_folder = left_consumer.into_folder();
+                left_folder = left_folder.consume(0);
+                left_folder.complete() // early return
+            },
+            || {
+                let mut right_folder = right_consumer.into_folder();
+                right_folder = right_folder.consume(2);
+                panic!("right consumer panic");
+            },
+        );
+    }
+    collect.complete();
+}

src/iter/mod.rs

@@ -155,6 +155,8 @@ mod sum;
 pub use self::cloned::Cloned;
 mod inspect;
 pub use self::inspect::Inspect;
+mod panic_fuse;
+pub use self::panic_fuse::PanicFuse;
 mod while_some;
 pub use self::while_some::WhileSome;
 mod extend;
@@ -1730,6 +1732,40 @@ pub trait ParallelIterator: Sized + Send {
         while_some::new(self)
     }
 
+    /// Wraps an iterator with a fuse in case of panics, to halt all threads
+    /// as soon as possible.
+    ///
+    /// Panics within parallel iterators are always propagated to the caller,
+    /// but they don't always halt the rest of the iterator right away, due to
+    /// the internal semantics of [`join`]. This adaptor makes a greater effort
+    /// to stop processing other items sooner, with the cost of additional
+    /// synchronization overhead, which may also inhibit some optimizations.
+    ///
+    /// [`join`]: ../fn.join.html#panics
+    ///
+    /// # Examples
+    ///
+    /// If this code didn't use `panic_fuse()`, it would continue processing
+    /// many more items in other threads (with long sleep delays) before the
+    /// panic is finally propagated.
+    ///
+    /// ```should_panic
+    /// use rayon::prelude::*;
+    /// use std::{thread, time};
+    ///
+    /// (0..1_000_000)
+    ///     .into_par_iter()
+    ///     .panic_fuse()
+    ///     .for_each(|i| {
+    ///         // simulate some work
+    ///         thread::sleep(time::Duration::from_secs(1));
+    ///         assert!(i > 0); // oops!
+    ///     });
+    /// ```
+    fn panic_fuse(self) -> PanicFuse<Self> {
+        panic_fuse::new(self)
+    }
+
     /// Create a fresh collection containing all the element produced
     /// by this parallel iterator.
     ///