GoToSocial/vendor/codeberg.org/gruf/go-cache/v3/ttl/ttl.go

package ttl

import (
	"sync"
	"time"

	"codeberg.org/gruf/go-maps"
)

// Entry represents an item in the cache, with it's currently calculated Expiry time.
type Entry[Key comparable, Value any] struct {
	Key    Key
	Value  Value
	Expiry time.Time
}

// Cache is the underlying Cache implementation, providing both the base Cache interface and unsafe access to underlying map to allow flexibility in building your own.
type Cache[Key comparable, Value any] struct {
	// TTL is the cache item TTL.
	TTL time.Duration

	// Evict is the hook that is called when an item is evicted from the cache.
	Evict func(*Entry[Key, Value])

	// Invalid is the hook that is called when an item's data in the cache is invalidated, includes Add/Set.
	Invalid func(*Entry[Key, Value])

	// Cache is the underlying hashmap used for this cache.
	Cache maps.LRUMap[Key, *Entry[Key, Value]]

	// stop is the eviction routine cancel func.
	stop func()

	// pool is a memory pool of entry objects.
	pool []*Entry[Key, Value]

	// Embedded mutex.
	sync.Mutex
}

// New returns a new initialized Cache with given initial length, maximum capacity and item TTL.
func New[K comparable, V any](len, cap int, ttl time.Duration) *Cache[K, V] {
	c := new(Cache[K, V])
	c.Init(len, cap, ttl)
	return c
}

// Init will initialize this cache with given initial length, maximum capacity and item TTL.
func (c *Cache[K, V]) Init(len, cap int, ttl time.Duration) {
	if ttl <= 0 {
		// Default duration
		ttl = time.Second * 5
	}
	c.TTL = ttl
	c.SetEvictionCallback(nil)
	c.SetInvalidateCallback(nil)
	c.Cache.Init(len, cap)
}

// Start: implements cache.Cache's Start().
func (c *Cache[K, V]) Start(freq time.Duration) (ok bool) {
	// Nothing to start
	if freq <= 0 {
		return false
	}

	// Safely start
	c.Lock()

	if ok = c.stop == nil; ok {
		// Not yet running, schedule us
		c.stop = schedule(c.Sweep, freq)
	}

	// Done with lock
	c.Unlock()

	return
}

// Stop: implements cache.Cache's Stop().
func (c *Cache[K, V]) Stop() (ok bool) {
	// Safely stop
	c.Lock()

	if ok = c.stop != nil; ok {
		// We're running, cancel evicts
		c.stop()
		c.stop = nil
	}

	// Done with lock
	c.Unlock()

	return
}

// Sweep attempts to evict expired items (with callback!) from cache.
func (c *Cache[K, V]) Sweep(now time.Time) {
	var after int

	// Sweep within lock
	c.Lock()
	defer c.Unlock()

	// Sentinel value
	after = -1

	// The cache will be ordered by expiry date, we iterate until we reach the index of
	// the youngest item that hsa expired, as all succeeding items will also be expired.
	c.Cache.RangeIf(0, c.Cache.Len(), func(i int, _ K, item *Entry[K, V]) bool {
		if now.After(item.Expiry) {
			after = i

			// All older than this (including) can be dropped
			return false
		}

		// Continue looping
		return true
	})

	if after == -1 {
		// No Truncation needed
		return
	}

	// Truncate items, calling eviction hook
	c.truncate(c.Cache.Len()-after, c.Evict)
}

// SetEvictionCallback: implements cache.Cache's SetEvictionCallback().
func (c *Cache[K, V]) SetEvictionCallback(hook func(*Entry[K, V])) {
	// Ensure non-nil hook
	if hook == nil {
		hook = func(*Entry[K, V]) {}
	}

	// Update within lock
	c.Lock()
	defer c.Unlock()

	// Update hook
	c.Evict = hook
}

// SetInvalidateCallback: implements cache.Cache's SetInvalidateCallback().
func (c *Cache[K, V]) SetInvalidateCallback(hook func(*Entry[K, V])) {
	// Ensure non-nil hook
	if hook == nil {
		hook = func(*Entry[K, V]) {}
	}

	// Update within lock
	c.Lock()
	defer c.Unlock()

	// Update hook
	c.Invalid = hook
}

// SetTTL: implements cache.Cache's SetTTL().
func (c *Cache[K, V]) SetTTL(ttl time.Duration, update bool) {
	if ttl < 0 {
		panic("ttl must be greater than zero")
	}

	// Update within lock
	c.Lock()
	defer c.Unlock()

	// Set updated TTL
	diff := ttl - c.TTL
	c.TTL = ttl

	if update {
		// Update existing cache entries with new expiry time
		c.Cache.Range(0, c.Cache.Len(), func(i int, _ K, item *Entry[K, V]) {
			item.Expiry = item.Expiry.Add(diff)
		})
	}
}

// Get: implements cache.Cache's Get().
func (c *Cache[K, V]) Get(key K) (V, bool) {
	// Read within lock
	c.Lock()
	defer c.Unlock()

	// Check for item in cache
	item, ok := c.Cache.Get(key)
	if !ok {
		var value V
		return value, false
	}

	// Update item expiry and return
	item.Expiry = time.Now().Add(c.TTL)
	return item.Value, true
}

