Allow cancellation in HNSW healing (#10426)

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

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 {