add EstimateSize method on cache

cache.go

@@ -79,6 +79,11 @@ func (c *Cache[K, V]) Size() int {
 	return c.store.Size()
 }
 
+// EstimateSize returns an approximate used size of the cache.
+func (c *Cache[K, V]) EstimateSize() int {
+	return c.store.EstimateSize()
+}
+
 // Close closes all goroutines created by cache.
 func (c *Cache[K, V]) Close() {
 	c.store.Close()
@@ -136,6 +141,11 @@ func (c *LoadingCache[K, V]) Size() int {
 	return c.store.Size()
 }
 
+// EstimateSize returns an approximate used size of the cache.
+func (c *LoadingCache[K, V]) EstimateSize() int {
+	return c.store.EstimateSize()
+}
+
 // SaveCache save cache data to writer.
 func (c *LoadingCache[K, V]) SaveCache(version uint64, writer io.Writer) error {
 	return c.store.Persist(version, writer)

cache_test.go

@@ -1,6 +1,7 @@
 package theine_test
 
 import (
+	"context"
 	"fmt"
 	"math/rand"
 	"strconv"
@@ -11,6 +12,7 @@ import (
 
 	"github.com/Yiling-J/theine-go"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
 )
 
 func TestMaxsizeZero(t *testing.T) {
@@ -227,6 +229,7 @@ func TestCost(t *testing.T) {
 	time.Sleep(time.Second)
 	require.True(t, client.Len() == 25)
 	require.True(t, client.Size() == 500)
+	require.True(t, client.EstimateSize() == 500)
 
 	// test cost func
 	builder := theine.NewBuilder[string, string](500)
@@ -245,6 +248,7 @@ func TestCost(t *testing.T) {
 	time.Sleep(time.Second)
 	require.True(t, client.Len() == 25)
 	require.True(t, client.Size() == 500)
+	require.True(t, client.EstimateSize() == 500)
 	client.Close()
 }
 
@@ -259,16 +263,51 @@ func TestCostUpdate(t *testing.T) {
 	time.Sleep(time.Second)
 	require.True(t, client.Len() == 25)
 	require.True(t, client.Size() == 500)
+	require.True(t, client.EstimateSize() == 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)
+	require.True(t, client.EstimateSize() == 15*20+200)
 	client.Close()
 }
 
+func TestEstimateSize(t *testing.T) {
+	client, err := theine.NewBuilder[int, int](500).Build()
+	require.Nil(t, err)
+	ctx, cfn := context.WithCancel(context.Background())
+	defer cfn()
+	wg, ctx := errgroup.WithContext(ctx)
+	wg.Go(func() error {
+		tkr := time.NewTicker(time.Nanosecond)
+		defer tkr.Stop()
+		for {
+			select {
+			case <-tkr.C:
+				client.EstimateSize()
+			case <-ctx.Done():
+				return nil
+			}
+		}
+	})
+	wg.Go(func() error {
+		defer cfn()
+		for i := 0; i < 10000000; i++ {
+			if i%2 == 0 {
+				client.Set(i, 1, 1)
+			} else {
+				client.Get(i)
+			}
+		}
+		return nil
+	})
+	require.Nil(t, wg.Wait())
+	require.True(t, client.EstimateSize() == client.Size())
+}
+
 func TestDoorkeeper(t *testing.T) {
 	builder := theine.NewBuilder[string, string](500)
 	builder.Doorkeeper(true)

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/gammazero/deque v0.2.1
 	github.com/stretchr/testify v1.8.2
 	github.com/zeebo/xxh3 v1.0.2
+	golang.org/x/sync v0.8.0
 	golang.org/x/sys v0.8.0
 )
 

go.sum

@@ -24,6 +24,8 @@ github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o
 github.com/zeebo/assert v1.3.0 h1:g7C04CbJuIDKNPFHmsk4hwZDO5O+kntRxzaUoNXj+IQ=
 github.com/zeebo/xxh3 v1.0.2 h1:xZmwmqxHZA8AI603jOQ0tMqmBr9lPeFwGg6d+xy9DC0=
 github.com/zeebo/xxh3 v1.0.2/go.mod h1:5NWz9Sef7zIDm2JHfFlcQvNekmcEl9ekUZQQKCYaDcA=
+golang.org/x/sync v0.8.0 h1:3NFvSEYkUoMifnESzZl15y791HH1qU2xm6eCJU5ZPXQ=
+golang.org/x/sync v0.8.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.8.0 h1:EBmGv8NaZBZTWvrbjNoL6HVt+IVy3QDQpJs7VRIw3tU=
 golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

internal/store.go

@@ -484,6 +484,16 @@ func (s *Store[K, V]) Size() int {
 	return total
 }
 
+func (s *Store[K, V]) EstimateSize() int {
+	total := s.policy.slru.protected.len + s.policy.slru.probation.len
+	for _, s := range s.shards {
+		tk := s.mu.RLock()
+		total += s.qlen
+		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)