Add support for resetting cache

2 files changed, 23 insertions, 3 deletions

diff --git a/cache/cache.go b/cache/cache.go
index 087ba24..b6c03dc 100644
--- a/cache/cache.go
+++ b/cache/cache.go
@@ -124,6 +124,7 @@ func (c *Cache) getValue(k uint64) (*Value, bool) {
 func (c *Cache) List(n int) []Value {
 	values := make([]Value, 0, n)
 	c.mu.RLock()
+	defer c.mu.RUnlock()
 	for i := len(c.keys) - 1; i >= 0; i-- {
 		if len(values) == n {
 			break
@@ -134,7 +135,6 @@ func (c *Cache) List(n int) []Value {
 		}
 		values = append(values, *v)
 	}
-	c.mu.RUnlock()
 	return values
 }
 
@@ -149,17 +149,26 @@ func (c *Cache) Set(k uint64, msg *dns.Msg) {
 	}
 	now := c.now()
 	c.mu.Lock()
+	defer c.mu.Unlock()
 	if len(c.values) == c.capacity && c.capacity > 0 {
 		delete(c.values, c.keys[0])
 		c.keys = c.keys[1:]
 	}
 	c.values[k] = &Value{CreatedAt: now, msg: msg}
 	c.keys = append(c.keys, k)
-	c.mu.Unlock()
+}
+
+// Reset removes all values contained in cache c.
+func (c *Cache) Reset() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.values = make(map[uint64]*Value)
+	c.keys = nil
 }
 
 func (c *Cache) evictExpired() {
 	c.mu.Lock()
+	defer c.mu.Unlock()
 	evictedKeys := make(map[uint64]bool)
 	for k, v := range c.values {
 		if c.isExpired(v) {
@@ -178,7 +187,6 @@ func (c *Cache) evictExpired() {
 		}
 		c.keys = keys
 	}
-	c.mu.Unlock()
 }
 
 func (c *Cache) isExpired(v *Value) bool {
diff --git a/cache/cache_test.go b/cache/cache_test.go
index 77f9f15..33f3a6b 100644
--- a/cache/cache_test.go
+++ b/cache/cache_test.go
@@ -218,6 +218,18 @@ func TestCacheList(t *testing.T) {
 	}
 }
 
+func TestReset(t *testing.T) {
+	c := New(10)
+	c.Set(uint64(1), &dns.Msg{})
+	c.Reset()
+	if got, want := len(c.values), 0; got != want {
+		t.Errorf("len(values) = %d, want %d", got, want)
+	}
+	if got, want := len(c.keys), 0; got != want {
+		t.Errorf("len(keys) = %d, want %d", got, want)
+	}
+}
+
 func BenchmarkNewKey(b *testing.B) {
 	for n := 0; n < b.N; n++ {
 		NewKey("key", 1, 1)