Introduce an atomic to catch panics in thread pools

milli/src/update/index_documents/mod.rs

@@ -8,6 +8,7 @@ use std::collections::{HashMap, HashSet};
 use std::io::{Read, Seek};
 use std::num::NonZeroU32;
 use std::result::Result as StdResult;
+use std::sync::atomic::Ordering;
 use std::sync::Arc;
 
 use crossbeam_channel::{Receiver, Sender};
@@ -296,20 +297,24 @@ where
         let settings_diff = Arc::new(settings_diff);
 
         let backup_pool;
+        let pool_catched_panic = self.indexer_config.pool_panic_catched.clone();
         let pool = match self.indexer_config.thread_pool {
             Some(ref pool) => pool,
-            #[cfg(not(test))]
             None => {
-                // We initialize a bakcup pool with the default
+                // We initialize a backup pool with the default
                 // settings if none have already been set.
-                backup_pool = rayon::ThreadPoolBuilder::new().build()?;
-                &backup_pool
-            }
-            #[cfg(test)]
-            None => {
-                // We initialize a bakcup pool with the default
-                // settings if none have already been set.
-                backup_pool = rayon::ThreadPoolBuilder::new().num_threads(1).build()?;
+                let mut pool_builder = rayon::ThreadPoolBuilder::new();
+                pool_builder = pool_builder.panic_handler({
+                    let catched_panic = pool_catched_panic.clone();
+                    move |_result| catched_panic.store(true, Ordering::SeqCst)
+                });
+
+                #[cfg(test)]
+                {
+                    pool_builder = pool_builder.num_threads(1);
+                }
+
+                backup_pool = pool_builder.build()?;
                 &backup_pool
             }
         };

milli/src/update/indexer_config.rs

@@ -1,3 +1,6 @@
+use std::sync::atomic::AtomicBool;
+use std::sync::Arc;
+
 use grenad::CompressionType;
 use rayon::ThreadPool;
 
@@ -10,6 +13,9 @@ pub struct IndexerConfig {
     pub chunk_compression_type: CompressionType,
     pub chunk_compression_level: Option<u32>,
     pub thread_pool: Option<ThreadPool>,
+    /// Set to true if the thread pool catched a panic
+    /// and we must abort the task
+    pub pool_panic_catched: Arc<AtomicBool>,
     pub max_positions_per_attributes: Option<u32>,
     pub skip_index_budget: bool,
 }
@@ -24,6 +30,7 @@ impl Default for IndexerConfig {
             chunk_compression_type: CompressionType::None,
             chunk_compression_level: None,
             thread_pool: None,
+            pool_panic_catched: Arc::default(),
             max_positions_per_attributes: None,
             skip_index_budget: false,
         }