refactor(cubestore): Replace invalidate_tables_cache flag with generic post-commit callback API

Introduce a generic `set_post_commit_callback` on `BatchPipe<S>` that fires after a
successful RocksDB commit on the RW-loop thread. This replaces the hard-coded
`invalidate_tables_cache` boolean flag, which was a hack that also ran before the
actual commit.

- `BatchPipe` is now generic over a store type `S` (defaults to `()`), allowing
  callbacks to receive a typed store reference
- `RocksMetaStore::write_operation` passes `&RocksMetaStore` to callbacks, giving
  them direct access to `cached_tables` and other metastore state
- Moved `cached_tables` from `RocksStore` to `RocksMetaStore` where it belongs
- Cache invalidation now only happens when the commit actually succeeds
- `RocksTable` methods are generic over `S` so they work with any `BatchPipe<S>`

Author: ovr

rust/cubestore/cubestore/src/cachestore/cache_rocksstore.rs

@@ -498,14 +498,14 @@ impl RocksCacheStore {
         f: F,
     ) -> Result<R, CubeError>
     where
-        F: for<'a> FnOnce(DbTableRef<'a>, &'a mut BatchPipe) -> Result<R, CubeError>
+        F: for<'a> FnOnce(DbTableRef<'a>, &mut BatchPipe<'a>) -> Result<R, CubeError>
             + Send
             + Sync
             + 'static,
         R: Send + Sync + 'static,
     {
         self.store
-            .write_operation_impl::<F, R>(&self.rw_loop_queue_cf, op_name, f)
+            .write_operation_impl::<F, R, ()>(&self.rw_loop_queue_cf, op_name, f, None)
             .await
     }
 

rust/cubestore/cubestore/src/metastore/mod.rs

@@ -26,6 +26,7 @@ use cuberockstore::rocksdb::{BlockBasedOptions, Cache, Env, MergeOperands, Optio
 use log::info;
 use serde::{Deserialize, Serialize};
 use std::hash::Hash;
+use std::sync::Mutex;
 use std::{env, io::Cursor, sync::Arc};
 
 use crate::config::injection::DIService;
@@ -1355,8 +1356,10 @@ impl RocksStoreDetails for RocksMetaStoreDetails {
     }
 }
 
+#[derive(Clone)]
 pub struct RocksMetaStore {
     store: Arc<RocksStore>,
+    cached_tables: Arc<Mutex<Option<Arc<Vec<TablePath>>>>>,
     disk_space_cache: Arc<RwLock<Option<(HashMap<String, u64>, SystemTime)>>>,
     upload_loop: Arc<WorkerLoop>,
 }
