From 2e104a8053fa31d76ba7f568fdc7ef20e38a5bf7 Mon Sep 17 00:00:00 2001
From: Nick Pillitteri <56quarters@users.noreply.github.com>
Date: Wed, 9 Oct 2024 10:11:03 -0400
Subject: [PATCH] cache: Add ability to manually advance "now" for mock cache
 (#601)

Add an `.Advance()` method to MockCache and InstrumentedMockCache to
allow the time considered "now" to be moved without needing to actually
sleep. This is useful for testing when items are set with a TTL and you
would like for them to actually expire as they would in a real cache.

Part of https://github.com/grafana/mimir/issues/9386

Signed-off-by: Nick Pillitteri <nick.pillitteri@grafana.com>
---
 CHANGELOG.md  |  1 +
 cache/mock.go | 30 +++++++++++++++++++++++-------
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 53a23358e..3288acdeb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -232,6 +232,7 @@
 * [ENHANCEMENT] grpcclient: Support custom gRPC compressors. #583
 * [ENHANCEMENT] Adapt `metrics.SendSumOfGaugesPerTenant` to use `metrics.MetricOption`. #584
 * [ENHANCEMENT] Cache: Add `.Add()` and `.Set()` methods to cache clients. #591
+* [ENHANCEMENT] Cache: Add `.Advance()` methods to mock cache clients for easier testing of TTLs. #601
 * [CHANGE] Backoff: added `Backoff.ErrCause()` which is like `Backoff.Err()` but returns the context cause if backoff is terminated because the context has been canceled. #538
 * [BUGFIX] spanlogger: Support multiple tenant IDs. #59
 * [BUGFIX] Memberlist: fixed corrupted packets when sending compound messages with more than 255 messages or messages bigger than 64KB. #85
diff --git a/cache/mock.go b/cache/mock.go
index 4a5dae962..15d95419a 100644
--- a/cache/mock.go
+++ b/cache/mock.go
@@ -19,10 +19,11 @@ var (
 type MockCache struct {
 	mu    sync.Mutex
 	cache map[string]Item
+	now   time.Time
 }
 
 func NewMockCache() *MockCache {
-	c := &MockCache{}
+	c := &MockCache{now: time.Now()}
 	c.Flush()
 	return c
 }
@@ -30,14 +31,14 @@ func NewMockCache() *MockCache {
 func (m *MockCache) SetAsync(key string, value []byte, ttl time.Duration) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
-	m.cache[key] = Item{Data: value, ExpiresAt: time.Now().Add(ttl)}
+	m.cache[key] = Item{Data: value, ExpiresAt: m.now.Add(ttl)}
 }
 
 func (m *MockCache) SetMultiAsync(data map[string][]byte, ttl time.Duration) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	exp := time.Now().Add(ttl)
+	exp := m.now.Add(ttl)
 	for key, val := range data {
 		m.cache[key] = Item{Data: val, ExpiresAt: exp}
 	}
@@ -46,7 +47,7 @@ func (m *MockCache) SetMultiAsync(data map[string][]byte, ttl time.Duration) {
 func (m *MockCache) Set(_ context.Context, key string, value []byte, ttl time.Duration) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
-	m.cache[key] = Item{Data: value, ExpiresAt: time.Now().Add(ttl)}
+	m.cache[key] = Item{Data: value, ExpiresAt: m.now.Add(ttl)}
 	return nil
 }
 
@@ -54,11 +55,11 @@ func (m *MockCache) Add(_ context.Context, key string, value []byte, ttl time.Du
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	if _, ok := m.cache[key]; ok {
+	if i, ok := m.cache[key]; ok && i.ExpiresAt.After(m.now) {
 		return ErrNotStored
 	}
 
-	m.cache[key] = Item{Data: value, ExpiresAt: time.Now().Add(ttl)}
+	m.cache[key] = Item{Data: value, ExpiresAt: m.now.Add(ttl)}
 	return nil
 }
 
@@ -68,7 +69,7 @@ func (m *MockCache) GetMulti(_ context.Context, keys []string, _ ...Option) map[
 
 	found := make(map[string][]byte, len(keys))
 
-	now := time.Now()
+	now := m.now
 	for _, k := range keys {
 		v, ok := m.cache[k]
 		if ok && now.Before(v.ExpiresAt) {
@@ -107,6 +108,7 @@ func (m *MockCache) Delete(_ context.Context, key string) error {
 	return nil
 }
 
+// Flush removes all entries from the cache
 func (m *MockCache) Flush() {
 	m.mu.Lock()
 	defer m.mu.Unlock()
@@ -114,6 +116,14 @@ func (m *MockCache) Flush() {
 	m.cache = map[string]Item{}
 }
 
+// Advance changes "now" by the given duration
+func (m *MockCache) Advance(d time.Duration) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	m.now = m.now.Add(d)
+}
+
 // InstrumentedMockCache is a mocked cache implementation which also tracks the number
 // of times its functions are called.
 type InstrumentedMockCache struct {
@@ -172,10 +182,16 @@ func (m *InstrumentedMockCache) GetItems() map[string]Item {
 	return m.cache.GetItems()
 }
 
+// Flush removes all entries from the cache
 func (m *InstrumentedMockCache) Flush() {
 	m.cache.Flush()
 }
 
+// Advance changes "now" by the given duration
+func (m *InstrumentedMockCache) Advance(d time.Duration) {
+	m.cache.Advance(d)
+}
+
 func (m *InstrumentedMockCache) CountStoreCalls() int {
 	return int(m.storeCount.Load())
 }