Move a modify_indexes::Db::update_items_count into a db actor

This is a part of #3.

The change moves update_items_count method from modify_indexes into a db
actor. Next paches will move other methods into a db actor.

Author: ewienik

src/db.rs

@@ -9,6 +9,7 @@ use crate::Dimensions;
 use crate::ExpansionAdd;
 use crate::ExpansionSearch;
 use crate::IndexId;
+use crate::IndexItemsCount;
 use crate::IndexMetadata;
 use crate::IndexVersion;
 use crate::KeyspaceName;
@@ -64,6 +65,11 @@ pub(crate) enum Db {
         #[allow(clippy::type_complexity)]
         tx: oneshot::Sender<anyhow::Result<Option<(Connectivity, ExpansionAdd, ExpansionSearch)>>>,
     },
+
+    UpdateItemsCount {
+        id: IndexId,
+        items_count: IndexItemsCount,
+    },
 }
 
 pub(crate) trait DbExt {
@@ -90,6 +96,12 @@ pub(crate) trait DbExt {
         &self,
         id: IndexId,
     ) -> anyhow::Result<Option<(Connectivity, ExpansionAdd, ExpansionSearch)>>;
+
+    async fn update_items_count(
+        &self,
+        id: IndexId,
+        items_count: IndexItemsCount,
+    ) -> anyhow::Result<()>;
 }
 
 impl DbExt for mpsc::Sender<Db> {
@@ -151,6 +163,15 @@ impl DbExt for mpsc::Sender<Db> {
         self.send(Db::GetIndexParams { id, tx }).await?;
         rx.await?
     }
+
+    async fn update_items_count(
+        &self,
+        id: IndexId,
+        items_count: IndexItemsCount,
+    ) -> anyhow::Result<()> {
+        self.send(Db::UpdateItemsCount { id, items_count }).await?;
+        Ok(())
+    }
 }
 
 pub(crate) async fn new(db_session: Arc<Session>) -> anyhow::Result<mpsc::Sender<Db>> {
@@ -207,6 +228,13 @@ async fn process(statements: Arc<Statements>, msg: Db) {
         Db::GetIndexParams { id, tx } => tx
             .send(statements.get_index_params(id).await)
             .unwrap_or_else(|_| warn!("db::process: Db::GetIndexParams: unable to send response")),
+
+        Db::UpdateItemsCount { id, items_count } => {
+            statements
+                .update_items_count(id, items_count)
+                .await
+                .unwrap_or_else(|err| warn!("db::process: Db::UpdateItemsCount: {err}"));
+        }
     }
 }
 
@@ -232,6 +260,7 @@ struct Statements {
     st_get_index_target_type: PreparedStatement,
     re_get_index_target_type: Regex,
     st_get_index_params: PreparedStatement,
+    st_update_items_count: PreparedStatement,
 }
 
 impl Statements {
@@ -265,6 +294,11 @@ impl Statements {
                 .await
                 .context("ST_GET_INDEX_PARAMS")?,
 
+            st_update_items_count: session
+                .prepare(Self::ST_UPDATE_ITEMS_COUNT)
+                .await
+                .context("ST_UPDATE_ITEMS_COUNT")?,
+
             session,
         })
     }
@@ -410,4 +444,21 @@ impl Statements {
             .try_next()
             .await?)
     }
+
+    const ST_UPDATE_ITEMS_COUNT: &str = "
+        UPDATE vector_benchmark.vector_indexes
+            SET indexed_elements_count = ?
+            WHERE id = ?
+        ";
+
+    async fn update_items_count(
+        &self,
+        id: IndexId,
+        items_count: IndexItemsCount,
+    ) -> anyhow::Result<()> {
+        self.session
+            .execute_unpaged(&self.st_update_items_count, (items_count, id))
+            .await?;
+        Ok(())
+    }
 }

src/engine.rs

