lru_cache.go

// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package cache borrows heavily from SmallLRUCache (originally by Nathan
// Schrenk). The object maintains a doubly-linked list of elements in the
// When an element is accessed it is promoted to the head of the list, and when
// space is needed the element at the tail of the list (the least recently used
// element) is evicted.
package cache

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

// LRUCache is static and TTL based implementation of LRU cache
type LRUCache[T any] struct {
	mu sync.Mutex

	// is this a ttl based evict lru?
	timed bool

	EvictCallback func(string, T)

	// list & table of *entry objects
	list  *list.List
	table map[string]*list.Element

	// How many items we are limiting the cache to.
	capacity uint64
}

// Item is a container for items listing
type Item[T any] struct {
	Key   string
	Value T
}

type entry[T any] struct {
	key          string
	value        T
	timeAccessed time.Time
}

// NewLRUCache make a new instance of static LRU cache where eviction is based
// on number of items in the cache
func NewLRUCache[T any](capacity uint64) *LRUCache[T] {
	return &LRUCache[T]{
		list:     list.New(),
		table:    make(map[string]*list.Element),
		capacity: capacity,
	}
}

// NewTimeEvictLRU creates a new Lru, which evicts units after a time period of non use
//
//	@ttl is time an entry exists before it is evicted
//	@checkEvery is duration to check every x
//
//	l := NewTimeEvictLru(10000, 200*time.Millisecond, 50*time.Millisecond)
//	var evictCt int = 0
//	l.EvictCallback = func(key string, v Value) {
//		evictCt++
//	}
func NewTimeEvictLRU[T any](capacity uint64, ttl, checkEvery time.Duration) *LRUCache[T] {

	lru := LRUCache[T]{
		timed:    true,
		list:     list.New(),
		table:    make(map[string]*list.Element),
		capacity: capacity,
	}

	timer := time.NewTicker(checkEvery)
	go func() {
		for now := range timer.C {
			TimeEvict(&lru, now, ttl)
		}
	}()

	return &lru
}

// TimeEvict this time eviction is called by a scheduler for background, periodic evictions
func TimeEvict[T any](lru *LRUCache[T], now time.Time, ttl time.Duration) {
	lru.mu.Lock()
	defer lru.mu.Unlock()

	if lru.list.Len() < 1 {
		return
	}

	var n *entry[T]
	var el *list.Element
	hasEvictCallback := lru.EvictCallback != nil

	el = lru.list.Back()

	for {
		if el == nil {
			break
		}
		n = el.Value.(*entry[T])
		if now.Sub(n.timeAccessed) > ttl {
			// the Difference is greater than max TTL so we need to evict this entry
			// first grab the next entry (as this is about to dissappear)
			el = el.Prev()
			lru.remove(n.key)
			if hasEvictCallback {
				lru.EvictCallback(n.key, n.value)
			}
		} else {
			// since they are stored in time order, as soon as we find
			// first item newer than ttl window, we are safe to bail
			break
		}
	}
}

func zero[T any]() T {
	var result T
	return result
}

// Get gets the value identified by a key and ok flag telling if it was found.
// It updates the access time of that item.
func (lru *LRUCache[T]) Get(key string) (v T, ok bool) {
	lru.mu.Lock()
	defer lru.mu.Unlock()

	element := lru.table[key]
	if element == nil {
		return zero[T](), false
	}
	lru.moveToFront(element)
	return element.Value.(*entry[T]).value, true
}

// Peek gets the value identified by a key and ok flag telling if it was found.
// It doesnt' update the access time of that item.
func (lru *LRUCache[T]) Peek(key string) (v T, ok bool) {
	lru.mu.Lock()
	defer lru.mu.Unlock()

	element := lru.table[key]
	if element == nil {
		return zero[T](), false
	}
	return element.Value.(*entry[T]).value, true
}

// Set stores value under a given key
func (lru *LRUCache[T]) Set(key string, value T) {
	lru.mu.Lock()
	defer lru.mu.Unlock()

	if element := lru.table[key]; element != nil {
		lru.updateInplace(element, value)
	} else {
		lru.addNew(key, value)
	}
}

