[Enhancement] Adjust calculation method for page cache mem usage metrics (#37740)

Signed-off-by: Zaorang Yang <[email protected]>
(cherry picked from commit 65b0fa9)

Author: zaorangyang

be/src/storage/page_cache.cpp

@@ -39,6 +39,7 @@
 #include "runtime/current_thread.h"
 #include "runtime/mem_tracker.h"
 #include "util/defer_op.h"
+#include "util/lru_cache.h"
 #include "util/metrics.h"
 #include "util/starrocks_metrics.h"
 
@@ -81,7 +82,7 @@ static void init_metrics() {
 }
 
 StoragePageCache::StoragePageCache(MemTracker* mem_tracker, size_t capacity)
-        : _mem_tracker(mem_tracker), _cache(new_lru_cache(capacity)) {
+        : _mem_tracker(mem_tracker), _cache(new_lru_cache(capacity, ChargeMode::MEMSIZE)) {
     init_metrics();
 }
 
@@ -123,8 +124,10 @@ bool StoragePageCache::lookup(const CacheKey& key, PageCacheHandle* handle) {
 }
 
 void StoragePageCache::insert(const CacheKey& key, const Slice& data, PageCacheHandle* handle, bool in_memory) {
+    // mem size should equals to data size when running UT
+    int64_t mem_size = data.size;
 #ifndef BE_TEST
-    int64_t mem_size = malloc_usable_size(data.data);
+    mem_size = malloc_usable_size(data.data);
     tls_thread_status.mem_release(mem_size);
     SCOPED_THREAD_LOCAL_MEM_TRACKER_SETTER(_mem_tracker);
     tls_thread_status.mem_consume(mem_size);
@@ -136,8 +139,9 @@ void StoragePageCache::insert(const CacheKey& key, const Slice& data, PageCacheH
     if (in_memory) {
         priority = CachePriority::DURABLE;
     }
-
-    auto* lru_handle = _cache->insert(key.encode(), data.data, data.size, deleter, priority);
+    // Use mem size managed by memory allocator as this record charge size. At the same time, we should record this record size
+    // for data fetching when lookup.
+    auto* lru_handle = _cache->insert(key.encode(), data.data, mem_size, deleter, priority, data.size);
     *handle = PageCacheHandle(_cache.get(), lru_handle);
 }
 

be/src/util/lru_cache.cpp

@@ -193,6 +193,10 @@ void LRUCache::set_capacity(size_t capacity) {
     }
 }
 
+void LRUCache::set_charge_mode(ChargeMode charge_mode) {
+    _charge_mode = charge_mode;
+}
+
 uint64_t LRUCache::get_lookup_count() const {
     std::lock_guard l(_mutex);
     return _lookup_count;
@@ -295,7 +299,8 @@ void LRUCache::_evict_one_entry(LRUHandle* e) {
 }
 
 Cache::Handle* LRUCache::insert(const CacheKey& key, uint32_t hash, void* value, size_t charge,
-                                void (*deleter)(const CacheKey& key, void* value), CachePriority priority) {
+                                void (*deleter)(const CacheKey& key, void* value), CachePriority priority,
+                                size_t value_size) {
     auto* e = reinterpret_cast<LRUHandle*>(malloc(sizeof(LRUHandle) - 1 + key.size()));
     e->value = value;
     e->deleter = deleter;
@@ -306,6 +311,7 @@ Cache::Handle* LRUCache::insert(const CacheKey& key, uint32_t hash, void* value,
     e->next = e->prev = nullptr;
     e->in_cache = true;
     e->priority = priority;
+    e->value_size = value_size;
     memcpy(e->key_data, key.data(), key.size());
     std::vector<LRUHandle*> last_ref_list;
     {
@@ -395,10 +401,12 @@ uint32_t ShardedLRUCache::_shard(uint32_t hash) {
     return hash >> (32 - kNumShardBits);
 }
 
-ShardedLRUCache::ShardedLRUCache(size_t capacity) : _last_id(0), _capacity(capacity) {
+ShardedLRUCache::ShardedLRUCache(size_t capacity, ChargeMode charge_mode)
+        : _last_id(0), _capacity(capacity), _charge_mode(charge_mode) {
     const size_t per_shard = (_capacity + (kNumShards - 1)) / kNumShards;
     for (auto& _shard : _shards) {
         _shard.set_capacity(per_shard);
+        _shard.set_charge_mode(_charge_mode);
     }
 }
 
@@ -427,9 +435,10 @@ bool ShardedLRUCache::adjust_capacity(int64_t delta, size_t min_capacity) {
 }
 
 Cache::Handle* ShardedLRUCache::insert(const CacheKey& key, void* value, size_t charge,
-                                       void (*deleter)(const CacheKey& key, void* value), CachePriority priority) {
+                                       void (*deleter)(const CacheKey& key, void* value), CachePriority priority,
+                                       size_t value_size) {
     const uint32_t hash = _hash_slice(key);
-    return _shards[_shard(hash)].insert(key, hash, value, charge, deleter, priority);
+    return _shards[_shard(hash)].insert(key, hash, value, charge, deleter, priority, value_size);
 }
 
 Cache::Handle* ShardedLRUCache::lookup(const CacheKey& key) {
@@ -453,7 +462,8 @@ void* ShardedLRUCache::value(Handle* handle) {
 
 Slice ShardedLRUCache::value_slice(Handle* handle) {
     auto lru_handle = reinterpret_cast<LRUHandle*>(handle);
-    return {(char*)lru_handle->value, lru_handle->charge};
+    size_t record_size = _charge_mode == ChargeMode::VALUESIZE ? lru_handle->charge : lru_handle->value_size;
+    return {(char*)lru_handle->value, record_size};
 }
 
 uint64_t ShardedLRUCache::new_id() {
@@ -526,8 +536,8 @@ void ShardedLRUCache::get_cache_status(rapidjson::Document* document) {
     }
 }
 
-Cache* new_lru_cache(size_t capacity) {
-    return new ShardedLRUCache(capacity);
+Cache* new_lru_cache(size_t capacity, ChargeMode charge_mode) {
+    return new ShardedLRUCache(capacity, charge_mode);
 }
 
 } // namespace starrocks

be/src/util/lru_cache.h

@@ -20,9 +20,16 @@ namespace starrocks {
 class Cache;
 class CacheKey;
 
+enum class ChargeMode {
+    // use value size as charge
+    VALUESIZE = 0,
+    // use allocator tracking size as charge
+    MEMSIZE = 1
+};
+
 // Create a new cache with a fixed size capacity.  This implementation
 // of Cache uses a least-recently-used eviction policy.
-extern Cache* new_lru_cache(size_t capacity);
+extern Cache* new_lru_cache(size_t capacity, ChargeMode charge_mode = ChargeMode::VALUESIZE);
 
 class CacheKey {
 public:
@@ -134,7 +141,7 @@ class Cache {
     // value will be passed to "deleter".
     virtual Handle* insert(const CacheKey& key, void* value, size_t charge,
                            void (*deleter)(const CacheKey& key, void* value),
-                           CachePriority priority = CachePriority::NORMAL) = 0;
+                           CachePriority priority = CachePriority::NORMAL, size_t value_size = 0) = 0;
 
     // If the cache has no mapping for "key", returns NULL.
     //
@@ -206,6 +213,7 @@ typedef struct LRUHandle {
     uint32_t refs;
     uint32_t hash; // Hash of key(); used for fast sharding and comparisons
     CachePriority priority = CachePriority::NORMAL;
+    size_t value_size;
     char key_data[1]; // Beginning of key
 
     CacheKey key() const {
@@ -266,10 +274,12 @@ class LRUCache {
     // Separate from constructor so caller can easily make an array of LRUCache
     void set_capacity(size_t capacity);
 
+    void set_charge_mode(ChargeMode charge_mode);
+
     // Like Cache methods, but with an extra "hash" parameter.
     Cache::Handle* insert(const CacheKey& key, uint32_t hash, void* value, size_t charge,
                           void (*deleter)(const CacheKey& key, void* value),
-                          CachePriority priority = CachePriority::NORMAL);
+                          CachePriority priority = CachePriority::NORMAL, size_t value_size = 0);
     Cache::Handle* lookup(const CacheKey& key, uint32_t hash);
     void release(Cache::Handle* handle);
     void erase(const CacheKey& key, uint32_t hash);
@@ -290,6 +300,8 @@ class LRUCache {
     // Initialized before use.
     size_t _capacity{0};
 
+    ChargeMode _charge_mode;
+
     // _mutex protects the following state.
     mutable std::mutex _mutex;
     size_t _usage{0};
@@ -310,10 +322,10 @@ static const int kNumShards = 1 << kNumShardBits;
 
 class ShardedLRUCache : public Cache {
 public:
-    explicit ShardedLRUCache(size_t capacity);
+    explicit ShardedLRUCache(size_t capacity, ChargeMode charge_mode = ChargeMode::VALUESIZE);
     ~ShardedLRUCache() override = default;
     Handle* insert(const CacheKey& key, void* value, size_t charge, void (*deleter)(const CacheKey& key, void* value),
-                   CachePriority priority = CachePriority::NORMAL) override;
+                   CachePriority priority = CachePriority::NORMAL, size_t value_size = 0) override;
     Handle* lookup(const CacheKey& key) override;
     void release(Handle* handle) override;
     void erase(const CacheKey& key) override;
@@ -339,6 +351,7 @@ class ShardedLRUCache : public Cache {
     std::mutex _mutex;
     uint64_t _last_id;
     size_t _capacity;
+    ChargeMode _charge_mode;
 };
 
 } // namespace starrocks