feat: improve lru.GetOrPrepare

lru.go

@@ -3,6 +3,7 @@ package rueidis
 import (
 	"container/list"
 	"sync"
+	"sync/atomic"
 	"time"
 	"unsafe"
 )
@@ -13,6 +14,8 @@ const (
 	stringSSize = int(unsafe.Sizeof(""))
 
 	entryMinSize = entrySize + elementSize + stringSSize*2 + messageStructSize
+
+	moveThreshold = uint64(1024 - 1)
 )
 
 type cache interface {
@@ -38,12 +41,14 @@ func (e *entry) Wait() RedisMessage {
 type keyCache struct {
 	cache map[string]*list.Element
 	ttl   time.Time
+	hits  uint64
+	miss  uint64
 }
 
 var _ cache = (*lru)(nil)
 
 type lru struct {
-	mu sync.Mutex
+	mu sync.RWMutex
 
 	store map[string]*keyCache
 	list  *list.List
@@ -61,23 +66,57 @@ func newLRU(max int) *lru {
 }
 
 func (c *lru) GetOrPrepare(key, cmd string, ttl time.Duration) (v RedisMessage, e *entry) {
+	var ok bool
+	var store *keyCache
+	var now = time.Now()
+	var storeTTL time.Time
+	var ele, back *list.Element
+
+	c.mu.RLock()
+	if store, ok = c.store[key]; ok {
+		storeTTL = store.ttl
+		if ele, ok = store.cache[cmd]; ok {
+			e = ele.Value.(*entry)
+			v = e.val
+			back = c.list.Back()
+		}
+	}
+	c.mu.RUnlock()
+
+	if e != nil && (v.typ == 0 || storeTTL.After(now)) {
+		hits := atomic.AddUint64(&store.hits, 1)
+		if ele != back && hits&moveThreshold == 0 {
+			c.mu.Lock()
+			c.list.MoveToBack(ele)
+			c.mu.Unlock()
+		}
+		return v, e
+	}
+
+	v = RedisMessage{}
+	e = nil
+
 	c.mu.Lock()
-	store, ok := c.store[key]
-	if !ok {
-		store = &keyCache{cache: make(map[string]*list.Element), ttl: time.Now().Add(ttl)}
-		c.store[key] = store
+	if store == nil {
+		if store, ok = c.store[key]; !ok {
+			store = &keyCache{cache: make(map[string]*list.Element), ttl: now.Add(ttl)}
+			c.store[key] = store
+		}
 	}
-	if ele, ok := store.cache[cmd]; ok {
-		if e = ele.Value.(*entry); e.val.typ == 0 || store.ttl.After(time.Now()) {
+	if ele, ok = store.cache[cmd]; ok {
+		if e = ele.Value.(*entry); e.val.typ == 0 || store.ttl.After(now) {
+			atomic.AddUint64(&store.hits, 1)
 			v = e.val
 			c.list.MoveToBack(ele)
 		} else {
-			e = nil
 			c.list.Remove(ele)
-			store.ttl = time.Now().Add(ttl)
+			c.size -= e.size
+			store.ttl = now.Add(ttl)
+			e = nil
 		}
 	}
 	if e == nil && c.list != nil {
+		atomic.AddUint64(&store.miss, 1)
 		c.list.PushBack(&entry{
 			key: key,
 			cmd: cmd,

lru_test.go

@@ -2,6 +2,8 @@ package rueidis
 
 import (
 	"strconv"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 )
@@ -17,7 +19,7 @@ func TestLRU(t *testing.T) {
 		if v, entry := lru.GetOrPrepare("0", "GET", TTL); v.typ != 0 || entry != nil {
 			t.Fatalf("got unexpected value from the first GetOrPrepare: %v %v", v, entry)
 		}
-		lru.Update("0", "GET", RedisMessage{typ: '+', string: "0"}, PTTL)
+		lru.Update("0", "GET", RedisMessage{typ: '+', string: "0", values: []RedisMessage{{}}}, PTTL)
 		return lru
 	}
 
@@ -42,10 +44,38 @@ func TestLRU(t *testing.T) {
 		}
 	})
 
-	t.Run("Cache Miss", func(t *testing.T) {
+	t.Run("Cache Miss Suppress", func(t *testing.T) {
+		count := 5000
 		lru := setup(t)
-		if v, _ := lru.GetOrPrepare("1", "GET", TTL); v.typ != 0 {
-			t.Fatalf("got unexpected value from the first GetOrPrepare: %v", v)
+		wg := sync.WaitGroup{}
+		wg.Add(count)
+		for i := 0; i < count; i++ {
+			go func() {
+				defer wg.Done()
+				if v, _ := lru.GetOrPrepare("1", "GET", TTL); v.typ != 0 {
+					t.Errorf("got unexpected value from the first GetOrPrepare: %v", v)
+				}
+				if v, _ := lru.GetOrPrepare("2", "GET", TTL); v.typ != 0 {
+					t.Errorf("got unexpected value from the first GetOrPrepare: %v", v)
+				}
+			}()
+		}
+		wg.Wait()
+		lru.mu.RLock()
+		store1 := lru.store["1"]
+		store2 := lru.store["2"]
+		lru.mu.RUnlock()
+		if miss := atomic.LoadUint64(&store1.miss); miss != 1 {
+			t.Fatalf("unexpected miss count %v", miss)
+		}
+		if hits := atomic.LoadUint64(&store1.hits); hits != uint64(count-1) {
+			t.Fatalf("unexpected hits count %v", hits)
+		}
+		if miss := atomic.LoadUint64(&store2.miss); miss != 1 {
+			t.Fatalf("unexpected miss count %v", miss)
+		}
+		if hits := atomic.LoadUint64(&store2.hits); hits != uint64(count-1) {
+			t.Fatalf("unexpected hits count %v", hits)
 		}
 	})