// SetIfAbsent set the value to the cache if it doesn't exist.
// If it exists it just refreshes the access time.
func (lru *LRUCache[T]) SetIfAbsent(key string, value T) {
	lru.mu.Lock()
	defer lru.mu.Unlock()

	if element := lru.table[key]; element != nil {
		lru.moveToFront(element)
	} else {
		lru.addNew(key, value)
	}
}

func (lru *LRUCache[T]) remove(key string) bool {
	element := lru.table[key]
	if element == nil {
		return false
	}

	lru.list.Remove(element)
	delete(lru.table, key)
	return true
}

// Delete removes the value identified by key from the cache and
// returns the flag if it was found.
func (lru *LRUCache[T]) Delete(key string) bool {
	lru.mu.Lock()
	defer lru.mu.Unlock()
	return lru.remove(key)
}

// Clear clears the cache
func (lru *LRUCache[T]) Clear() {
	lru.mu.Lock()
	defer lru.mu.Unlock()

	lru.list.Init()
	lru.table = make(map[string]*list.Element)
}

// SetCapacity sets the capacity of the cache
func (lru *LRUCache[T]) SetCapacity(capacity uint64) {
	lru.mu.Lock()
	defer lru.mu.Unlock()

	lru.capacity = capacity
	lru.checkCapacity()
}

// Stats gets the cache statistic (length, capacity, access time of oldest item)
func (lru *LRUCache[T]) Stats() (length, capacity uint64, oldest time.Time) {
	lru.mu.Lock()
	defer lru.mu.Unlock()
	if lastElem := lru.list.Back(); lastElem != nil {
		oldest = lastElem.Value.(*entry[T]).timeAccessed
	}
	return uint64(lru.list.Len()), lru.capacity, oldest
}

// StatsJSON gets the statistic in JSON format
func (lru *LRUCache[T]) StatsJSON() string {
	if lru == nil {
		return "{}"
	}
	l, c, o := lru.Stats()
	return fmt.Sprintf("{\"length\": %v, \"capacity\": %v, \"oldestAccess\": \"%v\"}", l, c, o)
}

// Keys gets the keys as a slice of string
func (lru *LRUCache[T]) Keys() []string {
	lru.mu.Lock()
	defer lru.mu.Unlock()

	keys := make([]string, 0, lru.list.Len())
	for e := lru.list.Front(); e != nil; e = e.Next() {
		keys = append(keys, e.Value.(*entry[T]).key)
	}
	return keys
}

// Len gets the current number of items stored in the cache
func (lru *LRUCache[T]) Len() int {
	lru.mu.Lock()
	defer lru.mu.Unlock()

	return lru.list.Len()
}

// Items provides view of Items stored in the cache
func (lru *LRUCache[T]) Items() []Item[T] {
	lru.mu.Lock()
	defer lru.mu.Unlock()

	items := make([]Item[T], 0, lru.list.Len())
	for e := lru.list.Front(); e != nil; e = e.Next() {
		v := e.Value.(*entry[T])
		items = append(items, Item[T]{Key: v.key, Value: v.value})
	}
	return items
}

func (lru *LRUCache[T]) updateInplace(element *list.Element, value T) {
	element.Value.(*entry[T]).value = value
	lru.moveToFront(element)
	lru.checkCapacity()
}

func (lru *LRUCache[T]) moveToFront(element *list.Element) {
	lru.list.MoveToFront(element)
	element.Value.(*entry[T]).timeAccessed = time.Now()
}

func (lru *LRUCache[T]) addNew(key string, value T) {
	newEntry := &entry[T]{key, value, time.Now()}
	element := lru.list.PushFront(newEntry)
	lru.table[key] = element
	lru.checkCapacity()
}

func (lru *LRUCache[T]) checkCapacity() {
	cl := uint64(lru.list.Len())

	for cl > lru.capacity {
		delElem := lru.list.Back()
		delValue := delElem.Value.(*entry[T])
		lru.list.Remove(delElem)
		delete(lru.table, delValue.key)
		cl--
	}
}