// Add: implements cache.Cache's Add().
func (c *Cache[K, V]) Add(key K, value V) bool {
	// Write within lock
	c.Lock()
	defer c.Unlock()

	// Check if already exists
	item, ok := c.Cache.Get(key)
	if ok {
		return false
	}

	// Alloc new item
	item = c.alloc()
	item.Key = key
	item.Value = value
	item.Expiry = time.Now().Add(c.TTL)

	var hook func(K, *Entry[K, V])

	if c.Evict != nil {
		// Pass evicted entry to user hook
		hook = func(_ K, item *Entry[K, V]) {
			c.Evict(item)
		}
	}

	// Place new item in the map with hook
	c.Cache.SetWithHook(key, item, hook)

	if c.Invalid != nil {
		// invalidate old
		c.Invalid(item)
	}

	return true
}

// Set: implements cache.Cache's Set().
func (c *Cache[K, V]) Set(key K, value V) {
	// Write within lock
	c.Lock()
	defer c.Unlock()

	// Check if already exists
	item, ok := c.Cache.Get(key)

	if !ok {
		var hook func(K, *Entry[K, V])

		// Allocate new item
		item = c.alloc()
		item.Key = key

		if c.Evict != nil {
			// Pass evicted entry to user hook
			hook = func(_ K, item *Entry[K, V]) {
				c.Evict(item)
			}
		}

		// Place new item in the map with hook
		c.Cache.SetWithHook(key, item, hook)
	}

	if c.Invalid != nil {
		// invalidate old
		c.Invalid(item)
	}

	// Update the item value + expiry
	item.Expiry = time.Now().Add(c.TTL)
	item.Value = value
}

// CAS: implements cache.Cache's CAS().
func (c *Cache[K, V]) CAS(key K, old V, new V, cmp func(V, V) bool) bool {
	// CAS within lock
	c.Lock()
	defer c.Unlock()

	// Check for item in cache
	item, ok := c.Cache.Get(key)
	if !ok || !cmp(item.Value, old) {
		return false
	}

	if c.Invalid != nil {
		// invalidate old
		c.Invalid(item)
	}

	// Update item + Expiry
	item.Value = new
	item.Expiry = time.Now().Add(c.TTL)

	return ok
}

// Swap: implements cache.Cache's Swap().
func (c *Cache[K, V]) Swap(key K, swp V) V {
	// Swap within lock
	c.Lock()
	defer c.Unlock()

	// Check for item in cache
	item, ok := c.Cache.Get(key)
	if !ok {
		var value V
		return value
	}

	if c.Invalid != nil {
		// invalidate old
		c.Invalid(item)
	}

	old := item.Value

	// update item + Expiry
	item.Value = swp
	item.Expiry = time.Now().Add(c.TTL)

	return old
}

// Has: implements cache.Cache's Has().
func (c *Cache[K, V]) Has(key K) bool {
	c.Lock()
	ok := c.Cache.Has(key)
	c.Unlock()
	return ok
}

// Invalidate: implements cache.Cache's Invalidate().
func (c *Cache[K, V]) Invalidate(key K) bool {
	// Delete within lock
	c.Lock()
	defer c.Unlock()

	// Check if we have item with key
	item, ok := c.Cache.Get(key)
	if !ok {
		return false
	}

	// Remove from cache map
	_ = c.Cache.Delete(key)

	if c.Invalid != nil {
		// Invalidate item
		c.Invalid(item)
	}

	// Return item to pool
	c.free(item)

	return true
}

// Clear: implements cache.Cache's Clear().
func (c *Cache[K, V]) Clear() {
	c.Lock()
	defer c.Unlock()
	c.truncate(c.Cache.Len(), c.Invalid)
}

// Len: implements cache.Cache's Len().
func (c *Cache[K, V]) Len() int {
	c.Lock()
	l := c.Cache.Len()
	c.Unlock()
	return l
}

// Cap: implements cache.Cache's Cap().
func (c *Cache[K, V]) Cap() int {
	c.Lock()
	l := c.Cache.Cap()
	c.Unlock()
	return l
}

// truncate will call Cache.Truncate(sz), and if provided a hook will temporarily store deleted items before passing them to the hook. This is required in order to prevent cache writes during .Truncate().
func (c *Cache[K, V]) truncate(sz int, hook func(*Entry[K, V])) {
	if hook == nil {
		// No hook was provided, we can simply truncate and free items immediately.
		c.Cache.Truncate(sz, func(_ K, item *Entry[K, V]) { c.free(item) })
		return
	}

	// Store list of deleted items for later callbacks
	deleted := make([]*Entry[K, V], 0, sz)

	// Truncate and store list of deleted items
	c.Cache.Truncate(sz, func(_ K, item *Entry[K, V]) {
		deleted = append(deleted, item)
	})

	// Pass each deleted to hook, then free
	for _, item := range deleted {
		hook(item)
		c.free(item)
	}
}

// alloc will acquire cache entry from pool, or allocate new.
func (c *Cache[K, V]) alloc() *Entry[K, V] {
	if len(c.pool) == 0 {
		return &Entry[K, V]{}
	}
	idx := len(c.pool) - 1
	e := c.pool[idx]
	c.pool = c.pool[:idx]
	return e
}

// free will reset entry fields and place back in pool.
func (c *Cache[K, V]) free(e *Entry[K, V]) {
	var (
		zk K
		zv V
	)
	e.Key = zk
	e.Value = zv
	e.Expiry = time.Time{}
	c.pool = append(c.pool, e)
}