improved expiration items

Author: Code-Hex

cache.go

@@ -46,9 +46,13 @@ type Item[K comparable, V any] struct {
 	InitialReferenceCount int
 }
 
+func (item *Item[K, V]) hasExpiration() bool {
+	return !item.Expiration.IsZero()
+}
+
 // Expired returns true if the item has expired.
 func (item *Item[K, V]) Expired() bool {
-	if item.Expiration.IsZero() {
+	if !item.hasExpiration() {
 		return false
 	}
 	return nowFunc().After(item.Expiration)
@@ -107,8 +111,9 @@ func newItem[K comparable, V any](key K, val V, opts ...ItemOption) *Item[K, V]
 type Cache[K comparable, V any] struct {
 	cache Interface[K, *Item[K, V]]
 	// mu is used to do lock in some method process.
-	mu      sync.Mutex
-	janitor *janitor
+	mu         sync.Mutex
+	janitor    *janitor
+	expManager *expirationManager[K]
 }
 
 // Option is an option for cache.
@@ -190,15 +195,16 @@ func NewContext[K comparable, V any](ctx context.Context, opts ...Option[K, V])
 		optFunc(o)
 	}
 	cache := &Cache[K, V]{
-		cache:   o.cache,
-		janitor: newJanitor(ctx, o.janitorInterval),
+		cache:      o.cache,
+		janitor:    newJanitor(ctx, o.janitorInterval),
+		expManager: newExpirationManager[K](),
 	}
 	cache.janitor.run(cache.DeleteExpired)
 	return cache
 }
 
 // Get looks up a key's value from the cache.
