querycache.go

package sql

import (
	"context"
	"slices"
	"sync"

	"github.com/hashicorp/golang-lru/v2/simplelru"
)

const defaultLRUCacheSize = 100

type (
	QueryCacheKind   string
	inGetValueCtxKey struct{}
)

var NullQueryCache QueryCache = (*queryCache)(nil)

type queryCacheKey struct {
	Kind QueryCacheKind
	Key  string
}

// QueryCacheKey creates a key for QueryCache.
func QueryCacheKey(kind QueryCacheKind, key string) queryCacheKey {
	return queryCacheKey{Kind: kind, Key: key}
}

// QueryCacheSubKey denotes a cache subkey. The empty subkey refers to the main
// key. All other subkeys are cleared by UpdateSlice for the key. The subkeys
// are intended to store data derived from the query results, such as serialized
// responses.
type QueryCacheSubKey string

const (
	// When UpdateSlice method or AppendToCachedSlice function is
	// called with mainSubKey, all other subkeys for the CacheKey
	// are invalidated.
	mainSubKey QueryCacheSubKey = ""
)

type (
	// UntypedRetrieveFunc retrieves a value to be cached.
	UntypedRetrieveFunc func(ctx context.Context) (any, error)
	// SliceAppender modifies slice value stored in the cache, appending the
	// specified item to it and returns the updated slice.
	SliceAppender func(s any) any
)

// QueryCache stores results of SQL queries and data derived from these results.
// Presently, the cached entries are never removed, but eventually, it might
// become an LRU cache.
type QueryCache interface {
	// GetValue retrieves the specified key+subKey value from the cache. If
	// the entry is absent from cache, it's populated by calling retrieve func.
	// Note that the retrieve func should never cause UpdateSlice to be
	// called for this cache.
	GetValue(
		ctx context.Context,
		key queryCacheKey,
		subKey QueryCacheSubKey,
		retrieve UntypedRetrieveFunc,
	) (any, error)
	// UpdateSlice updates the slice stored in the cache by invoking the
	// specified SliceAppender. If the entry is not cached, the method does
	// nothing.
	UpdateSlice(key queryCacheKey, update SliceAppender)
	// ClearCache empties the cache
	ClearCache()
}

// RetrieveFunc retrieves a value to be stored in the cache.
type RetrieveFunc[T any] func() (T, error)

// WithCachedValue retrieves the specified value from the cache. If the entry is
// absent from the cache, it's populated by calling retrieve func. Note that the
// retrieve func should never cause UpdateSlice to be called.
func WithCachedValue[T any](
	ctx context.Context,
	db any,
	key queryCacheKey,
	retrieve func(ctx context.Context) (T, error),
) (T, error) {
	return WithCachedSubKey(ctx, db, key, mainSubKey, retrieve)
}

// WithCachedValue retrieves the specified value identified by the key and
// subKey from the cache. If the entry is absent from the cache, it's populated
// by calling retrieve func. Note that the retrieve func should never cause
// UpdateSlice to be called.
func WithCachedSubKey[T any](
	ctx context.Context,
	db any,
	key queryCacheKey,
	subKey QueryCacheSubKey,
	retrieve func(ctx context.Context) (T, error),
) (T, error) {
	cache, ok := db.(QueryCache)
	if !ok {
		return retrieve(ctx)
	}

	v, err := cache.GetValue(
		ctx, key, subKey,
		func(ctx context.Context) (any, error) {
			return retrieve(ctx)
		})
	if err != nil {
		var r T
		return r, err
	}
	return v.(T), err
}

// AppendToCachedSlice adds a value to the slice stored in the cache by invoking
// the specified SliceAppender. If the entry is not cached, the function does
// nothing.
func AppendToCachedSlice[T any](db any, key queryCacheKey, v T) {
	if cache, ok := db.(QueryCache); ok {
		cache.UpdateSlice(key, func(s any) any {
			if s == nil {
				return []T{v}
			}
			return append(s.([]T), v)
		})
	}
}

type lruCacheKey struct {
	key    string
	subKey QueryCacheSubKey
}

type lru = simplelru.LRU[lruCacheKey, any]

