Add Size method on cache

README.md

@@ -124,6 +124,9 @@ client.Range(func(key, value int) bool {
 	return true
 })
 
+// returns an estimation of the cache size usage
+client.EstimatedSize()
+
 // close client, set hashmaps in shard to nil and close all goroutines
 client.Close()
 

cache.go

@@ -74,6 +74,11 @@ func (c *Cache[K, V]) Len() int {
 	return c.store.Len()
 }
 
+// EstimatedSize returns an approximate used size of the cache.
+func (c *Cache[K, V]) EstimatedSize() int {
+	return c.store.EstimatedSize()
+}
+
 // 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()
 }
 
+// EstimatedSize returns an approximate used size of the cache.
+func (c *LoadingCache[K, V]) EstimatedSize() int {
+	return c.store.EstimatedSize()
+}
+
 // 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) {
@@ -226,6 +228,7 @@ func TestCost(t *testing.T) {
 	}
 	time.Sleep(time.Second)
 	require.True(t, client.Len() == 25)
+	require.True(t, client.EstimatedSize() == 500)
 
 	// test cost func
 	builder := theine.NewBuilder[string, string](500)
@@ -243,6 +246,7 @@ func TestCost(t *testing.T) {
 	}
 	time.Sleep(time.Second)
 	require.True(t, client.Len() == 25)
+	require.True(t, client.EstimatedSize() == 500)
 	client.Close()
 }
 
@@ -256,15 +260,49 @@ func TestCostUpdate(t *testing.T) {
 	}
 	time.Sleep(time.Second)
 	require.True(t, client.Len() == 25)
+	require.True(t, client.EstimatedSize() == 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.EstimatedSize() == 15*20+200)
 	client.Close()
 }
 
+func TestEstimatedSize(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.EstimatedSize()
+			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())
+}
+
 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/list.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"strings"