-func (c *Cache[K, V]) Get(key K) (value V, ok bool) {
+func (c *Cache[K, V]) Get(key K) (zero V, ok bool) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	item, ok := c.cache.Get(key)
@@ -210,7 +216,7 @@ func (c *Cache[K, V]) Get(key K) (value V, ok bool) {
 	// Returns nil if the item has been expired.
 	// Do not delete here and leave it to an external process such as Janitor.
 	if item.Expired() {
-		return value, false
+		return zero, false
 	}
 
 	return item.Value, true
@@ -219,17 +225,30 @@ func (c *Cache[K, V]) Get(key K) (value V, ok bool) {
 // DeleteExpired all expired items from the cache.
 func (c *Cache[K, V]) DeleteExpired() {
 	c.mu.Lock()
-	keys := c.cache.Keys()
+	l := c.expManager.len()
 	c.mu.Unlock()
 
-	for _, key := range keys {
-		c.mu.Lock()
+	evict := func() bool {
+		key := c.expManager.pop()
 		// if is expired, delete it and return nil instead
 		item, ok := c.cache.Get(key)
-		if ok && item.Expired() {
-			c.cache.Delete(key)
+		if ok {
+			if item.Expired() {
+				c.cache.Delete(key)
+				return false
+			}
+			c.expManager.update(key, item.Expiration)
 		}
+		return true
+	}
+
+	for i := 0; i < l; i++ {
+		c.mu.Lock()
+		shouldBreak := evict()
 		c.mu.Unlock()
+		if shouldBreak {
+			break
+		}
 	}
 }
 
@@ -238,6 +257,9 @@ func (c *Cache[K, V]) Set(key K, val V, opts ...ItemOption) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	item := newItem(key, val, opts...)
+	if item.hasExpiration() {
+		c.expManager.update(key, item.Expiration)
+	}
 	c.cache.Set(key, item)
 }
 
@@ -253,6 +275,7 @@ func (c *Cache[K, V]) Delete(key K) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	c.cache.Delete(key)
+	c.expManager.remove(key)
 }
 
 // Len returns the number of items in the cache.

cache_internal_test.go

@@ -26,3 +26,108 @@ func TestDeletedCache(t *testing.T) {
 		t.Fatal("want false")
 	}
 }
+
+func TestDeleteExpired(t *testing.T) {
+	now := time.Now()
+	restore := func() {
+		nowFunc = time.Now
+	}
+
+	t.Run("normal", func(t *testing.T) {
+		defer restore()
+		c := New[string, int]()
+
+		c.Set("0", 0)
+		c.Set("1", 10, WithExpiration(10*time.Millisecond))
+		c.Set("2", 20, WithExpiration(20*time.Millisecond))
+		c.Set("3", 30, WithExpiration(30*time.Millisecond))
+		c.Set("4", 40, WithExpiration(40*time.Millisecond))
+		c.Set("5", 50)
+
+		maxItems := c.Len()
+
+		expItems := 2
+
+		for i := 0; i <= maxItems; i++ {
+			nowFunc = func() time.Time {
+				// Advance time to expire some items
+				advanced := time.Duration(i * 10)
+				return now.Add(advanced * time.Millisecond).Add(time.Millisecond)
+			}
+
+			c.DeleteExpired()
+
+			got := c.Len()
+			want := max(maxItems-i, expItems)
+			if want != got {
+				t.Errorf("want %d items but got %d", want, got)
+			}
+		}
+	})
+
+	t.Run("with remove", func(t *testing.T) {
+		defer restore()
+		c := New[string, int]()
+
+		c.Set("0", 0)
+		c.Set("1", 10, WithExpiration(10*time.Millisecond))
+		c.Set("2", 20, WithExpiration(20*time.Millisecond))
+
+		c.Delete("1")
+
+		nowFunc = func() time.Time {
+			return now.Add(30 * time.Millisecond).Add(time.Millisecond)
+		}
+
+		c.DeleteExpired()
+
+		keys := c.Keys()
+		want := 1
+		if want != len(keys) {
+			t.Errorf("want %d items but got %d", want, len(keys))
+		}
+	})
+
+	t.Run("with update", func(t *testing.T) {
+		defer restore()
+		c := New[string, int]()
+
+		c.Set("0", 0)
+		c.Set("1", 10, WithExpiration(10*time.Millisecond))
+		c.Set("2", 20, WithExpiration(20*time.Millisecond))
+		c.Set("1", 30, WithExpiration(30*time.Millisecond)) // update
+
+		maxItems := c.Len()
+
+		nowFunc = func() time.Time {
+			return now.Add(10 * time.Millisecond).Add(time.Millisecond)
+		}
+
+		c.DeleteExpired()
+
+		got1 := c.Len()
+		want1 := maxItems
+		if want1 != got1 {
+			t.Errorf("want1 %d items but got1 %d", want1, got1)
+		}
+
+		nowFunc = func() time.Time {
+			return now.Add(30 * time.Millisecond).Add(time.Millisecond)
+		}
+
+		c.DeleteExpired()
+
+		got2 := c.Len()
+		want2 := 1
+		if want2 != got2 {
+			t.Errorf("want2 %d items but got2 %d", want2, got2)
+		}
+	})
+}
+
+func max(x, y int) int {
+	if x < y {
+		return y
+	}
+	return x
+}

expiration.go

@@ -0,0 +1,91 @@
+package cache
+
+import (
+	"container/heap"
+	"time"
+)
+
+type expirationManager[K comparable] struct {
+	queue   expirationQueue[K]
+	mapping map[K]*expirationKey[K]
+}
+
+func newExpirationManager[K comparable]() *expirationManager[K] {
+	q := make(expirationQueue[K], 0)
+	heap.Init(&q)
+	return &expirationManager[K]{
+		queue:   q,
+		mapping: make(map[K]*expirationKey[K]),
+	}
+}
+
+func (m *expirationManager[K]) update(key K, expiration time.Time) {
+	if e, ok := m.mapping[key]; ok {
+		heap.Fix(&m.queue, e.index)
+	} else {
+		v := &expirationKey[K]{
+			key:        key,
+			expiration: expiration,
+		}
+		heap.Push(&m.queue, v)
+		m.mapping[key] = v
+	}
+}
+
+func (m *expirationManager[K]) len() int {
+	return m.queue.Len()
+}
+
+func (m *expirationManager[K]) pop() K {
+	v := heap.Pop(&m.queue)
+	key := v.(*expirationKey[K]).key
+	delete(m.mapping, key)
+	return key
+}
+
+func (m *expirationManager[K]) remove(key K) {
+	if e, ok := m.mapping[key]; ok {
+		heap.Remove(&m.queue, e.index)
+		delete(m.mapping, key)
+	}
+}
+
+type expirationKey[K comparable] struct {
+	key        K
+	expiration time.Time
+	index      int
+}
+
+// expirationQueue implements heap.Interface and holds CacheItems.
+type expirationQueue[K comparable] []*expirationKey[K]
+
+var _ heap.Interface = (*expirationQueue[int])(nil)
+
+func (pq expirationQueue[K]) Len() int { return len(pq) }
+
+func (pq expirationQueue[K]) Less(i, j int) bool {
+	// We want Pop to give us the least based on expiration time, not the greater
+	return pq[i].expiration.Before(pq[j].expiration)
+}
+
+func (pq expirationQueue[K]) Swap(i, j int) {
+	pq[i], pq[j] = pq[j], pq[i]
+	pq[i].index = i
+	pq[j].index = j
+}
+
+func (pq *expirationQueue[K]) Push(x interface{}) {
+	n := len(*pq)
+	item := x.(*expirationKey[K])
+	item.index = n
+	*pq = append(*pq, item)
+}
+
+func (pq *expirationQueue[K]) Pop() interface{} {
+	old := *pq
+	n := len(old)
+	item := old[n-1]
+	item.index = -1 // For safety
+	*pq = old[0 : n-1]
+	return item
+}