feat: use entry size as weight in block and meta cache, export cache …

…capacity config (#1512)

rust/bench/ss_bench/main.rs

@@ -125,6 +125,8 @@ async fn main() {
         data_directory: "hummock_001".to_string(),
         async_checkpoint_enabled: true,
         write_conflict_detection_enabled: false,
+        block_cache_capacity: 256 << 20,
+        meta_cache_capacity: 64 << 20,
     });
 
     let mock_hummock_meta_service = Arc::new(MockHummockMetaService::new());

rust/common/src/config.rs

@@ -110,6 +110,14 @@ pub struct StorageConfig {
     /// Whether to enable write conflict detection
     #[serde(default = "default::write_conflict_detection_enabled")]
     pub write_conflict_detection_enabled: bool,
+
+    /// Capacity of sstable block cache.
+    #[serde(default = "default::block_cache_capacity")]
+    pub block_cache_capacity: usize,
+
+    /// Capacity of sstable meta cache.
+    #[serde(default = "default::meta_cache_capacity")]
+    pub meta_cache_capacity: usize,
 }
 
 impl Default for StorageConfig {
@@ -181,6 +189,16 @@ mod default {
     pub fn write_conflict_detection_enabled() -> bool {
         cfg!(debug_assertions)
     }
+
+    pub fn block_cache_capacity() -> usize {
+        // 256 MB
+        268435456
+    }
+
+    pub fn meta_cache_capacity() -> usize {
+        // 64 MB
+        67108864
+    }
 }
 
 #[cfg(test)]

rust/config/risingwave.toml

@@ -12,4 +12,6 @@ sstable_size = 268435456
 block_size = 4096
 bloom_false_positive = 0.1
 data_directory = "hummock_001"
-async_checkpoint_enabled = true
+async_checkpoint_enabled = true
+block_cache_capacity = 268435456
+meta_cache_capacity = 67108864

rust/storage/src/hummock/block_cache.rs

@@ -18,17 +18,20 @@ use bytes::{BufMut, Bytes, BytesMut};
 use futures::Future;
 use moka::future::Cache;
 
-use super::{Block, HummockError, HummockResult};
+use super::{Block, HummockError, HummockResult, DEFAULT_ENTRY_SIZE};
 
 pub struct BlockCache {
     inner: Cache<Bytes, Arc<Block>>,
 }
 
 impl BlockCache {
     pub fn new(capacity: usize) -> Self {
-        Self {
-            inner: Cache::new(capacity as u64),
-        }
+        let cache: Cache<Bytes, Arc<Block>> = Cache::builder()
+            .weigher(|_k, v: &Arc<Block>| v.len() as u32)
+            .initial_capacity(capacity / DEFAULT_ENTRY_SIZE)
+            .max_capacity(capacity as u64)
+            .build();
+        Self { inner: cache }
     }
 
     // TODO: Optimize for concurrent get https://github.com/singularity-data/risingwave/pull/627#discussion_r817354730.

rust/storage/src/hummock/iterator/test_utils.rs

@@ -52,6 +52,8 @@ pub fn mock_sstable_store_with_object_store(object_store: ObjectStoreRef) -> Sst
         object_store,
         path,
         Arc::new(StateStoreMetrics::unused()),
+        64 << 20,
+        64 << 20,
     ))
 }
 

rust/storage/src/hummock/shared_buffer/shared_buffer_manager.rs

@@ -244,6 +244,8 @@ mod tests {
             obj_client,
             remote_dir.to_string(),
             Arc::new(StateStoreMetrics::unused()),
+            64 << 20,
+            64 << 20,
         ));
         let vm = Arc::new(LocalVersionManager::new(sstable_store.clone()));
         let mock_hummock_meta_client = Arc::new(MockHummockMetaClient::new(Arc::new(

rust/storage/src/hummock/snapshot_tests.rs

@@ -45,6 +45,8 @@ async fn gen_and_upload_table(
         object_store,
         remote_dir.to_string(),
         Arc::new(StateStoreMetrics::unused()),
+        64 << 20,
+        64 << 20,
     ));
     let sst = gen_test_sstable(
         default_builder_opt_for_test(),
@@ -153,6 +155,8 @@ async fn test_snapshot() {
         object_store.clone(),
         remote_dir.to_string(),
         Arc::new(StateStoreMetrics::unused()),
+        64 << 20,
+        64 << 20,
     ));
     let vm = Arc::new(LocalVersionManager::new(sstable_store.clone()));
     let mock_hummock_meta_service = Arc::new(MockHummockMetaService::new());
@@ -230,6 +234,8 @@ async fn test_snapshot_range_scan() {
         object_store.clone(),
         remote_dir.to_string(),
         Arc::new(StateStoreMetrics::unused()),
+        64 << 20,
+        64 << 20,
     ));
     let vm = Arc::new(LocalVersionManager::new(sstable_store.clone()));
     let mock_hummock_meta_service = Arc::new(MockHummockMetaService::new());
@@ -286,6 +292,8 @@ async fn test_snapshot_reverse_range_scan() {
         object_store.clone(),
         remote_dir.to_string(),
         Arc::new(StateStoreMetrics::unused()),
+        64 << 20,
+        64 << 20,
     ));
     let vm = Arc::new(LocalVersionManager::new(sstable_store.clone()));
     let mock_hummock_meta_service = Arc::new(MockHummockMetaService::new());

