github.com/ava-labs/avalanchego@v1.11.11/cache/lru_sized_cache.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package cache
     5  
     6  import (
     7  	"sync"
     8  
     9  	"github.com/ava-labs/avalanchego/utils"
    10  	"github.com/ava-labs/avalanchego/utils/linked"
    11  )
    12  
    13  var _ Cacher[struct{}, any] = (*sizedLRU[struct{}, any])(nil)
    14  
    15  // sizedLRU is a key value store with bounded size. If the size is attempted to
    16  // be exceeded, then elements are removed from the cache until the bound is
    17  // honored, based on evicting the least recently used value.
    18  type sizedLRU[K comparable, V any] struct {
    19  	lock        sync.Mutex
    20  	elements    *linked.Hashmap[K, V]
    21  	maxSize     int
    22  	currentSize int
    23  	size        func(K, V) int
    24  }
    25  
    26  func NewSizedLRU[K comparable, V any](maxSize int, size func(K, V) int) Cacher[K, V] {
    27  	return &sizedLRU[K, V]{
    28  		elements: linked.NewHashmap[K, V](),
    29  		maxSize:  maxSize,
    30  		size:     size,
    31  	}
    32  }
    33  
    34  func (c *sizedLRU[K, V]) Put(key K, value V) {
    35  	c.lock.Lock()
    36  	defer c.lock.Unlock()
    37  
    38  	c.put(key, value)
    39  }
    40  
    41  func (c *sizedLRU[K, V]) Get(key K) (V, bool) {
    42  	c.lock.Lock()
    43  	defer c.lock.Unlock()
    44  
    45  	return c.get(key)
    46  }
    47  
    48  func (c *sizedLRU[K, V]) Evict(key K) {
    49  	c.lock.Lock()
    50  	defer c.lock.Unlock()
    51  
    52  	c.evict(key)
    53  }
    54  
    55  func (c *sizedLRU[K, V]) Flush() {
    56  	c.lock.Lock()
    57  	defer c.lock.Unlock()
    58  
    59  	c.flush()
    60  }
    61  
    62  func (c *sizedLRU[_, _]) Len() int {
    63  	c.lock.Lock()
    64  	defer c.lock.Unlock()
    65  
    66  	return c.len()
    67  }
    68  
    69  func (c *sizedLRU[_, _]) PortionFilled() float64 {
    70  	c.lock.Lock()
    71  	defer c.lock.Unlock()
    72  
    73  	return c.portionFilled()
    74  }
    75  
    76  func (c *sizedLRU[K, V]) put(key K, value V) {
    77  	newEntrySize := c.size(key, value)
    78  	if newEntrySize > c.maxSize {
    79  		c.flush()
    80  		return
    81  	}
    82  
    83  	if oldValue, ok := c.elements.Get(key); ok {
    84  		c.currentSize -= c.size(key, oldValue)
    85  	}
    86  
    87  	// Remove elements until the size of elements in the cache <= [c.maxSize].
    88  	for c.currentSize > c.maxSize-newEntrySize {
    89  		oldestKey, oldestValue, _ := c.elements.Oldest()
    90  		c.elements.Delete(oldestKey)
    91  		c.currentSize -= c.size(oldestKey, oldestValue)
    92  	}
    93  
    94  	c.elements.Put(key, value)
    95  	c.currentSize += newEntrySize
    96  }
    97  
    98  func (c *sizedLRU[K, V]) get(key K) (V, bool) {
    99  	value, ok := c.elements.Get(key)
   100  	if !ok {
   101  		return utils.Zero[V](), false
   102  	}
   103  
   104  	c.elements.Put(key, value) // Mark [k] as MRU.
   105  	return value, true
   106  }
   107  
   108  func (c *sizedLRU[K, _]) evict(key K) {
   109  	if value, ok := c.elements.Get(key); ok {
   110  		c.elements.Delete(key)
   111  		c.currentSize -= c.size(key, value)
   112  	}
   113  }
   114  
   115  func (c *sizedLRU[K, V]) flush() {
   116  	c.elements.Clear()
   117  	c.currentSize = 0
   118  }
   119  
   120  func (c *sizedLRU[_, _]) len() int {
   121  	return c.elements.Len()
   122  }
   123  
   124  func (c *sizedLRU[_, _]) portionFilled() float64 {
   125  	return float64(c.currentSize) / float64(c.maxSize)
   126  }