chore(cubestore): Allow to configure checksum_type/cache_capacity for RocksDB (#6335)

Author: ovr

rust/cubestore/Cargo.lock

@@ -53,7 +53,8 @@ tar = "0.4.38"
 
 [dev-dependencies]
 criterion = { version = "0.4.0", features = ["async_tokio", "html_reports"] }
-rocksdb = { version = "0.20.1", default-features = false, features = ["bzip2"] }
+# Awaiting 0.20.2
+rocksdb = { git = "https://github.com/rust-rocksdb/rust-rocksdb", rev = "44dc84171adefbbe75a25b72c35f773a643655a0", default-features = false, features = ["bzip2"] }
 
 [[bench]]
 name = "in_process"

rust/cubestore/cubestore-sql-tests/Cargo.toml

@@ -41,7 +41,8 @@ simple_logger = "1.7.0"
 async-trait = "0.1.36"
 actix-rt = "2.7.0"
 regex = "1.3.9"
-rocksdb = { version = "0.20.1", default-features = false, features = ["bzip2"] }
+# Awaiting 0.20.2
+rocksdb = { git = "https://github.com/rust-rocksdb/rust-rocksdb", rev = "44dc84171adefbbe75a25b72c35f773a643655a0", default-features = false, features = ["bzip2"] }
 uuid = { version = "0.8", features = ["serde", "v4"] }
 num = "0.3.0"
 enum_primitive = "0.1.1"

rust/cubestore/cubestore/Cargo.toml

@@ -24,7 +24,7 @@ use crate::CubeError;
 use async_trait::async_trait;
 
 use futures_timer::Delay;
-use rocksdb::{BlockBasedOptions, Options, DB};
+use rocksdb::{BlockBasedOptions, Cache, Options, DB};
 
 use crate::cachestore::compaction::CompactionPreloadedState;
 use crate::cachestore::listener::RocksCacheStoreListener;
@@ -67,7 +67,8 @@ impl RocksCacheStoreDetails {
 }
 
 impl RocksStoreDetails for RocksCacheStoreDetails {
-    fn open_db(&self, path: &Path) -> Result<DB, CubeError> {
+    fn open_db(&self, path: &Path, config: &Arc<dyn ConfigObj>) -> Result<DB, CubeError> {
+        let rocksdb_config = config.cachestore_rocksdb_config();
         let compaction_state = Arc::new(Mutex::new(Some(
             RocksCacheStoreDetails::get_compaction_state(),
         )));
@@ -81,9 +82,17 @@ impl RocksStoreDetails for RocksCacheStoreDetails {
         // Disable automatic compaction before migration, will be enabled later in after_migration
         opts.set_disable_auto_compactions(true);
 
-        let mut block_opts = BlockBasedOptions::default();
-        // https://github.com/facebook/rocksdb/blob/v7.9.2/include/rocksdb/table.h#L524
-        block_opts.set_format_version(5);
+        let block_opts = {
+            let mut block_opts = BlockBasedOptions::default();
+            // https://github.com/facebook/rocksdb/blob/v7.9.2/include/rocksdb/table.h#L524
+            block_opts.set_format_version(5);
+            block_opts.set_checksum_type(rocksdb_config.checksum_type.as_rocksdb_enum());
+
+            let cache = Cache::new_lru_cache(rocksdb_config.cache_capacity)?;
+            block_opts.set_block_cache(&cache);
+
+            block_opts
+        };
 
         opts.set_block_based_table_factory(&block_opts);
 

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

@@ -14,7 +14,9 @@ use crate::config::processing_loop::ProcessingLoop;
 use crate::http::HttpServer;
 use crate::import::limits::ConcurrencyLimits;
 use crate::import::{ImportService, ImportServiceImpl};
-use crate::metastore::{BaseRocksStoreFs, MetaStore, MetaStoreRpcClient, RocksMetaStore};
+use crate::metastore::{
+    BaseRocksStoreFs, MetaStore, MetaStoreRpcClient, RocksMetaStore, RocksStoreConfig,
+};
 use crate::mysql::{MySqlServer, SqlAuthDefaultImpl, SqlAuthService};
 use crate::queryplanner::query_executor::{QueryExecutor, QueryExecutorImpl};
 use crate::queryplanner::{QueryPlanner, QueryPlannerImpl};
@@ -392,8 +394,12 @@ pub trait ConfigObj: DIService {
 
     fn metastore_bind_address(&self) -> &Option<String>;
 
+    fn metastore_rocksdb_config(&self) -> &RocksStoreConfig;
+
     fn metastore_remote_address(&self) -> &Option<String>;
 
+    fn cachestore_rocksdb_config(&self) -> &RocksStoreConfig;
+
     fn download_concurrency(&self) -> u64;
 
     fn upload_concurrency(&self) -> u64;
@@ -478,6 +484,8 @@ pub struct ConfigObjImpl {
     pub worker_bind_address: Option<String>,
     pub metastore_bind_address: Option<String>,
     pub metastore_remote_address: Option<String>,
+    pub metastore_rocks_store_config: RocksStoreConfig,
+    pub cachestore_rocks_store_config: RocksStoreConfig,
     pub upload_concurrency: u64,
     pub download_concurrency: u64,
     pub connection_timeout: u64,
@@ -626,6 +634,14 @@ impl ConfigObj for ConfigObjImpl {
         &self.metastore_remote_address
     }
 
+    fn metastore_rocksdb_config(&self) -> &RocksStoreConfig {
+        &self.metastore_rocks_store_config
+    }
+
+    fn cachestore_rocksdb_config(&self) -> &RocksStoreConfig {
+        &self.cachestore_rocks_store_config
+    }
+
     fn download_concurrency(&self) -> u64 {
         self.download_concurrency
     }
@@ -881,6 +897,8 @@ impl Config {
                     env_optparse::<u16>("CUBESTORE_META_PORT").map(|v| format!("0.0.0.0:{}", v))
                 }),
                 metastore_remote_address: env::var("CUBESTORE_META_ADDR").ok(),
+                metastore_rocks_store_config: RocksStoreConfig::metastore_default(),
+                cachestore_rocks_store_config: RocksStoreConfig::cachestore_default(),
                 upload_concurrency: env_parse("CUBESTORE_MAX_ACTIVE_UPLOADS", 4),
                 download_concurrency: env_parse("CUBESTORE_MAX_ACTIVE_DOWNLOADS", 8),
                 max_ingestion_data_frames: env_parse("CUBESTORE_MAX_DATA_FRAMES", 4),
@@ -981,6 +999,8 @@ impl Config {
                 worker_bind_address: None,
                 metastore_bind_address: None,
                 metastore_remote_address: None,
+                metastore_rocks_store_config: RocksStoreConfig::metastore_default(),
+                cachestore_rocks_store_config: RocksStoreConfig::cachestore_default(),
                 upload_concurrency: 4,
                 download_concurrency: 8,
                 max_ingestion_data_frames: 4,

rust/cubestore/cubestore/src/config/mod.rs

@@ -21,7 +21,7 @@ pub use rocks_table::*;
 use crate::cluster::node_name_by_partition;
 use async_trait::async_trait;
 use log::info;
-use rocksdb::{BlockBasedOptions, Env, MergeOperands, Options, DB};
+use rocksdb::{BlockBasedOptions, Cache, Env, MergeOperands, Options, DB};
 use serde::{Deserialize, Serialize};
 use std::hash::Hash;
 use std::{env, io::Cursor, sync::Arc};
@@ -1146,15 +1146,24 @@ fn meta_store_merge(
 struct RocksMetaStoreDetails {}
 
 impl RocksStoreDetails for RocksMetaStoreDetails {
-    fn open_db(&self, path: &Path) -> Result<DB, CubeError> {
+    fn open_db(&self, path: &Path, config: &Arc<dyn ConfigObj>) -> Result<DB, CubeError> {
+        let rocksdb_config = config.metastore_rocksdb_config();
         let mut opts = Options::default();
         opts.create_if_missing(true);
         opts.set_prefix_extractor(rocksdb::SliceTransform::create_fixed_prefix(13));
         opts.set_merge_operator_associative("meta_store merge", meta_store_merge);
 
-        let mut block_opts = BlockBasedOptions::default();
-        // https://github.com/facebook/rocksdb/blob/v7.9.2/include/rocksdb/table.h#L524
-        block_opts.set_format_version(5);
+        let block_opts = {
+            let mut block_opts = BlockBasedOptions::default();
+            // https://github.com/facebook/rocksdb/blob/v7.9.2/include/rocksdb/table.h#L524
+            block_opts.set_format_version(5);
+            block_opts.set_checksum_type(rocksdb_config.checksum_type.as_rocksdb_enum());
+
+            let cache = Cache::new_lru_cache(rocksdb_config.cache_capacity)?;
+            block_opts.set_block_cache(&cache);
+
+            block_opts
+        };
 
         opts.set_block_based_table_factory(&block_opts);
 

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

@@ -491,8 +491,51 @@ macro_rules! meta_store_table_impl {
     };
 }
 
+#[derive(Debug, Clone)]
+pub enum RocksStoreChecksumType {
+    NoChecksum = 0,
+    CRC32c = 1,
+    XXHash = 2,
+    XXHash64 = 3,
+    XXH3 = 4, // Supported since RocksDB 6.27
+}
+
+impl RocksStoreChecksumType {
+    pub fn as_rocksdb_enum(&self) -> rocksdb::ChecksumType {
+        match &self {
+            RocksStoreChecksumType::NoChecksum => rocksdb::ChecksumType::NoChecksum,
+            RocksStoreChecksumType::CRC32c => rocksdb::ChecksumType::CRC32c,
+            RocksStoreChecksumType::XXHash => rocksdb::ChecksumType::XXHash,
+            RocksStoreChecksumType::XXHash64 => rocksdb::ChecksumType::XXHash64,
+            RocksStoreChecksumType::XXH3 => rocksdb::ChecksumType::XXH3,
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct RocksStoreConfig {
+    pub checksum_type: RocksStoreChecksumType,
+    pub cache_capacity: usize,
+}
+
+impl RocksStoreConfig {
+    pub fn metastore_default() -> Self {
+        Self {
+            checksum_type: RocksStoreChecksumType::XXHash,
+            cache_capacity: 1024 * 8,
+        }
+    }
+
+    pub fn cachestore_default() -> Self {
+        Self {
+            checksum_type: RocksStoreChecksumType::XXHash,
+            cache_capacity: 1024 * 8,
+        }
+    }
+}
+
 pub trait RocksStoreDetails: Send + Sync {
-    fn open_db(&self, path: &Path) -> Result<DB, CubeError>;
+    fn open_db(&self, path: &Path, config: &Arc<dyn ConfigObj>) -> Result<DB, CubeError>;
 
     fn migrate(&self, table_ref: DbTableRef) -> Result<(), CubeError>;
 
@@ -554,7 +597,7 @@ impl RocksStore {
         config: Arc<dyn ConfigObj>,
         details: Arc<dyn RocksStoreDetails>,
     ) -> Result<Self, CubeError> {
-        let db = details.open_db(path)?;
+        let db = details.open_db(path, &config)?;
         let db_arc = Arc::new(db);
 
         let (rw_loop_tx, mut rw_loop_rx) = tokio::sync::mpsc::channel::<