Update ordered list of keys on eviction

2 files changed, 28 insertions, 5 deletions

diff --git a/cache/cache.go b/cache/cache.go
index dec3cfc..7dc654f 100644
--- a/cache/cache.go
+++ b/cache/cache.go
@@ -144,7 +144,7 @@ func (c *Cache) Set(k uint64, msg *dns.Msg) {
 	if c.capacity == 0 {
 		return
 	}
-	if !isCacheable(msg) {
+	if !canCache(msg) {
 		return
 	}
 	now := c.now()
@@ -158,13 +158,26 @@ func (c *Cache) Set(k uint64, msg *dns.Msg) {
 	c.mu.Unlock()
 }
 
-func (c *Cache) deleteExpired() {
+func (c *Cache) evictExpired() {
 	c.mu.Lock()
+	evictedKeys := make(map[uint64]bool)
 	for k, v := range c.values {
 		if c.isExpired(v) {
 			delete(c.values, k)
+			evictedKeys[k] = true
 		}
 	}
+	if len(evictedKeys) > 0 {
+		// At least one entry was evicted. The ordered list of keys must be updated.
+		var keys []uint64
+		for _, k := range c.keys {
+			if _, ok := evictedKeys[k]; ok {
+				continue
+			}
+			keys = append(keys, k)
+		}
+		c.keys = keys
+	}
 	c.mu.Unlock()
 }
 
@@ -196,7 +209,7 @@ func minTTL(m *dns.Msg) time.Duration {
 	return time.Duration(ttl) * time.Second
 }
 
-func isCacheable(m *dns.Msg) bool {
+func canCache(m *dns.Msg) bool {
 	if minTTL(m) == 0 {
 		return false
 	}
diff --git a/cache/cache_test.go b/cache/cache_test.go
index cec3ec8..df05721 100644
--- a/cache/cache_test.go
+++ b/cache/cache_test.go
@@ -95,7 +95,7 @@ func TestCache(t *testing.T) {
 	}{
 		{msg, createdAt, createdAt, true, &Value{CreatedAt: createdAt, msg: msg}},                       // Not expired when query time == create time
 		{msg, createdAt, createdAt.Add(30 * time.Second), true, &Value{CreatedAt: createdAt, msg: msg}}, // Not expired when below TTL
-		{msg, createdAt, createdAt.Add(60 * time.Second), true, &Value{CreatedAt: createdAt, msg: msg}}, //,  Not expired until TTL exceeds
+		{msg, createdAt, createdAt.Add(60 * time.Second), true, &Value{CreatedAt: createdAt, msg: msg}}, // Not expired until TTL exceeds
 		{msg, createdAt, createdAt.Add(61 * time.Second), false, nil},                                   // Expired
 		{msgWithZeroTTL, createdAt, createdAt, false, nil},                                              // 0 TTL is not cached
 		{msgFailure, createdAt, createdAt, false, nil},                                                  // Non-cacheable rcode
@@ -115,7 +115,17 @@ func TestCache(t *testing.T) {
 			awaitExpiry(t, c, k)
 		}
 		if _, ok := c.values[k]; ok != tt.ok {
-			t.Errorf("#%d: Cache[%d] = %t, want %t", i, k, ok, tt.ok)
+			t.Errorf("#%d: values[%d] = %t, want %t", i, k, ok, tt.ok)
+		}
+		keyIdx := -1
+		for i, key := range c.keys {
+			if key == k {
+				keyIdx = i
+				break
+			}
+		}
+		if (keyIdx != -1) != tt.ok {
+			t.Errorf("#%d: keys[%d] = %d, should not exist", i, keyIdx, k)
 		}
 	}
 }