github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/pkg/cache/cache_ristretto.go (about)

     1  //go:build !wasm
     2  // +build !wasm
     3  
     4  package cache
     5  
     6  import (
     7  	"time"
     8  
     9  	"github.com/outcaste-io/ristretto"
    10  	"github.com/outcaste-io/ristretto/z"
    11  	"github.com/rs/zerolog"
    12  
    13  	"github.com/authzed/spicedb/internal/dispatch/keys"
    14  )
    15  
    16  func ristrettoConfig(config *Config) *ristretto.Config {
    17  	return &ristretto.Config{
    18  		NumCounters: config.NumCounters,
    19  		MaxCost:     config.MaxCost,
    20  		BufferItems: 64, // Recommended constant by Ristretto authors.
    21  		KeyToHash: func(key any) (uint64, uint64) {
    22  			dispatchCacheKey, ok := key.(keys.DispatchCacheKey)
    23  			if !ok {
    24  				return z.KeyToHash(key)
    25  			}
    26  			return dispatchCacheKey.AsUInt64s()
    27  		},
    28  	}
    29  }
    30  
    31  // NewCacheWithMetrics creates a new ristretto cache from the given config
    32  // that also reports metrics to the default Prometheus registry.
    33  func NewCacheWithMetrics(name string, config *Config) (Cache, error) {
    34  	cfg := ristrettoConfig(config)
    35  	cfg.Metrics = true
    36  
    37  	rcache, err := ristretto.NewCache(cfg)
    38  	if err != nil {
    39  		return nil, err
    40  	}
    41  
    42  	cache := wrapped{name, config, config.DefaultTTL, rcache}
    43  	mustRegisterCache(name, cache)
    44  	return &cache, nil
    45  }
    46  
    47  // NewCache creates a new ristretto cache from the given config.
    48  func NewCache(config *Config) (Cache, error) {
    49  	rcache, err := ristretto.NewCache(ristrettoConfig(config))
    50  	return &wrapped{"", config, config.DefaultTTL, rcache}, err
    51  }
    52  
    53  type wrapped struct {
    54  	name       string
    55  	config     *Config
    56  	defaultTTL time.Duration
    57  	*ristretto.Cache
    58  }
    59  
    60  func (w wrapped) Set(key, entry any, cost int64) bool {
    61  	if w.defaultTTL <= 0 {
    62  		return w.Cache.Set(key, entry, cost)
    63  	}
    64  	return w.Cache.SetWithTTL(key, entry, cost, w.defaultTTL)
    65  }
    66  
    67  var _ Cache = (*wrapped)(nil)
    68  
    69  func (w wrapped) GetMetrics() Metrics                   { return w.Cache.Metrics }
    70  func (w wrapped) MarshalZerologObject(e *zerolog.Event) { e.EmbedObject(w.config) }
    71  
    72  func (w wrapped) Close() {
    73  	w.Cache.Close()
    74  	unregisterCache(w.name)
    75  }