44package merkledb
55
66import (
7+ "errors"
78 "sync"
89
910 "github.com/ava-labs/avalanchego/utils/linkedhashmap"
1011 "github.com/ava-labs/avalanchego/utils/wrappers"
1112)
1213
14+ var errEmptyCacheTooLarge = errors .New ("cache is empty yet still too large" )
15+
1316// A cache that calls [onEviction] on the evicted element.
1417type onEvictCache [K comparable , V any ] struct {
15- lock sync.RWMutex
16- maxSize int
17- fifo linkedhashmap.LinkedHashmap [K , V ]
18+ lock sync.RWMutex
19+ maxSize int
20+ currentSize int
21+ fifo linkedhashmap.LinkedHashmap [K , V ]
22+ size func (K , V ) int
1823 // Must not call any method that grabs [c.lock]
1924 // because this would cause a deadlock.
2025 onEviction func (K , V ) error
2126}
2227
23- func newOnEvictCache [K comparable , V any ](maxSize int , onEviction func (K , V ) error ) onEvictCache [K , V ] {
28+ func newOnEvictCache [K comparable , V any ](
29+ maxSize int ,
30+ size func (K , V ) int ,
31+ onEviction func (K , V ) error ,
32+ ) onEvictCache [K , V ] {
2433 return onEvictCache [K , V ]{
2534 maxSize : maxSize ,
2635 fifo : linkedhashmap .New [K , V ](),
36+ size : size ,
2737 onEviction : onEviction ,
2838 }
2939}
3040
31- // removeOldest returns and removes the oldest element from this cache.
32- // Assumes [c.lock] is held.
33- func (c * onEvictCache [K , V ]) removeOldest () (K , V , bool ) {
34- k , v , exists := c .fifo .Oldest ()
35- if exists {
36- c .fifo .Delete (k )
37- }
38- return k , v , exists
39- }
40-
4141// Get an element from this cache.
4242func (c * onEvictCache [K , V ]) Get (key K ) (V , bool ) {
4343 c .lock .RLock ()
@@ -53,14 +53,14 @@ func (c *onEvictCache[K, V]) Put(key K, value V) error {
5353 c .lock .Lock ()
5454 defer c .lock .Unlock ()
5555
56+ if oldValue , replaced := c .fifo .Get (key ); replaced {
57+ c .currentSize -= c .size (key , oldValue )
58+ }
59+
60+ c .currentSize += c .size (key , value )
5661 c .fifo .Put (key , value ) // Mark as MRU
5762
58- if c .fifo .Len () > c .maxSize {
59- oldestKey , oldestVal , _ := c .fifo .Oldest ()
60- c .fifo .Delete (oldestKey )
61- return c .onEviction (oldestKey , oldestVal )
62- }
63- return nil
63+ return c .resize (c .maxSize )
6464}
6565
6666// Flush removes all elements from the cache.
@@ -74,16 +74,37 @@ func (c *onEvictCache[K, V]) Flush() error {
7474 c .lock .Unlock ()
7575 }()
7676
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 {
7797 // Note that we can't use [c.fifo]'s iterator because [c.onEviction]
7898 // modifies [c.fifo], which violates the iterator's invariant.
7999 var errs wrappers.Errs
80- for {
81- key , value , exists := c .removeOldest ()
100+ for c . currentSize > target {
101+ k , v , exists := c .removeOldest ()
82102 if ! exists {
83- // The cache is empty.
84- return errs .Err
103+ // This should really never happen unless the size of an entry
104+ // changed or the target size is negative.
105+ return errEmptyCacheTooLarge
85106 }
86-
87- errs .Add (c .onEviction (key , value ))
107+ errs .Add (c .onEviction (k , v ))
88108 }
109+ return errs .Err
89110}
0 commit comments