github.com/koko1123/flow-go-1@v0.29.6/storage/badger/cache.go (about)

     1  package badger
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  
     7  	"github.com/dgraph-io/badger/v3"
     8  	lru "github.com/hashicorp/golang-lru"
     9  
    10  	"github.com/koko1123/flow-go-1/module"
    11  	"github.com/koko1123/flow-go-1/storage"
    12  	"github.com/koko1123/flow-go-1/storage/badger/transaction"
    13  )
    14  
    15  func withLimit(limit uint) func(*Cache) {
    16  	return func(c *Cache) {
    17  		c.limit = limit
    18  	}
    19  }
    20  
    21  type storeFunc func(key interface{}, val interface{}) func(*transaction.Tx) error
    22  
    23  const DefaultCacheSize = uint(1000)
    24  
    25  func withStore(store storeFunc) func(*Cache) {
    26  	return func(c *Cache) {
    27  		c.store = store
    28  	}
    29  }
    30  
    31  func noStore(key interface{}, val interface{}) func(*transaction.Tx) error {
    32  	return func(tx *transaction.Tx) error {
    33  		return fmt.Errorf("no store function for cache put available")
    34  	}
    35  }
    36  
    37  func noopStore(key interface{}, val interface{}) func(*transaction.Tx) error {
    38  	return func(tx *transaction.Tx) error {
    39  		return nil
    40  	}
    41  }
    42  
    43  type retrieveFunc func(key interface{}) func(*badger.Txn) (interface{}, error)
    44  
    45  func withRetrieve(retrieve retrieveFunc) func(*Cache) {
    46  	return func(c *Cache) {
    47  		c.retrieve = retrieve
    48  	}
    49  }
    50  
    51  func noRetrieve(key interface{}) func(*badger.Txn) (interface{}, error) {
    52  	return func(tx *badger.Txn) (interface{}, error) {
    53  		return nil, fmt.Errorf("no retrieve function for cache get available")
    54  	}
    55  }
    56  
    57  type Cache struct {
    58  	metrics  module.CacheMetrics
    59  	limit    uint
    60  	store    storeFunc
    61  	retrieve retrieveFunc
    62  	resource string
    63  	cache    *lru.Cache
    64  }
    65  
    66  func newCache(collector module.CacheMetrics, resourceName string, options ...func(*Cache)) *Cache {
    67  	c := Cache{
    68  		metrics:  collector,
    69  		limit:    1000,
    70  		store:    noStore,
    71  		retrieve: noRetrieve,
    72  		resource: resourceName,
    73  	}
    74  	for _, option := range options {
    75  		option(&c)
    76  	}
    77  	c.cache, _ = lru.New(int(c.limit))
    78  	c.metrics.CacheEntries(c.resource, uint(c.cache.Len()))
    79  	return &c
    80  }
    81  
    82  // Get will try to retrieve the resource from cache first, and then from the
    83  // injected
    84  func (c *Cache) Get(key interface{}) func(*badger.Txn) (interface{}, error) {
    85  	return func(tx *badger.Txn) (interface{}, error) {
    86  
    87  		// check if we have it in the cache
    88  		resource, cached := c.cache.Get(key)
    89  		if cached {
    90  			c.metrics.CacheHit(c.resource)
    91  			return resource, nil
    92  		}
    93  
    94  		// get it from the database
    95  		resource, err := c.retrieve(key)(tx)
    96  		if err != nil {
    97  			if errors.Is(err, storage.ErrNotFound) {
    98  				c.metrics.CacheNotFound(c.resource)
    99  			}
   100  			return nil, fmt.Errorf("could not retrieve resource: %w", err)
   101  		}
   102  
   103  		c.metrics.CacheMiss(c.resource)
   104  
   105  		// cache the resource and eject least recently used one if we reached limit
   106  		evicted := c.cache.Add(key, resource)
   107  		if !evicted {
   108  			c.metrics.CacheEntries(c.resource, uint(c.cache.Len()))
   109  		}
   110  
   111  		return resource, nil
   112  	}
   113  }
   114  
   115  func (c *Cache) Remove(key interface{}) {
   116  	c.cache.Remove(key)
   117  }
   118  
   119  // Insert will add a resource directly to the cache with the given ID
   120  func (c *Cache) Insert(key interface{}, resource interface{}) {
   121  	// cache the resource and eject least recently used one if we reached limit
   122  	evicted := c.cache.Add(key, resource)
   123  	if !evicted {
   124  		c.metrics.CacheEntries(c.resource, uint(c.cache.Len()))
   125  	}
   126  }
   127  
   128  // PutTx will return tx which adds a resource to the cache with the given ID.
   129  func (c *Cache) PutTx(key interface{}, resource interface{}) func(*transaction.Tx) error {
   130  	storeOps := c.store(key, resource) // assemble DB operations to store resource (no execution)
   131  
   132  	return func(tx *transaction.Tx) error {
   133  		err := storeOps(tx) // execute operations to store resource
   134  		if err != nil {
   135  			return fmt.Errorf("could not store resource: %w", err)
   136  		}
   137  
   138  		tx.OnSucceed(func() {
   139  			c.Insert(key, resource)
   140  		})
   141  
   142  		return nil
   143  	}
   144  }