milvus/internal/storage/lru_cache.go

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package storage

import (
	"container/list"
	"errors"
)

type LRU struct {
	evictList *list.List
	items     map[interface{}]*list.Element
	size      int
	onEvicted func(k, v interface{})
}

type entry struct {
	key   interface{}
	value interface{}
}

func NewLRU(size int, onEvicted func(k, v interface{})) (*LRU, error) {
	if size <= 0 {
		return nil, errors.New("cache size must be positive")
	}
	c := &LRU{
		size:      size,
		evictList: list.New(),
		items:     make(map[interface{}]*list.Element),
		onEvicted: onEvicted,
	}
	return c, nil
}

func (c *LRU) Add(key, value interface{}) {
	if e, ok := c.items[key]; ok {
		c.evictList.MoveToFront(e)
		e.Value.(*entry).value = value
		return
	}
	e := &entry{key: key, value: value}
	listE := c.evictList.PushFront(e)
	c.items[key] = listE

	if c.evictList.Len() > c.size {
		c.RemoveOldest()
	}
}

func (c *LRU) Get(key interface{}) (value interface{}, ok bool) {
	if e, ok := c.items[key]; ok {
		c.evictList.MoveToFront(e)
		kv := e.Value.(*entry)
		return kv.value, true
	}
	return nil, false
}

func (c *LRU) Remove(key interface{}) {
	if e, ok := c.items[key]; ok {
		c.evictList.Remove(e)
		kv := e.Value.(*entry)
		delete(c.items, kv.key)
		if c.onEvicted != nil {
			c.onEvicted(kv.key, kv.value)
		}
	}
}

func (c *LRU) Contains(key interface{}) bool {
	_, ok := c.items[key]
	return ok
}

// Peek get value but not update the recently-used list.
func (c *LRU) Peek(key interface{}) (value interface{}, ok bool) {
	if e, ok := c.items[key]; ok {
		kv := e.Value.(*entry)
		return kv.value, true
	}
	return nil, false
}

func (c *LRU) Keys() []interface{} {
	keys := make([]interface{}, len(c.items))
	i := 0
	for ent := c.evictList.Back(); ent != nil; ent = ent.Prev() {
		keys[i] = ent.Value.(*entry).key
		i++
	}
	return keys
}

func (c *LRU) Len() int {
	return c.evictList.Len()
}

func (c *LRU) Purge() {
	for k, v := range c.items {
		if c.onEvicted != nil {
			c.onEvicted(k, v.Value.(*entry).value)
		}
		delete(c.items, k)
	}
	c.evictList.Init()
}

func (c *LRU) Resize(size int) int {
	c.size = size
	if size >= c.evictList.Len() {
		return 0
	}
	diff := c.evictList.Len() - c.size
	for i := 0; i < diff; i++ {
		c.RemoveOldest()
	}
	return diff
}

func (c *LRU) GetOldest() (interface{}, interface{}, bool) {
	ent := c.evictList.Back()
	if ent != nil {
		kv := ent.Value.(*entry)
		return kv.key, kv.value, true
	}
	return nil, nil, false
}

func (c *LRU) RemoveOldest() {
	e := c.evictList.Back()
	if e != nil {
		c.evictList.Remove(e)
		kv := e.Value.(*entry)
		delete(c.items, kv.key)
		if c.onEvicted != nil {
			c.onEvicted(kv.key, kv.value)
		}
	}
}