add test for expirationMap cleanup

ttl_test.go

@@ -0,0 +1,85 @@
+package ristretto
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+// TestExpirationMapCleanup tests the cleanup functionality of the expiration map.
+// It verifies that expired items are correctly evicted from the store and that
+// non-expired items remain in the store.
+func TestExpirationMapCleanup(t *testing.T) {
+	// Create a new expiration map
+	em := newExpirationMap[int]()
+	// Create a new store
+	s := newShardedMap[int]()
+	// Create a new policy
+	p := newDefaultPolicy[int](100, 10)
+
+	// Add items to the store and expiration map
+	now := time.Now()
+	i1 := &Item[int]{Key: 1, Conflict: 1, Value: 100, Expiration: now.Add(1 * time.Second)}
+	s.Set(i1)
+	em.add(i1.Key, i1.Conflict, i1.Expiration)
+
+	i2 := &Item[int]{Key: 2, Conflict: 2, Value: 200, Expiration: now.Add(3 * time.Second)}
+	s.Set(i2)
+	em.add(i2.Key, i2.Conflict, i2.Expiration)
+
+	// Wait for the first item to expire
+	time.Sleep(2 * time.Second)
+
+	// Cleanup the expiration map
+	evictedItems := make(map[uint64]int)
+	em.cleanup(s, p, func(item *Item[int]) {
+		evictedItems[item.Key] = item.Value
+	})
+
+	// Check that the first item was evicted
+	require.Equal(t, 1, len(evictedItems), "evictedItems should have 1 item")
+	require.Equal(t, 100, evictedItems[1], "evictedItems should have the first item")
+	_, ok := s.Get(i1.Key, i1.Conflict)
+	require.False(t, ok, "i1 should have been evicted")
+
+	// Check that the second item is still in the store
+	_, ok = s.Get(i2.Key, i2.Conflict)
+	require.True(t, ok, "i2 should still be in the store")
+
+	// Wait for the second item to expire
+	time.Sleep(2 * time.Second)
+
+	// Cleanup the expiration map
+	em.cleanup(s, p, func(item *Item[int]) {
+		evictedItems[item.Key] = item.Value
+	})
+
+	// Check that the second item was evicted
+	require.Equal(t, 2, len(evictedItems), "evictedItems should have 2 items")
+	require.Equal(t, 200, evictedItems[2], "evictedItems should have the second item")
+	_, ok = s.Get(i2.Key, i2.Conflict)
+	require.False(t, ok, "i2 should have been evicted")
+
+	t.Run("Miscalculation of buckets does not cause memory leaks", func(t *testing.T) {
+		// Break lastCleanedBucketNum, this can happen if the system time is changed.
+		em.lastCleanedBucketNum = storageBucket(now.AddDate(-1, 0, 0))
+
+		ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+		defer cancel()
+
+		done := make(chan struct{})
+		go func() {
+			em.cleanup(s, p, func(item *Item[int]) {})
+			close(done)
+		}()
+
+		select {
+		case <-done:
+			// cleanup completed!
+		case <-ctx.Done():
+			require.Fail(t, "cleanup method hangs / there may be a memory leak!")
+		}
+	})
+}