github.com/MetalBlockchain/metalgo@v1.11.9/x/merkledb/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 merkledb
     5  
     6  import (
     7  	"errors"
     8  	"sync"
     9  
    10  	"github.com/MetalBlockchain/metalgo/utils/linked"
    11  	"github.com/MetalBlockchain/metalgo/utils/wrappers"
    12  )
    13  
    14  var errEmptyCacheTooLarge = errors.New("cache is empty yet still too large")
    15  
    16  // A cache that calls [onEviction] on the evicted element.
    17  type onEvictCache[K comparable, V any] struct {
    18  	lock        sync.RWMutex
    19  	maxSize     int
    20  	currentSize int
    21  	fifo        *linked.Hashmap[K, V]
    22  	size        func(K, V) int
    23  	// Must not call any method that grabs [c.lock]
    24  	// because this would cause a deadlock.
    25  	onEviction func(K, V) error
    26  }
    27  
    28  // [size] must always return a positive number.
    29  func newOnEvictCache[K comparable, V any](
    30  	maxSize int,
    31  	size func(K, V) int,
    32  	onEviction func(K, V) error,
    33  ) onEvictCache[K, V] {
    34  	return onEvictCache[K, V]{
    35  		maxSize:    maxSize,
    36  		fifo:       linked.NewHashmap[K, V](),
    37  		size:       size,
    38  		onEviction: onEviction,
    39  	}
    40  }
    41  
    42  // Get an element from this cache.
    43  func (c *onEvictCache[K, V]) Get(key K) (V, bool) {
    44  	c.lock.RLock()
    45  	defer c.lock.RUnlock()
    46  
    47  	return c.fifo.Get(key)
    48  }
    49  
    50  // Put an element into this cache. If this causes an element
    51  // to be evicted, calls [c.onEviction] on the evicted element
    52  // and returns the error from [c.onEviction]. Otherwise, returns nil.
    53  func (c *onEvictCache[K, V]) Put(key K, value V) error {
    54  	c.lock.Lock()
    55  	defer c.lock.Unlock()
    56  
    57  	if oldValue, replaced := c.fifo.Get(key); replaced {
    58  		c.currentSize -= c.size(key, oldValue)
    59  	}
    60  
    61  	c.currentSize += c.size(key, value)
    62  	c.fifo.Put(key, value) // Mark as MRU
    63  
    64  	return c.resize(c.maxSize)
    65  }
    66  
    67  // Flush removes all elements from the cache.
    68  //
    69  // Returns the first non-nil error returned by [c.onEviction], if any.
    70  //
    71  // If [c.onEviction] errors, it will still be called for any subsequent elements
    72  // and the cache will still be emptied.
    73  func (c *onEvictCache[K, V]) Flush() error {
    74  	c.lock.Lock()
    75  	defer c.lock.Unlock()
    76  
    77  	return c.resize(0)
    78  }
    79  
    80  // removeOldest returns and removes the oldest element from this cache.
    81  //
    82  // Assumes [c.lock] is held.
    83  func (c *onEvictCache[K, V]) removeOldest() (K, V, bool) {
    84  	k, v, exists := c.fifo.Oldest()
    85  	if exists {
    86  		c.currentSize -= c.size(k, v)
    87  		c.fifo.Delete(k)
    88  	}
    89  	return k, v, exists
    90  }
    91  
    92  // resize removes the oldest elements from the cache until the cache is not
    93  // larger than the provided target.
    94  //
    95  // Assumes [c.lock] is held.
    96  func (c *onEvictCache[K, V]) resize(target int) error {
    97  	// Note that we can't use [c.fifo]'s iterator because [c.onEviction]
    98  	// modifies [c.fifo], which violates the iterator's invariant.
    99  	var errs wrappers.Errs
   100  	for c.currentSize > target {
   101  		k, v, exists := c.removeOldest()
   102  		if !exists {
   103  			// This should really never happen unless the size of an entry
   104  			// changed or the target size is negative.
   105  			return errEmptyCacheTooLarge
   106  		}
   107  		errs.Add(c.onEviction(k, v))
   108  	}
   109  	return errs.Err
   110  }