+	"sync/atomic"
 )
 
 const (
@@ -17,9 +18,9 @@ const (
 // List represents a doubly linked list.
 // The zero value for List is an empty list ready to use.
 type List[K comparable, V any] struct {
-	root     Entry[K, V] // sentinel list element, only &root, root.prev, and root.next are used
-	len      int         // current list length(sum of costs) excluding (this) sentinel element
-	count    int         // count of entries in list
+	root     Entry[K, V]  // sentinel list element, only &root, root.prev, and root.next are used
+	len      atomic.Int64 // current list length(sum of costs) excluding (this) sentinel element
+	count    int          // count of entries in list
 	capacity uint
 	bounded  bool
 	listType uint8 // 1 tinylfu list, 2 timerwheel list
@@ -31,7 +32,7 @@ func NewList[K comparable, V any](size uint, listType uint8) *List[K, V] {
 	l.root.root = true
 	l.root.setNext(&l.root, l.listType)
 	l.root.setPrev(&l.root, l.listType)
-	l.len = 0
+	l.len = atomic.Int64{}
 	l.capacity = size
 	if size > 0 {
 		l.bounded = true
@@ -42,12 +43,12 @@ func NewList[K comparable, V any](size uint, listType uint8) *List[K, V] {
 func (l *List[K, V]) Reset() {
 	l.root.setNext(&l.root, l.listType)
 	l.root.setPrev(&l.root, l.listType)
-	l.len = 0
+	l.len.Store(0)
 }
 
 // Len returns the number of elements of list l.
 // The complexity is O(1).
-func (l *List[K, V]) Len() int { return l.len }
+func (l *List[K, V]) Len() int { return int(l.len.Load()) }
 
 func (l *List[K, V]) display() string {
 	var s []string
@@ -86,7 +87,7 @@ func (l *List[K, V]) Back() *Entry[K, V] {
 // insert inserts e after at, increments l.len, and evicted entry if capacity exceed
 func (l *List[K, V]) insert(e, at *Entry[K, V]) *Entry[K, V] {
 	var evicted *Entry[K, V]
-	if l.bounded && l.len >= int(l.capacity) {
+	if l.bounded && l.len.Load() >= int64(l.capacity) {
 		evicted = l.PopTail()
 	}
 	if l.listType != WHEEL_LIST {
@@ -97,7 +98,8 @@ func (l *List[K, V]) insert(e, at *Entry[K, V]) *Entry[K, V] {
 	e.prev(l.listType).setNext(e, l.listType)
 	e.next(l.listType).setPrev(e, l.listType)
 	if l.bounded {
-		l.len += int(e.cost.Load())
+		l.len.Add(e.cost.Load())
+		// l.len += int(e.cost.Load())
 		l.count += 1
 	}
 	return evicted
@@ -123,7 +125,7 @@ func (l *List[K, V]) remove(e *Entry[K, V]) {
 		e.list = 0
 	}
 	if l.bounded {
-		l.len -= int(e.cost.Load())
+		l.len.Add(-e.cost.Load())
 		l.count -= 1
 	}
 }

internal/list_test.go

@@ -15,13 +15,13 @@ func TestList(t *testing.T) {
 		evicted := l.PushFront(NewEntry(fmt.Sprintf("%d", i), "", 1, 0))
 		require.Nil(t, evicted)
 	}
-	require.Equal(t, 5, l.len)
+	require.Equal(t, 5, int(l.len.Load()))
 	require.Equal(t, "4/3/2/1/0", l.display())
 	require.Equal(t, "0/1/2/3/4", l.displayReverse())
 
 	evicted := l.PushFront(NewEntry("5", "", 1, 0))
 	require.Equal(t, "0", evicted.key)
-	require.Equal(t, 5, l.len)
+	require.Equal(t, 5, int(l.len.Load()))
 	require.Equal(t, "5/4/3/2/1", l.display())
 	require.Equal(t, "1/2/3/4/5", l.displayReverse())
 
@@ -63,13 +63,13 @@ func TestListCountCost(t *testing.T) {
 		evicted := l.PushFront(NewEntry(fmt.Sprintf("%d", i), "", 20, 0))
 		require.Nil(t, evicted)
 	}
-	require.Equal(t, 100, l.len)
+	require.Equal(t, 100, int(l.len.Load()))
 	require.Equal(t, 5, l.count)
 	for i := 0; i < 3; i++ {
 		entry := l.PopTail()
 		require.NotNil(t, entry)
 	}
-	require.Equal(t, 40, l.len)
+	require.Equal(t, 40, int(l.len.Load()))
 	require.Equal(t, 2, l.count)
 }
 
@@ -81,13 +81,13 @@ func TestWheelList(t *testing.T) {
 		evicted := l.PushFront(NewEntry(fmt.Sprintf("%d", i), "", 1, 0))
 		require.Nil(t, evicted)
 	}
-	require.Equal(t, 5, l.len)
+	require.Equal(t, 5, int(l.len.Load()))
 	require.Equal(t, "4/3/2/1/0", l.display())
 	require.Equal(t, "0/1/2/3/4", l.displayReverse())
 
 	evicted := l.PushFront(NewEntry("5", "", 1, 0))
 	require.Equal(t, "0", evicted.key)
-	require.Equal(t, 5, l.len)
+	require.Equal(t, 5, int(l.len.Load()))
 	require.Equal(t, "5/4/3/2/1", l.display())
 	require.Equal(t, "1/2/3/4/5", l.displayReverse())
 

internal/slru.go

@@ -55,8 +55,8 @@ func (s *Slru[K, V]) remove(entry *Entry[K, V]) {
 func (s *Slru[K, V]) updateCost(entry *Entry[K, V], delta int64) {
 	switch entry.list {
 	case LIST_PROBATION:
-		s.probation.len += int(delta)
+		s.probation.len.Add(delta)
 	case LIST_PROTECTED:
-		s.protected.len += int(delta)
+		s.protected.len.Add(delta)
 	}
 }

internal/store.go

@@ -472,6 +472,16 @@ func (s *Store[K, V]) Len() int {
 	return total
 }
 
+func (s *Store[K, V]) EstimatedSize() 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
+}
+
 func (s *Store[K, V]) index(key K) (uint64, int) {
 	base := s.hasher.hash(key)
 	return base, int(base & uint64(s.shardCount-1))
@@ -829,7 +839,7 @@ func (s *Store[K, V]) Recover(version uint64, reader io.Reader) error {
 					continue
 				}
 				l := s.policy.slru.protected
-				if l.len < int(l.capacity) {
+				if l.len.Load() < int64(l.capacity) {
 					entry := pentry.entry()
 					l.PushBack(entry)
 					s.insertSimple(entry)
@@ -852,7 +862,7 @@ func (s *Store[K, V]) Recover(version uint64, reader io.Reader) error {
 				}
 				l1 := s.policy.slru.protected
 				l2 := s.policy.slru.probation
-				if l1.len+l2.len < int(s.policy.slru.maxsize) {
+				if l1.len.Load()+l2.len.Load() < int64(s.policy.slru.maxsize) {
 					entry := pentry.entry()
 					l2.PushBack(entry)
 					s.insertSimple(entry)