vector-store: fix fullscan in progress status (#288)

ewienik · web-flow · commit 6b97f6d9edab · 2025-11-19T12:13:46.000Z
The are issues with wrong index status during fullscan - some vectors
are still in the pipeline when status is changing to SERVING meaning
fullscan finished. The patch introduces marker for asynchronous
operation in pipeline - now fullscan waits for all add operation to
finish.
diff --git a/crates/vector-store/src/db.rs b/crates/vector-store/src/db.rs
@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0
  */
 
+use crate::AsyncInProgress;
 use crate::ColumnName;
 use crate::Connectivity;
 use crate::Credentials;
@@ -53,7 +54,10 @@ use tracing::trace;
 use tracing::warn;
 use uuid::Uuid;
 
-type GetDbIndexR = anyhow::Result<(mpsc::Sender<DbIndex>, mpsc::Receiver<DbEmbedding>)>;
+type GetDbIndexR = anyhow::Result<(
+    mpsc::Sender<DbIndex>,
+    mpsc::Receiver<(DbEmbedding, Option<AsyncInProgress>)>,
+)>;
 pub(crate) type LatestSchemaVersionR = anyhow::Result<Option<CqlTimeuuid>>;
 type GetIndexesR = anyhow::Result<Vec<DbCustomIndex>>;
 type GetIndexVersionR = anyhow::Result<Option<IndexVersion>>;
diff --git a/crates/vector-store/src/db_index.rs b/crates/vector-store/src/db_index.rs
@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0
  */
 
+use crate::AsyncInProgress;
 use crate::ColumnName;
 use crate::DbEmbedding;
 use crate::IndexMetadata;
@@ -112,7 +113,10 @@ pub(crate) async fn new(
     db_session: Arc<Session>,
     metadata: IndexMetadata,
     node_state: Sender<NodeState>,
-) -> anyhow::Result<(mpsc::Sender<DbIndex>, mpsc::Receiver<DbEmbedding>)> {
+) -> anyhow::Result<(
+    mpsc::Sender<DbIndex>,
+    mpsc::Receiver<(DbEmbedding, Option<AsyncInProgress>)>,
+)> {
     let id = metadata.id();
 
     // TODO: The value of channel size was taken from initial benchmarks. Needs more testing
@@ -313,7 +317,7 @@ impl Statements {
     /// to send read embeddings into the pipeline.
     async fn initial_scan(
         &self,
-        tx: mpsc::Sender<DbEmbedding>,
+        tx: mpsc::Sender<(DbEmbedding, Option<AsyncInProgress>)>,
         completed_scan_length: Arc<AtomicU64>,
     ) {
         let semaphore_capacity = self.nr_parallel_queries().get();
@@ -324,14 +328,27 @@ impl Statements {
 
             let range_scan = self.preform_range_scan(begin, end).await;
             if let Ok(embeddings) = range_scan {
-                let tx: Sender<DbEmbedding> = tx.clone();
+                let tx = tx.clone();
                 let scan_length = completed_scan_length.clone();
                 tokio::spawn(async move {
+                    let (tx_in_progress, mut rx_in_progress) = mpsc::channel(1);
                     embeddings
-                        .for_each(|embedding| async {
-                            _ = tx.send(embedding).await;
+                        .for_each(move |embedding| {
+                            let tx = tx.clone();
+                            let tx_in_progress = tx_in_progress.clone();
+                            async move {
+                                _ = tx
+                                    .send((embedding, Some(AsyncInProgress(tx_in_progress))))
+                                    .await;
+                            }
                         })
                         .await;
+
+                    // wait until all in-progress markers are dropped
+                    while rx_in_progress.recv().await.is_some() {
+                        rx_in_progress.len();
+                    }
+
                     //Safety: end > begin, and the range fits into u64
                     scan_length.fetch_add(
                         end.value().abs_diff(begin.value() - 1),
@@ -493,7 +510,7 @@ impl Statements {
 struct CdcConsumerData {
     primary_key_columns: Vec<ColumnName>,
     target_column: ColumnName,
-    tx: mpsc::Sender<DbEmbedding>,
+    tx: mpsc::Sender<(DbEmbedding, Option<AsyncInProgress>)>,
     gregorian_epoch: OffsetDateTime,
 }
 
@@ -563,11 +580,14 @@ impl Consumer for CdcConsumer {
         _ = self
             .0
             .tx
-            .send(DbEmbedding {
-                primary_key,
-                embedding,
-                timestamp,
-            })
+            .send((
+                DbEmbedding {
+                    primary_key,
+                    embedding,
+                    timestamp,
+                },
+                None,
+            ))
             .await;
         Ok(())
     }
@@ -586,7 +606,7 @@ impl CdcConsumerFactory {
     fn new(
         session: Arc<Session>,
         metadata: &IndexMetadata,
-        tx: mpsc::Sender<DbEmbedding>,
+        tx: mpsc::Sender<(DbEmbedding, Option<AsyncInProgress>)>,
     ) -> anyhow::Result<Self> {
         let cluster_state = session.get_cluster_state();
         let table = cluster_state
diff --git a/crates/vector-store/src/index/actor.rs b/crates/vector-store/src/index/actor.rs
@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0
  */
 
+use crate::AsyncInProgress;
 use crate::Distance;
 use crate::Limit;
 use crate::PrimaryKey;
@@ -17,9 +18,11 @@ pub enum Index {
     AddOrReplace {
         primary_key: PrimaryKey,
         embedding: Vector,
+        in_progress: Option<AsyncInProgress>,
     },
     Remove {
         primary_key: PrimaryKey,
+        in_progress: Option<AsyncInProgress>,
     },
     Ann {
         embedding: Vector,
@@ -32,26 +35,40 @@ pub enum Index {
 }
 
 pub(crate) trait IndexExt {
-    async fn add_or_replace(&self, primary_key: PrimaryKey, embedding: Vector);
-    async fn remove(&self, primary_key: PrimaryKey);
+    async fn add_or_replace(
+        &self,
+        primary_key: PrimaryKey,
+        embedding: Vector,
+        in_progress: Option<AsyncInProgress>,
+    );
+    async fn remove(&self, primary_key: PrimaryKey, in_progress: Option<AsyncInProgress>);
     async fn ann(&self, embedding: Vector, limit: Limit) -> AnnR;
     async fn count(&self) -> CountR;
 }
 
 impl IndexExt for mpsc::Sender<Index> {
-    async fn add_or_replace(&self, primary_key: PrimaryKey, embedding: Vector) {
+    async fn add_or_replace(
+        &self,
+        primary_key: PrimaryKey,
+        embedding: Vector,
+        in_progress: Option<AsyncInProgress>,
+    ) {
         self.send(Index::AddOrReplace {
             primary_key,
             embedding,
+            in_progress,
         })
         .await
         .expect("internal actor should receive request");
     }
 
-    async fn remove(&self, primary_key: PrimaryKey) {
-        self.send(Index::Remove { primary_key })
-            .await
-            .expect("internal actor should receive request");
+    async fn remove(&self, primary_key: PrimaryKey, in_progress: Option<AsyncInProgress>) {
+        self.send(Index::Remove {
+            primary_key,
+            in_progress,
+        })
+        .await
+        .expect("internal actor should receive request");
     }
 
     async fn ann(&self, embedding: Vector, limit: Limit) -> AnnR {
diff --git a/crates/vector-store/src/index/opensearch.rs b/crates/vector-store/src/index/opensearch.rs
@@ -255,8 +255,12 @@ async fn process(
         Index::AddOrReplace {
             primary_key,
             embedding,
+            in_progress: _in_progress,
         } => add_or_replace(id, keys, opensearch_key, primary_key, embedding, client).await,
-        Index::Remove { primary_key } => remove(id, keys, primary_key, client).await,
+        Index::Remove {
+            primary_key,
+            in_progress: _in_progress,
+        } => remove(id, keys, primary_key, client).await,
         Index::Ann {
             embedding,
             limit,
diff --git a/crates/vector-store/src/index/usearch.rs b/crates/vector-store/src/index/usearch.rs
@@ -158,11 +158,15 @@ fn process(
         Index::AddOrReplace {
             primary_key,
             embedding,
+            in_progress: _in_progress,
         } => {
             add_or_replace(idx, keys, usearch_key, primary_key, embedding, allocate);
         }
 
-        Index::Remove { primary_key } => {
+        Index::Remove {
+            primary_key,
+            in_progress: _in_progress,
+        } => {
             remove(idx, keys, primary_key);
         }
 
@@ -339,18 +343,21 @@ mod tests {
             .add_or_replace(
                 vec![CqlValue::Int(1), CqlValue::Text("one".to_string())].into(),
                 vec![1., 1., 1.].into(),
+                None,
             )
             .await;
         actor
             .add_or_replace(
                 vec![CqlValue::Int(2), CqlValue::Text("two".to_string())].into(),
                 vec![2., -2., 2.].into(),
+                None,
             )
             .await;
         actor
             .add_or_replace(
                 vec![CqlValue::Int(3), CqlValue::Text("three".to_string())].into(),
                 vec![3., 3., 3.].into(),
+                None,
             )
             .await;
 
@@ -380,6 +387,7 @@ mod tests {
             .add_or_replace(
                 vec![CqlValue::Int(3), CqlValue::Text("three".to_string())].into(),
                 vec![2.1, -2.1, 2.1].into(),
+                None,
             )
             .await;
 
@@ -403,7 +411,10 @@ mod tests {
         .unwrap();
 
         actor
-            .remove(vec![CqlValue::Int(3), CqlValue::Text("three".to_string())].into())
+            .remove(
+                vec![CqlValue::Int(3), CqlValue::Text("three".to_string())].into(),
+                None,
+            )
             .await;
 
         time::timeout(Duration::from_secs(10), async {
diff --git a/crates/vector-store/src/lib.rs b/crates/vector-store/src/lib.rs
@@ -42,6 +42,7 @@ use tokio::runtime::Builder;
 use tokio::runtime::Handle;
 use tokio::signal;
 use tokio::sync::Semaphore;
+use tokio::sync::mpsc;
 use tokio::sync::mpsc::Sender;
 use tokio::task;
 use utoipa::PartialSchema;
@@ -494,6 +495,11 @@ impl From<SocketAddr> for HttpServerConfig {
     }
 }
 
+#[derive(Clone)]
+/// Marker struct to indicate that an async operation is in progress.
+#[allow(dead_code)]
+pub struct AsyncInProgress(mpsc::Sender<()>);
+
 pub fn block_on<Output>(threads: Option<usize>, f: impl AsyncFnOnce() -> Output) -> Output {
     let mut builder = match threads {
         Some(0) | None => Builder::new_multi_thread(),
diff --git a/crates/vector-store/src/monitor_items.rs b/crates/vector-store/src/monitor_items.rs
diff --git a/crates/vector-store/tests/integration/db_basic.rs b/crates/vector-store/tests/integration/db_basic.rs
