add onEvict callback for TTLCache

map_ttl.go

@@ -23,16 +23,20 @@ func zero[T any]() T {
 	return z
 }
 
+type onEvictFunc[K comparable, V any] func(key K, value V)
+
 // MapTTLCache is the thread-safe map-based cache with TTL cache invalidation support.
 // MapTTLCache uses double linked list to maintain FIFO order of inserted values.
 type MapTTLCache[K comparable, V any] struct {
-	data map[K]ttlRec[K, V]
-	mux  sync.RWMutex
-	ttl  time.Duration
-	now  func() time.Time
-	tail K
-	head K
-	zero K
+	data    map[K]ttlRec[K, V]
+	mux     sync.RWMutex
+	ttl     time.Duration
+	// TODO: replace with sync.Test
+	now     func() time.Time
+	onEvict onEvictFunc[K, V]
+	tail    K
+	head    K
+	zero    K
 }
 
 // NewMapTTLCache creates MapTTLCache instance and spawns background
@@ -70,6 +74,15 @@ func NewMapTTLCache[K comparable, V any](
 	return &c
 }
 
+// OnEvict sets a callback function that will be called when an entry is evicted from the cache
+// due to TTL expiration. The callback receives the key and value of the evicted entry.
+// Note that the eviction callback is not called for Del operation.
+func (c *MapTTLCache[K, V]) OnEvict(f onEvictFunc[K, V]) {
+	c.mux.Lock()
+	c.onEvict = f
+	c.mux.Unlock()
+}
+
 func (c *MapTTLCache[K, V]) Set(key K, value V) {
 	c.mux.Lock()
 	defer c.mux.Unlock()
@@ -145,10 +158,22 @@ func (c *MapTTLCache[K, V]) Del(key K) error {
 	return nil
 }
 
-// cleanup removes outdated records.
+// cleanup removes outdated records
+// and calls eviction callbacks.
 func (c *MapTTLCache[K, V]) cleanup() error {
+	var (
+		evicted map[K]V
+		onEvict onEvictFunc[K, V]
+	)
+
 	c.mux.Lock()
-	defer c.mux.Unlock()
+
+	// Preallocate a small map for evicted records
+	// if eviction callback is set.
+	if c.onEvict != nil {
+		onEvict = c.onEvict
+		evicted = make(map[K]V, 16)
+	}
 
 	key := c.head
 	for {
@@ -164,9 +189,13 @@ func (c *MapTTLCache[K, V]) cleanup() error {
 		c.head = rec.next
 		delete(c.data, key)
 
+		if onEvict != nil {
+			evicted[key] = rec.value
+		}
+
 		if key == c.tail {
 			c.tail = c.zero
-			return nil
+			break
 		}
 
 		next, ok := c.data[rec.next]
@@ -176,6 +205,12 @@ func (c *MapTTLCache[K, V]) cleanup() error {
 		}
 		key = rec.next
 	}
+	c.mux.Unlock()
+
+	// Call eviction callbacks outside of the lock.
+	for k, v := range evicted {
+		onEvict(k, v)
+	}
 
 	return nil
 }

ttl_test.go

@@ -332,3 +332,193 @@ func TestSetIfPresentResetsTTL(t *testing.T) {
 		t.Errorf("value was not updated by SetIfPresent, expected %v, but got %v", "value2", v)
 	}
 }
+
+func TestOnEvict(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	c := NewMapTTLCache[string, string](ctx, time.Second, time.Hour)
+	ts := time.Now()
+
+	evicted := make(map[string]string)
+	var mu sync.Mutex
+
+	c.OnEvict(func(key string, value string) {
+		mu.Lock()
+		evicted[key] = value
+		mu.Unlock()
+	})
+
+	c.Set("key1", "value1")
+	c.Set("key2", "value2")
+	c.Set("key3", "value3")
+
+	// Override now to simulate time passing
+	c.mux.Lock()
+	c.now = func() time.Time { return ts.Add(2 * time.Second) }
+	c.mux.Unlock()
+
+	// Manually trigger cleanup
+	if err := c.cleanup(); err != nil {
+		t.Errorf("unexpected error in cleanup: %v", err)
+	}
+
+	mu.Lock()
+	defer mu.Unlock()
+
+	if len(evicted) != 3 {
+		t.Errorf("expected 3 evictions, got %d", len(evicted))
+	}
+
+	expected := map[string]string{
+		"key1": "value1",
+		"key2": "value2",
+		"key3": "value3",
+	}
+
+	for k, v := range expected {
+		if evicted[k] != v {
+			t.Errorf("expected evicted[%q] = %q, got %q", k, v, evicted[k])
+		}
+	}
+}
+
+func TestOnEvictNotCalledForDel(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	c := NewMapTTLCache[string, string](ctx, time.Second, time.Hour)
+	ts := time.Now()
+
+	evicted := make(map[string]string)
+	var mu sync.Mutex
+
+	c.OnEvict(func(key string, value string) {
+		mu.Lock()
+		evicted[key] = value
+		mu.Unlock()
+	})
+
+	c.Set("key1", "value1")
+	c.Set("key2", "value2")
+
+	// Delete key1 explicitly
+	if err := c.Del("key1"); err != nil {
+		t.Errorf("unexpected error in Del: %v", err)
+	}
+
+	mu.Lock()
+	if len(evicted) != 0 {
+		t.Errorf("expected no evictions from Del, got %d", len(evicted))
+	}
+	mu.Unlock()
+
+	// Override now to simulate TTL expiration
+	c.mux.Lock()
+	c.now = func() time.Time { return ts.Add(2 * time.Second) }
+	c.mux.Unlock()
+
+	// Manually trigger cleanup
+	if err := c.cleanup(); err != nil {
+		t.Errorf("unexpected error in cleanup: %v", err)
+	}
+
+	mu.Lock()
+	defer mu.Unlock()
+
+	if len(evicted) != 1 {
+		t.Errorf("expected 1 eviction from TTL, got %d", len(evicted))
+	}
+
+	if evicted["key2"] != "value2" {
+		t.Errorf("expected evicted[key2] = value2, got %q", evicted["key2"])
+	}
+
+	if _, ok := evicted["key1"]; ok {
+		t.Errorf("key1 should not be in evicted map as it was deleted with Del()")
+	}
+}
+
+func TestOnEvictPartialCleanup(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	c := NewMapTTLCache[string, string](ctx, time.Second, time.Hour)
+	ts := time.Now()
+
+	evicted := make(map[string]string)
+	var mu sync.Mutex
+
+	c.OnEvict(func(key string, value string) {
+		mu.Lock()
+		evicted[key] = value
+		mu.Unlock()
+	})
+
+	// Override now to time T for first batch
+	c.mux.Lock()
+	c.now = func() time.Time { return ts }
+	c.mux.Unlock()
+
+	// Add first batch at time T
+	c.Set("key1", "value1")
+	c.Set("key2", "value2")
+
+	// Override now to time T+700ms for second batch
+	c.mux.Lock()
+	c.now = func() time.Time { return ts.Add(700 * time.Millisecond) }
+	c.mux.Unlock()
+
+	// Add second batch at time T+700ms
+	c.Set("key3", "value3")
+	c.Set("key4", "value4")
+
+	// Override now to T+1.5s - first batch expires but second doesn't
+	c.mux.Lock()
+	c.now = func() time.Time { return ts.Add(1500 * time.Millisecond) }
+	c.mux.Unlock()
+
+	// Manually trigger cleanup - should evict first batch only
+	if err := c.cleanup(); err != nil {
+		t.Errorf("unexpected error in cleanup: %v", err)
+	}
+
+	mu.Lock()
+	if len(evicted) != 2 {
+		t.Errorf("expected 2 evictions, got %d", len(evicted))
+	}
+
+	if evicted["key1"] != "value1" {
+		t.Errorf("expected evicted[key1] = value1, got %q", evicted["key1"])
+	}
+
+	if evicted["key2"] != "value2" {
+		t.Errorf("expected evicted[key2] = value2, got %q", evicted["key2"])
+	}
+
+	if _, ok := evicted["key3"]; ok {
+		t.Errorf("key3 should not be evicted yet")
+	}
+
+	if _, ok := evicted["key4"]; ok {
+		t.Errorf("key4 should not be evicted yet")
+	}
+	mu.Unlock()
+
+	// Override now to T+2s - second batch now expires
+	c.mux.Lock()
+	c.now = func() time.Time { return ts.Add(2 * time.Second) }
+	c.mux.Unlock()
+
+	// Manually trigger cleanup - should evict second batch
+	if err := c.cleanup(); err != nil {
+		t.Errorf("unexpected error in cleanup: %v", err)
+	}
+
+	mu.Lock()
+	defer mu.Unlock()
+
+	if len(evicted) != 4 {
+		t.Errorf("expected 4 evictions total, got %d", len(evicted))
+	}
+}