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 grafana/mimir#9386

Signed-off-by: Nick Pillitteri <[email protected]>

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

cache/mock.go

@@ -19,25 +19,26 @@ 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
 }
 
 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,19 +47,19 @@ 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
 }
 
 func (m *MockCache) Add(_ context.Context, key string, value []byte, ttl time.Duration) error {
 	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,13 +108,22 @@ 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()
 
 	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())
 }