Fix the way we spawn the rayon threadpool

crates/index-scheduler/src/batch.rs

@@ -1351,7 +1351,10 @@ impl IndexScheduler {
                 let pool = match &indexer_config.thread_pool {
                     Some(pool) => pool,
                     None => {
-                        local_pool = ThreadPoolNoAbortBuilder::new().build().unwrap();
+                        local_pool = ThreadPoolNoAbortBuilder::new()
+                            .thread_name(|i| format!("indexing-thread-{i}"))
+                            .build()
+                            .unwrap();
                         &local_pool
                     }
                 };
@@ -1399,10 +1402,10 @@ impl IndexScheduler {
                 }
 
                 if tasks.iter().any(|res| res.error.is_none()) {
-                    pool.install(|| {
                     indexer::index(
                         index_wtxn,
                         index,
+                        pool,
                         indexer_config.grenad_parameters(),
                         &db_fields_ids_map,
                         new_fields_ids_map,
@@ -1411,9 +1414,7 @@ impl IndexScheduler {
                         embedders,
                         &|| must_stop_processing.get(),
                         &send_progress,
-                        )
+                    )?;
-                    })
-                    .unwrap()?;
 
                     tracing::info!(indexing_result = ?addition, processed_in = ?started_processing_at.elapsed(), "document indexing done");
                 }
@@ -1489,21 +1490,25 @@ impl IndexScheduler {
                     let pool = match &indexer_config.thread_pool {
                         Some(pool) => pool,
                         None => {
-                            local_pool = ThreadPoolNoAbortBuilder::new().build().unwrap();
+                            local_pool = ThreadPoolNoAbortBuilder::new()
+                                .thread_name(|i| format!("indexing-thread-{i}"))
+                                .build()
+                                .unwrap();
                             &local_pool
                         }
                     };
 
-                    pool.install(|| {
+                    let indexer = UpdateByFunction::new(candidates, context.clone(), code.clone());
-                        let indexer =
+                    let document_changes =
-                            UpdateByFunction::new(candidates, context.clone(), code.clone());
+                        pool.install(|| indexer.into_changes(&primary_key)).unwrap()?;
-                        let document_changes = indexer.into_changes(&primary_key)?;
+
                     let embedders = index.embedding_configs(index_wtxn)?;
                     let embedders = self.embedders(embedders)?;
 
                     indexer::index(
                         index_wtxn,
                         index,
+                        pool,
                         indexer_config.grenad_parameters(),
                         &db_fields_ids_map,
                         new_fields_ids_map,
@@ -1514,10 +1519,6 @@ impl IndexScheduler {
                         &send_progress,
                     )?;
 
-                        Result::Ok(())
-                    })
-                    .unwrap()?;
-
                     // tracing::info!(indexing_result = ?addition, processed_in = ?started_processing_at.elapsed(), "document indexing done");
                 }
 
@@ -1641,7 +1642,10 @@ impl IndexScheduler {
                     let pool = match &indexer_config.thread_pool {
                         Some(pool) => pool,
                         None => {
-                            local_pool = ThreadPoolNoAbortBuilder::new().build().unwrap();
+                            local_pool = ThreadPoolNoAbortBuilder::new()
+                                .thread_name(|i| format!("indexing-thread-{i}"))
+                                .build()
+                                .unwrap();
                             &local_pool
                         }
                     };
@@ -1652,10 +1656,10 @@ impl IndexScheduler {
                     let embedders = index.embedding_configs(index_wtxn)?;
                     let embedders = self.embedders(embedders)?;
 
-                    pool.install(|| {
                     indexer::index(
                         index_wtxn,
                         index,
+                        pool,
                         indexer_config.grenad_parameters(),
                         &db_fields_ids_map,
                         new_fields_ids_map,
@@ -1664,9 +1668,7 @@ impl IndexScheduler {
                         embedders,
                         &|| must_stop_processing.get(),
                         &send_progress,
-                        )
+                    )?;
-                    })
-                    .unwrap()?;
 
                     // tracing::info!(indexing_result = ?addition, processed_in = ?started_processing_at.elapsed(), "document indexing done");
                 }

crates/milli/src/update/new/channel.rs

@@ -55,6 +55,12 @@ pub fn extractor_writer_bbqueue(
 
     let producers = ThreadLocal::with_capacity(bbbuffers.len());
     let consumers = rayon::broadcast(|bi| {
+        eprintln!(
+            "hello thread #{:?} (#{:?}, #{:?})",
+            bi.index(),
+            std::thread::current().name(),
+            std::thread::current().id(),
+        );
         let bbqueue = &bbbuffers[bi.index()];
         let (producer, consumer) = bbqueue.try_split_framed().unwrap();
         producers.get_or(|| FullySend(RefCell::new(producer)));
@@ -399,7 +405,15 @@ impl<'b> ExtractorBbqueueSender<'b> {
 
     fn delete_vector(&self, docid: DocumentId) -> crate::Result<()> {
         let capacity = self.capacity;
-        let refcell = self.producers.get().unwrap();
+        let refcell = match self.producers.get() {
+            Some(refcell) => refcell,
+            None => panic!(
+                "hello thread #{:?} (#{:?}, #{:?})",
+                rayon::current_thread_index(),
+                std::thread::current().name(),
+                std::thread::current().id()
+            ),
+        };
         let mut producer = refcell.0.borrow_mut_or_yield();
 
         let payload_header = EntryHeader::ArroyDeleteVector(ArroyDeleteVector { docid });
@@ -438,7 +452,15 @@ impl<'b> ExtractorBbqueueSender<'b> {
         embedding: &[f32],
     ) -> crate::Result<()> {
         let capacity = self.capacity;
-        let refcell = self.producers.get().unwrap();
+        let refcell = match self.producers.get() {
+            Some(refcell) => refcell,
+            None => panic!(
+                "hello thread #{:?} (#{:?}, #{:?})",
+                rayon::current_thread_index(),
+                std::thread::current().name(),
+                std::thread::current().id()
+            ),
+        };
         let mut producer = refcell.0.borrow_mut_or_yield();
 
         let payload_header =
@@ -496,7 +518,15 @@ impl<'b> ExtractorBbqueueSender<'b> {
         F: FnOnce(&mut [u8]) -> crate::Result<()>,
     {
         let capacity = self.capacity;
-        let refcell = self.producers.get().unwrap();
+        let refcell = match self.producers.get() {
+            Some(refcell) => refcell,
+            None => panic!(
+                "hello thread #{:?} (#{:?}, #{:?})",
+                rayon::current_thread_index(),
+                std::thread::current().name(),
+                std::thread::current().id()
+            ),
+        };
         let mut producer = refcell.0.borrow_mut_or_yield();
 
         let operation = DbOperation { database, key_length: Some(key_length) };

crates/milli/src/update/new/indexer/mod.rs

@@ -62,6 +62,7 @@ mod update_by_function;
 pub fn index<'pl, 'indexer, 'index, DC, MSP, SP>(
     wtxn: &mut RwTxn,
     index: &'index Index,
+    pool: &ThreadPoolNoAbort,
     grenad_parameters: GrenadParameters,
     db_fields_ids_map: &'indexer FieldsIdsMap,
     new_fields_ids_map: FieldsIdsMap,
@@ -77,10 +78,15 @@ where
     SP: Fn(Progress) + Sync,
 {
     /// TODO restrict memory and remove this memory from the extractors bump allocators
-    let bbbuffers: Vec<_> = (0..rayon::current_num_threads())
+    let bbbuffers: Vec<_> = pool
+        .install(|| {
+            (0..rayon::current_num_threads())
                 .map(|_| bbqueue::BBBuffer::new(100 * 1024 * 1024)) // 100 MiB by thread
-        .collect();
+                .collect()
-    let (extractor_sender, mut writer_receiver) = extractor_writer_bbqueue(&bbbuffers, 1000);
+        })
+        .unwrap();
+    let (extractor_sender, mut writer_receiver) =
+        pool.install(|| extractor_writer_bbqueue(&bbbuffers, 1000)).unwrap();
     let finished_extraction = AtomicBool::new(false);
 
     let metadata_builder = MetadataBuilder::from_index(index, wtxn)?;
@@ -112,6 +118,7 @@ where
         let field_distribution = &mut field_distribution;
         let document_ids = &mut document_ids;
         let extractor_handle = Builder::new().name(S("indexer-extractors")).spawn_scoped(s, move || {
+            pool.install(move || {
                 let span = tracing::trace_span!(target: "indexing::documents", parent: &indexer_span, "extract");
                 let _entered = span.enter();
 
@@ -359,6 +366,7 @@ where
                 finished_extraction.store(true, std::sync::atomic::Ordering::Relaxed);
 
                 Result::Ok((facet_field_ids_delta, index_embeddings))
+            }).unwrap()
         })?;
 
         let global_fields_ids_map = GlobalFieldsIdsMap::new(&new_fields_ids_map);