diff --git a/crates/milli/src/update/new/extract/documents.rs b/crates/milli/src/update/new/extract/documents.rs index 79e1a2462..2c93a5def 100644 --- a/crates/milli/src/update/new/extract/documents.rs +++ b/crates/milli/src/update/new/extract/documents.rs @@ -6,9 +6,8 @@ use hashbrown::HashMap; use super::DelAddRoaringBitmap; use crate::update::new::channel::DocumentsSender; use crate::update::new::document::{write_to_obkv, Document as _}; -use crate::update::new::indexer::document_changes::{ - DocumentChangeContext, Extractor, FullySend, RefCellExt as _, -}; +use crate::update::new::indexer::document_changes::{DocumentChangeContext, Extractor, FullySend}; +use crate::update::new::ref_cell_ext::RefCellExt as _; use crate::update::new::DocumentChange; use crate::vector::EmbeddingConfigs; use crate::Result; diff --git a/crates/milli/src/update/new/extract/faceted/extract_facets.rs b/crates/milli/src/update/new/extract/faceted/extract_facets.rs index 1aaae1cb8..11dc8f3c7 100644 --- a/crates/milli/src/update/new/extract/faceted/extract_facets.rs +++ b/crates/milli/src/update/new/extract/faceted/extract_facets.rs @@ -13,8 +13,9 @@ use crate::facet::value_encoding::f64_into_bytes; use crate::update::new::extract::DocidsExtractor; use crate::update::new::indexer::document_changes::{ extract, DocumentChangeContext, DocumentChanges, Extractor, FullySend, IndexingContext, - Progress, RefCellExt, ThreadLocal, + Progress, ThreadLocal, }; +use crate::update::new::ref_cell_ext::RefCellExt as _; use crate::update::new::DocumentChange; use crate::update::GrenadParameters; use crate::{DocumentId, FieldId, Index, Result, MAX_FACET_VALUE_LENGTH}; diff --git a/crates/milli/src/update/new/extract/searchable/extract_word_docids.rs b/crates/milli/src/update/new/extract/searchable/extract_word_docids.rs index cadea7251..89583bd93 100644 --- a/crates/milli/src/update/new/extract/searchable/extract_word_docids.rs +++ b/crates/milli/src/update/new/extract/searchable/extract_word_docids.rs @@ -12,8 +12,9 @@ use crate::update::new::extract::cache::BalancedCaches; use crate::update::new::extract::perm_json_p::contained_in; use crate::update::new::indexer::document_changes::{ extract, DocumentChangeContext, DocumentChanges, Extractor, FullySend, IndexingContext, - MostlySend, Progress, RefCellExt, ThreadLocal, + MostlySend, Progress, ThreadLocal, }; +use crate::update::new::ref_cell_ext::RefCellExt as _; use crate::update::new::DocumentChange; use crate::update::GrenadParameters; use crate::{bucketed_position, DocumentId, FieldId, Index, Result, MAX_POSITION_PER_ATTRIBUTE}; diff --git a/crates/milli/src/update/new/extract/searchable/extract_word_pair_proximity_docids.rs b/crates/milli/src/update/new/extract/searchable/extract_word_pair_proximity_docids.rs index 796b8c943..7f9fff38f 100644 --- a/crates/milli/src/update/new/extract/searchable/extract_word_pair_proximity_docids.rs +++ b/crates/milli/src/update/new/extract/searchable/extract_word_pair_proximity_docids.rs @@ -9,7 +9,8 @@ use super::SearchableExtractor; use crate::proximity::{index_proximity, MAX_DISTANCE}; use crate::update::new::document::Document; use crate::update::new::extract::cache::BalancedCaches; -use crate::update::new::indexer::document_changes::{DocumentChangeContext, RefCellExt}; +use crate::update::new::indexer::document_changes::DocumentChangeContext; +use crate::update::new::ref_cell_ext::RefCellExt as _; use crate::update::new::DocumentChange; use crate::{FieldId, GlobalFieldsIdsMap, Index, Result}; diff --git a/crates/milli/src/update/new/indexer/document_changes.rs b/crates/milli/src/update/new/indexer/document_changes.rs index 0a9155b1c..b9bf79e47 100644 --- a/crates/milli/src/update/new/indexer/document_changes.rs +++ b/crates/milli/src/update/new/indexer/document_changes.rs @@ -1,4 +1,4 @@ -use std::cell::{Cell, Ref, RefCell, RefMut}; +use std::cell::{Cell, RefCell}; use std::sync::{Arc, RwLock}; use bumpalo::Bump; @@ -10,57 +10,6 @@ use crate::fields_ids_map::metadata::FieldIdMapWithMetadata; use crate::update::new::parallel_iterator_ext::ParallelIteratorExt as _; use crate::{FieldsIdsMap, GlobalFieldsIdsMap, Index, InternalError, Result}; -pub trait RefCellExt { - fn try_borrow_or_yield(&self) -> std::result::Result, std::cell::BorrowError>; - fn try_borrow_mut_or_yield( - &self, - ) -> std::result::Result, std::cell::BorrowMutError>; - - fn borrow_or_yield(&self) -> Ref<'_, T> { - self.try_borrow_or_yield().unwrap() - } - - fn borrow_mut_or_yield(&self) -> RefMut<'_, T> { - self.try_borrow_mut_or_yield().unwrap() - } -} - -impl RefCellExt for RefCell { - fn try_borrow_or_yield(&self) -> std::result::Result, std::cell::BorrowError> { - /// TODO: move this trait and impl elsewhere - loop { - match self.try_borrow() { - Ok(borrow) => break Ok(borrow), - Err(error) => { - tracing::warn!("dynamic borrow failed, yielding to local tasks"); - match rayon::yield_local() { - Some(rayon::Yield::Executed) => continue, - _ => return Err(error), - } - } - } - } - } - - fn try_borrow_mut_or_yield( - &self, - ) -> std::result::Result, std::cell::BorrowMutError> { - loop { - match self.try_borrow_mut() { - Ok(borrow) => break Ok(borrow), - Err(error) => { - tracing::warn!("dynamic borrow failed, yielding to local tasks"); - - match rayon::yield_local() { - Some(rayon::Yield::Executed) => continue, - _ => return Err(error), - } - } - } - } - } -} - /// A trait for types that are **not** [`Send`] only because they would then allow concurrent access to a type that is not [`Sync`]. /// /// The primary example of such a type is `&T`, with `T: !Sync`. diff --git a/crates/milli/src/update/new/indexer/partial_dump.rs b/crates/milli/src/update/new/indexer/partial_dump.rs index 4984eb354..2da047824 100644 --- a/crates/milli/src/update/new/indexer/partial_dump.rs +++ b/crates/milli/src/update/new/indexer/partial_dump.rs @@ -3,10 +3,11 @@ use std::ops::DerefMut; use rayon::iter::IndexedParallelIterator; use serde_json::value::RawValue; -use super::document_changes::{DocumentChangeContext, DocumentChanges, MostlySend, RefCellExt}; +use super::document_changes::{DocumentChangeContext, DocumentChanges, MostlySend}; use crate::documents::PrimaryKey; use crate::update::concurrent_available_ids::ConcurrentAvailableIds; use crate::update::new::document::Versions; +use crate::update::new::ref_cell_ext::RefCellExt as _; use crate::update::new::{DocumentChange, Insertion}; use crate::{Error, InternalError, Result, UserError}; diff --git a/crates/milli/src/update/new/indexer/update_by_function.rs b/crates/milli/src/update/new/indexer/update_by_function.rs index 3eb0cc306..eb7252445 100644 --- a/crates/milli/src/update/new/indexer/update_by_function.rs +++ b/crates/milli/src/update/new/indexer/update_by_function.rs @@ -4,12 +4,13 @@ use rayon::slice::ParallelSlice as _; use rhai::{Dynamic, Engine, OptimizationLevel, Scope, AST}; use roaring::RoaringBitmap; -use super::document_changes::{DocumentChangeContext, MostlySend, RefCellExt}; +use super::document_changes::{DocumentChangeContext, MostlySend}; use super::DocumentChanges; use crate::documents::Error::InvalidDocumentFormat; use crate::documents::PrimaryKey; use crate::error::{FieldIdMapMissingEntry, InternalError}; use crate::update::new::document::Versions; +use crate::update::new::ref_cell_ext::RefCellExt as _; use crate::update::new::{Deletion, DocumentChange, KvReaderFieldId, Update}; use crate::{all_obkv_to_json, Error, FieldsIdsMap, Object, Result, UserError}; diff --git a/crates/milli/src/update/new/mod.rs b/crates/milli/src/update/new/mod.rs index 36a9a3f92..7a749228e 100644 --- a/crates/milli/src/update/new/mod.rs +++ b/crates/milli/src/update/new/mod.rs @@ -16,6 +16,7 @@ mod fst_merger_builder; pub mod indexer; mod merger; mod parallel_iterator_ext; +mod ref_cell_ext; mod top_level_map; pub mod vector_document; mod word_fst_builder; diff --git a/crates/milli/src/update/new/ref_cell_ext.rs b/crates/milli/src/update/new/ref_cell_ext.rs new file mode 100644 index 000000000..b147c00e5 --- /dev/null +++ b/crates/milli/src/update/new/ref_cell_ext.rs @@ -0,0 +1,52 @@ +use std::cell::{Ref, RefCell, RefMut}; + +pub trait RefCellExt { + fn try_borrow_or_yield(&self) -> std::result::Result, std::cell::BorrowError>; + fn try_borrow_mut_or_yield( + &self, + ) -> std::result::Result, std::cell::BorrowMutError>; + + fn borrow_or_yield(&self) -> Ref<'_, T> { + self.try_borrow_or_yield().unwrap() + } + + fn borrow_mut_or_yield(&self) -> RefMut<'_, T> { + self.try_borrow_mut_or_yield().unwrap() + } +} + +impl RefCellExt for RefCell { + fn try_borrow_or_yield(&self) -> std::result::Result, std::cell::BorrowError> { + /// TODO: move this trait and impl elsewhere + loop { + match self.try_borrow() { + Ok(borrow) => break Ok(borrow), + Err(error) => { + tracing::warn!("dynamic borrow failed, yielding to local tasks"); + match rayon::yield_local() { + Some(rayon::Yield::Executed) => continue, + _ => return Err(error), + } + } + } + } + } + + fn try_borrow_mut_or_yield( + &self, + ) -> std::result::Result, std::cell::BorrowMutError> { + loop { + match self.try_borrow_mut() { + Ok(borrow) => break Ok(borrow), + Err(error) => { + tracing::warn!("dynamic borrow failed, yielding to local tasks"); + + match rayon::yield_local() { + Some(rayon::Yield::Executed) => continue, + _ => return Err(error), + } + } + } + } + } +} diff --git a/crates/milli/src/update/new/words_prefix_docids.rs b/crates/milli/src/update/new/words_prefix_docids.rs index edc09c5f3..ffc0c5048 100644 --- a/crates/milli/src/update/new/words_prefix_docids.rs +++ b/crates/milli/src/update/new/words_prefix_docids.rs @@ -10,7 +10,7 @@ use roaring::MultiOps; use tempfile::tempfile; use thread_local::ThreadLocal; -use super::indexer::document_changes::RefCellExt; +use super::ref_cell_ext::RefCellExt as _; use crate::heed_codec::StrBEU16Codec; use crate::{CboRoaringBitmapCodec, Index, Prefix, Result};