123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294 |
- package lru
- // Modified by https://github.com/die-net/lrucache
- import (
- "sync"
- "time"
- list "github.com/bahlo/generic-list-go"
- "github.com/samber/lo"
- )
- // Option is part of Functional Options Pattern
- type Option[K comparable, V any] func(*LruCache[K, V])
- // EvictCallback is used to get a callback when a cache entry is evicted
- type EvictCallback[K comparable, V any] func(key K, value V)
- // WithEvict set the evict callback
- func WithEvict[K comparable, V any](cb EvictCallback[K, V]) Option[K, V] {
- return func(l *LruCache[K, V]) {
- l.onEvict = cb
- }
- }
- // WithUpdateAgeOnGet update expires when Get element
- func WithUpdateAgeOnGet[K comparable, V any]() Option[K, V] {
- return func(l *LruCache[K, V]) {
- l.updateAgeOnGet = true
- }
- }
- // WithAge defined element max age (second)
- 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
- 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`.
- 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.
- type LruCache[K comparable, V any] struct {
- maxAge int64
- maxSize int
- mu sync.Mutex
- cache map[K]*list.Element[*entry[K, V]]
- lru *list.List[*entry[K, V]] // Front is least-recent
- updateAgeOnGet bool
- staleReturn bool
- onEvict EvictCallback[K, V]
- }
- // New creates an LruCache
- func New[K comparable, V any](options ...Option[K, V]) *LruCache[K, V] {
- lc := &LruCache[K, V]{
- lru: list.New[*entry[K, V]](),
- cache: make(map[K]*list.Element[*entry[K, V]]),
- }
- for _, option := range options {
- option(lc)
- }
- return lc
- }
- // Get returns any representation of a cached response and a bool
- // set to true if the key was found.
- func (c *LruCache[K, V]) Get(key K) (V, bool) {
- c.mu.Lock()
- defer c.mu.Unlock()
- el := c.get(key)
- if el == nil {
- return lo.Empty[V](), false
- }
- value := el.value
- return value, true
- }
- 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
- }
- // GetWithExpire returns 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.
- func (c *LruCache[K, V]) GetWithExpire(key K) (V, time.Time, bool) {
- c.mu.Lock()
- defer c.mu.Unlock()
- el := c.get(key)
- if el == nil {
- return lo.Empty[V](), time.Time{}, false
- }
- return el.value, time.Unix(el.expires, 0), true
- }
- // Exist returns if key exist in cache but not put item to the head of linked list
- func (c *LruCache[K, V]) Exist(key K) bool {
- c.mu.Lock()
- defer c.mu.Unlock()
- _, ok := c.cache[key]
- return ok
- }
- // Set stores any representation of a response for a given key.
- func (c *LruCache[K, V]) Set(key K, value V) {
- 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
- }
- c.setWithExpire(key, value, time.Unix(expires, 0))
- }
- // SetWithExpire stores any representation of a response for a given key and given expires.
- // The expires time will round to second.
- func (c *LruCache[K, V]) SetWithExpire(key K, value V, expires time.Time) {
- c.mu.Lock()
- defer c.mu.Unlock()
- 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)
- e := le.Value
- e.value = value
- e.expires = expires.Unix()
- } else {
- e := &entry[K, V]{key: key, value: value, expires: expires.Unix()}
- c.cache[key] = c.lru.PushBack(e)
- if c.maxSize > 0 {
- if elLen := c.lru.Len(); elLen > c.maxSize {
- c.deleteElement(c.lru.Front())
- }
- }
- }
- c.maybeDeleteOldest()
- }
- // CloneTo clone and overwrite elements to another LruCache
- func (c *LruCache[K, V]) CloneTo(n *LruCache[K, V]) {
- c.mu.Lock()
- defer c.mu.Unlock()
- n.mu.Lock()
- defer n.mu.Unlock()
- 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() {
- elm := e.Value
- n.cache[elm.key] = n.lru.PushBack(elm)
- }
- }
- func (c *LruCache[K, V]) get(key K) *entry[K, V] {
- le, ok := c.cache[key]
- if !ok {
- return nil
- }
- if !c.staleReturn && c.maxAge > 0 && le.Value.expires <= time.Now().Unix() {
- c.deleteElement(le)
- c.maybeDeleteOldest()
- return nil
- }
- c.lru.MoveToBack(le)
- el := le.Value
- if c.maxAge > 0 && c.updateAgeOnGet {
- el.expires = time.Now().Unix() + c.maxAge
- }
- return el
- }
- // Delete removes the value associated with a key.
- func (c *LruCache[K, V]) Delete(key K) {
- c.mu.Lock()
- defer c.mu.Unlock()
- c.delete(key)
- }
- func (c *LruCache[K, V]) delete(key K) {
- if le, ok := c.cache[key]; ok {
- c.deleteElement(le)
- }
- }
- func (c *LruCache[K, V]) maybeDeleteOldest() {
- if !c.staleReturn && c.maxAge > 0 {
- now := time.Now().Unix()
- for le := c.lru.Front(); le != nil && le.Value.expires <= now; le = c.lru.Front() {
- c.deleteElement(le)
- }
- }
- }
- func (c *LruCache[K, V]) deleteElement(le *list.Element[*entry[K, V]]) {
- c.lru.Remove(le)
- e := le.Value
- delete(c.cache, e.key)
- if c.onEvict != nil {
- c.onEvict(e.key, e.value)
- }
- }
- func (c *LruCache[K, V]) Clear() error {
- c.mu.Lock()
- defer c.mu.Unlock()
- c.cache = make(map[K]*list.Element[*entry[K, V]])
- return nil
- }
- // Compute either sets the computed new value for the key or deletes
- // the value for the key. When the delete result of the valueFn function
- // is set to true, the value will be deleted, if it exists. When delete
- // is set to false, the value is updated to the newValue.
- // The ok result indicates whether value was computed and stored, thus, is
- // present in the map. The actual result contains the new value in cases where
- // the value was computed and stored.
- func (c *LruCache[K, V]) Compute(
- key K,
- valueFn func(oldValue V, loaded bool) (newValue V, delete bool),
- ) (actual V, ok bool) {
- c.mu.Lock()
- defer c.mu.Unlock()
- if el := c.get(key); el != nil {
- actual, ok = el.value, true
- }
- if newValue, del := valueFn(actual, ok); del {
- if ok { // data not in cache, so needn't delete
- c.delete(key)
- }
- return lo.Empty[V](), false
- } else {
- c.set(key, newValue)
- return newValue, true
- }
- }
- type entry[K comparable, V any] struct {
- key K
- value V
- expires int64
- }
|