From c687481db84d3f07e7922b95eb95dd0817c72da8 Mon Sep 17 00:00:00 2001
From: HappenLee <happenlee@hotmail.com>
Date: Wed, 25 Sep 2024 20:23:42 +0800
Subject: [PATCH] [Exec)(cache) add element count in LRU cache (#41199)

add element_count() in LRU cache
---
 be/src/olap/lru_cache.cpp                   | 8 ++++++++
 be/src/olap/lru_cache.h                     | 4 ++++
 be/src/pipeline/query_cache/query_cache.cpp | 2 +-
 be/src/runtime/memory/lru_cache_policy.h    | 2 ++
 4 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/be/src/olap/lru_cache.cpp b/be/src/olap/lru_cache.cpp
index 6e5bb2fa31578f..b0ad59b6c8d15c 100644
--- a/be/src/olap/lru_cache.cpp
+++ b/be/src/olap/lru_cache.cpp
@@ -700,6 +700,14 @@ int64_t ShardedLRUCache::get_usage() {
     return total_usage;
 }
 
+size_t ShardedLRUCache::get_element_count() {
+    size_t total_element_count = 0;
+    for (int i = 0; i < _num_shards; i++) {
+        total_element_count += _shards[i]->get_element_count();
+    }
+    return total_element_count;
+}
+
 void ShardedLRUCache::update_cache_metrics() const {
     size_t capacity = 0;
     size_t total_usage = 0;
diff --git a/be/src/olap/lru_cache.h b/be/src/olap/lru_cache.h
index de7084382d7398..ba2dd2b5c52c56 100644
--- a/be/src/olap/lru_cache.h
+++ b/be/src/olap/lru_cache.h
@@ -230,6 +230,8 @@ class Cache {
     virtual PrunedInfo set_capacity(size_t capacity) = 0;
     virtual size_t get_capacity() = 0;
 
+    virtual size_t get_element_count() = 0;
+
 private:
     DISALLOW_COPY_AND_ASSIGN(Cache);
 };
@@ -404,6 +406,7 @@ class ShardedLRUCache : public Cache {
     PrunedInfo prune() override;
     PrunedInfo prune_if(CachePrunePredicate pred, bool lazy_mode = false) override;
     int64_t get_usage() override;
+    size_t get_element_count() override;
     PrunedInfo set_capacity(size_t capacity) override;
     size_t get_capacity() override;
 
@@ -467,6 +470,7 @@ class DummyLRUCache : public Cache {
     int64_t get_usage() override { return 0; };
     PrunedInfo set_capacity(size_t capacity) override { return {0, 0}; };
     size_t get_capacity() override { return 0; };
+    size_t get_element_count() override { return 0; };
 };
 
 } // namespace doris
diff --git a/be/src/pipeline/query_cache/query_cache.cpp b/be/src/pipeline/query_cache/query_cache.cpp
index 20e342e140f156..b69e202d67baca 100644
--- a/be/src/pipeline/query_cache/query_cache.cpp
+++ b/be/src/pipeline/query_cache/query_cache.cpp
@@ -66,4 +66,4 @@ bool QueryCache::lookup(const CacheKey& key, int64_t version, doris::QueryCacheH
     return false;
 }
 
-} // namespace doris
\ No newline at end of file
+} // namespace doris
diff --git a/be/src/runtime/memory/lru_cache_policy.h b/be/src/runtime/memory/lru_cache_policy.h
index 7b5a8ab9fec6d9..f65a1e23e1a3af 100644
--- a/be/src/runtime/memory/lru_cache_policy.h
+++ b/be/src/runtime/memory/lru_cache_policy.h
@@ -132,6 +132,8 @@ class LRUCachePolicy : public CachePolicy {
 
     int64_t get_usage() { return _cache->get_usage(); }
 
+    size_t get_element_count() { return _cache->get_element_count(); }
+
     size_t get_capacity() override { return _cache->get_capacity(); }
 
     uint64_t new_id() { return _cache->new_id(); };