diff --git a/internal/storage/vector_chunk_manager.go b/internal/storage/vector_chunk_manager.go index 90a5bb6925..d111210543 100644 --- a/internal/storage/vector_chunk_manager.go +++ b/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) { - r := v.(*mmap.ReaderAt) - size := r.Len() - err := r.Close() + c, err := cache.NewLRU(defaultLocalCacheSize, func(k string, v *mmap.ReaderAt) { + size := v.Len() + err := v.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, at.Len()) - _, err := at.ReadAt(p, 0) + p := make([]byte, r.Len()) + _, err := r.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 } diff --git a/internal/util/cache/lru_cache.go b/internal/util/cache/lru_cache.go index 9070c5ab93..5694f3e5c8 100644 --- a/internal/util/cache/lru_cache.go +++ b/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 Value interface { +type entry[K comparable, V any] struct { + key K + value V } -type entry struct { - key Key - value Value -} - -func NewLRU(capacity int, onEvicted func(k Key, v Value)) (*LRU, error) { +func NewLRU[K comparable, V any](capacity int, onEvicted func(k K, v V)) (*LRU[K, V], 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 } diff --git a/internal/util/cache/lru_cache_test.go b/internal/util/cache/lru_cache_test.go index 63a13e6324..700bf74df0 100644 --- a/internal/util/cache/lru_cache_test.go +++ b/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.Nil(t, v) + assert.Empty(t, k) + 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)