feature add lru cache evict (#30)

ecodeclub · Jan 11, 2024 · 7567983 · 7567983
1 parent 37fe769
commit 7567983
Show file tree

Hide file tree

Showing 6 changed files with 850 additions and 160 deletions.
diff --git a/go.mod b/go.mod
@@ -13,7 +13,6 @@ require (
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
-	github.com/hashicorp/golang-lru/v2 v2.0.6
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	golang.org/x/sync v0.5.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect

diff --git a/go.sum b/go.sum
@@ -9,8 +9,6 @@ github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/r
 github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
 github.com/ecodeclub/ekit v0.0.8-0.20230925161647-c5bfbd460261 h1:FunYsaj58DVk4iIBXeU8hwdbvlGS1hc7ZbWXOx/+Vj0=
 github.com/ecodeclub/ekit v0.0.8-0.20230925161647-c5bfbd460261/go.mod h1:OqTojKeKFTxeeAAUwNIPKu339SRkX6KAuoK/8A5BCEs=
-github.com/hashicorp/golang-lru/v2 v2.0.6 h1:3xi/Cafd1NaoEnS/yDssIiuVeDVywU0QdFGl3aQaQHM=
-github.com/hashicorp/golang-lru/v2 v2.0.6/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e h1:fD57ERR4JtEqsWbfPhv4DMiApHyliiK5xCTNVSPiaAs=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=

diff --git a/memory/lru/cache.go b/memory/lru/cache.go
@@ -27,53 +27,202 @@ import (
 
 	"github.com/ecodeclub/ecache"
 	"github.com/ecodeclub/ecache/internal/errs"
-	"github.com/hashicorp/golang-lru/v2/simplelru"
 )
 
 var (
 	_ ecache.Cache = (*Cache)(nil)
 )
 
+type entry struct {
+	key       string
+	value     any
+	expiresAt time.Time
+}
+
+func (e entry) isExpired() bool {
+	return !e.expiresAt.IsZero() && e.expiresAt.Before(time.Now())
+}
+
+type EvictCallback func(key string, value any)
+
+type Option func(l *Cache)
+
+func WithEvictCallback(callback func(k string, v any)) Option {
+	return func(l *Cache) {
+		l.callback = callback
+	}
+}
+
+func WithCycleInterval(interval time.Duration) Option {
+	return func(l *Cache) {
+		l.cycleInterval = interval
+	}
+}
+
 type Cache struct {
-	lock   sync.RWMutex
-	client simplelru.LRUCache[string, any]
+	lock          sync.RWMutex
+	capacity      int
+	list          *linkedList[entry]
+	data          map[string]*element[entry]
+	callback      EvictCallback
+	cycleInterval time.Duration
 }
 
-func NewCache(client simplelru.LRUCache[string, any]) *Cache {
-	return &Cache{
-		lock:   sync.RWMutex{},
-		client: client,
+func NewCache(capacity int, options ...Option) *Cache {
+	res := &Cache{
+		list:          newLinkedList[entry](),
+		data:          make(map[string]*element[entry], capacity),
+		capacity:      capacity,
+		cycleInterval: time.Second * 10,
+	}
+	for _, opt := range options {
+		opt(res)
 	}
+	res.cleanCycle()
+	return res
+}
+
+func (c *Cache) cleanCycle() {
+	go func() {
+		ticker := time.NewTicker(c.cycleInterval)
+		for range ticker.C {
+			cnt := 0
+			c.lock.Lock()
+			limit := c.list.len() / 3
+			for elem, i := c.list.back(), 0; i < c.list.len(); i++ {
+				if elem.Value.isExpired() {
+					c.removeElement(elem)
+				}
+				elem = elem.prev
+				cnt++
+				if cnt >= limit {
+					break
+				}
+			}
+			c.lock.Unlock()
+		}
+	}()
+}
+
+func (c *Cache) pushEntry(key string, ent entry) bool {
+	if len(c.data) >= c.capacity && c.len() >= c.capacity {
+		if elem, ok := c.data[key]; ok {
+			elem.Value = ent
+			c.list.moveToFront(elem)
+			return false
+		}
+		c.removeOldest()
+	}
+	if elem, ok := c.data[key]; ok {
+		elem.Value = ent
+		c.list.moveToFront(elem)
+		return false
+	}
+	elem := c.list.pushFront(ent)
+	c.data[key] = elem
+	return true
+}
+
+func (c *Cache) addTTL(key string, value any, expiration time.Duration) bool {
+	ent := entry{key: key, value: value,
+		expiresAt: time.Now().Add(expiration)}
+	return c.pushEntry(key, ent)
+}
+
+func (c *Cache) add(key string, value any) bool {
+	ent := entry{key: key, value: value}
+	return c.pushEntry(key, ent)
+}
+
+func (c *Cache) get(key string) (value any, ok bool) {
+	if elem, exist := c.data[key]; exist {
+		ent := elem.Value
+		if ent.isExpired() {
+			c.removeElement(elem)
+			return
+		}
+		c.list.moveToFront(elem)
+		return ent.value, true
+	}
+	return
+}
+
+func (c *Cache) removeOldest() {
+	if elem := c.list.back(); elem != nil {
+		c.removeElement(elem)
+	}
+}
+
+func (c *Cache) removeElement(elem *element[entry]) {
+	c.list.removeElem(elem)
+	ent := elem.Value
+	c.delete(ent.key)
+	if c.callback != nil {
+		c.callback(ent.key, ent.value)
+	}
+}
+
+func (c *Cache) remove(key string) bool {
+	if elem, ok := c.data[key]; ok {
+		c.removeElement(elem)
+		return !elem.Value.isExpired()
+	}
+	return false
+}
+
+func (c *Cache) contains(key string) (ok bool) {
+	elem, ok := c.data[key]
+	if ok {
+		if elem.Value.isExpired() {
+			c.removeElement(elem)
+			return false
+		}
+	}
+	return ok
+}
+
+func (c *Cache) delete(key string) {
+	delete(c.data, key)
+}
+
+func (c *Cache) len() int {
+	var length int
+	for elem, i := c.list.back(), 0; i < c.list.len(); i++ {
+		if elem.Value.isExpired() {
+			c.removeElement(elem)
+			continue
+		}
+		elem = elem.prev
+		length++
+	}
+	return length
 }
 
-// Set expiration 无效 由lru 统一控制过期时间
 func (c *Cache) Set(ctx context.Context, key string, val any, expiration time.Duration) error {
 	c.lock.Lock()
 	defer c.lock.Unlock()
-
-	c.client.Add(key, val)
+	c.addTTL(key, val, expiration)
 	return nil
 }
 
-// SetNX expiration 无效 由lru 统一控制过期时间
 func (c *Cache) SetNX(ctx context.Context, key string, val any, expiration time.Duration) (bool, error) {
 	c.lock.Lock()
 	defer c.lock.Unlock()
 
-	if c.client.Contains(key) {
+	if c.contains(key) {
 		return false, nil
 	}
 
-	c.client.Add(key, val)
+	c.addTTL(key, val, expiration)
 
 	return true, nil
 }
 
 func (c *Cache) Get(ctx context.Context, key string) (val ecache.Value) {
-	c.lock.RLock()
-	defer c.lock.RUnlock()
+	c.lock.Lock()
+	defer c.lock.Unlock()
 	var ok bool
-	val.Val, ok = c.client.Get(key)
+	val.Val, ok = c.get(key)
 	if !ok {
 		val.Err = errs.ErrKeyNotExist
 	}
@@ -86,12 +235,12 @@ func (c *Cache) GetSet(ctx context.Context, key string, val string) (result ecac
 	defer c.lock.Unlock()
 
 	var ok bool
-	result.Val, ok = c.client.Get(key)
+	result.Val, ok = c.get(key)
 	if !ok {
 		result.Err = errs.ErrKeyNotExist
 	}
 
-	c.client.Add(key, val)
+	c.add(key, val)
 
 	return
 }
@@ -105,11 +254,11 @@ func (c *Cache) Delete(ctx context.Context, key ...string) (int64, error) {
 		if ctx.Err() != nil {
 			return n, ctx.Err()
 		}
-		_, ok := c.client.Get(k)
+		_, ok := c.get(k)
 		if !ok {
 			continue
 		}
-		if c.client.Remove(k) {
+		if c.remove(k) {
 			n++
 		} else {
 			return n, fmt.Errorf("%w: key = %s", errs.ErrDeleteKeyFailed, k)
@@ -137,12 +286,12 @@ func (c *Cache) LPush(ctx context.Context, key string, val ...any) (int64, error
 		ok     bool
 		result = ecache.Value{}
 	)
-	result.Val, ok = c.client.Get(key)
+	result.Val, ok = c.get(key)
 	if !ok {
 		l := &list.ConcurrentList[ecache.Value]{
 			List: list.NewLinkedListOf[ecache.Value](c.anySliceToValueSlice(val...)),
 		}
-		c.client.Add(key, l)
+		c.add(key, l)
 		return int64(l.Len()), nil
 	}
 
@@ -156,7 +305,7 @@ func (c *Cache) LPush(ctx context.Context, key string, val ...any) (int64, error
 		return 0, err
 	}
 
-	c.client.Add(key, data)
+	c.add(key, data)
 	return int64(data.Len()), nil
 }
 
@@ -167,7 +316,7 @@ func (c *Cache) LPop(ctx context.Context, key string) (val ecache.Value) {
 	var (
 		ok bool
 	)
-	val.Val, ok = c.client.Get(key)
+	val.Val, ok = c.get(key)
 	if !ok {
 		val.Err = errs.ErrKeyNotExist
 		return
@@ -197,7 +346,7 @@ func (c *Cache) SAdd(ctx context.Context, key string, members ...any) (int64, er
 		ok     bool
 		result = ecache.Value{}
 	)
-	result.Val, ok = c.client.Get(key)
+	result.Val, ok = c.get(key)
 	if !ok {
 		result.Val = set.NewMapSet[any](8)
 	}
@@ -210,7 +359,7 @@ func (c *Cache) SAdd(ctx context.Context, key string, members ...any) (int64, er
 	for _, value := range members {
 		s.Add(value)
 	}
-	c.client.Add(key, s)
+	c.add(key, s)
 
 	return int64(len(s.Keys())), nil
 }
@@ -219,7 +368,7 @@ func (c *Cache) SRem(ctx context.Context, key string, members ...any) (int64, er
 	c.lock.Lock()
 	defer c.lock.Unlock()
 
-	result, ok := c.client.Get(key)
+	result, ok := c.get(key)
 	if !ok {
 		return 0, errs.ErrKeyNotExist
 	}
@@ -247,9 +396,9 @@ func (c *Cache) IncrBy(ctx context.Context, key string, value int64) (int64, err
 		ok     bool
 		result = ecache.Value{}
 	)
-	result.Val, ok = c.client.Get(key)
+	result.Val, ok = c.get(key)
 	if !ok {
-		c.client.Add(key, value)
+		c.add(key, value)
 		return value, nil
 	}
 
@@ -259,7 +408,7 @@ func (c *Cache) IncrBy(ctx context.Context, key string, value int64) (int64, err
 	}
 
 	newVal := incr + value
-	c.client.Add(key, newVal)
+	c.add(key, newVal)
 
 	return newVal, nil
 }
@@ -272,9 +421,9 @@ func (c *Cache) DecrBy(ctx context.Context, key string, value int64) (int64, err
 		ok     bool
 		result = ecache.Value{}
 	)
-	result.Val, ok = c.client.Get(key)
+	result.Val, ok = c.get(key)
 	if !ok {
-		c.client.Add(key, -value)
+		c.add(key, -value)
 		return -value, nil
 	}
 
@@ -284,7 +433,7 @@ func (c *Cache) DecrBy(ctx context.Context, key string, value int64) (int64, err
 	}
 
 	newVal := decr - value
-	c.client.Add(key, newVal)
+	c.add(key, newVal)
 
 	return newVal, nil
 }
@@ -297,9 +446,9 @@ func (c *Cache) IncrByFloat(ctx context.Context, key string, value float64) (flo
 		ok     bool
 		result = ecache.Value{}
 	)
-	result.Val, ok = c.client.Get(key)
+	result.Val, ok = c.get(key)
 	if !ok {
-		c.client.Add(key, value)
+		c.add(key, value)
 		return value, nil
 	}
 
@@ -309,7 +458,7 @@ func (c *Cache) IncrByFloat(ctx context.Context, key string, value float64) (flo
 	}
 
 	newVal := val + value
-	c.client.Add(key, newVal)
+	c.add(key, newVal)
 
 	return newVal, nil
 }