From 038e9321c69f72070dc7ab454d367dd3abfd16bd Mon Sep 17 00:00:00 2001
From: Nicolas Lach <dev@nlach.fr>
Date: Wed, 7 Aug 2024 11:15:54 +0200
Subject: [PATCH] add Size method on cache

---
 cache.go          | 10 ++++++++++
 cache_test.go     |  4 ++++
 internal/store.go | 13 +++++++++++++
 3 files changed, 27 insertions(+)

diff --git a/cache.go b/cache.go
index fb7a2d5..5e65910 100644
--- a/cache.go
+++ b/cache.go
@@ -74,6 +74,11 @@ func (c *Cache[K, V]) Len() int {
 	return c.store.Len()
 }
 
+// Size returns the used size of the cache.
+func (c *Cache[K, V]) Size() int {
+	return c.store.Size()
+}
+
 // Close closes all goroutines created by cache.
 func (c *Cache[K, V]) Close() {
 	c.store.Close()
@@ -126,6 +131,11 @@ func (c *LoadingCache[K, V]) Len() int {
 	return c.store.Len()
 }
 
+// Size returns the used size of the cache.
+func (c *LoadingCache[K, V]) Size() int {
+	return c.store.Size()
+}
+
 // SaveCache save cache data to writer.
 func (c *LoadingCache[K, V]) SaveCache(version uint64, writer io.Writer) error {
 	return c.store.Persist(version, writer)
diff --git a/cache_test.go b/cache_test.go
index 1ca9c9f..fead8e8 100644
--- a/cache_test.go
+++ b/cache_test.go
@@ -226,6 +226,7 @@ func TestCost(t *testing.T) {
 	}
 	time.Sleep(time.Second)
 	require.True(t, client.Len() == 25)
+	require.True(t, client.Size() == 500)
 
 	// test cost func
 	builder := theine.NewBuilder[string, string](500)
@@ -243,6 +244,7 @@ func TestCost(t *testing.T) {
 	}
 	time.Sleep(time.Second)
 	require.True(t, client.Len() == 25)
+	require.True(t, client.Size() == 500)
 	client.Close()
 }
 
@@ -256,12 +258,14 @@ func TestCostUpdate(t *testing.T) {
 	}
 	time.Sleep(time.Second)
 	require.True(t, client.Len() == 25)
+	require.True(t, client.Size() == 500)
 	// update cost
 	success := client.Set("key:10", "", 200)
 	require.True(t, success)
 	time.Sleep(time.Second)
 	// 15 * 20 + 200
 	require.True(t, client.Len() == 16)
+	require.True(t, client.Size() == 15*20+200)
 	client.Close()
 }
 
diff --git a/internal/store.go b/internal/store.go
index 5450393..325040b 100644
--- a/internal/store.go
+++ b/internal/store.go
@@ -472,6 +472,19 @@ func (s *Store[K, V]) Len() int {
 	return total
 }
 
+func (s *Store[K, V]) Size() int {
+	total := 0
+	for _, s := range s.shards {
+		tk := s.mu.RLock()
+		for _, entry := range s.hashmap {
+			total += int(entry.cost.Load())
+		}
+		s.mu.RUnlock(tk)
+	}
+	return total
+}
+
+// spread hash before get index
 func (s *Store[K, V]) index(key K) (uint64, int) {
 	base := s.hasher.hash(key)
 	return base, int(base & uint64(s.shardCount-1))