@@ -1379,13 +1382,14 @@ impl RocksMetaStore {
     fn new_from_store(store: Arc<RocksStore>) -> Arc<Self> {
         Arc::new(Self {
             store,
+            cached_tables: Arc::new(Mutex::new(None)),
             disk_space_cache: Arc::new(RwLock::new(None)),
             upload_loop: Arc::new(WorkerLoop::new("Metastore upload")),
         })
     }
 
     pub fn reset_cached_tables(&self) {
-        *self.store.cached_tables.lock().unwrap() = None;
+        *self.cached_tables.lock().unwrap() = None;
     }
 
     pub async fn load_from_dump(
@@ -1512,19 +1516,29 @@ impl RocksMetaStore {
     #[inline(always)]
     pub async fn write_operation<F, R>(&self, op_name: &'static str, f: F) -> Result<R, CubeError>
     where
-        F: for<'a> FnOnce(DbTableRef<'a>, &'a mut BatchPipe) -> Result<R, CubeError>
+        F: for<'a> FnOnce(
+                DbTableRef<'a>,
+                &mut BatchPipe<'a, RocksMetaStore>,
+            ) -> Result<R, CubeError>
             + Send
             + Sync
             + 'static,
         R: Send + Sync + 'static,
     {
-        self.store.write_operation(op_name, f).await
+        self.store
+            .write_operation_impl(
+                &self.store.rw_loop_default_cf,
+                op_name,
+                f,
+                Some(self.clone()),
+            )
+            .await
     }
 
-    fn drop_table_impl(
+    fn drop_table_impl<S>(
         table_id: u64,
         db_ref: DbTableRef,
-        batch_pipe: &mut BatchPipe,
+        batch_pipe: &mut BatchPipe<'_, S>,
     ) -> Result<IdRow<Table>, CubeError> {
         let tables_table = TableRocksTable::new(db_ref.clone());
         let indexes_table = IndexRocksTable::new(db_ref.clone());
@@ -1554,8 +1568,8 @@ impl RocksMetaStore {
 }
 
 impl RocksMetaStore {
-    fn add_index(
-        batch_pipe: &mut BatchPipe,
+    fn add_index<S>(
+        batch_pipe: &mut BatchPipe<'_, S>,
         rocks_index: &IndexRocksTable,
         rocks_partition: &PartitionRocksTable,
         table_cols: &Vec<Column>,
@@ -1585,8 +1599,8 @@ impl RocksMetaStore {
             ),
         }
     }
-    fn add_regular_index(
-        batch_pipe: &mut BatchPipe,
+    fn add_regular_index<S>(
+        batch_pipe: &mut BatchPipe<'_, S>,
         rocks_index: &IndexRocksTable,
         rocks_partition: &PartitionRocksTable,
         table_cols: &Vec<Column>,
@@ -1733,8 +1747,8 @@ impl RocksMetaStore {
         Ok(index_id)
     }
 
-    fn add_aggregate_index(
-        batch_pipe: &mut BatchPipe,
+    fn add_aggregate_index<S>(
+        batch_pipe: &mut BatchPipe<'_, S>,
         rocks_index: &IndexRocksTable,
         rocks_partition: &PartitionRocksTable,
         table_cols: &Vec<Column>,
@@ -1848,9 +1862,9 @@ impl RocksMetaStore {
     }
 
     // Must be run under write_operation(). Returns activated row count.
-    fn activate_chunks_impl(
+    fn activate_chunks_impl<S>(
         db_ref: DbTableRef,
-        batch_pipe: &mut BatchPipe,
+        batch_pipe: &mut BatchPipe<'_, S>,
         uploaded_chunk_ids: &[(u64, Option<u64>)],
         replay_handle_id: Option<u64>,
     ) -> Result<(u64, HashMap</*partition_id*/ u64, /*rows*/ u64>), CubeError> {
@@ -1907,7 +1921,9 @@ impl MetaStore for RocksMetaStore {
         if_not_exists: bool,
     ) -> Result<IdRow<Schema>, CubeError> {
         self.write_operation("create_schema", move |db_ref, batch_pipe| {
-            batch_pipe.invalidate_tables_cache();
+            batch_pipe.set_post_commit_callback(|metastore| {
+                *metastore.cached_tables.lock().unwrap() = None;
+            });
             let table = SchemaRocksTable::new(db_ref.clone());
             if if_not_exists {
                 let rows = table.get_rows_by_index(&schema_name, &SchemaRocksIndex::Name)?;
@@ -1968,7 +1984,9 @@ impl MetaStore for RocksMetaStore {
         new_schema_name: String,
     ) -> Result<IdRow<Schema>, CubeError> {
         self.write_operation("rename_schema", move |db_ref, batch_pipe| {
-            batch_pipe.invalidate_tables_cache();
+            batch_pipe.set_post_commit_callback(|metastore| {
+                *metastore.cached_tables.lock().unwrap() = None;
+            });
             let table = SchemaRocksTable::new(db_ref.clone());
             let existing_keys =
                 table.get_row_ids_by_index(&old_schema_name, &SchemaRocksIndex::Name)?;
@@ -1992,7 +2010,9 @@ impl MetaStore for RocksMetaStore {
         new_schema_name: String,
     ) -> Result<IdRow<Schema>, CubeError> {
         self.write_operation("rename_schema_by_id", move |db_ref, batch_pipe| {
-            batch_pipe.invalidate_tables_cache();
+            batch_pipe.set_post_commit_callback(|metastore| {
+                *metastore.cached_tables.lock().unwrap() = None;
+            });
             let table = SchemaRocksTable::new(db_ref.clone());
 
             let old_schema = table.get_row(schema_id)?.unwrap();
@@ -2008,7 +2028,9 @@ impl MetaStore for RocksMetaStore {
     #[tracing::instrument(level = "trace", skip(self))]
     async fn delete_schema(&self, schema_name: String) -> Result<(), CubeError> {
         self.write_operation("delete_schema", move |db_ref, batch_pipe| {
-            batch_pipe.invalidate_tables_cache();
+            batch_pipe.set_post_commit_callback(|metastore| {
+                *metastore.cached_tables.lock().unwrap() = None;
+            });
             let table = SchemaRocksTable::new(db_ref.clone());
             let existing_keys =
                 table.get_row_ids_by_index(&schema_name, &SchemaRocksIndex::Name)?;
@@ -2035,7 +2057,9 @@ impl MetaStore for RocksMetaStore {
     #[tracing::instrument(level = "trace", skip(self))]
     async fn delete_schema_by_id(&self, schema_id: u64) -> Result<(), CubeError> {
         self.write_operation("delete_schema_by_id", move |db_ref, batch_pipe| {
-            batch_pipe.invalidate_tables_cache();
+            batch_pipe.set_post_commit_callback(|metastore| {
+                *metastore.cached_tables.lock().unwrap() = None;
+            });
             let tables = TableRocksTable::new(db_ref.clone()).all_rows()?;
             if tables
                 .into_iter()
@@ -2100,7 +2124,9 @@ impl MetaStore for RocksMetaStore {
         extension: Option<String>,
     ) -> Result<IdRow<Table>, CubeError> {
         self.write_operation("create_table", move |db_ref, batch_pipe| {
-            batch_pipe.invalidate_tables_cache();
+            batch_pipe.set_post_commit_callback(|metastore| {
+                *metastore.cached_tables.lock().unwrap() = None;
+            });
             if drop_if_exists {
                 if let Ok(exists_table) = get_table_impl(db_ref.clone(), schema_name.clone(), table_name.clone()) {
                     RocksMetaStore::drop_table_impl(exists_table.get_id(), db_ref.clone(), batch_pipe)?;
@@ -2295,7 +2321,9 @@ impl MetaStore for RocksMetaStore {
     #[tracing::instrument(level = "trace", skip(self))]
     async fn table_ready(&self, id: u64, is_ready: bool) -> Result<IdRow<Table>, CubeError> {
         self.write_operation("table_ready", move |db_ref, batch_pipe| {
-            batch_pipe.invalidate_tables_cache();
+            batch_pipe.set_post_commit_callback(|metastore| {
+                *metastore.cached_tables.lock().unwrap() = None;
+            });
             let rocks_table = TableRocksTable::new(db_ref.clone());
             Ok(rocks_table.update_with_fn(id, |r| r.update_is_ready(is_ready), batch_pipe)?)
         })
@@ -2305,7 +2333,9 @@ impl MetaStore for RocksMetaStore {
     #[tracing::instrument(level = "trace", skip(self))]
     async fn seal_table(&self, id: u64) -> Result<IdRow<Table>, CubeError> {
         self.write_operation("seal_table", move |db_ref, batch_pipe| {
-            batch_pipe.invalidate_tables_cache();
+            batch_pipe.set_post_commit_callback(|metastore| {
+                *metastore.cached_tables.lock().unwrap() = None;
+            });
             let rocks_table = TableRocksTable::new(db_ref.clone());
             Ok(rocks_table.update_with_fn(id, |r| r.update_sealed(true), batch_pipe)?)
         })
@@ -2334,13 +2364,16 @@ impl MetaStore for RocksMetaStore {
         self.write_operation(
             "update_location_download_size",
             move |db_ref, batch_pipe| {
-                batch_pipe.invalidate_tables_cache();
+                batch_pipe.set_post_commit_callback(|metastore| {
+                    *metastore.cached_tables.lock().unwrap() = None;
+                });
+
                 let rocks_table = TableRocksTable::new(db_ref.clone());
-                Ok(rocks_table.update_with_res_fn(
+                rocks_table.update_with_res_fn(
                     id,
                     |r| r.update_location_download_size(&location, download_size),
                     batch_pipe,
-                )?)
+                )
             },
         )
         .await
@@ -2393,14 +2426,14 @@ impl MetaStore for RocksMetaStore {
             })
             .await
         } else {
-            let cache = self.store.cached_tables.clone();
+            let cache = self.cached_tables.clone();
 
             if let Some(t) = cube_ext::spawn_blocking(move || cache.lock().unwrap().clone()).await?
             {
                 return Ok(t);
             }
 
-            let cache = self.store.cached_tables.clone();
+            let cache = self.cached_tables.clone();
             // Can't do read_operation_out_of_queue as we need to update cache on the same thread where it's dropped
             self.read_operation("get_tables_with_path", move |db_ref| {
                 let cached_tables = { cache.lock().unwrap().clone() };
@@ -2465,7 +2498,9 @@ impl MetaStore for RocksMetaStore {
     #[tracing::instrument(level = "trace", skip(self))]
     async fn drop_table(&self, table_id: u64) -> Result<IdRow<Table>, CubeError> {
         self.write_operation("drop_table", move |db_ref, batch_pipe| {
-            batch_pipe.invalidate_tables_cache();
+            batch_pipe.set_post_commit_callback(|metastore| {
+                *metastore.cached_tables.lock().unwrap() = None;
+            });
             RocksMetaStore::drop_table_impl(table_id, db_ref, batch_pipe)
         })
         .await
@@ -4817,9 +4852,9 @@ fn get_default_index_impl(db_ref: DbTableRef, table_id: u64) -> Result<IdRow<Ind
 /// Note that [current_active] and [new_active] are snapshots at some older point in time. The
 /// relevant partitions might be dropped or changed by the time this function runs. Implementation
 /// must take great care to avoid inconsistencies caused by this.
-fn swap_active_partitions_impl(
+fn swap_active_partitions_impl<S>(
     db_ref: DbTableRef,
-    batch_pipe: &mut BatchPipe,
+    batch_pipe: &mut BatchPipe<'_, S>,
     current_active: &[(IdRow<Partition>, Vec<IdRow<Chunk>>)],
     new_active: &[(IdRow<Partition>, u64)],
     mut update_new_partition_stats: impl FnMut(/*index*/ usize, &Partition) -> Partition,
@@ -7132,11 +7167,11 @@ mod tests {
 }
 
 impl RocksMetaStore {
-    fn swap_chunks_impl(
+    fn swap_chunks_impl<S>(
         deactivate_ids: Vec<u64>,
         uploaded_ids_and_sizes: Vec<(u64, Option<u64>)>,
         db_ref: DbTableRef,
-        batch_pipe: &mut BatchPipe,
+        batch_pipe: &mut BatchPipe<'_, S>,
         check_rows: bool,
         new_replay_handle_id: Option<u64>,
     ) -> Result<(), CubeError> {
@@ -7234,9 +7269,9 @@ impl RocksMetaStore {
 }
 
 impl RocksMetaStore {
-    fn drop_index(
+    fn drop_index<S>(
         db: DbTableRef,
-        pipe: &mut BatchPipe,
+        pipe: &mut BatchPipe<'_, S>,
         index_id: u64,
         update_multi_partitions: bool,
     ) -> Result<(), CubeError> {