fix: 修复memory cache并发问题 (#889)

cache/memory.go

@@ -5,53 +5,64 @@ import (
 	"time"
 )
 
-// Memory struct contains *memcache.Client
+// Memory is an concurrent-safe in-process cache backed by a plain map.
 type Memory struct {
-	sync.Mutex
-
+	mu   sync.RWMutex
 	data map[string]*data
 }
 
+// data holds a single cached value together with its expiry time.
 type data struct {
 	Data    interface{}
 	Expired time.Time
 }
 
-// NewMemory create new memcache
+// NewMemory returns a ready-to-use in-memory cache.
 func NewMemory() *Memory {
 	return &Memory{
 		data: map[string]*data{},
 	}
 }
 
-// Get return cached value
+// Get returns the cached value for key.
+// Returns nil if the key does not exist or has expired.
 func (mem *Memory) Get(key string) interface{} {
-	if ret, ok := mem.data[key]; ok {
-		if ret.Expired.Before(time.Now()) {
-			mem.deleteKey(key)
-			return nil
-		}
-		return ret.Data
+	mem.mu.RLock()
+	ret, ok := mem.data[key]
+	expired := ok && ret.Expired.Before(time.Now())
+	mem.mu.RUnlock()
+
+	if !ok {
+		return nil
 	}
-	return nil
+	if expired {
+		mem.deleteExpiredKey(key)
+		return nil
+	}
+	return ret.Data
 }
 
 // IsExist check value exists in memcache.
 func (mem *Memory) IsExist(key string) bool {
-	if ret, ok := mem.data[key]; ok {
-		if ret.Expired.Before(time.Now()) {
-			mem.deleteKey(key)
-			return false
-		}
-		return true
+	mem.mu.RLock()
+	ret, ok := mem.data[key]
+	expired := ok && ret.Expired.Before(time.Now())
+	mem.mu.RUnlock()
+
+	if !ok {
+		return false
 	}
-	return false
+	if expired {
+		mem.deleteExpiredKey(key)
+		return false
+	}
+	return true
 }
 
 // Set cached value with key and expire time.
 func (mem *Memory) Set(key string, val interface{}, timeout time.Duration) (err error) {
-	mem.Lock()
-	defer mem.Unlock()
+	mem.mu.Lock()
+	defer mem.mu.Unlock()
 
 	mem.data[key] = &data{
 		Data:    val,
@@ -62,13 +73,18 @@ func (mem *Memory) Set(key string, val interface{}, timeout time.Duration) (err
 
 // Delete delete value in memcache.
 func (mem *Memory) Delete(key string) error {
-	mem.deleteKey(key)
+	mem.mu.Lock()
+	defer mem.mu.Unlock()
+	delete(mem.data, key)
 	return nil
 }
 
-// deleteKey
-func (mem *Memory) deleteKey(key string) {
-	mem.Lock()
-	defer mem.Unlock()
-	delete(mem.data, key)
+// deleteExpiredKey deletes a key only if it has expired.
+// Caller must NOT hold any lock.
+func (mem *Memory) deleteExpiredKey(key string) {
+	mem.mu.Lock()
+	defer mem.mu.Unlock()
+	if d, ok := mem.data[key]; ok && d.Expired.Before(time.Now()) {
+		delete(mem.data, key)
+	}
 }

cache/memory_test.go

@@ -0,0 +1,116 @@
+// 运行测试：go test -race -v ./cache/ -run "TestMemory" -count=1
+package cache
+
+import (
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMemoryGet(t *testing.T) {
+	mem := NewMemory()
+
+	err := mem.Set("username", "silenceper", 10*time.Second)
+	assert.NoError(t, err)
+
+	val := mem.Get("username")
+	assert.Equal(t, "silenceper", val)
+
+	val = mem.Get("unknown-key")
+	assert.Nil(t, val)
+}
+
+func TestMemoryIsExist(t *testing.T) {
+	mem := NewMemory()
+
+	err := mem.Set("username", "silenceper", 10*time.Second)
+	assert.NoError(t, err)
+
+	assert.True(t, mem.IsExist("username"))
+	assert.False(t, mem.IsExist("unknown-key"))
+}
+
+func TestMemoryDelete(t *testing.T) {
+	mem := NewMemory()
+
+	err := mem.Set("username", "silenceper", 10*time.Second)
+	assert.NoError(t, err)
+
+	err = mem.Delete("username")
+	assert.NoError(t, err)
+
+	// delete 不存在的 key 不应报错
+	err = mem.Delete("unknown-key")
+	assert.NoError(t, err)
+}
+
+func TestMemoryExpire(t *testing.T) {
+	mem := NewMemory()
+
+	err := mem.Set("username", "silenceper", 10*time.Millisecond)
+	assert.NoError(t, err)
+
+	assert.True(t, mem.IsExist("username"))
+
+	time.Sleep(20 * time.Millisecond)
+
+	assert.False(t, mem.IsExist("username"))
+	assert.Nil(t, mem.Get("username"))
+}
+
+// TestMemoryConcurrentOps 验证 Get/IsExist/Set/Delete 并发执行不产生数据竞争
+func TestMemoryConcurrentOps(t *testing.T) {
+	mem := NewMemory()
+	_ = mem.Set("key", "value", time.Minute)
+
+	var wg sync.WaitGroup
+	for i := 0; i < 100; i++ {
+		wg.Add(4)
+		// 并发读
+		go func() {
+			defer wg.Done()
+			_ = mem.Get("key")
+		}()
+		// 并发 IsExist
+		go func() {
+			defer wg.Done()
+			_ = mem.IsExist("key")
+		}()
+		// 并发写
+		go func(i int) {
+			defer wg.Done()
+			_ = mem.Set("key", i, time.Minute)
+		}(i)
+		// 并发删除
+		go func() {
+			defer wg.Done()
+			_ = mem.Delete("key")
+		}()
+	}
+	wg.Wait()
+}
+
+// TestMemoryConcurrentExpireAndRead 验证过期惰性删除在并发场景下不会死锁
+func TestMemoryConcurrentExpireAndRead(t *testing.T) {
+	mem := NewMemory()
+
+	var wg sync.WaitGroup
+	for i := 0; i < 50; i++ {
+		wg.Add(2)
+		go func(i int) {
+			defer wg.Done()
+			key := "expire-key"
+			_ = mem.Set(key, i, 5*time.Millisecond)
+		}(i)
+		go func() {
+			defer wg.Done()
+			time.Sleep(3 * time.Millisecond)
+			// 此时 key 可能已过期，触发 deleteKey，验证不死锁
+			_ = mem.Get("expire-key")
+			_ = mem.IsExist("expire-key")
+		}()
+	}
+	wg.Wait()
+}