rust/storage/src/hummock/sstable/mod.rs

@@ -69,6 +69,11 @@ impl Sstable {
     pub fn block_count(&self) -> usize {
         self.meta.block_metas.len()
     }
+
+    #[inline]
+    pub fn encoded_size(&self) -> usize {
+        8 /* id */ + self.meta.encoded_size()
+    }
 }
 
 #[derive(Clone, PartialEq, Eq, Debug)]
@@ -100,6 +105,11 @@ impl BlockMeta {
             len,
         }
     }
+
+    #[inline]
+    pub fn encoded_size(&self) -> usize {
+        12 /* offset + len + key len */ + self.smallest_key.len()
+    }
 }
 
 #[derive(Clone, PartialEq, Eq, Debug)]
@@ -189,6 +199,27 @@ impl SstableMeta {
             version,
         })
     }
+
+    #[inline]
+    pub fn encoded_size(&self) -> usize {
+        4 // block meta count
+            + self
+            .block_metas
+            .iter()
+            .map(|block_meta| block_meta.encoded_size())
+            .sum::<usize>()
+            + 4 // bloom filter len
+            + self.bloom_filter.len()
+            + 4 // estimated size
+            + 4 // key count
+            + 4 // key len
+            + self.smallest_key.len()
+            + 4 // key len
+            + self.largest_key.len()
+            + 8 // checksum
+            + 4 // version
+            + 4 // magic
+    }
 }
 
 #[cfg(test)]

rust/storage/src/hummock/sstable_store.rs

@@ -22,6 +22,8 @@ use crate::hummock::{HummockError, HummockResult};
 use crate::monitor::StateStoreMetrics;
 use crate::object::{BlockLocation, ObjectStoreRef};
 
+const DEFAULT_META_CACHE_INIT_CAPACITY: usize = 1024;
+
 // TODO: Define policy based on use cases (read / compaction / ...).
 pub enum CachePolicy {
     Disable,
@@ -33,18 +35,30 @@ pub struct SstableStore {
     path: String,
     store: ObjectStoreRef,
     block_cache: BlockCache,
-    sstable_cache: Cache<u64, Arc<Sstable>>,
+    meta_cache: Cache<u64, Arc<Sstable>>,
     /// Statistics.
     stats: Arc<StateStoreMetrics>,
 }
 
 impl SstableStore {
-    pub fn new(store: ObjectStoreRef, path: String, stats: Arc<StateStoreMetrics>) -> Self {
+    pub fn new(
+        store: ObjectStoreRef,
+        path: String,
+        stats: Arc<StateStoreMetrics>,
+        block_cache_capacity: usize,
+        meta_cache_capacity: usize,
+    ) -> Self {
+        let meta_cache: Cache<u64, Arc<Sstable>> = Cache::builder()
+            .weigher(|_k, v: &Arc<Sstable>| v.encoded_size() as u32)
+            .initial_capacity(DEFAULT_META_CACHE_INIT_CAPACITY)
+            .max_capacity(meta_cache_capacity as u64)
+            .build();
+
         Self {
             path,
             store,
-            block_cache: BlockCache::new(65536),
-            sstable_cache: Cache::new(1024),
+            block_cache: BlockCache::new(block_cache_capacity),
+            meta_cache,
             stats,
         }
     }
@@ -152,7 +166,7 @@ impl SstableStore {
             Ok::<_, TracedHummockError>(sst)
         };
 
-        self.sstable_cache
+        self.meta_cache
             .try_get_with(sst_id, fetch)
             .await
             .map_err(|e| HummockError::Other(e.to_string()).into())

rust/storage/src/hummock/test_utils.rs

@@ -39,6 +39,8 @@ pub fn default_config_for_test() -> StorageConfig {
         data_directory: "hummock_001".to_string(),
         async_checkpoint_enabled: true,
         write_conflict_detection_enabled: true,
+        block_cache_capacity: 64 << 20,
+        meta_cache_capacity: 64 << 20,
     }
 }
 

rust/storage/src/hummock/vacuum.rs

@@ -72,6 +72,8 @@ mod tests {
             Arc::new(InMemObjectStore::new()),
             String::from("test_dir"),
             Arc::new(StateStoreMetrics::unused()),
+            64 << 20,
+            64 << 20,
         ));
 
         // Put some SSTs to object store

rust/storage/src/store_impl.rs

@@ -118,6 +118,8 @@ impl StateStoreImpl {
                     object_store,
                     config.data_directory.to_string(),
                     state_store_stats.clone(),
+                    config.block_cache_capacity,
+                    config.meta_cache_capacity,
                 ));
                 let inner = HummockStorage::new(
                     config.clone(),