diff --git a/ttl_test.go b/ttl_test.go new file mode 100644 index 00000000..dcd2417c --- /dev/null +++ b/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!") + } + }) +}