generic.go

package loadingcache

import (
	"context"
	"sync"
	"time"

	"github.com/devzero-inc/loadingcache/stats"
	"github.com/pkg/errors"
	"go.uber.org/atomic"
)

// genericCache is an implementation of a cache where keys and values are
// of type interface{}
type genericCache[Key comparable, Value any] struct {
	cacheOptions[Key, Value]

	data     map[Key]*cacheEntry[Key, Value]
	dataLock sync.RWMutex

	done         chan struct{}
	backgroundWg sync.WaitGroup

	stats *stats.InternalStats
}

func (g *genericCache[K, V]) isExpired(entry *cacheEntry[K, V]) bool {
	lastRead := entry.lastRead.Load()
	lastWrite := entry.lastWrite.Load()

	if g.expiresAfterRead() && lastRead.Add(g.expireAfterRead).Before(g.clock.Now()) {
		return true
	}
	if g.expiresAfterWrite() && lastWrite.Add(g.expireAfterWrite).Before(g.clock.Now()) {
		return true
	}
	if entry.isTombstone() && lastWrite.Add(entry.tombstoneTTL).Before(g.clock.Now()) {
		return true
	}

	return false
}

func (g *genericCache[K, V]) Get(ctx context.Context, key K) (V, error) {
	g.dataLock.RLock()

	entry, exists := g.data[key]
	if !exists {
		g.dataLock.RUnlock()

		val, err := g.load(ctx, key)
		return val, errors.Wrap(err, "")
	}
	// Create a copy of the value to return to avoid concurrent updates
	toReturn := entry.value

	g.dataLock.RUnlock()

	if g.isExpired(entry) {
		g.concurrentEvict(key, RemovalReasonExpired)
		val, err := g.load(ctx, key)
		return val, errors.Wrap(err, "")
	}

	if entry.isTombstone() {
		return toReturn, errors.Wrap(ErrTombstone, "")
	}

	entry.lastRead.Store(g.clock.Now())
	g.stats.Hit()
	return toReturn, nil
}

func (g *genericCache[K, V]) load(ctx context.Context, key K) (V, error) {
	var Nil V

	g.dataLock.Lock()
	defer g.dataLock.Unlock()

	// It is possible that another call loaded the value for this key.
	// Let's do a double check if that was the case, since we have
	// the lock.
	if val, exists := g.data[key]; exists {
		if val.isTombstone() {
			g.stats.Miss()
			return Nil, errors.Wrap(ErrTombstone, "")
		}

		g.stats.Hit()
		return val.value, nil
	}

	if g.loadFunc == nil {
		g.stats.Miss()
		return Nil, errors.Wrap(ErrKeyNotFound, "")
	}

	loadStartTime := g.clock.Now()
	val, err := g.loadFunc(ctx, key)
	if err != nil {
		switch err.(type) {
		case LoadingError:
			g.stats.LoadError()
			e := err.(loadingError)
			g.internalPutTombstone(key, e.CacheFor())
			return Nil, errors.Wrapf(err, "failed to load key %v", key)
		default:
			g.stats.LoadError()
			return Nil, errors.Wrapf(err, "failed to load key %v", key)
		}
	}

	g.stats.LoadTime(g.clock.Now().Sub(loadStartTime))
	g.stats.LoadSuccess()
	g.internalPut(key, val)
	return val, nil
}

func (g *genericCache[K, V]) concurrentEvict(key K, reason RemovalReason) {
	g.dataLock.Lock()
	defer g.dataLock.Unlock()
	g.evict(key, reason)
}

func (g *genericCache[K, V]) runBackgroundEvict() {
	ticker := g.clock.Ticker(g.backgroundEvictFrequency)
	defer ticker.Stop()
	defer g.backgroundWg.Done()
	for {
		select {
		case <-g.done:
			return
		case <-ticker.C:
			g.backgroundEvict()
		}
	}
}

