cache.go

// Package loadingcache provides a way for clients to create a cache capable
// of loading values on demand, should they get cache misses.
//
// You can configure the cache to expire entries after a certain amount elapses
// since the last write and/or read.
//
// This project is heavily inspired by Guava Cache (https://github.com/google/guava/wiki/CachesExplained).
//
// All errors are wrapped by github.com/pkg/errors.Wrap. If you which to check
// the type of it, please use github.com/pkg/errors.Is.
package loadingcache

import (
	"context"
	"time"

	"github.com/benbjohnson/clock"
	"github.com/pkg/errors"
	"go.uber.org/atomic"

	"github.com/devzero-inc/loadingcache/stats"
)

// ErrKeyNotFound represents an error indicating that the key was not found
var ErrKeyNotFound error = errors.New("Key not found")

// ErrTombstone represents an error indicating that the key is a "tombstone"
var ErrTombstone error = errors.New("tombstone")

// RemovalReason is an enum describing the causes for an entry to
// be removed from the cache.
type RemovalReason string

const (
	// RemovalReasonExplicit means the entry was explicitly invalidated
	RemovalReasonExplicit RemovalReason = "EXPLICIT"

	// RemovalReasonReplaced means the entry was replaced by a new one
	RemovalReasonReplaced RemovalReason = "REPLACED"

	// RemovalReasonExpired means the entry expired, e.g. too much time
	// since last read/write.
	RemovalReasonExpired RemovalReason = "EXPIRED"

	// RemovalReasonSize means the entry was removed due to the cache size.
	RemovalReasonSize RemovalReason = "SIZE"
)

// RemovalNotification is passed to listeners everytime an entry is removed
type RemovalNotification[Key comparable, Value any] struct {
	Key    Key
	Value  Value
	Reason RemovalReason
}

// RemovalListener represents a removal listener
type RemovalListener[Key comparable, Value any] func(RemovalNotification[Key, Value])

// Cache describe the base interface to interact with a generic cache.
//
// This interface reduces all keys and values to a generic interface{}.
type Cache[Key comparable, Value any] interface {
	// Get returns the value associated with a given key. If no entry exists for
	// the provided key, loadingcache.ErrKeyNotFound is returned.
	Get(ctx context.Context, key Key) (Value, error)

	// Put adds a value to the cache identified by a key.
	// If a value already exists associated with that key, it
	// is replaced.
	Put(key Key, value Value)

	// Invalidate removes keys from the cache. If a key does not exists it is a noop.
	Invalidate(keys ...Key)

	// InvalidateAll invalidates all keys
	InvalidateAll()

	// Close cleans up any resources used by the cache
	Close()

	// Stats returns the curret stats
	Stats() Stats
}

// CacheOptions available options to initialize the cache
type cacheOptions[Key comparable, Value any] struct {
	// Clock allows passing a custom clock to be used with the cache.
	//
	// This is useful for testing, where controlling time is important.
	clock clock.Clock

	// ExpireAfterWrite configures the cache to expire entries after
	// a given duration after the last write was made to a given key.
	// For example, setting this value to 24 hours will mean that any given
	// value will expire 24 hours after it was last written to.
	expireAfterWrite time.Duration

	// ExpireAfterRead configures the cache to expire entries after
	// a given duration after the last read was made. For example, setting
	// this value to 24 hours will mean that any given value will expire if
	// 24 hours has elapsed since this key was last read.
	expireAfterRead time.Duration

	// Load configures a loading function
	loadFunc LoadFunc[Key, Value]

	// MaxSize limits the number of entries allowed in the cache.
	// If the limit is achieved, an eviction process will take place,
	// this means that eviction policies will be executed such as write
	// time, read time or random entry if no evection policy frees up
	// space.
	//
	// If the cache is sharded, MaxSize is applied to each shard,
	// meaning that the overall capacity will be MaxSize * ShardCount.
	maxSize int32

	// RemovalListeners configures a removal listeners
	removalListeners []RemovalListener[Key, Value]

	// ShardCount indicates how many shards will be used by the cache.
	// This allows some degree of parallelism in read and writing to the cache.
	//
	// If the shard count is greater than 1, then HashCodeFunc must be provided
	// otherwise the constructor will panic.
	shardCount int

	// HashCodeFunc is a function that produces a hashcode of the key.
	//
	// See https://docs.oracle.com/en/java/javase/15/docs/api/java.base/java/lang/Object.html#hashCode()
	// for best practices surrounding hash code functions.
	hashCodeFunc func(key Key) uint64

	// BackgroundEvictFrequency controls if a background go routine should be created
	// which automatically evicts entries that have expired. If not speficied
	// no background goroutine will be created.
	//
	// The background go routine runs with the provided frequency.
	// To avoid go routine leaks, use the close function when you're done with the cache.
	backgroundEvictFrequency time.Duration

	// Name is the name of the cache. This is used for metrics and logging.
	name string
}

func (c cacheOptions[K, V]) expiresAfterRead() bool {
	return c.expireAfterRead > 0
}

func (c cacheOptions[K, V]) expiresAfterWrite() bool {
	return c.expireAfterWrite > 0
}

// CacheOption describes an option that can configure the cache
type CacheOption[Key comparable, Value any] func(cacheOptions[Key, Value]) cacheOptions[Key, Value]

// LoadFunc represents a function that given a key, it returns a value or an error.
type LoadFunc[Key comparable, Value any] func(context.Context, Key) (Value, error)

type cacheEntry[Key comparable, Value any] struct {
	key          Key
	value        Value
	tombstone    bool
	tombstoneTTL time.Duration
	lastRead     *atomic.Time
	lastWrite    *atomic.Time
}

func (c cacheEntry[Key, Value]) isTombstone() bool {
	return c.tombstone && c.tombstoneTTL != 0
}

// New instantiates a new cache
func New[Key comparable, Value any](opts ...CacheOption[Key, Value]) Cache[Key, Value] {
	options := cacheOptions[Key, Value]{}

	for _, opt := range opts {
		options = opt(options)
	}

	return newCache(options)
}

func newCache[Key comparable, Value any](options cacheOptions[Key, Value]) Cache[Key, Value] {
	if options.clock == nil {
		options.clock = clock.New()
	}

	if options.shardCount < 0 {
		panic("shard count must be non-negative")
	}

	switch options.shardCount {
	case 0, 1:
		c := &genericCache[Key, Value]{
			cacheOptions: options,
			data:         map[Key]*cacheEntry[Key, Value]{},
			done:         make(chan struct{}),
			stats:        &stats.InternalStats{},
		}
		if options.backgroundEvictFrequency > 0 {
			c.backgroundWg.Add(1)
			go c.runBackgroundEvict()
		}
		return c
	default:
		if options.hashCodeFunc == nil {
			panic("cannot have a sharded cache without a hashcode function")
		}
		singleShardOptions := options
		singleShardOptions.shardCount = 1
		s := &shardedCache[Key, Value]{
			cacheOptions: options,
			shards:       make([]Cache[Key, Value], options.shardCount),
		}
		for i := range options.shardCount {
			s.shards[i] = newCache(singleShardOptions)
		}
		return s
	}
}