Fix some memory leaks in TTL implementation

store.go

@@ -111,6 +111,7 @@ func (sm *shardedMap[V]) Clear(onEvict func(item *Item[V])) {
 	for i := uint64(0); i < numShards; i++ {
 		sm.shards[i].Clear(onEvict)
 	}
+	sm.expiryMap.clear()
 }
 
 type lockedMap[V any] struct {

ttl.go

@@ -42,12 +42,14 @@ type bucket map[uint64]uint64
 // expirationMap is a map of bucket number to the corresponding bucket.
 type expirationMap[V any] struct {
 	sync.RWMutex
-	buckets map[int64]bucket
+	buckets              map[int64]bucket
+	lastCleanedBucketNum int64
 }
 
 func newExpirationMap[V any]() *expirationMap[V] {
 	return &expirationMap[V]{
-		buckets: make(map[int64]bucket),
+		buckets:              make(map[int64]bucket),
+		lastCleanedBucketNum: cleanupBucket(time.Now()),
 	}
 }
 
@@ -87,6 +89,11 @@ func (m *expirationMap[_]) update(key, conflict uint64, oldExpTime, newExpTime t
 		delete(oldBucket, key)
 	}
 
+	// Items that don't expire don't need to be in the expiration map.
+	if newExpTime.IsZero() {
+		return
+	}
+
 	newBucketNum := storageBucket(newExpTime)
 	newBucket, ok := m.buckets[newBucketNum]
 	if !ok {
@@ -121,29 +128,50 @@ func (m *expirationMap[V]) cleanup(store store[V], policy policy[V], onEvict fun
 
 	m.Lock()
 	now := time.Now()
-	bucketNum := cleanupBucket(now)
-	keys := m.buckets[bucketNum]
-	delete(m.buckets, bucketNum)
+	currentBucketNum := cleanupBucket(now)
+	// Clean up all buckets up to and including currentBucketNum, starting from
+	// (but not including) the last one that was cleaned up
+	var buckets []bucket
+	for bucketNum := m.lastCleanedBucketNum + 1; bucketNum <= currentBucketNum; bucketNum++ {
+		buckets = append(buckets, m.buckets[bucketNum])
+		delete(m.buckets, bucketNum)
+	}
+	m.lastCleanedBucketNum = currentBucketNum
 	m.Unlock()
 
-	for key, conflict := range keys {
-		expr := store.Expiration(key)
-		// Sanity check. Verify that the store agrees that this key is expired.
-		if expr.After(now) {
-			continue
+	for _, keys := range buckets {
+		for key, conflict := range keys {
+			expr := store.Expiration(key)
+			// Sanity check. Verify that the store agrees that this key is expired.
+			if store.Expiration(key).After(now) {
+				continue
+			}
+
+			cost := policy.Cost(key)
+			policy.Del(key)
+			_, value := store.Del(key, conflict)
+
+			if onEvict != nil {
+				onEvict(&Item[V]{Key: key,
+					Conflict:   conflict,
+					Value:      value,
+					Cost:       cost,
+					Expiration: expr,
+				})
+			}
 		}
+	}
+}
 
-		cost := policy.Cost(key)
-		policy.Del(key)
-		_, value := store.Del(key, conflict)
-
-		if onEvict != nil {
-			onEvict(&Item[V]{Key: key,
-				Conflict:   conflict,
-				Value:      value,
-				Cost:       cost,
-				Expiration: expr,
-			})
-		}
+// clear clears the expirationMap, the caller is responsible for properly
+// evicting the referenced items
+func (m *expirationMap[V]) clear() {
+	if m == nil {
+		return
 	}
+
+	m.Lock()
+	m.buckets = make(map[int64]bucket)
+	m.lastCleanedBucketNum = cleanupBucket(time.Now())
+	m.Unlock()
 }