// backgroundEvict performs a scan of the cache in search for expired entries
// and evicts them
func (g *genericCache[K, V]) backgroundEvict() {
	g.dataLock.Lock()
	defer g.dataLock.Unlock()
	for key := range g.data {
		entry := g.data[key]
		if g.isExpired(entry) {
			// TODO: There's a possibility that we want to evict
			// in a go routine so we can get through
			// all expired entries as fast as possible without
			// having to sequentially wait for removal listeners.
			g.evict(entry.key, RemovalReasonExpired)
		}
	}
}

func (g *genericCache[K, V]) evict(key K, reason RemovalReason) {
	val, exists := g.data[key]
	if !exists {
		return
	}
	g.stats.Eviction()
	delete(g.data, key)

	if len(g.removalListeners) == 0 {
		return
	}
	notification := RemovalNotification[K, V]{
		Key:    key,
		Value:  val.value,
		Reason: reason,
	}
	// Each removal listener is called on its own goroutine
	// so a slow one does not affect the others.
	// This could potentially be early optimization, but seems
	// simple enough.
	var listenerWg sync.WaitGroup
	listenerWg.Add(len(g.removalListeners))
	for i := range g.removalListeners {
		listener := g.removalListeners[i]
		go func() {
			defer listenerWg.Done()
			listener(notification)
		}()
	}
	listenerWg.Wait()
}

// internalPut actually saves the values into the internal structures.
// It does not handle any synchronization, leaving that to the caller.
func (g *genericCache[K, V]) internalPut(key K, value V) {
	if g.maxSize > 0 && int32(len(g.data)) >= g.maxSize {
		// If eviction is needed it currently removes a random entry,
		// since maps do not have a deterministic order.
		// TODO: Apply smarter eviction policies if available
		for toEvict := range g.data {
			g.evict(toEvict, RemovalReasonSize)
			break
		}
	}
	g.data[key] = &cacheEntry[K, V]{
		key:       key,
		value:     value,
		lastRead:  atomic.NewTime(g.clock.Now()),
		lastWrite: atomic.NewTime(g.clock.Now()),
	}
}

// internalPutTombstone actually saves the values into the internal structures.
// It does not handle any synchronization, leaving that to the caller.
func (g *genericCache[K, V]) internalPutTombstone(key K, ttl time.Duration) {
	if g.maxSize > 0 && int32(len(g.data)) >= g.maxSize {
		// If eviction is needed it currently removes a random entry,
		// since maps do not have a deterministic order.
		// TODO: Apply smarter eviction policies if available
		for toEvict := range g.data {
			g.evict(toEvict, RemovalReasonSize)
			break
		}
	}
	g.data[key] = &cacheEntry[K, V]{
		key:          key,
		tombstone:    true,
		tombstoneTTL: ttl,
		lastRead:     atomic.NewTime(g.clock.Now()),
		lastWrite:    atomic.NewTime(g.clock.Now()),
	}
}

// preWriteCleanup does a pass through all entries to assess if any are expired
// and should be removed.
//
// If background cleanup os enabled, this becomes a noop.
func (g *genericCache[K, V]) preWriteCleanup() {
	if g.backgroundEvictFrequency > 0 {
		return
	}
	for key := range g.data {
		if g.isExpired(g.data[key]) {
			g.evict(key, RemovalReasonExpired)
		}
	}
}

func (g *genericCache[K, V]) Put(key K, value V) {
	g.dataLock.Lock()
	defer g.dataLock.Unlock()
	g.preWriteCleanup()
	if _, exists := g.data[key]; exists {
		g.evict(key, RemovalReasonReplaced)
	}
	g.internalPut(key, value)
}

func (g *genericCache[K, V]) Invalidate(keys ...K) {
	g.dataLock.Lock()
	defer g.dataLock.Unlock()
	for _, k := range keys {
		delete(g.data, k)
	}
}

func (g *genericCache[K, V]) InvalidateAll() {
	g.dataLock.Lock()
	defer g.dataLock.Unlock()
	for key := range g.data {
		delete(g.data, key)
	}
}

func (g *genericCache[K, V]) Close() {
	close(g.done)
	// Ensure that we wait for all background tasks to complete.
	g.backgroundWg.Wait()
}

func (g *genericCache[K, V]) Stats() Stats {
	return g.stats
}