go/kvmemcache/cache.go

package kvmemcache

import (
	"container/list"
	"sync"
	"time"

	"git.crumpington.com/lib/keyedmutex"
)

type Cache[K comparable, V any] struct {
	updateLock keyedmutex.KeyedMutex[K]
	src        func(K) (V, error)
	ttl        time.Duration
	maxSize    int

	// Lock protects variables below.
	lock  sync.Mutex
	cache map[K]*list.Element
	ll    *list.List
	stats Stats
}

type lruItem[K comparable, V any] struct {
	key       K
	createdAt time.Time
	value     V
	err       error
}

type Config[K comparable, V any] struct {
	MaxSize int
	TTL     time.Duration // Zero to ignore.
	Src     func(K) (V, error)
}

func New[K comparable, V any](conf Config[K, V]) *Cache[K, V] {
	return &Cache[K, V]{
		updateLock: keyedmutex.New[K](),
		src:        conf.Src,
		ttl:        conf.TTL,
		maxSize:    conf.MaxSize,
		lock:       sync.Mutex{},
		cache:      make(map[K]*list.Element, conf.MaxSize+1),
		ll:         list.New(),
	}
}

func (c *Cache[K, V]) Get(key K) (V, error) {
	ok, val, err := c.get(key)
	if ok {
		return val, err
	}

	return c.load(key)
}

func (c *Cache[K, V]) Evict(key K) {
	c.lock.Lock()
	defer c.lock.Unlock()
	c.evict(key)
}

func (c *Cache[K, V]) Stats() Stats {
	c.lock.Lock()
	defer c.lock.Unlock()
	return c.stats
}

func (c *Cache[K, V]) put(key K, value V, err error) {
	c.lock.Lock()
	defer c.lock.Unlock()

	c.stats.Misses++

	c.cache[key] = c.ll.PushFront(lruItem[K, V]{
		key:       key,
		createdAt: time.Now(),
		value:     value,
		err:       err,
	})

	if c.maxSize != 0 && len(c.cache) > c.maxSize {
		li := c.ll.Back()
		c.ll.Remove(li)
		delete(c.cache, li.Value.(lruItem[K, V]).key)
	}
}

func (c *Cache[K, V]) evict(key K) {
	elem := c.cache[key]
	if elem != nil {
		delete(c.cache, key)
		c.ll.Remove(elem)
	}
}

func (c *Cache[K, V]) get(key K) (ok bool, val V, err error) {
	c.lock.Lock()
	defer c.lock.Unlock()

	li := c.cache[key]
	if li == nil {
		return false, val, nil
	}

	item := li.Value.(lruItem[K, V])
	// Maybe evict.
	if c.ttl != 0 && time.Since(item.createdAt) > c.ttl {
		c.evict(key)
		return false, val, nil
	}

	c.stats.Hits++

	c.ll.MoveToFront(li)
	return true, item.value, item.err
}

func (c *Cache[K, V]) load(key K) (V, error) {
	c.updateLock.Lock(key)
	defer c.updateLock.Unlock(key)

	// Check again in case we lost the update race.
	ok, val, err := c.get(key)
	if ok {
		return val, err
	}

	// Won the update race.
	val, err = c.src(key)
	c.put(key, val, err)
	return val, err
}