diff --git a/map_ttl.go b/map_ttl.go index c81480b..1b082ad 100644 --- a/map_ttl.go +++ b/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 } diff --git a/ttl_test.go b/ttl_test.go index ecd4c91..331439e 100644 --- a/ttl_test.go +++ b/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)) + } +}