mihomo/common/lru/lrucache.go

263 lines
5.9 KiB
Go
Raw Normal View History

2023-12-02 17:07:36 +08:00
package lru
// Modified by https://github.com/die-net/lrucache
import (
"sync"
"time"
2022-04-24 02:07:57 +08:00
list "github.com/bahlo/generic-list-go"
2023-08-26 21:19:53 +08:00
"github.com/samber/lo"
)
// Option is part of Functional Options Pattern
2022-04-05 20:23:16 +08:00
type Option[K comparable, V any] func(*LruCache[K, V])
2019-12-05 00:17:24 +08:00
// EvictCallback is used to get a callback when a cache entry is evicted
2022-04-24 02:07:57 +08:00
type EvictCallback[K comparable, V any] func(key K, value V)
2019-12-05 00:17:24 +08:00
// WithEvict set the evict callback
2022-04-24 02:07:57 +08:00
func WithEvict[K comparable, V any](cb EvictCallback[K, V]) Option[K, V] {
2022-04-05 20:23:16 +08:00
return func(l *LruCache[K, V]) {
2019-12-05 00:17:24 +08:00
l.onEvict = cb
}
}
// WithUpdateAgeOnGet update expires when Get element
2022-04-05 20:23:16 +08:00
func WithUpdateAgeOnGet[K comparable, V any]() Option[K, V] {
return func(l *LruCache[K, V]) {
l.updateAgeOnGet = true
}
}
// WithAge defined element max age (second)
2022-04-05 20:23:16 +08:00
func WithAge[K comparable, V any](maxAge int64) Option[K, V] {
return func(l *LruCache[K, V]) {
l.maxAge = maxAge
}
}
// WithSize defined max length of LruCache
2022-04-05 20:23:16 +08:00
func WithSize[K comparable, V any](maxSize int) Option[K, V] {
return func(l *LruCache[K, V]) {
l.maxSize = maxSize
}
}
// WithStale decide whether Stale return is enabled.
// If this feature is enabled, element will not get Evicted according to `WithAge`.
2022-04-05 20:23:16 +08:00
func WithStale[K comparable, V any](stale bool) Option[K, V] {
return func(l *LruCache[K, V]) {
l.staleReturn = stale
}
}
// LruCache is a thread-safe, in-memory lru-cache that evicts the
// least recently used entries from memory when (if set) the entries are
// older than maxAge (in seconds). Use the New constructor to create one.
2022-04-05 20:23:16 +08:00
type LruCache[K comparable, V any] struct {
maxAge int64
maxSize int
mu sync.Mutex
2022-04-24 02:07:57 +08:00
cache map[K]*list.Element[*entry[K, V]]
lru *list.List[*entry[K, V]] // Front is least-recent
updateAgeOnGet bool
staleReturn bool
2022-04-24 02:07:57 +08:00
onEvict EvictCallback[K, V]
}
2022-08-17 11:43:13 +08:00
// New creates an LruCache
func New[K comparable, V any](options ...Option[K, V]) *LruCache[K, V] {
2022-04-05 20:23:16 +08:00
lc := &LruCache[K, V]{
2022-04-24 02:07:57 +08:00
lru: list.New[*entry[K, V]](),
cache: make(map[K]*list.Element[*entry[K, V]]),
}
for _, option := range options {
option(lc)
}
return lc
}
2022-03-16 12:10:13 +08:00
// Get returns the any representation of a cached response and a bool
// set to true if the key was found.
2022-04-05 20:23:16 +08:00
func (c *LruCache[K, V]) Get(key K) (V, bool) {
2023-08-03 23:07:30 +08:00
c.mu.Lock()
defer c.mu.Unlock()
2022-04-05 20:23:16 +08:00
el := c.get(key)
if el == nil {
2023-08-26 21:19:53 +08:00
return lo.Empty[V](), false
}
2022-04-05 20:23:16 +08:00
value := el.value
return value, true
}
2023-08-03 23:07:30 +08:00
func (c *LruCache[K, V]) GetOrStore(key K, constructor func() V) (V, bool) {
c.mu.Lock()
defer c.mu.Unlock()
el := c.get(key)
if el == nil {
value := constructor()
c.set(key, value)
return value, false
}
value := el.value
return value, true
}
2022-03-16 12:10:13 +08:00
// GetWithExpire returns the any representation of a cached response,
// a time.Time Give expected expires,
// and a bool set to true if the key was found.
// This method will NOT check the maxAge of element and will NOT update the expires.
2022-04-05 20:23:16 +08:00
func (c *LruCache[K, V]) GetWithExpire(key K) (V, time.Time, bool) {
2023-08-03 23:07:30 +08:00
c.mu.Lock()
defer c.mu.Unlock()
2022-04-05 20:23:16 +08:00
el := c.get(key)
if el == nil {
2023-08-26 21:19:53 +08:00
return lo.Empty[V](), time.Time{}, false
}
2022-04-05 20:23:16 +08:00
return el.value, time.Unix(el.expires, 0), true
}
2019-10-11 14:01:16 +08:00
// Exist returns if key exist in cache but not put item to the head of linked list
2022-04-05 20:23:16 +08:00
func (c *LruCache[K, V]) Exist(key K) bool {
2019-10-11 14:01:16 +08:00
c.mu.Lock()
defer c.mu.Unlock()
_, ok := c.cache[key]
return ok
}
2022-03-16 12:10:13 +08:00
// Set stores the any representation of a response for a given key.
2022-04-05 20:23:16 +08:00
func (c *LruCache[K, V]) Set(key K, value V) {
2023-08-03 23:07:30 +08:00
c.mu.Lock()
defer c.mu.Unlock()
c.set(key, value)
}
func (c *LruCache[K, V]) set(key K, value V) {
expires := int64(0)
if c.maxAge > 0 {
expires = time.Now().Unix() + c.maxAge
}
2023-08-03 23:07:30 +08:00
c.setWithExpire(key, value, time.Unix(expires, 0))
}
2022-03-16 12:10:13 +08:00
// SetWithExpire stores the any representation of a response for a given key and given expires.
// The expires time will round to second.
2022-04-05 20:23:16 +08:00
func (c *LruCache[K, V]) SetWithExpire(key K, value V, expires time.Time) {
c.mu.Lock()
defer c.mu.Unlock()
2023-08-03 23:07:30 +08:00
c.setWithExpire(key, value, expires)
}
func (c *LruCache[K, V]) setWithExpire(key K, value V, expires time.Time) {
if le, ok := c.cache[key]; ok {
c.lru.MoveToBack(le)
2022-04-24 02:07:57 +08:00
e := le.Value
e.value = value
e.expires = expires.Unix()
} else {
2022-04-05 20:23:16 +08:00
e := &entry[K, V]{key: key, value: value, expires: expires.Unix()}
c.cache[key] = c.lru.PushBack(e)
if c.maxSize > 0 {
2022-04-05 20:23:16 +08:00
if elLen := c.lru.Len(); elLen > c.maxSize {
c.deleteElement(c.lru.Front())
}
}
}
c.maybeDeleteOldest()
}
// CloneTo clone and overwrite elements to another LruCache
2022-04-05 20:23:16 +08:00
func (c *LruCache[K, V]) CloneTo(n *LruCache[K, V]) {
c.mu.Lock()
defer c.mu.Unlock()
n.mu.Lock()
defer n.mu.Unlock()
2022-04-24 02:07:57 +08:00
n.lru = list.New[*entry[K, V]]()
n.cache = make(map[K]*list.Element[*entry[K, V]])
for e := c.lru.Front(); e != nil; e = e.Next() {
2022-04-24 02:07:57 +08:00
elm := e.Value
n.cache[elm.key] = n.lru.PushBack(elm)
}
}
2022-04-05 20:23:16 +08:00
func (c *LruCache[K, V]) get(key K) *entry[K, V] {
le, ok := c.cache[key]
if !ok {
return nil
}
2022-04-24 02:07:57 +08:00
if !c.staleReturn && c.maxAge > 0 && le.Value.expires <= time.Now().Unix() {
c.deleteElement(le)
c.maybeDeleteOldest()
return nil
}
c.lru.MoveToBack(le)
2022-04-24 02:07:57 +08:00
el := le.Value
if c.maxAge > 0 && c.updateAgeOnGet {
2022-04-05 20:23:16 +08:00
el.expires = time.Now().Unix() + c.maxAge
}
2022-04-05 20:23:16 +08:00
return el
}
// Delete removes the value associated with a key.
2022-04-05 20:23:16 +08:00
func (c *LruCache[K, V]) Delete(key K) {
c.mu.Lock()
2023-08-03 23:07:30 +08:00
defer c.mu.Unlock()
if le, ok := c.cache[key]; ok {
c.deleteElement(le)
}
}
2022-04-05 20:23:16 +08:00
func (c *LruCache[K, V]) maybeDeleteOldest() {
if !c.staleReturn && c.maxAge > 0 {
now := time.Now().Unix()
2022-04-24 02:07:57 +08:00
for le := c.lru.Front(); le != nil && le.Value.expires <= now; le = c.lru.Front() {
c.deleteElement(le)
}
}
}
2022-04-24 02:07:57 +08:00
func (c *LruCache[K, V]) deleteElement(le *list.Element[*entry[K, V]]) {
c.lru.Remove(le)
2022-04-24 02:07:57 +08:00
e := le.Value
delete(c.cache, e.key)
2019-12-05 00:17:24 +08:00
if c.onEvict != nil {
c.onEvict(e.key, e.value)
}
}
2022-04-05 20:23:16 +08:00
func (c *LruCache[K, V]) Clear() error {
2022-03-23 01:05:43 +08:00
c.mu.Lock()
2023-08-03 23:07:30 +08:00
defer c.mu.Unlock()
2022-03-23 01:05:43 +08:00
2022-04-24 02:07:57 +08:00
c.cache = make(map[K]*list.Element[*entry[K, V]])
2022-03-23 01:05:43 +08:00
return nil
}
2022-04-05 20:23:16 +08:00
type entry[K comparable, V any] struct {
key K
value V
expires int64
}