type queryCache struct {
	sync.Mutex
	updateMtx        sync.RWMutex
	subKeyMap        map[queryCacheKey][]QueryCacheSubKey
	cacheSizesByKind map[QueryCacheKind]int
	caches           map[QueryCacheKind]*lru
}

var _ QueryCache = &queryCache{}

func (c *queryCache) ensureLRU(kind QueryCacheKind) *lru {
	if lruForKind, found := c.caches[kind]; found {
		return lruForKind
	}
	size, found := c.cacheSizesByKind[kind]
	if !found || size <= 0 {
		size = defaultLRUCacheSize
	}
	lruForKind, err := simplelru.NewLRU[lruCacheKey, any](size, func(k lruCacheKey, v any) {
		if k.subKey == mainSubKey {
			c.clearSubKeys(queryCacheKey{Kind: kind, Key: k.key})
		}
	})
	if err != nil {
		panic("NewLRU failed: " + err.Error())
	}
	if c.caches == nil {
		c.caches = make(map[QueryCacheKind]*lru)
	}
	c.caches[kind] = lruForKind
	return lruForKind
}

func (c *queryCache) clearSubKeys(key queryCacheKey) {
	lru, found := c.caches[key.Kind]
	if !found {
		return
	}
	for _, sk := range c.subKeyMap[key] {
		lru.Remove(lruCacheKey{
			key:    key.Key,
			subKey: sk,
		})
	}
}

func (c *queryCache) get(key queryCacheKey, subKey QueryCacheSubKey) (any, bool) {
	c.Lock()
	defer c.Unlock()
	lru, found := c.caches[key.Kind]
	if !found {
		return nil, false
	}

	return lru.Get(lruCacheKey{
		key:    key.Key,
		subKey: subKey,
	})
}

func (c *queryCache) set(key queryCacheKey, subKey QueryCacheSubKey, v any) {
	c.Lock()
	defer c.Unlock()
	if subKey != mainSubKey {
		sks := c.subKeyMap[key]
		if slices.Index(sks, subKey) < 0 {
			if c.subKeyMap == nil {
				c.subKeyMap = make(map[queryCacheKey][]QueryCacheSubKey)
			}
			c.subKeyMap[key] = append(sks, subKey)
		}
	}
	lru := c.ensureLRU(key.Kind)
	lru.Add(lruCacheKey{key: key.Key, subKey: subKey}, v)
}

func (c *queryCache) GetValue(
	ctx context.Context,
	key queryCacheKey,
	subKey QueryCacheSubKey,
	retrieve UntypedRetrieveFunc,
) (any, error) {
	if c == nil {
		return retrieve(ctx)
	}
	// Avoid recursive locking from within retrieve()
	if ctx.Value(inGetValueCtxKey{}) == nil {
		c.updateMtx.RLock()
		defer c.updateMtx.RUnlock()
	}
	v, found := c.get(key, subKey)
	var err error
	if !found {
		// This may seem like a race, but at worst, retrieve() will be
		// called several times when populating this cached entry.
		// That's better than locking for the duration of retrieve(),
		// which can also refer to this cache
		v, err = retrieve(context.WithValue(ctx, inGetValueCtxKey{}, true))
		if err == nil {
			c.set(key, subKey, v)
		}
	}
	return v, err
}

func (c *queryCache) UpdateSlice(key queryCacheKey, update SliceAppender) {
	if c == nil {
		return
	}

	// Here we lock for the call b/c we can't have conflicting updates for
	// the slice at the same time
	c.updateMtx.Lock()
	c.Lock()
	defer func() {
		c.Unlock()
		c.updateMtx.Unlock()
	}()
	lru, found := c.caches[key.Kind]
	if !found {
		return
	}

	k := lruCacheKey{key: key.Key, subKey: mainSubKey}
	if v, found := lru.Get(k); found {
		lru.Add(k, update(v))
		c.clearSubKeys(key)
	}
}

func (c *queryCache) ClearCache() {
	if c == nil {
		return
	}
	c.updateMtx.Lock()
	c.Lock()
	defer func() {
		c.Unlock()
		c.updateMtx.Unlock()
	}()
	// No need to clear c.subKeyMap as it's only used to keep track of possible subkeys for each key
	c.caches = nil
}