180 lines
4.0 KiB
Go
180 lines
4.0 KiB
Go
// Package cache provides a bounded, TTL-expiring LRU map.
|
|
//
|
|
// It exists for the one place that genuinely needs eviction: the failed-auth
|
|
// rate limiter, whose key space is attacker-controlled and must not be allowed
|
|
// to grow without bound. The other caches in this server deliberately do not
|
|
// use it — the project registry and deployment manifests are copy-on-write
|
|
// snapshots read without any lock at all, and adding an LRU there would put a
|
|
// contended mutex on the hot serving path to solve a problem that does not
|
|
// exist.
|
|
package cache
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type node[K comparable, V any] struct {
|
|
key K
|
|
val V
|
|
expires time.Time
|
|
prev, next *node[K, V]
|
|
}
|
|
|
|
// Cache maps K to V with a size cap and a per-entry TTL. It is safe for
|
|
// concurrent use. Entries are evicted when the cap is exceeded (least recently
|
|
// used first) or when their TTL passes, whichever comes first.
|
|
//
|
|
// The zero value is not usable; call New.
|
|
type Cache[K comparable, V any] struct {
|
|
mu sync.Mutex
|
|
m map[K]*node[K, V]
|
|
head *node[K, V] // most recently used
|
|
tail *node[K, V] // least recently used
|
|
max int
|
|
ttl time.Duration
|
|
|
|
// now is swapped out by tests. Production always uses time.Now.
|
|
now func() time.Time
|
|
}
|
|
|
|
// New returns a cache holding at most max entries for at most ttl each.
|
|
func New[K comparable, V any](max int, ttl time.Duration) *Cache[K, V] {
|
|
if max < 1 {
|
|
max = 1
|
|
}
|
|
return &Cache[K, V]{
|
|
m: make(map[K]*node[K, V]),
|
|
max: max,
|
|
ttl: ttl,
|
|
now: time.Now,
|
|
}
|
|
}
|
|
|
|
// Get returns the value for k, refreshing its recency. A value whose TTL has
|
|
// passed is reported as absent and dropped.
|
|
func (c *Cache[K, V]) Get(k K) (V, bool) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
|
|
n, ok := c.m[k]
|
|
if !ok {
|
|
var zero V
|
|
return zero, false
|
|
}
|
|
if !c.now().Before(n.expires) {
|
|
c.remove(n)
|
|
var zero V
|
|
return zero, false
|
|
}
|
|
c.moveToFront(n)
|
|
return n.val, true
|
|
}
|
|
|
|
// Put inserts or replaces the value for k and resets its TTL.
|
|
func (c *Cache[K, V]) Put(k K, v V) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
c.put(k, v)
|
|
}
|
|
|
|
// GetOrCreate returns the existing value for k, or stores and returns the one
|
|
// newVal produces.
|
|
//
|
|
// The whole operation happens under the lock, which is what makes it usable for
|
|
// the rate limiter: two concurrent requests from the same address must share
|
|
// one token bucket, and a Get-then-Put pair would hand each of them its own.
|
|
// newVal must not call back into the cache.
|
|
func (c *Cache[K, V]) GetOrCreate(k K, newVal func() V) V {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
|
|
if n, ok := c.m[k]; ok {
|
|
if c.now().Before(n.expires) {
|
|
c.moveToFront(n)
|
|
return n.val
|
|
}
|
|
c.remove(n)
|
|
}
|
|
v := newVal()
|
|
c.put(k, v)
|
|
return v
|
|
}
|
|
|
|
// Delete drops k if present.
|
|
func (c *Cache[K, V]) Delete(k K) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if n, ok := c.m[k]; ok {
|
|
c.remove(n)
|
|
}
|
|
}
|
|
|
|
// Len reports the number of entries, including any whose TTL has passed but
|
|
// that have not been touched since. It is meant for tests and diagnostics.
|
|
func (c *Cache[K, V]) Len() int {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
return len(c.m)
|
|
}
|
|
|
|
// ---------------------------------------------------------------- internals
|
|
// All of these require c.mu.
|
|
|
|
func (c *Cache[K, V]) put(k K, v V) {
|
|
if n, ok := c.m[k]; ok {
|
|
n.val = v
|
|
n.expires = c.now().Add(c.ttl)
|
|
c.moveToFront(n)
|
|
return
|
|
}
|
|
n := &node[K, V]{key: k, val: v, expires: c.now().Add(c.ttl)}
|
|
c.m[k] = n
|
|
c.pushFront(n)
|
|
for len(c.m) > c.max {
|
|
c.remove(c.tail)
|
|
}
|
|
}
|
|
|
|
func (c *Cache[K, V]) pushFront(n *node[K, V]) {
|
|
n.prev = nil
|
|
n.next = c.head
|
|
if c.head != nil {
|
|
c.head.prev = n
|
|
}
|
|
c.head = n
|
|
if c.tail == nil {
|
|
c.tail = n
|
|
}
|
|
}
|
|
|
|
func (c *Cache[K, V]) moveToFront(n *node[K, V]) {
|
|
if c.head == n {
|
|
return
|
|
}
|
|
c.unlink(n)
|
|
c.pushFront(n)
|
|
}
|
|
|
|
func (c *Cache[K, V]) remove(n *node[K, V]) {
|
|
if n == nil {
|
|
return
|
|
}
|
|
c.unlink(n)
|
|
delete(c.m, n.key)
|
|
}
|
|
|
|
func (c *Cache[K, V]) unlink(n *node[K, V]) {
|
|
if n.prev != nil {
|
|
n.prev.next = n.next
|
|
} else if c.head == n {
|
|
c.head = n.next
|
|
}
|
|
if n.next != nil {
|
|
n.next.prev = n.prev
|
|
} else if c.tail == n {
|
|
c.tail = n.prev
|
|
}
|
|
n.prev, n.next = nil, nil
|
|
}
|