tstor/pkg/maxcache/cache.go
2024-11-15 16:39:56 +03:00

114 lines
2.6 KiB
Go

package maxcache
import (
"context"
"sync"
"time"
"github.com/cespare/xxhash/v2"
"github.com/goware/singleflight"
lru "github.com/hashicorp/golang-lru/v2"
)
func NewCache[K comparable, V any](size int, freshFor, ttl time.Duration) *Cache[K, V] {
values, _ := lru.New[K, value[V]](size)
return &Cache[K, V]{
values: values,
}
}
type Cache[K comparable, V any] struct {
values *lru.Cache[K, value[V]]
mu sync.RWMutex
callGroup singleflight.Group[K, V]
}
func (c *Cache[K, V]) Get(ctx context.Context, key K, fn singleflight.DoFunc[V]) (V, error) {
return c.get(ctx, key, false, fn)
}
func (c *Cache[K, V]) GetFresh(ctx context.Context, key K, fn singleflight.DoFunc[V]) (V, error) {
return c.get(ctx, key, true, fn)
}
func (c *Cache[K, V]) Set(ctx context.Context, key K, fn singleflight.DoFunc[V]) (V, bool, error) {
v, err, shared := c.callGroup.Do(key, c.set(key, fn))
return v, shared, err
}
func (c *Cache[K, V]) get(ctx context.Context, key K, freshOnly bool, fn singleflight.DoFunc[V]) (V, error) {
c.mu.RLock()
val, ok := c.values.Get(key)
c.mu.RUnlock()
// value exists and is fresh - just return
if ok && val.IsFresh() {
return val.Value(), nil
}
// value exists and is stale, and we're OK with serving it stale while updating in the background
// note: stale means its still okay, but not fresh. but if its expired, then it means its useless.
if ok && !freshOnly && !val.IsExpired() {
// TODO: technically could be a stampede of goroutines here if the value is expired
// and we're OK with serving it stale
go c.Set(ctx, key, fn)
return val.Value(), nil
}
// value doesn't exist or is expired, or is stale and we need it fresh (freshOnly:true) - sync update
v, _, err := c.Set(ctx, key, fn)
return v, err
}
func (c *Cache[K, V]) set(key K, fn singleflight.DoFunc[V]) singleflight.DoFunc[V] {
return singleflight.DoFunc[V](func() (V, error) {
val, err := fn()
if err != nil {
return val, err
}
c.mu.Lock()
c.values.Add(key, value[V]{
v: val,
})
c.mu.Unlock()
return val, nil
})
}
type value[V any] struct {
v V
bestBefore time.Time // cache entry freshness cutoff
expiry time.Time // cache entry time to live cutoff
}
func (v *value[V]) IsFresh() bool {
return v.bestBefore.After(time.Now())
}
func (v *value[V]) IsExpired() bool {
return v.expiry.Before(time.Now())
}
func (v *value[V]) Value() V {
return v.v
}
func BytesToHash(b ...[]byte) uint64 {
d := xxhash.New()
for _, v := range b {
d.Write(v)
}
return d.Sum64()
}
func StringToHash(s ...string) uint64 {
d := xxhash.New()
for _, v := range s {
d.WriteString(v)
}
return d.Sum64()
}