From 4fda5d7357d2dad6d6ef31d476370c94ebda839d Mon Sep 17 00:00:00 2001
From: Croxx <mrcroxx@outlook.com>
Date: Mon, 13 May 2024 19:29:17 +0800
Subject: [PATCH] fix(storage): fix compactor oom by count the memory for cache
 in compactor (#16727)

Signed-off-by: MrCroxx <mrcroxx@outlook.com>
---
 src/storage/src/hummock/sstable_store.rs | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/storage/src/hummock/sstable_store.rs b/src/storage/src/hummock/sstable_store.rs
index 63581e3fb742..6dfbe5468426 100644
--- a/src/storage/src/hummock/sstable_store.rs
+++ b/src/storage/src/hummock/sstable_store.rs
@@ -140,6 +140,7 @@ impl SstableStore {
 
     /// For compactor, we do not need a high concurrency load for cache. Instead, we need the cache
     ///  can be evict more effective.
+    #[expect(clippy::borrowed_box)]
     pub async fn for_compactor(
         store: ObjectStoreRef,
         path: String,
@@ -149,6 +150,9 @@ impl SstableStore {
         let meta_cache_v2 = HybridCacheBuilder::new()
             .memory(meta_cache_capacity)
             .with_shards(1)
+            .with_weighter(|_: &HummockSstableObjectId, value: &Box<Sstable>| {
+                u64::BITS as usize / 8 + value.estimate_size()
+            })
             .storage()
             .build()
             .await
@@ -157,6 +161,10 @@ impl SstableStore {
         let block_cache_v2 = HybridCacheBuilder::new()
             .memory(block_cache_capacity)
             .with_shards(1)
+            .with_weighter(|_: &SstableBlockIndex, value: &Box<Block>| {
+                // FIXME(MrCroxx): Calculate block weight more accurately.
+                u64::BITS as usize * 2 / 8 + value.raw().len()
+            })
             .storage()
             .build()
             .await