@@ -106,7 +106,6 @@ pub(crate) async fn new(
 
                     let Ok(index_actor) = index::new(
                         id.clone(),
-                        modify_actor.clone(),
                         db.clone(),
                         metadata.dimensions,
                         metadata.connectivity,

src/index.rs

@@ -3,6 +3,8 @@
  * SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0
  */
 
+use crate::db::Db;
+use crate::db::DbExt;
 use crate::Connectivity;
 use crate::Dimensions;
 use crate::Distance;
@@ -13,25 +15,22 @@ use crate::IndexId;
 use crate::IndexItemsCount;
 use crate::Key;
 use crate::Limit;
-use crate::db::Db;
-use crate::modify_indexes::ModifyIndexes;
-use crate::modify_indexes::ModifyIndexesExt;
 use anyhow::anyhow;
 use std::num::NonZeroUsize;
-use std::sync::Arc;
-use std::sync::RwLock;
 use std::sync::atomic::AtomicU32;
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
+use std::sync::Arc;
+use std::sync::RwLock;
 use tokio::sync::mpsc;
 use tokio::sync::oneshot;
 use tokio::time;
-use tracing::Instrument;
 use tracing::debug;
 use tracing::debug_span;
 use tracing::error;
 use tracing::info;
 use tracing::warn;
+use tracing::Instrument;
 use usearch::IndexOptions;
 use usearch::ScalarKind;
 
@@ -89,8 +88,7 @@ impl IndexExt for mpsc::Sender<Index> {
 
 pub(crate) fn new(
     id: IndexId,
-    modify_actor: mpsc::Sender<ModifyIndexes>,
-    _db: mpsc::Sender<Db>,
+    db: mpsc::Sender<Db>,
     dimensions: Dimensions,
     connectivity: Connectivity,
     expansion_add: ExpansionAdd,
@@ -119,9 +117,11 @@ pub(crate) fn new(
             async move {
                 let mut items_count_db = IndexItemsCount(0);
                 let items_count = Arc::new(AtomicU32::new(0));
-                modify_actor
-                    .update_items_count(id.clone(), items_count_db)
-                    .await;
+
+                db.update_items_count(id.clone(), items_count_db)
+                    .await
+                    .unwrap_or_else(|err| warn!("index::new: unable update items count: {err}"));
+
                 let mut housekeeping_interval = time::interval(time::Duration::from_secs(1));
                 let idx_lock = Arc::new(RwLock::new(()));
                 let counter_add = Arc::new(AtomicUsize::new(0));
@@ -131,7 +131,7 @@ pub(crate) fn new(
                     tokio::select! {
                         _ = housekeeping_interval.tick() => {
                             housekeeping(
-                                &modify_actor,
+                                &db,
                                 id.clone(),
                                 &mut items_count_db,
                                 &items_count,
@@ -181,7 +181,7 @@ pub(crate) fn new(
 }
 
 async fn housekeeping(
-    modify_actor: &mpsc::Sender<ModifyIndexes>,
+    db: &mpsc::Sender<Db>,
     id: IndexId,
     items_count_db: &mut IndexItemsCount,
     items_count: &AtomicU32,
@@ -200,7 +200,9 @@ async fn housekeeping(
     if items != items_count_db.0 {
         debug!("housekeeping update items count: {items_count_db}",);
         items_count_db.0 = items;
-        modify_actor.update_items_count(id, *items_count_db).await;
+        db.update_items_count(id, *items_count_db)
+            .await
+            .unwrap_or_else(|err| warn!("index::housekeeping: unable update items count: {err}"));
     }
 }
 

src/modify_indexes.rs

@@ -3,9 +3,8 @@
  * SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0
  */
 
-use crate::IndexId;
-use crate::IndexItemsCount;
 use crate::db;
+use crate::IndexId;
 use anyhow::Context;
 use scylla::client::session::Session;
 use scylla::statement::prepared::PreparedStatement;
@@ -15,29 +14,14 @@ use tokio::sync::mpsc::Sender;
 use tracing::warn;
 
 pub(crate) enum ModifyIndexes {
-    UpdateItemsCount {
-        id: IndexId,
-        items_count: IndexItemsCount,
-    },
-    Del {
-        id: IndexId,
-    },
+    Del { id: IndexId },
 }
 
 pub(crate) trait ModifyIndexesExt {
-    async fn update_items_count(&self, id: IndexId, items_count: IndexItemsCount);
     async fn del(&self, id: IndexId);
 }
 
 impl ModifyIndexesExt for Sender<ModifyIndexes> {
-    async fn update_items_count(&self, id: IndexId, items_count: IndexItemsCount) {
-        self.send(ModifyIndexes::UpdateItemsCount { id, items_count })
-            .await
-            .unwrap_or_else(|err| {
-                warn!("ModifyIndexesExt::update_items_count: unable to send request: {err}")
-            });
-    }
-
     async fn del(&self, id: IndexId) {
         self.send(ModifyIndexes::Del { id })
             .await
@@ -58,12 +42,6 @@ pub(crate) async fn new(
     tokio::spawn(async move {
         while let Some(msg) = rx.recv().await {
             match msg {
-                ModifyIndexes::UpdateItemsCount { id, items_count } => db
-                    .update_items_count(id, items_count)
-                    .await
-                    .unwrap_or_else(|err| {
-                        warn!("modify_indexes: unable to update items count for index in db: {err}")
-                    }),
                 ModifyIndexes::Del { id } => db.remove_index(id).await.unwrap_or_else(|err| {
                     warn!("modify_indexes: unable to remove index from db: {err}")
                 }),
@@ -75,17 +53,12 @@ pub(crate) async fn new(
 
 struct Db {
     session: Arc<Session>,
-    st_update_items_count: PreparedStatement,
     st_remove_index: PreparedStatement,
 }
 
 impl Db {
     async fn new(session: Arc<Session>) -> anyhow::Result<Self> {
         Ok(Self {
-            st_update_items_count: session
-                .prepare(Self::UPDATE_ITEMS_COUNT)
-                .await
-                .context("UPDATE_ITEMS_COUNT")?,
             st_remove_index: session
                 .prepare(Self::REMOVE_INDEX)
                 .await
@@ -94,22 +67,6 @@ impl Db {
         })
     }
 
-    const UPDATE_ITEMS_COUNT: &str = "
-        UPDATE vector_benchmark.vector_indexes
-            SET indexed_elements_count = ?
-            WHERE id = ?
-        ";
-    async fn update_items_count(
-        &self,
-        id: IndexId,
-        items_count: IndexItemsCount,
-    ) -> anyhow::Result<()> {
-        self.session
-            .execute_unpaged(&self.st_update_items_count, (items_count, id))
-            .await?;
-        Ok(())
-    }
-
     const REMOVE_INDEX: &str = "DELETE FROM vector_benchmark.vector_indexes WHERE id = ?";
     async fn remove_index(&self, id: IndexId) -> anyhow::Result<()> {
         self.session