github.com/metacubex/mihomo@v1.18.5/common/lru/lrucache.go (about)

     1  package lru
     2  
     3  // Modified by https://github.com/die-net/lrucache
     4  
     5  import (
     6  	"sync"
     7  	"time"
     8  
     9  	list "github.com/bahlo/generic-list-go"
    10  	"github.com/samber/lo"
    11  )
    12  
    13  // Option is part of Functional Options Pattern
    14  type Option[K comparable, V any] func(*LruCache[K, V])
    15  
    16  // EvictCallback is used to get a callback when a cache entry is evicted
    17  type EvictCallback[K comparable, V any] func(key K, value V)
    18  
    19  // WithEvict set the evict callback
    20  func WithEvict[K comparable, V any](cb EvictCallback[K, V]) Option[K, V] {
    21  	return func(l *LruCache[K, V]) {
    22  		l.onEvict = cb
    23  	}
    24  }
    25  
    26  // WithUpdateAgeOnGet update expires when Get element
    27  func WithUpdateAgeOnGet[K comparable, V any]() Option[K, V] {
    28  	return func(l *LruCache[K, V]) {
    29  		l.updateAgeOnGet = true
    30  	}
    31  }
    32  
    33  // WithAge defined element max age (second)
    34  func WithAge[K comparable, V any](maxAge int64) Option[K, V] {
    35  	return func(l *LruCache[K, V]) {
    36  		l.maxAge = maxAge
    37  	}
    38  }
    39  
    40  // WithSize defined max length of LruCache
    41  func WithSize[K comparable, V any](maxSize int) Option[K, V] {
    42  	return func(l *LruCache[K, V]) {
    43  		l.maxSize = maxSize
    44  	}
    45  }
    46  
    47  // WithStale decide whether Stale return is enabled.
    48  // If this feature is enabled, element will not get Evicted according to `WithAge`.
    49  func WithStale[K comparable, V any](stale bool) Option[K, V] {
    50  	return func(l *LruCache[K, V]) {
    51  		l.staleReturn = stale
    52  	}
    53  }
    54  
    55  // LruCache is a thread-safe, in-memory lru-cache that evicts the
    56  // least recently used entries from memory when (if set) the entries are
    57  // older than maxAge (in seconds).  Use the New constructor to create one.
    58  type LruCache[K comparable, V any] struct {
    59  	maxAge         int64
    60  	maxSize        int
    61  	mu             sync.Mutex
    62  	cache          map[K]*list.Element[*entry[K, V]]
    63  	lru            *list.List[*entry[K, V]] // Front is least-recent
    64  	updateAgeOnGet bool
    65  	staleReturn    bool
    66  	onEvict        EvictCallback[K, V]
    67  }
    68  
    69  // New creates an LruCache
    70  func New[K comparable, V any](options ...Option[K, V]) *LruCache[K, V] {
    71  	lc := &LruCache[K, V]{
    72  		lru:   list.New[*entry[K, V]](),
    73  		cache: make(map[K]*list.Element[*entry[K, V]]),
    74  	}
    75  
    76  	for _, option := range options {
    77  		option(lc)
    78  	}
    79  
    80  	return lc
    81  }
    82  
    83  // Get returns any representation of a cached response and a bool
    84  // set to true if the key was found.
    85  func (c *LruCache[K, V]) Get(key K) (V, bool) {
    86  	c.mu.Lock()
    87  	defer c.mu.Unlock()
    88  
    89  	el := c.get(key)
    90  	if el == nil {
    91  		return lo.Empty[V](), false
    92  	}
    93  	value := el.value
    94  
    95  	return value, true
    96  }
    97  
    98  func (c *LruCache[K, V]) GetOrStore(key K, constructor func() V) (V, bool) {
    99  	c.mu.Lock()
   100  	defer c.mu.Unlock()
   101  
   102  	el := c.get(key)
   103  	if el == nil {
   104  		value := constructor()
   105  		c.set(key, value)
   106  		return value, false
   107  	}
   108  	value := el.value
   109  
   110  	return value, true
   111  }
   112  
   113  // GetWithExpire returns any representation of a cached response,
   114  // a time.Time Give expected expires,
   115  // and a bool set to true if the key was found.
   116  // This method will NOT check the maxAge of element and will NOT update the expires.
   117  func (c *LruCache[K, V]) GetWithExpire(key K) (V, time.Time, bool) {
   118  	c.mu.Lock()
   119  	defer c.mu.Unlock()
   120  
   121  	el := c.get(key)
   122  	if el == nil {
   123  		return lo.Empty[V](), time.Time{}, false
   124  	}
   125  
   126  	return el.value, time.Unix(el.expires, 0), true
   127  }
   128  
   129  // Exist returns if key exist in cache but not put item to the head of linked list
   130  func (c *LruCache[K, V]) Exist(key K) bool {
   131  	c.mu.Lock()
   132  	defer c.mu.Unlock()
   133  
   134  	_, ok := c.cache[key]
   135  	return ok
   136  }
   137  
   138  // Set stores any representation of a response for a given key.
   139  func (c *LruCache[K, V]) Set(key K, value V) {
   140  	c.mu.Lock()
   141  	defer c.mu.Unlock()
   142  
   143  	c.set(key, value)
   144  }
   145  
   146  func (c *LruCache[K, V]) set(key K, value V) {
   147  	expires := int64(0)
   148  	if c.maxAge > 0 {
   149  		expires = time.Now().Unix() + c.maxAge
   150  	}
   151  	c.setWithExpire(key, value, time.Unix(expires, 0))
   152  }
   153  
   154  // SetWithExpire stores any representation of a response for a given key and given expires.
   155  // The expires time will round to second.
   156  func (c *LruCache[K, V]) SetWithExpire(key K, value V, expires time.Time) {
   157  	c.mu.Lock()
   158  	defer c.mu.Unlock()
   159  
   160  	c.setWithExpire(key, value, expires)
   161  }
   162  
   163  func (c *LruCache[K, V]) setWithExpire(key K, value V, expires time.Time) {
   164  	if le, ok := c.cache[key]; ok {
   165  		c.lru.MoveToBack(le)
   166  		e := le.Value
   167  		e.value = value
   168  		e.expires = expires.Unix()
   169  	} else {
   170  		e := &entry[K, V]{key: key, value: value, expires: expires.Unix()}
   171  		c.cache[key] = c.lru.PushBack(e)
   172  
   173  		if c.maxSize > 0 {
   174  			if elLen := c.lru.Len(); elLen > c.maxSize {
   175  				c.deleteElement(c.lru.Front())
   176  			}
   177  		}
   178  	}
   179  
   180  	c.maybeDeleteOldest()
   181  }
   182  
   183  // CloneTo clone and overwrite elements to another LruCache
   184  func (c *LruCache[K, V]) CloneTo(n *LruCache[K, V]) {
   185  	c.mu.Lock()
   186  	defer c.mu.Unlock()
   187  
   188  	n.mu.Lock()
   189  	defer n.mu.Unlock()
   190  
   191  	n.lru = list.New[*entry[K, V]]()
   192  	n.cache = make(map[K]*list.Element[*entry[K, V]])
   193  
   194  	for e := c.lru.Front(); e != nil; e = e.Next() {
   195  		elm := e.Value
   196  		n.cache[elm.key] = n.lru.PushBack(elm)
   197  	}
   198  }
   199  
   200  func (c *LruCache[K, V]) get(key K) *entry[K, V] {
   201  	le, ok := c.cache[key]
   202  	if !ok {
   203  		return nil
   204  	}
   205  
   206  	if !c.staleReturn && c.maxAge > 0 && le.Value.expires <= time.Now().Unix() {
   207  		c.deleteElement(le)
   208  		c.maybeDeleteOldest()
   209  
   210  		return nil
   211  	}
   212  
   213  	c.lru.MoveToBack(le)
   214  	el := le.Value
   215  	if c.maxAge > 0 && c.updateAgeOnGet {
   216  		el.expires = time.Now().Unix() + c.maxAge
   217  	}
   218  	return el
   219  }
   220  
   221  // Delete removes the value associated with a key.
   222  func (c *LruCache[K, V]) Delete(key K) {
   223  	c.mu.Lock()
   224  	defer c.mu.Unlock()
   225  
   226  	if le, ok := c.cache[key]; ok {
   227  		c.deleteElement(le)
   228  	}
   229  }
   230  
   231  func (c *LruCache[K, V]) maybeDeleteOldest() {
   232  	if !c.staleReturn && c.maxAge > 0 {
   233  		now := time.Now().Unix()
   234  		for le := c.lru.Front(); le != nil && le.Value.expires <= now; le = c.lru.Front() {
   235  			c.deleteElement(le)
   236  		}
   237  	}
   238  }
   239  
   240  func (c *LruCache[K, V]) deleteElement(le *list.Element[*entry[K, V]]) {
   241  	c.lru.Remove(le)
   242  	e := le.Value
   243  	delete(c.cache, e.key)
   244  	if c.onEvict != nil {
   245  		c.onEvict(e.key, e.value)
   246  	}
   247  }
   248  
   249  func (c *LruCache[K, V]) Clear() error {
   250  	c.mu.Lock()
   251  	defer c.mu.Unlock()
   252  
   253  	c.cache = make(map[K]*list.Element[*entry[K, V]])
   254  
   255  	return nil
   256  }
   257  
   258  type entry[K comparable, V any] struct {
   259  	key     K
   260  	value   V
   261  	expires int64
   262  }