backend/internal/cache/lru_striped.go

package cache

import (
	"math"
	"runtime"
	"time"

	"github.com/cespare/xxhash/v2"
)

type LRUStriped struct {
	buckets []*LRU
	name    string
}

type LRUStripedOptions struct {
	Size          int
	DefaultExpiry time.Duration
	Name          string
	Buckets       int
}

func NewLRUStriped(opts *LRUStripedOptions) *LRUStriped {
	if opts.Buckets <= 0 {
		opts.Buckets = runtime.NumCPU()
	}
	if opts.Buckets < 4 {
		opts.Buckets = 4
	}
	if opts.Size < opts.Buckets {
		opts.Size = opts.Buckets * 100
	}

	opts.Size += int(math.Ceil(float64(opts.Size) * 10.0 / 100.0))
	bucketSize := (opts.Size / opts.Buckets) + (opts.Size % opts.Buckets)

	buckets := make([]*LRU, opts.Buckets)
	for i := 0; i < opts.Buckets; i++ {
		buckets[i] = NewLRU(&LRUOptions{
			Size:          bucketSize,
			DefaultExpiry: opts.DefaultExpiry,
			Name:          opts.Name,
		})
	}

	return &LRUStriped{
		buckets: buckets,
		name:    opts.Name,
	}
}

func (l *LRUStriped) hashKey(key string) uint64 {
	return xxhash.Sum64String(key)
}

func (l *LRUStriped) keyBucket(key string) *LRU {
	return l.buckets[l.hashKey(key)%uint64(len(l.buckets))]
}

func (l *LRUStriped) Set(key string, value any, ttl time.Duration) {
	l.keyBucket(key).Set(key, value, ttl)
}

func (l *LRUStriped) Get(key string) (any, bool) {
	return l.keyBucket(key).Get(key)
}

func (l *LRUStriped) Delete(key string) {
	l.keyBucket(key).Delete(key)
}

func (l *LRUStriped) DeleteByPrefix(prefix string) {
	for _, bucket := range l.buckets {
		bucket.DeleteByPrefix(prefix)
	}
}

func (l *LRUStriped) Clear() {
	for _, bucket := range l.buckets {
		bucket.Clear()
	}
}

func (l *LRUStriped) Exists(key string) bool {
	return l.keyBucket(key).Exists(key)
}

func (l *LRUStriped) Increment(key string) int64 {
	return l.keyBucket(key).Increment(key)
}

func (l *LRUStriped) IncrementBy(key string, value int64) int64 {
	return l.keyBucket(key).IncrementBy(key, value)
}

func (l *LRUStriped) Len() int {
	var total int
	for _, bucket := range l.buckets {
		total += bucket.Len()
	}
	return total
}

func (l *LRUStriped) Name() string {
	return l.name
}

func (l *LRUStriped) DeleteMulti(keys []string) error {
	for _, key := range keys {
		l.Delete(key)
	}
	return nil
}

func (l *LRUStriped) DeleteBatch(keys []string) {
	for _, key := range keys {
		l.Delete(key)
	}
}

func (l *LRUStriped) SetRaw(key string, raw []byte, ttl time.Duration) {
	l.keyBucket(key).SetRaw(key, raw, ttl)
}

func (l *LRUStriped) GetMulti(keys []string, values []any) []error {
	errs := make([]error, len(values))
	for i, key := range keys {
		if val, ok := l.Get(key); ok {
			values[i] = val
		} else {
			errs[i] = ErrKeyNotFound
		}
	}
	return errs
}

func (l *LRUStriped) SetMulti(items map[string]any, ttl time.Duration) error {
	for key, value := range items {
		l.Set(key, value, ttl)
	}
	return nil
}

func (l *LRUStriped) Keys() []string {
	keys := make([]string, 0)
	for _, bucket := range l.buckets {
		bucket.lock.RLock()
		for key, ent := range bucket.items {
			e := ent.Value.(*lruEntry)
			if e.generation == bucket.currentGeneration && (e.expires.IsZero() || !time.Now().After(e.expires)) {
				keys = append(keys, key)
			}
		}
		bucket.lock.RUnlock()
	}
	return keys
}

type LRUStripedInterface interface {
	Set(key string, value any, ttl time.Duration)
	Get(key string) (any, bool)
	Delete(key string)
	DeleteByPrefix(prefix string)
	Clear()
	Exists(key string) bool
	Increment(key string) int64
	IncrementBy(key string, value int64) int64
	Len() int
	Name() string
	DeleteMulti(keys []string) error
	GetMulti(keys []string, values []any) []error
	SetMulti(items map[string]any, ttl time.Duration) error
	Keys() []string
}

var _ LRUStripedInterface = (*LRUStriped)(nil)