diff --git a/lib/segment/src/index/hnsw_index/graph_layers_healer.rs b/lib/segment/src/index/hnsw_index/graph_layers_healer.rs index dbd9408443..45fb7e9394 100644 --- a/lib/segment/src/index/hnsw_index/graph_layers_healer.rs +++ b/lib/segment/src/index/hnsw_index/graph_layers_healer.rs @@ -6,7 +6,7 @@ use parking_lot::RwLock; use rayon::ThreadPool; use rayon::iter::{IntoParallelIterator as _, ParallelIterator as _}; -use crate::common::operation_error::OperationResult; +use crate::common::operation_error::{OperationResult, check_process_stopped}; use crate::index::hnsw_index::HnswM; use crate::index::hnsw_index::graph_layers::GraphLayers; use crate::index::hnsw_index::graph_layers_builder::{GraphLayersBuilder, LockedLayersContainer}; @@ -211,11 +211,14 @@ impl<'a> GraphLayersHealer<'a> { pool: &ThreadPool, vector_storage: &VectorStorageEnum, quantized_vectors: Option<&QuantizedVectors>, + stopped: &std::sync::atomic::AtomicBool, ) -> OperationResult<()> { pool.install(|| { std::mem::take(&mut self.to_heal) .into_par_iter() .try_for_each(|(offset, level)| { + check_process_stopped(stopped)?; + // Internal operation. No measurements needed. let internal_hardware_counter = HardwareCounterCell::disposable(); let query = vector_storage diff --git a/lib/segment/src/index/hnsw_index/hnsw/build.rs b/lib/segment/src/index/hnsw_index/hnsw/build.rs index 42c262f9f3..ea857a38d5 100644 --- a/lib/segment/src/index/hnsw_index/hnsw/build.rs +++ b/lib/segment/src/index/hnsw_index/hnsw/build.rs @@ -278,6 +278,8 @@ impl HNSWIndex { None }; + check_process_stopped(stopped)?; + if build_main_graph { let progress_main_graph = progress_main_graph.unwrap(); let progress_migrate = progress_migrate.unwrap(); @@ -300,7 +302,14 @@ impl HNSWIndex { ); let old_vector_storage = old_index.index.vector_storage.borrow(); let old_quantized_vectors = old_index.index.quantized_vectors.borrow(); - healer.heal(&pool, &old_vector_storage, old_quantized_vectors.as_ref())?; + + healer.heal( + &pool, + &old_vector_storage, + old_quantized_vectors.as_ref(), + stopped, + )?; + check_process_stopped(stopped)?; healer.save_into_builder(&graph_layers_builder); for vector_id in ids_iter {