Merge pull request #5670 from meilisearch/export-and-transfer-route

Introduce a new route to export indexes

crates/milli/src/thread_pool_no_abort.rs

@@ -1,7 +1,7 @@
 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use std::sync::Arc;
 
-use rayon::{ThreadPool, ThreadPoolBuilder};
+use rayon::{BroadcastContext, ThreadPool, ThreadPoolBuilder};
 use thiserror::Error;
 
 /// A rayon ThreadPool wrapper that can catch panics in the pool
@@ -32,6 +32,22 @@ impl ThreadPoolNoAbort {
         }
     }
 
+    pub fn broadcast<OP, R>(&self, op: OP) -> Result<Vec<R>, PanicCatched>
+    where
+        OP: Fn(BroadcastContext<'_>) -> R + Sync,
+        R: Send,
+    {
+        self.active_operations.fetch_add(1, Ordering::Relaxed);
+        let output = self.thread_pool.broadcast(op);
+        self.active_operations.fetch_sub(1, Ordering::Relaxed);
+        // While reseting the pool panic catcher we return an error if we catched one.
+        if self.pool_catched_panic.swap(false, Ordering::SeqCst) {
+            Err(PanicCatched)
+        } else {
+            Ok(output)
+        }
+    }
+
     pub fn current_num_threads(&self) -> usize {
         self.thread_pool.current_num_threads()
     }

crates/milli/src/update/index_documents/extract/mod.rs

@@ -213,7 +213,7 @@ fn run_extraction_task<FE, FS, M>(
     })
 }
 
-fn request_threads() -> &'static ThreadPoolNoAbort {
+pub fn request_threads() -> &'static ThreadPoolNoAbort {
     static REQUEST_THREADS: OnceLock<ThreadPoolNoAbort> = OnceLock::new();
 
     REQUEST_THREADS.get_or_init(|| {

crates/milli/src/update/index_documents/mod.rs

@@ -12,6 +12,7 @@ use std::sync::Arc;
 
 use crossbeam_channel::{Receiver, Sender};
 use enrich::enrich_documents_batch;
+pub use extract::request_threads;
 use grenad::{Merger, MergerBuilder};
 use hashbrown::HashMap;
 use heed::types::Str;

crates/milli/src/update/mod.rs

@@ -4,7 +4,7 @@ pub use self::clear_documents::ClearDocuments;
 pub use self::concurrent_available_ids::ConcurrentAvailableIds;
 pub use self::facet::bulk::FacetsUpdateBulk;
 pub use self::facet::incremental::FacetsUpdateIncrementalInner;
-pub use self::index_documents::*;
+pub use self::index_documents::{request_threads, *};
 pub use self::indexer_config::{default_thread_pool_and_threads, IndexerConfig};
 pub use self::new::ChannelCongestion;
 pub use self::settings::{validate_embedding_settings, Setting, Settings};