From dba8c2c043fa9fcf2140b5c49e579c73d262cc3e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Mon, 18 Nov 2024 12:36:15 +0800
Subject: [PATCH] Fix clear lru cache

---
 common/cache/lrucache.go      |  10 ++-
 common/cache/lrucache_test.go | 114 ++++++++++++++++++++++++++++++++++
 2 files changed, 122 insertions(+), 2 deletions(-)
 create mode 100644 common/cache/lrucache_test.go

diff --git a/common/cache/lrucache.go b/common/cache/lrucache.go
index a8912e8e..4b99c54c 100644
--- a/common/cache/lrucache.go
+++ b/common/cache/lrucache.go
@@ -261,9 +261,15 @@ func (c *LruCache[K, V]) Delete(key K) {
 func (c *LruCache[K, V]) Clear() {
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	for element := c.lru.Front(); element != nil; element = element.Next() {
-		c.deleteElement(element)
+
+	if c.onEvict != nil {
+		for le := c.lru.Front(); le != nil; le = le.Next() {
+			c.onEvict(le.Value.key, le.Value.value)
+		}
 	}
+
+	c.lru.Init()
+	c.cache = make(map[K]*list.Element[*entry[K, V]])
 }
 
 func (c *LruCache[K, V]) maybeDeleteOldest() {
diff --git a/common/cache/lrucache_test.go b/common/cache/lrucache_test.go
new file mode 100644
index 00000000..330dc140
--- /dev/null
+++ b/common/cache/lrucache_test.go
@@ -0,0 +1,114 @@
+package cache_test
+
+import (
+	"testing"
+	"time"
+
+	"github.com/sagernet/sing/common/cache"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestLRUCache(t *testing.T) {
+	t.Parallel()
+	t.Run("basic operations", func(t *testing.T) {
+		t.Parallel()
+		c := cache.New[string, int]()
+
+		c.Store("key1", 1)
+		value, exists := c.Load("key1")
+		require.True(t, exists)
+		require.Equal(t, 1, value)
+
+		value, exists = c.Load("missing")
+		require.False(t, exists)
+		require.Zero(t, value)
+
+		c.Delete("key1")
+		_, exists = c.Load("key1")
+		require.False(t, exists)
+	})
+
+	t.Run("max size", func(t *testing.T) {
+		t.Parallel()
+		c := cache.New[string, int](cache.WithSize[string, int](2))
+
+		c.Store("key1", 1)
+		c.Store("key2", 2)
+		c.Store("key3", 3)
+
+		_, exists := c.Load("key1")
+		require.False(t, exists)
+
+		value, exists := c.Load("key2")
+		require.True(t, exists)
+		require.Equal(t, 2, value)
+	})
+
+	t.Run("expiration", func(t *testing.T) {
+		t.Parallel()
+		c := cache.New[string, int](cache.WithAge[string, int](1))
+
+		c.Store("key1", 1)
+
+		value, exists := c.Load("key1")
+		require.True(t, exists)
+		require.Equal(t, 1, value)
+
+		time.Sleep(time.Second * 2)
+
+		value, exists = c.Load("key1")
+		require.False(t, exists)
+		require.Zero(t, value)
+	})
+
+	t.Run("clear", func(t *testing.T) {
+		t.Parallel()
+		evicted := make(map[string]int)
+		c := cache.New[string, int](
+			cache.WithEvict[string, int](func(key string, value int) {
+				evicted[key] = value
+			}),
+		)
+
+		c.Store("key1", 1)
+		c.Store("key2", 2)
+
+		c.Clear()
+
+		require.Equal(t, map[string]int{"key1": 1, "key2": 2}, evicted)
+		_, exists := c.Load("key1")
+		require.False(t, exists)
+	})
+
+	t.Run("load or store", func(t *testing.T) {
+		t.Parallel()
+		c := cache.New[string, int]()
+
+		value, loaded := c.LoadOrStore("key1", func() int { return 1 })
+		require.False(t, loaded)
+		require.Equal(t, 1, value)
+
+		value, loaded = c.LoadOrStore("key1", func() int { return 2 })
+		require.True(t, loaded)
+		require.Equal(t, 1, value)
+	})
+
+	t.Run("update age on get", func(t *testing.T) {
+		t.Parallel()
+		c := cache.New[string, int](
+			cache.WithAge[string, int](5),
+			cache.WithUpdateAgeOnGet[string, int](),
+		)
+
+		c.Store("key1", 1)
+
+		time.Sleep(time.Second * 3)
+		_, exists := c.Load("key1")
+		require.True(t, exists)
+
+		time.Sleep(time.Second * 3)
+		_, exists = c.Load("key1")
+		require.True(t, exists)
+	})
+}