chore(cubestore): Rockstore - use named threads per RWLoop/Store (#9906)

ovr · web-flow · commit eace9dbd50b6 · 2025-08-19T17:12:08.000+02:00
diff --git a/rust/cubestore/cubestore/benches/cachestore_queue.rs b/rust/cubestore/cubestore/benches/cachestore_queue.rs
@@ -43,12 +43,18 @@ fn generate_queue_path(queue_path: &str, queue_id: usize) -> String {
 }
 
 async fn do_insert(
+    cachestore_name: &str,
     cachestore: &Arc<RocksCacheStore>,
     total: usize,
     size_kb: usize,
     queue_path: &str,
     insert_id_padding: usize,
 ) {
+    println!(
+        "[Preparing] {}: Inserting {} items into queue: {}",
+        cachestore_name, total, queue_path
+    );
+
     for i in 0..total {
         let fut = cachestore.queue_add(QueueAddPayload {
             path: generate_queue_path(queue_path, i + insert_id_padding),
@@ -60,12 +66,13 @@ async fn do_insert(
         let res = fut.await;
         assert!(res.is_ok());
     }
+
+    println!("[Preparing] {}: Done", cachestore_name);
 }
 
 fn do_insert_bench(c: &mut Criterion, runtime: &Runtime, total: usize, size_kb: usize) {
-    let cachestore = runtime.block_on(async {
-        prepare_cachestore(&format!("cachestore_queue_add_{}", size_kb)).unwrap()
-    });
+    let cachestore_name = format!("cachestore_queue_add_{}", size_kb);
+    let cachestore = runtime.block_on(async { prepare_cachestore(&cachestore_name).unwrap() });
 
     c.bench_with_input(
         BenchmarkId::new(format!("queue_add queues:1, size:{} kb", size_kb), total),
@@ -78,6 +85,7 @@ fn do_insert_bench(c: &mut Criterion, runtime: &Runtime, total: usize, size_kb:
                 insert_id_padding += total;
 
                 do_insert(
+                    &cachestore_name,
                     &cachestore,
                     *total,
                     *size_kb,
@@ -116,18 +124,24 @@ fn do_list_bench(
     total: usize,
 ) {
     let cachestore = runtime.block_on(async {
-        let cachestore = prepare_cachestore(&format!(
+        let cachestore_name = format!(
             "cachestore_queue_list_{}_{}",
             format!("{:?}", status_filter).to_ascii_lowercase(),
             size_kb
-        ))
-        .unwrap();
-
-        do_insert(&cachestore, per_queue, size_kb, "STANDALONE#queue", 0).await;
-        do_insert(&cachestore, per_queue, size_kb, "STANDALONE#queue", 0).await;
-        do_insert(&cachestore, per_queue, size_kb, "STANDALONE#queue", 0).await;
-        do_insert(&cachestore, per_queue, size_kb, "STANDALONE#queue", 0).await;
-        do_insert(&cachestore, per_queue, size_kb, "STANDALONE#queue", 0).await;
+        );
+        let cachestore = prepare_cachestore(&cachestore_name).unwrap();
+
+        for idx in 0..5 {
+            do_insert(
+                &cachestore_name,
+                &cachestore,
+                per_queue,
+                size_kb,
+                &format!("STANDALONE#queue{}", idx + 1),
+                0,
+            )
+            .await;
+        }
 
         cachestore
     });
@@ -148,11 +162,11 @@ fn do_list_bench(
     );
 }
 
-async fn do_get(cachestore: &Arc<RocksCacheStore>, total: usize) {
+async fn do_get(cachestore: &Arc<RocksCacheStore>, total: usize, total_queues: usize) {
     for i in 0..total {
         let fut = cachestore.queue_get(QueueKey::ByPath(generate_queue_path(
-            "STANDALONE#queue",
-            i + ((i - 1) * 5),
+            &format!("STANDALONE#queue{}", (i % total_queues) + 1),
+            i,
         )));
 
         let res = fut.await;
@@ -166,15 +180,23 @@ fn do_get_bench(
     per_queue: usize,
     size_kb: usize,
     total: usize,
+    queues: usize,
 ) {
     let cachestore = runtime.block_on(async {
-        let cachestore = prepare_cachestore(&format!("cachestore_queue_get_{}", size_kb)).unwrap();
-
-        do_insert(&cachestore, per_queue, size_kb, "STANDALONE#queue", 0).await;
-        do_insert(&cachestore, per_queue, size_kb, "STANDALONE#queue", 0).await;
-        do_insert(&cachestore, per_queue, size_kb, "STANDALONE#queue", 0).await;
-        do_insert(&cachestore, per_queue, size_kb, "STANDALONE#queue", 0).await;
-        do_insert(&cachestore, per_queue, size_kb, "STANDALONE#queue", 0).await;
+        let cachestore_name = format!("cachestore_queue_get_{}", size_kb);
+        let cachestore = prepare_cachestore(&cachestore_name).unwrap();
+
+        for idx in 0..queues {
+            do_insert(
+                &cachestore_name,
+                &cachestore,
+                per_queue,
+                size_kb,
+                &format!("STANDALONE#queue{}", idx + 1),
+                0,
+            )
+            .await;
+        }
 
         cachestore
     });
@@ -189,7 +211,7 @@ fn do_get_bench(
         ),
         &total,
         |b, total| {
-            b.to_async(runtime).iter(|| do_get(&cachestore, *total));
+            b.to_async(runtime).iter(|| do_get(&cachestore, *total, 3));
         },
     );
 }
@@ -205,7 +227,7 @@ fn do_benches(c: &mut Criterion) {
     do_list_bench(c, &runtime, Some(QueueItemStatus::Pending), 1_000, 128, 128);
     do_list_bench(c, &runtime, Some(QueueItemStatus::Active), 1_000, 128, 128);
 
-    do_get_bench(c, &runtime, 10_000, 128, 128);
+    do_get_bench(c, &runtime, 2_500, 128, 128, 4na);
 }
 
 criterion_group!(benches, do_benches);
diff --git a/rust/cubestore/cubestore/benches/metastore.rs b/rust/cubestore/cubestore/benches/metastore.rs
@@ -140,41 +140,44 @@ fn do_get_tables_with_path_bench(
     );
 }
 
-async fn do_cold_cache_test(num_schemas: usize, tables_per_schema: usize) {
-    let fresh_metastore = prepare_metastore("cold_cache_fresh").unwrap();
-    populate_metastore(&fresh_metastore, num_schemas, tables_per_schema)
-        .await
-        .unwrap();
-    let result = fresh_metastore.get_tables_with_path(false).await;
-    assert!(result.is_ok());
-}
-
-async fn do_warm_cache_test(metastore: &Arc<RocksMetaStore>) {
-    let result = metastore.get_tables_with_path(false).await;
-    assert!(result.is_ok());
-}
-
 fn do_cold_vs_warm_cache_bench(
     c: &mut Criterion,
     runtime: &Runtime,
     num_schemas: usize,
     tables_per_schema: usize,
 ) {
-    let metastore = runtime.block_on(async {
-        let metastore = prepare_metastore("cold_warm_cache").unwrap();
+    let cold_metastore = runtime.block_on(async {
+        let metastore = prepare_metastore("warm_cache").unwrap();
+        populate_metastore(&metastore, num_schemas, tables_per_schema)
+            .await
+            .unwrap();
+        metastore
+    });
+
+    let warm_metastore = runtime.block_on(async {
+        let metastore = prepare_metastore("cold_cache").unwrap();
         populate_metastore(&metastore, num_schemas, tables_per_schema)
             .await
             .unwrap();
         metastore
     });
 
     c.bench_function("get_tables_with_path_cold_cache", |b| {
-        b.to_async(runtime)
-            .iter(|| do_cold_cache_test(num_schemas, tables_per_schema));
+        b.to_async(runtime).iter_batched(
+            || cold_metastore.reset_cached_tables(),
+            async |_| {
+                let result = cold_metastore.get_tables_with_path(false).await;
+                assert!(result.is_ok());
+            },
+            criterion::BatchSize::SmallInput,
+        );
     });
 
     c.bench_function("get_tables_with_path_warm_cache", |b| {
-        b.to_async(runtime).iter(|| do_warm_cache_test(&metastore));
+        b.to_async(runtime).iter(async || {
+            let result = warm_metastore.get_tables_with_path(false).await;
+            assert!(result.is_ok());
+        });
     });
 }
 
diff --git a/rust/cubestore/cubestore/src/cachestore/cache_rocksstore.rs b/rust/cubestore/cubestore/src/cachestore/cache_rocksstore.rs
@@ -223,7 +223,7 @@ impl RocksCacheStore {
             cache_eviction_manager,
             upload_loop: Arc::new(WorkerLoop::new("Cachestore upload")),
             metrics_loop: Arc::new(WorkerLoop::new("Cachestore metrics")),
-            rw_loop_queue_cf: RocksStoreRWLoop::new("queue"),
+            rw_loop_queue_cf: RocksStoreRWLoop::new("cachestore", "queue"),
         }))
     }
 
diff --git a/rust/cubestore/cubestore/src/metastore/mod.rs b/rust/cubestore/cubestore/src/metastore/mod.rs
@@ -1388,6 +1388,10 @@ impl RocksMetaStore {
         })
     }
 
+    pub fn reset_cached_tables(&self) {
+        *self.store.cached_tables.lock().unwrap() = None;
+    }
+
     pub async fn load_from_dump(
         path: &Path,
         dump_path: &Path,
diff --git a/rust/cubestore/cubestore/src/metastore/rocks_store.rs b/rust/cubestore/cubestore/src/metastore/rocks_store.rs
@@ -1,7 +1,6 @@
 use crate::config::ConfigObj;
 use crate::metastore::table::TablePath;
 use crate::metastore::{MetaStoreEvent, MetaStoreFs};
-use crate::util::aborting_join_handle::AbortingJoinHandle;
 use crate::util::time_span::warn_long;
 
 use crate::CubeError;
@@ -813,36 +812,42 @@ pub type RocksStoreRWLoopFn = Box<dyn FnOnce() -> Result<(), CubeError> + Send +
 pub struct RocksStoreRWLoop {
     name: &'static str,
     tx: tokio::sync::mpsc::Sender<RocksStoreRWLoopFn>,
-    _join_handle: Arc<AbortingJoinHandle<()>>,
 }
 
 impl RocksStoreRWLoop {
-    pub fn new(name: &'static str) -> Self {
+    pub fn new(store_name: &'static str, name: &'static str) -> Self {
         let (tx, mut rx) = tokio::sync::mpsc::channel::<RocksStoreRWLoopFn>(32_768);
 
-        let join_handle = cube_ext::spawn_blocking(move || loop {
-            if let Some(fun) = rx.blocking_recv() {
-                match std::panic::catch_unwind(std::panic::AssertUnwindSafe(fun)) {
-                    Err(panic_payload) => {
-                        let restore_error = CubeError::from_panic_payload(panic_payload);
-                        log::error!("Panic during read write loop execution: {}", restore_error);
-                    }
-                    Ok(res) => {
-                        if let Err(e) = res {
-                            log::error!("Error during read write loop execution: {}", e);
+        let thread_name = format!("{}-{}-rwloop", store_name, name);
+        std::thread::Builder::new()
+            .name(thread_name)
+            .spawn(move || loop {
+                if let Some(fun) = rx.blocking_recv() {
+                    match std::panic::catch_unwind(std::panic::AssertUnwindSafe(fun)) {
+                        Err(panic_payload) => {
+                            let restore_error = CubeError::from_panic_payload(panic_payload);
+                            log::error!(
+                                "Panic during read write loop execution: {}",
+                                restore_error
+                            );
+                        }
+                        Ok(res) => {
+                            if let Err(e) = res {
+                                log::error!("Error during read write loop execution: {}", e);
+                            }
                         }
                     }
+                } else {
+                    return;
                 }
-            } else {
-                return;
-            }
-        });
+            })
+            .expect(&format!(
+                "Failed to spawn RWLoop thread for store '{}', name '{}'",
+                store_name, name
+            ));
 
-        Self {
-            name,
-            tx,
-            _join_handle: Arc::new(AbortingJoinHandle::new(join_handle)),
-        }
+        // Thread handle is intentionally dropped - thread will exit when tx is dropped
+        Self { name, tx }
     }
 
     pub async fn schedule(&self, fun: RocksStoreRWLoopFn) -> Result<(), CubeError> {
@@ -928,7 +933,7 @@ impl RocksStore {
             snapshots_upload_stopped: Arc::new(AsyncMutex::new(false)),
             config,
             cached_tables: Arc::new(Mutex::new(None)),
-            rw_loop_default_cf: RocksStoreRWLoop::new("default"),
+            rw_loop_default_cf: RocksStoreRWLoop::new("metastore", "default"),
             details,
         };
 
diff --git a/rust/cubestore/cubestore/src/metastore/rocks_table.rs b/rust/cubestore/cubestore/src/metastore/rocks_table.rs
@@ -573,7 +573,7 @@ pub trait RocksTable: BaseRocksTable + Debug + Send + Sync {
     fn migration_check_table(&self) -> Result<(), CubeError> {
         let snapshot = self.snapshot();
 
-        let table_info = snapshot.get_pinned(
+        let table_info = snapshot.get(
             &RowKey::TableInfo {
                 table_id: Self::table_id(),
             }
@@ -633,7 +633,7 @@ pub trait RocksTable: BaseRocksTable + Debug + Send + Sync {
     fn migration_check_indexes(&self) -> Result<(), CubeError> {
         let snapshot = self.snapshot();
         for index in Self::indexes().into_iter() {
-            let index_info = snapshot.get_pinned(
+            let index_info = snapshot.get(
                 &RowKey::SecondaryIndexInfo {
                     index_id: Self::index_id(index.get_id()),
                 }
@@ -977,7 +977,7 @@ pub trait RocksTable: BaseRocksTable + Debug + Send + Sync {
             RowKey::SecondaryIndex(Self::index_id(index_id), secondary_key_hash, row_id);
         let secondary_index_key = secondary_index_row_key.to_bytes();
 
-        if let Some(secondary_key_bytes) = self.db().get_pinned(&secondary_index_key)? {
+        if let Some(secondary_key_bytes) = self.db().get(&secondary_index_key)? {
             let index_value_version = RocksSecondaryIndex::value_version(secondary_index);
             let new_value = match RocksSecondaryIndexValue::from_bytes(
                 &secondary_key_bytes,
@@ -1102,6 +1102,11 @@ pub trait RocksTable: BaseRocksTable + Debug + Send + Sync {
 
     fn get_row(&self, row_id: u64) -> Result<Option<IdRow<Self::T>>, CubeError> {
         let ref db = self.snapshot();
+        // Use pinned access to avoid double copying. While zero-copy deserialization would be ideal, but
+        // we're using flex buffers with serde, which copies String values during deserialization. There is a way
+        // to solve it by using &[u8] types, but it's not worth the effort right now.
+        //
+        // Let's avoid copying on lookup row, but doing copy on deserialization.
         let res = db.get_pinned(RowKey::Table(Self::table_id(), row_id).to_bytes())?;
 
         if let Some(buffer) = res {

Original file line number	Diff line number	Diff line change
`@@ -223,7 +223,7 @@ impl RocksCacheStore {`
`223`	`223`	`cache_eviction_manager,`
`224`	`224`	`upload_loop: Arc::new(WorkerLoop::new("Cachestore upload")),`
`225`	`225`	`metrics_loop: Arc::new(WorkerLoop::new("Cachestore metrics")),`
`226`		`- rw_loop_queue_cf: RocksStoreRWLoop::new("queue"),`
	`226`	`+ rw_loop_queue_cf: RocksStoreRWLoop::new("cachestore", "queue"),`
`227`	`227`	`}))`
`228`	`228`	`}`
`229`	`229`