Change LRU to a generic type (#20744)

Signed-off-by: yah01 <yang.cen@zilliz.com>

Signed-off-by: yah01 <yang.cen@zilliz.com>

internal/storage/vector_chunk_manager.go

@@ -40,7 +40,7 @@ var (
 type VectorChunkManager struct {
 	cacheStorage  ChunkManager
 	vectorStorage ChunkManager
-	cache         *cache.LRU
+	cache         *cache.LRU[string, *mmap.ReaderAt]
 
 	insertCodec *InsertCodec
 
@@ -68,14 +68,13 @@ func NewVectorChunkManager(ctx context.Context, cacheStorage ChunkManager, vecto
 		if cacheLimit <= 0 {
 			return nil, errors.New("cache limit must be positive if cacheEnable")
 		}
-		c, err := cache.NewLRU(defaultLocalCacheSize, func(k cache.Key, v cache.Value) {
+		c, err := cache.NewLRU(defaultLocalCacheSize, func(k string, v *mmap.ReaderAt) {
-			r := v.(*mmap.ReaderAt)
+			size := v.Len()
-			size := r.Len()
+			err := v.Close()
-			err := r.Close()
 			if err != nil {
 				log.Error("Unmmap file failed", zap.Any("file", k))
 			}
-			err = cacheStorage.Remove(ctx, k.(string))
+			err = cacheStorage.Remove(ctx, k)
 			if err != nil {
 				log.Error("cache storage remove file failed", zap.Any("file", k))
 			}
@@ -191,9 +190,8 @@ func (vcm *VectorChunkManager) readWithCache(ctx context.Context, filePath strin
 func (vcm *VectorChunkManager) Read(ctx context.Context, filePath string) ([]byte, error) {
 	if vcm.cacheEnable {
 		if r, ok := vcm.cache.Get(filePath); ok {
-			at := r.(*mmap.ReaderAt)
+			p := make([]byte, r.Len())
-			p := make([]byte, at.Len())
+			_, err := r.ReadAt(p, 0)
-			_, err := at.ReadAt(p, 0)
 			if err != nil {
 				return p, err
 			}
@@ -241,7 +239,7 @@ func (vcm *VectorChunkManager) ListWithPrefix(ctx context.Context, prefix string
 func (vcm *VectorChunkManager) Mmap(ctx context.Context, filePath string) (*mmap.ReaderAt, error) {
 	if vcm.cacheEnable && vcm.cache != nil {
 		if r, ok := vcm.cache.Get(filePath); ok {
-			return r.(*mmap.ReaderAt), nil
+			return r, nil
 		}
 	}
 	return nil, errors.New("the file mmap has not been cached")
@@ -255,9 +253,8 @@ func (vcm *VectorChunkManager) Reader(ctx context.Context, filePath string) (Fil
 func (vcm *VectorChunkManager) ReadAt(ctx context.Context, filePath string, off int64, length int64) ([]byte, error) {
 	if vcm.cacheEnable {
 		if r, ok := vcm.cache.Get(filePath); ok {
-			at := r.(*mmap.ReaderAt)
 			p := make([]byte, length)
-			_, err := at.ReadAt(p, off)
+			_, err := r.ReadAt(p, off)
 			if err != nil {
 				return nil, err
 			}

internal/util/cache/lru_cache.go

@@ -24,15 +24,15 @@ import (
 	"sync"
 )
 
-type LRU struct {
+type LRU[K comparable, V any] struct {
 	ctx       context.Context
 	cancel    context.CancelFunc
 	evictList *list.List
-	items     map[interface{}]*list.Element
+	items     map[K]*list.Element
 	capacity  int
-	onEvicted func(k Key, v Value)
+	onEvicted func(k K, v V)
 	m         sync.RWMutex
-	evictedCh chan *entry
+	evictedCh chan *entry[K, V]
 	stats     *Stats
 }
 
@@ -54,36 +54,31 @@ func (s *Stats) String() string {
 	return fmt.Sprintf("lru cache hit ratio = %f, evictedRatio = %f", hitRatio, evictedRatio)
 }
 
-type Key interface {
+type entry[K comparable, V any] struct {
-}
+	key   K
-type Value interface {
+	value V
 }
 
-type entry struct {
+func NewLRU[K comparable, V any](capacity int, onEvicted func(k K, v V)) (*LRU[K, V], error) {
-	key   Key
-	value Value
-}
-
-func NewLRU(capacity int, onEvicted func(k Key, v Value)) (*LRU, error) {
 	if capacity <= 0 {
 		return nil, errors.New("cache size must be positive")
 	}
 	ctx, cancel := context.WithCancel(context.Background())
-	c := &LRU{
+	c := &LRU[K, V]{
 		ctx:       ctx,
 		cancel:    cancel,
 		capacity:  capacity,
 		evictList: list.New(),
-		items:     make(map[interface{}]*list.Element),
+		items:     make(map[K]*list.Element),
 		onEvicted: onEvicted,
-		evictedCh: make(chan *entry, 16),
+		evictedCh: make(chan *entry[K, V], 16),
 		stats:     &Stats{},
 	}
 	go c.evictedWorker()
 	return c, nil
 }
 
-func (c *LRU) evictedWorker() {
+func (c *LRU[K, V]) evictedWorker() {
 	for {
 		select {
 		case <-c.ctx.Done():
@@ -98,16 +93,16 @@ func (c *LRU) evictedWorker() {
 	}
 }
 
-func (c *LRU) Add(key, value Value) {
+func (c *LRU[K, V]) Add(key K, value V) {
 	c.m.Lock()
 	defer c.m.Unlock()
 	c.stats.writeCount++
 	if e, ok := c.items[key]; ok {
 		c.evictList.MoveToFront(e)
-		e.Value.(*entry).value = value
+		e.Value.(*entry[K, V]).value = value
 		return
 	}
-	e := &entry{key: key, value: value}
+	e := &entry[K, V]{key: key, value: value}
 	listE := c.evictList.PushFront(e)
 	c.items[key] = listE
 
@@ -116,7 +111,7 @@ func (c *LRU) Add(key, value Value) {
 		oldestE := c.evictList.Back()
 		if oldestE != nil {
 			c.evictList.Remove(oldestE)
-			kv := oldestE.Value.(*entry)
+			kv := oldestE.Value.(*entry[K, V])
 			delete(c.items, kv.key)
 			if c.onEvicted != nil {
 				c.evictedCh <- kv
@@ -125,25 +120,27 @@ func (c *LRU) Add(key, value Value) {
 	}
 }
 
-func (c *LRU) Get(key Key) (value Value, ok bool) {
+func (c *LRU[K, V]) Get(key K) (value V, ok bool) {
 	c.m.RLock()
 	defer c.m.RUnlock()
 	c.stats.readCount++
 	if e, ok := c.items[key]; ok {
 		c.stats.hitCount++
 		c.evictList.MoveToFront(e)
-		kv := e.Value.(*entry)
+		kv := e.Value.(*entry[K, V])
 		return kv.value, true
 	}
-	return nil, false
+
+	var zeroV V
+	return zeroV, false
 }
 
-func (c *LRU) Remove(key Key) {
+func (c *LRU[K, V]) Remove(key K) {
 	c.m.Lock()
 	defer c.m.Unlock()
 	if e, ok := c.items[key]; ok {
 		c.evictList.Remove(e)
-		kv := e.Value.(*entry)
+		kv := e.Value.(*entry[K, V])
 		delete(c.items, kv.key)
 		if c.onEvicted != nil {
 			c.evictedCh <- kv
@@ -151,48 +148,48 @@ func (c *LRU) Remove(key Key) {
 	}
 }
 
-func (c *LRU) Contains(key Key) bool {
+func (c *LRU[K, V]) Contains(key K) bool {
 	c.m.RLock()
 	defer c.m.RUnlock()
 	_, ok := c.items[key]
 	return ok
 }
 
-func (c *LRU) Keys() []Key {
+func (c *LRU[K, V]) Keys() []K {
 	c.m.RLock()
 	defer c.m.RUnlock()
-	keys := make([]Key, len(c.items))
+	keys := make([]K, len(c.items))
 	i := 0
 	for ent := c.evictList.Back(); ent != nil; ent = ent.Prev() {
-		keys[i] = ent.Value.(*entry).key
+		keys[i] = ent.Value.(*entry[K, V]).key
 		i++
 	}
 	return keys
 }
 
-func (c *LRU) Len() int {
+func (c *LRU[K, V]) Len() int {
 	c.m.RLock()
 	defer c.m.RUnlock()
 	return c.evictList.Len()
 }
 
-func (c *LRU) Capacity() int {
+func (c *LRU[K, V]) Capacity() int {
 	return c.capacity
 }
 
-func (c *LRU) Purge() {
+func (c *LRU[K, V]) Purge() {
 	c.m.Lock()
 	defer c.m.Unlock()
 	for k, v := range c.items {
 		if c.onEvicted != nil {
-			c.evictedCh <- v.Value.(*entry)
+			c.evictedCh <- v.Value.(*entry[K, V])
 		}
 		delete(c.items, k)
 	}
 	c.evictList.Init()
 }
 
-func (c *LRU) Resize(capacity int) int {
+func (c *LRU[K, V]) Resize(capacity int) int {
 	c.m.Lock()
 	defer c.m.Unlock()
 	c.capacity = capacity
@@ -204,7 +201,7 @@ func (c *LRU) Resize(capacity int) int {
 		oldestE := c.evictList.Back()
 		if oldestE != nil {
 			c.evictList.Remove(oldestE)
-			kv := oldestE.Value.(*entry)
+			kv := oldestE.Value.(*entry[K, V])
 			delete(c.items, kv.key)
 			if c.onEvicted != nil {
 				c.evictedCh <- kv
@@ -214,18 +211,23 @@ func (c *LRU) Resize(capacity int) int {
 	return diff
 }
 
-func (c *LRU) GetOldest() (Key, Value, bool) {
+func (c *LRU[K, V]) GetOldest() (K, V, bool) {
 	c.m.RLock()
 	defer c.m.RUnlock()
 	ent := c.evictList.Back()
 	if ent != nil {
-		kv := ent.Value.(*entry)
+		kv := ent.Value.(*entry[K, V])
 		return kv.key, kv.value, true
 	}
-	return nil, nil, false
+
+	var (
+		zeroK K
+		zeroV V
+	)
+	return zeroK, zeroV, false
 }
 
-func (c *LRU) Close() {
+func (c *LRU[K, V]) Close() {
 	c.Purge()
 	c.cancel()
 	remain := len(c.evictedCh)
@@ -238,6 +240,6 @@ func (c *LRU) Close() {
 	close(c.evictedCh)
 }
 
-func (c *LRU) Stats() *Stats {
+func (c *LRU[K, V]) Stats() *Stats {
 	return c.stats
 }

internal/util/cache/lru_cache_test.go

@@ -25,22 +25,22 @@ import (
 )
 
 func TestNewLRU(t *testing.T) {
-	c, err := NewLRU(1, nil)
+	c, err := NewLRU[int, int](1, nil)
 	assert.Nil(t, err)
 	assert.NotNil(t, c)
 
-	c, err = NewLRU(0, nil)
+	c, err = NewLRU[int, int](0, nil)
 	assert.NotNil(t, err)
 	assert.Nil(t, c)
 
-	c, err = NewLRU(-1, nil)
+	c, err = NewLRU[int, int](-1, nil)
 	assert.NotNil(t, err)
 	assert.Nil(t, c)
 }
 
 func TestLRU_Add(t *testing.T) {
 	evicted := int32(0)
-	c, err := NewLRU(2, func(Key, Value) { atomic.AddInt32(&evicted, 1) })
+	c, err := NewLRU(2, func(string, string) { atomic.AddInt32(&evicted, 1) })
 	assert.Nil(t, err)
 
 	testKey1 := "test_key_1"
@@ -77,7 +77,7 @@ func TestLRU_Add(t *testing.T) {
 
 	v, ok = c.Get(testKey2)
 	assert.False(t, ok)
-	assert.Nil(t, v)
+	assert.Empty(t, v)
 
 	assert.Eventually(t, func() bool {
 		return atomic.LoadInt32(&evicted) == 1
@@ -86,7 +86,7 @@ func TestLRU_Add(t *testing.T) {
 
 func TestLRU_Contains(t *testing.T) {
 	evicted := int32(0)
-	c, err := NewLRU(1, func(Key, Value) { atomic.AddInt32(&evicted, 1) })
+	c, err := NewLRU(1, func(string, string) { atomic.AddInt32(&evicted, 1) })
 	assert.Nil(t, err)
 
 	testKey1 := "test_key_1"
@@ -112,7 +112,7 @@ func TestLRU_Contains(t *testing.T) {
 
 func TestLRU_Get(t *testing.T) {
 	evicted := int32(0)
-	c, err := NewLRU(1, func(Key, Value) { atomic.AddInt32(&evicted, 1) })
+	c, err := NewLRU(1, func(string, string) { atomic.AddInt32(&evicted, 1) })
 	assert.Nil(t, err)
 
 	testKey1 := "test_key_1"
@@ -132,7 +132,7 @@ func TestLRU_Get(t *testing.T) {
 
 	v, ok = c.Get(testKey1)
 	assert.False(t, ok)
-	assert.Nil(t, v)
+	assert.Empty(t, v)
 
 	assert.Eventually(t, func() bool {
 		return atomic.LoadInt32(&evicted) == 1
@@ -141,7 +141,7 @@ func TestLRU_Get(t *testing.T) {
 
 func TestLRU_GetOldest(t *testing.T) {
 	evicted := int32(0)
-	c, err := NewLRU(2, func(Key, Value) { atomic.AddInt32(&evicted, 1) })
+	c, err := NewLRU(2, func(string, string) { atomic.AddInt32(&evicted, 1) })
 	assert.Nil(t, err)
 
 	testKey1 := "test_key_1"
@@ -153,8 +153,8 @@ func TestLRU_GetOldest(t *testing.T) {
 
 	k, v, ok := c.GetOldest()
 	assert.False(t, ok)
-	assert.Nil(t, k)
+	assert.Empty(t, k)
-	assert.Nil(t, v)
+	assert.Empty(t, v)
 
 	c.Add(testKey1, testValue1)
 	k, v, ok = c.GetOldest()
@@ -190,7 +190,7 @@ func TestLRU_GetOldest(t *testing.T) {
 
 func TestLRU_Keys(t *testing.T) {
 	evicted := int32(0)
-	c, err := NewLRU(2, func(Key, Value) { atomic.AddInt32(&evicted, 1) })
+	c, err := NewLRU(2, func(string, string) { atomic.AddInt32(&evicted, 1) })
 	assert.Nil(t, err)
 
 	testKey1 := "test_key_1"
@@ -222,7 +222,7 @@ func TestLRU_Keys(t *testing.T) {
 }
 
 func TestLRU_Len(t *testing.T) {
-	c, err := NewLRU(2, nil)
+	c, err := NewLRU[string, string](2, nil)
 	assert.Nil(t, err)
 	assert.EqualValues(t, c.Len(), 0)
 
@@ -242,7 +242,7 @@ func TestLRU_Len(t *testing.T) {
 }
 
 func TestLRU_Capacity(t *testing.T) {
-	c, err := NewLRU(5, nil)
+	c, err := NewLRU[string, string](5, nil)
 	assert.Nil(t, err)
 	assert.EqualValues(t, c.Len(), 0)
 
@@ -264,7 +264,7 @@ func TestLRU_Capacity(t *testing.T) {
 
 func TestLRU_Purge(t *testing.T) {
 	evicted := int32(0)
-	c, err := NewLRU(2, func(Key, Value) { atomic.AddInt32(&evicted, 1) })
+	c, err := NewLRU(2, func(string, string) { atomic.AddInt32(&evicted, 1) })
 	assert.Nil(t, err)
 	assert.EqualValues(t, c.Len(), 0)
 
@@ -291,7 +291,7 @@ func TestLRU_Purge(t *testing.T) {
 
 func TestLRU_Remove(t *testing.T) {
 	evicted := int32(0)
-	c, err := NewLRU(2, func(Key, Value) { atomic.AddInt32(&evicted, 1) })
+	c, err := NewLRU(2, func(string, string) { atomic.AddInt32(&evicted, 1) })
 	assert.Nil(t, err)
 	assert.EqualValues(t, c.Len(), 0)
 
@@ -315,7 +315,7 @@ func TestLRU_Remove(t *testing.T) {
 
 func TestLRU_RemoveOldest(t *testing.T) {
 	evicted := int32(0)
-	c, err := NewLRU(2, func(Key, Value) { atomic.AddInt32(&evicted, 1) })
+	c, err := NewLRU(2, func(string, string) { atomic.AddInt32(&evicted, 1) })
 	assert.Nil(t, err)
 	assert.EqualValues(t, c.Len(), 0)
 
@@ -341,11 +341,11 @@ func TestLRU_RemoveOldest(t *testing.T) {
 
 	v, ok = c.Get(testKey1)
 	assert.False(t, ok)
-	assert.Nil(t, v)
+	assert.Empty(t, v)
 
 	v, ok = c.Get(testKey2)
 	assert.False(t, ok)
-	assert.Nil(t, v)
+	assert.Empty(t, v)
 
 	assert.EqualValues(t, c.Len(), 0)
 	assert.Eventually(t, func() bool {
@@ -356,7 +356,7 @@ func TestLRU_RemoveOldest(t *testing.T) {
 
 func TestLRU_Resize(t *testing.T) {
 	evicted := int32(0)
-	c, err := NewLRU(2, func(Key, Value) { atomic.AddInt32(&evicted, 1) })
+	c, err := NewLRU(2, func(string, string) { atomic.AddInt32(&evicted, 1) })
 	assert.Nil(t, err)
 	assert.EqualValues(t, c.Len(), 0)
 
@@ -381,7 +381,7 @@ func TestLRU_Resize(t *testing.T) {
 
 	v, ok = c.Get(testKey1)
 	assert.False(t, ok)
-	assert.Nil(t, v)
+	assert.Empty(t, v)
 
 	v, ok = c.Get(testKey2)
 	assert.True(t, ok)