feat: New WithMaxCost option for custom cache capacity management strategies

cache.go

@@ -7,6 +7,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/DmitriyVTitov/size"
 	"golang.org/x/sync/singleflight"
 )
 
@@ -15,6 +16,7 @@ const (
 	EvictionReasonDeleted EvictionReason = iota + 1
 	EvictionReasonCapacityReached
 	EvictionReasonExpired
+	EvictionReasonMaxMemorySizeExceeded
 )
 
 // EvictionReason is used to specify why a certain item was
@@ -36,6 +38,7 @@ type Cache[K comparable, V any] struct {
 
 		timerCh chan time.Duration
 	}
+	sizeInBytes uint64
 
 	metricsMu sync.RWMutex
 	metrics   Metrics
@@ -137,7 +140,23 @@ func (c *Cache[K, V]) set(key K, value V, ttl time.Duration) *Item[K, V] {
 	if elem != nil {
 		// update/overwrite an existing item
 		item := elem.Value.(*Item[K, V])
+		oldValue := item.value
 		item.update(value, ttl)
+
+		if c.options.sizeInBytes != 0 {
+			oldSize := size.Of(oldValue)
+			newSize := size.Of(value)
+
+			// size.Of returns -1 on errors
+			if oldSize != -1 && newSize != -1 {
+				c.sizeInBytes = c.sizeInBytes - uint64(oldSize) + uint64(newSize)
+			}
+
+			for c.sizeInBytes > c.options.sizeInBytes {
+				c.evict(EvictionReasonMaxMemorySizeExceeded, c.items.lru.Back())
+			}
+		}
+
 		c.updateExpirations(false, elem)
 
 		return item
@@ -158,6 +177,17 @@ func (c *Cache[K, V]) set(key K, value V, ttl time.Duration) *Item[K, V] {
 	c.items.values[key] = elem
 	c.updateExpirations(true, elem)
 
+	if c.options.sizeInBytes != 0 {
+		itemSize := size.Of(item)
+		if itemSize != -1 {
+			c.sizeInBytes += uint64(itemSize)
+		}
+
+		for c.sizeInBytes > c.options.sizeInBytes {
+			c.evict(EvictionReasonMaxMemorySizeExceeded, c.items.lru.Back())
+		}
+	}
+
 	c.metricsMu.Lock()
 	c.metrics.Insertions++
 	c.metricsMu.Unlock()
@@ -258,6 +288,14 @@ func (c *Cache[K, V]) evict(reason EvictionReason, elems ...*list.Element) {
 		for i := range elems {
 			item := elems[i].Value.(*Item[K, V])
 			delete(c.items.values, item.key)
+
+			if c.options.sizeInBytes != 0 {
+				itemSize := size.Of(item)
+				if itemSize != -1 {
+					c.sizeInBytes -= uint64(itemSize)
+				}
+			}
+
 			c.items.lru.Remove(elems[i])
 			c.items.expQueue.remove(elems[i])
 

go.mod

@@ -3,6 +3,7 @@ module github.com/jellydator/ttlcache/v3
 go 1.18
 
 require (
+	github.com/DmitriyVTitov/size v1.5.0
 	github.com/stretchr/testify v1.9.0
 	go.uber.org/goleak v1.3.0
 	golang.org/x/sync v0.8.0

go.sum

@@ -1,6 +1,10 @@
+github.com/DmitriyVTitov/size v1.5.0 h1:/PzqxYrOyOUX1BXj6J9OuVRVGe+66VL4D9FlUaW515g=
+github.com/DmitriyVTitov/size v1.5.0/go.mod h1:le6rNI4CoLQV1b9gzp1+3d7hMAD/uu2QcJ+aYbNgiU0=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da h1:oI5xCqsCo564l8iNU+DwB5epxmsaqB+rhGL0m5jtYqE=
+github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
 github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=

options.go

@@ -18,6 +18,7 @@ func (fn optionFunc[K, V]) apply(opts *options[K, V]) {
 // options holds all available cache configuration options.
 type options[K comparable, V any] struct {
 	capacity              uint64
+	sizeInBytes           uint64
 	ttl                   time.Duration
 	loader                Loader[K, V]
 	disableTouchOnHit     bool
@@ -75,3 +76,11 @@ func WithDisableTouchOnHit[K comparable, V any]() Option[K, V] {
 		opts.disableTouchOnHit = true
 	})
 }
+
+// WithMemorySize sets the maximum memory size the cache is allowed to grow.
+// If used together with WithCapacity, WithMemorySize overrules the maximum capacity.
+func WithMemorySize[K comparable, V any](s uint64) Option[K, V] {
+	return optionFunc[K, V](func(opts *options[K, V]) {
+		opts.sizeInBytes = s
+	})
+}

options_test.go

@@ -68,3 +68,10 @@ func Test_WithDisableTouchOnHit(t *testing.T) {
 	WithDisableTouchOnHit[string, string]().apply(&opts)
 	assert.True(t, opts.disableTouchOnHit)
 }
+
+func Test_WithMemorySize(t *testing.T) {
+	var opts options[string, string]
+
+	WithMemorySize[string, string](1024).apply(&opts)
+	assert.Equal(t, uint64(1024), opts.sizeInBytes)
+}