github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/storage/pebble/cache.go (about)

     1  package pebble
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  
     7  	lru "github.com/hashicorp/golang-lru/v2"
     8  
     9  	"github.com/onflow/flow-go/model/flow"
    10  	"github.com/onflow/flow-go/module"
    11  	"github.com/onflow/flow-go/storage"
    12  )
    13  
    14  const DefaultCacheSize = uint(10_000)
    15  
    16  type CacheType int
    17  
    18  const (
    19  	CacheTypeLRU CacheType = iota + 1
    20  	CacheTypeTwoQueue
    21  )
    22  
    23  func ParseCacheType(s string) (CacheType, error) {
    24  	switch s {
    25  	case CacheTypeLRU.String():
    26  		return CacheTypeLRU, nil
    27  	case CacheTypeTwoQueue.String():
    28  		return CacheTypeTwoQueue, nil
    29  	default:
    30  		return 0, errors.New("invalid cache type")
    31  	}
    32  }
    33  
    34  func (m CacheType) String() string {
    35  	switch m {
    36  	case CacheTypeLRU:
    37  		return "lru"
    38  	case CacheTypeTwoQueue:
    39  		return "2q"
    40  	default:
    41  		return ""
    42  	}
    43  }
    44  
    45  type CacheBackend interface {
    46  	Get(key string) (value flow.RegisterValue, ok bool)
    47  	Add(key string, value flow.RegisterValue)
    48  	Contains(key string) bool
    49  	Len() int
    50  	Remove(key string)
    51  }
    52  
    53  // wrapped is a wrapper around lru.Cache to implement CacheBackend
    54  // this is needed because the standard lru cache implementation provides additional features that
    55  // the 2Q cache do not. This standardizes the interface to allow swapping between types.
    56  type wrapped struct {
    57  	cache *lru.Cache[string, flow.RegisterValue]
    58  }
    59  
    60  func (c *wrapped) Get(key string) (value flow.RegisterValue, ok bool) {
    61  	return c.cache.Get(key)
    62  }
    63  func (c *wrapped) Add(key string, value flow.RegisterValue) {
    64  	_ = c.cache.Add(key, value)
    65  }
    66  func (c *wrapped) Contains(key string) bool {
    67  	return c.cache.Contains(key)
    68  }
    69  func (c *wrapped) Len() int {
    70  	return c.cache.Len()
    71  }
    72  func (c *wrapped) Remove(key string) {
    73  	_ = c.cache.Remove(key)
    74  }
    75  
    76  type ReadCache struct {
    77  	metrics  module.CacheMetrics
    78  	resource string
    79  	cache    CacheBackend
    80  	retrieve func(key string) (flow.RegisterValue, error)
    81  }
    82  
    83  func newReadCache(
    84  	collector module.CacheMetrics,
    85  	resourceName string,
    86  	cacheType CacheType,
    87  	cacheSize uint,
    88  	retrieve func(key string) (flow.RegisterValue, error),
    89  ) (*ReadCache, error) {
    90  	cache, err := getCache(cacheType, int(cacheSize))
    91  	if err != nil {
    92  		return nil, fmt.Errorf("could not create cache: %w", err)
    93  	}
    94  
    95  	c := ReadCache{
    96  		metrics:  collector,
    97  		resource: resourceName,
    98  		cache:    cache,
    99  		retrieve: retrieve,
   100  	}
   101  	c.metrics.CacheEntries(c.resource, uint(c.cache.Len()))
   102  
   103  	return &c, nil
   104  }
   105  
   106  func getCache(cacheType CacheType, size int) (CacheBackend, error) {
   107  	switch cacheType {
   108  	case CacheTypeLRU:
   109  		cache, err := lru.New[string, flow.RegisterValue](size)
   110  		if err != nil {
   111  			return nil, err
   112  		}
   113  		return &wrapped{cache: cache}, nil
   114  	case CacheTypeTwoQueue:
   115  		return lru.New2Q[string, flow.RegisterValue](size)
   116  	default:
   117  		return nil, fmt.Errorf("unknown cache type: %d", cacheType)
   118  	}
   119  }
   120  
   121  // IsCached returns true if the key exists in the cache.
   122  // It DOES NOT check whether the key exists in the underlying data store.
   123  func (c *ReadCache) IsCached(key string) bool {
   124  	return c.cache.Contains(key)
   125  }
   126  
   127  // Get will try to retrieve the resource from cache first, and then from the
   128  // injected. During normal operations, the following error returns are expected:
   129  //   - `storage.ErrNotFound` if key is unknown.
   130  func (c *ReadCache) Get(key string) (flow.RegisterValue, error) {
   131  	resource, cached := c.cache.Get(key)
   132  	if cached {
   133  		c.metrics.CacheHit(c.resource)
   134  		if resource == nil {
   135  			return nil, storage.ErrNotFound
   136  		}
   137  		return resource, nil
   138  	}
   139  
   140  	// get it from the database
   141  	resource, err := c.retrieve(key)
   142  	if err != nil {
   143  		if errors.Is(err, storage.ErrNotFound) {
   144  			c.cache.Add(key, nil)
   145  			c.metrics.CacheEntries(c.resource, uint(c.cache.Len()))
   146  			c.metrics.CacheNotFound(c.resource)
   147  		}
   148  		return nil, fmt.Errorf("could not retrieve resource: %w", err)
   149  	}
   150  
   151  	c.metrics.CacheMiss(c.resource)
   152  
   153  	c.cache.Add(key, resource)
   154  	c.metrics.CacheEntries(c.resource, uint(c.cache.Len()))
   155  
   156  	return resource, nil
   157  }
   158  
   159  func (c *ReadCache) Remove(key string) {
   160  	c.cache.Remove(key)
   161  }
   162  
   163  // Insert will add a resource directly to the cache with the given ID
   164  func (c *ReadCache) Insert(key string, resource flow.RegisterValue) {
   165  	c.cache.Add(key, resource)
   166  	c.metrics.CacheEntries(c.resource, uint(c.cache.Len()))
   167  }