feat: add reason to OnEvictCallback

Author: costela

s3fifo/s3fifo.go

@@ -3,9 +3,10 @@ package s3fifo
 import (
 	"container/list"
 	"context"
-	"github.com/scalalang2/golang-fifo/types"
 	"sync"
 	"time"
+
+	"github.com/scalalang2/golang-fifo/types"
 )
 
 const numberOfBuckets = 100
@@ -150,7 +151,7 @@ func (s *S3FIFO[K, V]) Remove(key K) (ok bool) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if e, ok := s.items[key]; ok {
-		s.removeEntry(e)
+		s.removeEntry(e, types.EvictReasonRemoved)
 		return true
 	}
 
@@ -209,9 +210,9 @@ func (s *S3FIFO[K, V]) Close() {
 	s.mu.Unlock()
 }
 
-func (s *S3FIFO[K, V]) removeEntry(e *entry[K, V]) {
+func (s *S3FIFO[K, V]) removeEntry(e *entry[K, V], reason types.EvictReason) {
 	if s.callback != nil {
-		s.callback(e.key, e.value)
+		s.callback(e.key, e.value, reason)
 	}
 
 	if s.ghost.contains(e.key) {
@@ -256,7 +257,7 @@ func (s *S3FIFO[K, V]) deleteExpired() {
 	}
 
 	for _, e := range bucket.entries {
-		s.removeEntry(e)
+		s.removeEntry(e, types.EvictReasonExpired)
 	}
 
 	s.mu.Unlock()
@@ -292,7 +293,7 @@ func (s *S3FIFO[K, V]) evictFromSmall() {
 				s.evictFromMain()
 			}
 		} else {
-			s.removeEntry(el)
+			s.removeEntry(el, types.EvictReasonEvicted)
 			s.ghost.add(key)
 			evicted = true
 			delete(s.items, key)
@@ -314,7 +315,7 @@ func (s *S3FIFO[K, V]) evictFromMain() {
 			s.items[key].freq -= 1
 			s.items[key].element = s.main.PushFront(el.key)
 		} else {
-			s.removeEntry(el)
+			s.removeEntry(el, types.EvictReasonEvicted)
 			evicted = true
 			delete(s.items, key)
 		}

s3fifo/s3fifo_test.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"fortio.org/assert"
+	"github.com/scalalang2/golang-fifo/types"
 )
 
 const noEvictionTTL = 0
@@ -192,7 +193,7 @@ func TestEvictionCallback(t *testing.T) {
 	cache := New[int, int](10, noEvictionTTL)
 	evicted := make(map[int]int)
 
-	cache.SetOnEvicted(func(key int, value int) {
+	cache.SetOnEvicted(func(key int, value int, _ types.EvictReason) {
 		evicted[key] = value
 	})
 
@@ -216,7 +217,7 @@ func TestEvictionCallbackWithTTL(t *testing.T) {
 	var mu sync.Mutex
 	cache := New[int, int](10, time.Second)
 	evicted := make(map[int]int)
-	cache.SetOnEvicted(func(key int, value int) {
+	cache.SetOnEvicted(func(key int, value int, _ types.EvictReason) {
 		mu.Lock()
 		evicted[key] = value
 		mu.Unlock()

sieve/sieve.go

@@ -81,23 +81,25 @@ func New[K comparable, V any](size int, ttl time.Duration) *Sieve[K, V] {
 	}
 
 	if ttl != 0 {
-		go func(ctx context.Context) {
-			ticker := time.NewTicker(ttl / numberOfBuckets)
-			defer ticker.Stop()
-			for {
-				select {
-				case <-ctx.Done():
-					return
-				case <-ticker.C:
-					cache.deleteExpired()
-				}
-			}
-		}(cache.ctx)
+		go cache.cleanup(cache.ctx)
 	}
 
 	return cache
 }
 
+func (s *Sieve[K, V]) cleanup(ctx context.Context) {
+	ticker := time.NewTicker(s.ttl / numberOfBuckets)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			s.deleteExpired()
+		}
+	}
+}
+
 func (s *Sieve[K, V]) Set(key K, value V) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -147,7 +149,7 @@ func (s *Sieve[K, V]) Remove(key K) (ok bool) {
 			s.hand = s.hand.Prev()
 		}
 
-		s.removeEntry(e)
+		s.removeEntry(e, types.EvictReasonRemoved)
 		return true
 	}
 
@@ -191,7 +193,7 @@ func (s *Sieve[K, V]) Purge() {
 	defer s.mu.Unlock()
 
 	for _, e := range s.items {
-		s.removeEntry(e)
+		s.removeEntry(e, types.EvictReasonRemoved)
 	}
 
 	for i := range s.buckets {
@@ -213,9 +215,9 @@ func (s *Sieve[K, V]) Close() {
 	s.mu.Unlock()
 }
 
-func (s *Sieve[K, V]) removeEntry(e *entry[K, V]) {
+func (s *Sieve[K, V]) removeEntry(e *entry[K, V], reason types.EvictReason) {
 	if s.callback != nil {
-		s.callback(e.key, e.value)
+		s.callback(e.key, e.value, reason)
 	}
 
 	s.ll.Remove(e.element)
@@ -249,7 +251,7 @@ func (s *Sieve[K, V]) evict() {
 	}
 
 	s.hand = o.Prev()
-	s.removeEntry(el)
+	s.removeEntry(el, types.EvictReasonEvicted)
 }
 
 func (s *Sieve[K, V]) addToBucket(e *entry[K, V]) {
@@ -285,7 +287,7 @@ func (s *Sieve[K, V]) deleteExpired() {
 	}
 
 	for _, e := range bucket.entries {
-		s.removeEntry(e)
+		s.removeEntry(e, types.EvictReasonExpired)
 	}
 
 	s.mu.Unlock()

sieve/sieve_test.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"fortio.org/assert"
+	"github.com/scalalang2/golang-fifo/types"
 )
 
 const noEvictionTTL = 0
@@ -146,7 +147,7 @@ func TestEvictionCallback(t *testing.T) {
 	cache := New[int, int](10, noEvictionTTL)
 	evicted := make(map[int]int)
 
-	cache.SetOnEvicted(func(key int, value int) {
+	cache.SetOnEvicted(func(key int, value int, _ types.EvictReason) {
 		evicted[key] = value
 	})
 
@@ -170,7 +171,7 @@ func TestEvictionCallbackWithTTL(t *testing.T) {
 	var mu sync.Mutex
 	cache := New[int, int](10, time.Second)
 	evicted := make(map[int]int)
-	cache.SetOnEvicted(func(key int, value int) {
+	cache.SetOnEvicted(func(key int, value int, _ types.EvictReason) {
 		mu.Lock()
 		evicted[key] = value
 		mu.Unlock()

types/types.go

@@ -1,6 +1,19 @@
 package types
 
-type OnEvictCallback[K comparable, V any] func(key K, value V)
+// EvictReason is the reason for an entry to be evicted from the cache.
+// It is used in the [OnEvictCallback] function.
+type EvictReason int
+
+const (
+	// EvictReasonExpired is used when an item is removed because its TTL has expired.
+	EvictReasonExpired = iota
+	// EvictReasonEvicted is used when an item is removed because the cache size limit was exceeded.
+	EvictReasonEvicted
+	// EvictReasonRemoved is used when an item is explicitly deleted.
+	EvictReasonRemoved
+)
+
+type OnEvictCallback[K comparable, V any] func(key K, value V, reason EvictReason)
 
 // Cache is the interface for a cache.
 type Cache[K comparable, V any] interface {