backend/internal/cache/lru.go

package cache

import (
	"container/list"
	"encoding/json"
	"sync"
	"time"
)

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

type LRU struct {
	lock              sync.RWMutex
	size              int
	len               int
	currentGeneration int64
	evictList         *list.List
	items             map[string]*list.Element
	defaultExpiry     time.Duration
	name              string
}

type lruEntry struct {
	key        string
	value      []byte
	expires    time.Time
	generation int64
}

func NewLRU(opts *LRUOptions) *LRU {
	if opts.Size <= 0 {
		opts.Size = 1000
	}
	return &LRU{
		name:          opts.Name,
		size:          opts.Size,
		evictList:     list.New(),
		items:         make(map[string]*list.Element, opts.Size),
		defaultExpiry: opts.DefaultExpiry,
	}
}

func (l *LRU) Set(key string, value any, ttl time.Duration) {
	if ttl <= 0 {
		ttl = l.defaultExpiry
	}
	l.set(key, value, ttl)
}

func (l *LRU) set(key string, value any, ttl time.Duration) {
	var expires time.Time
	if ttl > 0 {
		expires = time.Now().Add(ttl)
	}

	buf, err := json.Marshal(value)
	if err != nil {
		return
	}

	l.lock.Lock()
	defer l.lock.Unlock()

	if ent, ok := l.items[key]; ok {
		l.evictList.MoveToFront(ent)
		e := ent.Value.(*lruEntry)
		e.value = buf
		e.expires = expires
		if e.generation != l.currentGeneration {
			e.generation = l.currentGeneration
			l.len++
		}
		return
	}

	ent := &lruEntry{key, buf, expires, l.currentGeneration}
	entry := l.evictList.PushFront(ent)
	l.items[key] = entry
	l.len++

	if l.evictList.Len() > l.size {
		l.removeElement(l.evictList.Back())
	}
}

func (l *LRU) Get(key string) (any, bool) {
	val, err := l.getItem(key)
	if err != nil {
		return nil, false
	}
	return val, true
}

func (l *LRU) getItem(key string) ([]byte, error) {
	l.lock.Lock()
	defer l.lock.Unlock()

	ent, ok := l.items[key]
	if !ok {
		return nil, ErrKeyNotFound
	}
	e := ent.Value.(*lruEntry)
	if e.generation != l.currentGeneration || (!e.expires.IsZero() && time.Now().After(e.expires)) {
		l.removeElement(ent)
		return nil, ErrKeyNotFound
	}
	l.evictList.MoveToFront(ent)
	return e.value, nil
}

func (l *LRU) Delete(key string) {
	l.lock.Lock()
	defer l.lock.Unlock()

	if ent, ok := l.items[key]; ok {
		l.removeElement(ent)
	}
}

func (l *LRU) DeleteByPrefix(prefix string) {
	l.lock.Lock()
	defer l.lock.Unlock()

	for key, ent := range l.items {
		if len(key) >= len(prefix) && key[:len(prefix)] == prefix {
			l.removeElement(ent)
		}
	}
}

func (l *LRU) Clear() {
	l.lock.Lock()
	defer l.lock.Unlock()

	l.len = 0
	l.currentGeneration++
}

func (l *LRU) Exists(key string) bool {
	l.lock.RLock()
	defer l.lock.RUnlock()

	ent, ok := l.items[key]
	if !ok {
		return false
	}
	e := ent.Value.(*lruEntry)
	return e.generation == l.currentGeneration && (e.expires.IsZero() || !time.Now().After(e.expires))
}

func (l *LRU) Increment(key string) int64 {
	return l.IncrementBy(key, 1)
}

func (l *LRU) IncrementBy(key string, value int64) int64 {
	l.lock.Lock()
	defer l.lock.Unlock()

	var current int64
	if ent, ok := l.items[key]; ok {
		e := ent.Value.(*lruEntry)
		if e.generation == l.currentGeneration && (e.expires.IsZero() || !time.Now().After(e.expires)) {
			json.Unmarshal(e.value, &current)
		}
	}

	current += value
	buf, _ := json.Marshal(current)

	if ent, ok := l.items[key]; ok {
		e := ent.Value.(*lruEntry)
		e.value = buf
		l.evictList.MoveToFront(ent)
	} else {
		ent := &lruEntry{key, buf, time.Time{}, l.currentGeneration}
		entry := l.evictList.PushFront(ent)
		l.items[key] = entry
		l.len++
	}

	return current
}

func (l *LRU) Len() int {
	l.lock.RLock()
	defer l.lock.RUnlock()
	return l.len
}

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

func (l *LRU) removeElement(e *list.Element) {
	l.evictList.Remove(e)
	kv := e.Value.(*lruEntry)
	if kv.generation == l.currentGeneration {
		l.len--
	}
	delete(l.items, kv.key)
}