fancy-dns/internal/cache/lru.go

package cache

import (
	"container/list"
	"hash/maphash"
	"sync"
	"sync/atomic"
	"time"
)

// DefaultShardCount is used when the config does not pick one.
const DefaultShardCount = 32

// LRU is a sharded, size-bounded LRU keyed by Key.
//
// Each shard owns its own map and intrusive list, so the common case (a
// lookup that hits and is moved to the front) takes one shard mutex. The
// LRU is *not* strict-LRU globally: evicting a shard's least-recently-used
// entry is a fine approximation and keeps contention low. The global entry
// budget is divided evenly between shards.
type LRU struct {
	shards []*shard
	mask   uint64

	// perShard is the per-shard entry budget. It is atomic because a config
	// reload resizes the cache from the API goroutine while queries are
	// inserting in every shard.
	perShard atomic.Int64
	// maxBytes bounds the whole LRU by accounted bytes, divided between the
	// shards the same way perShard is. 0 means unbounded. This is the budget
	// that actually protects the process: perShard alone cannot, because an
	// entry's size is set by the answer, not by the count.
	maxBytes atomic.Int64

	// counters are the hot-path stats, kept as atomics so Get does not need
	// the lock just to record a hit.
	hits      atomic.Uint64
	misses    atomic.Uint64
	staleHits atomic.Uint64
	inserts   atomic.Uint64
	evictions atomic.Uint64
	// expired drops entries found already past their stale window.
	expired atomic.Uint64
	bytes   atomic.Int64
	entries atomic.Int64

	seed maphash.Seed

	// prefetch delays revalidation without blocking Get. It is a semaphore of
	// "refreshes in flight" so a burst of stale hits triggers one fetch.
	prefetchGate chan struct{}
}

type shard struct {
	mu    sync.Mutex
	items map[Key]*list.Element
	ll    *list.List
	bytes int
}

type lruItem struct {
	key   Key
	entry *Entry
}

// NewLRU builds an LRU holding at most size entries across shards shards
// (rounded up to a power of two, minimum 1).
// SetMaxBytes bounds the LRU by accounted bytes, dividing the budget between
// shards the same way the entry budget is. Zero removes the bound. Existing
// entries are kept; a reload that tightens the budget reclaims on the next
// insert, exactly like the entry budget.
func (l *LRU) SetMaxBytes(n int64) {
	if n < 0 {
		n = 0
	}
	l.maxBytes.Store(n)
}

// MaxBytes reports the whole-LRU byte budget, or 0 when unbounded.
func (l *LRU) MaxBytes() int64 { return l.maxBytes.Load() }

func NewLRU(size, shards int) *LRU {
	if shards <= 0 {
		shards = DefaultShardCount
	}
	shards = roundPow2(shards)
	if size < shards {
		size = shards
	}
	l := &LRU{
		shards:       make([]*shard, shards),
		mask:         uint64(shards - 1),
		seed:         maphash.MakeSeed(),
		prefetchGate: make(chan struct{}, 64),
	}
	l.perShard.Store(int64(max(1, size/shards)))
	for i := range l.shards {
		l.shards[i] = &shard{
			items: make(map[Key]*list.Element, 1024),
			ll:    list.New(),
		}
	}
	return l
}

func roundPow2(n int) int {
	if n <= 1 {
		return 1
	}
	p := 1
	for p < n {
		p <<= 1
	}
	return p
}

func (l *LRU) shardFor(k Key) *shard {
	h := maphash.String(l.seed, k.Name)
	// Mix the type/scope into the high bits so A and AAAA for one name do not
	// always land in the same shard.
	h ^= uint64(k.Type)<<40 | uint64(k.Scope)<<56
	return l.shards[h&l.mask]
}

// Get returns the entry for k and reports its kind. The returned entry is
// shared and must be treated as read-only.
//
// stale controls whether an expired-but-not-yet-dead entry is returned. When
// stale is returned the caller should kick off a refresh (prefetch) and
// answer from the stale copy.
func (l *LRU) Get(k Key, now time.Time) (e *Entry, kind HitKind) {
	s := l.shardFor(k)
	s.mu.Lock()
	el, ok := s.items[k]
	if !ok {
		s.mu.Unlock()
		l.misses.Add(1)
		return nil, Miss
	}
	it := el.Value.(*lruItem)
	if now.After(it.entry.StaleUntil) {
		// Fully dead: drop it now so the map does not grow unbounded.
		s.ll.Remove(el)
		delete(s.items, k)
		s.bytes -= it.entry.Bytes
		s.mu.Unlock()
		l.bytes.Add(-int64(it.entry.Bytes))
		l.entries.Add(-1)
		l.expired.Add(1)
		l.misses.Add(1)
		return nil, Miss
	}
	s.ll.MoveToFront(el)
	e = it.entry
	kind = Fresh
	if e.Expired(now) {
		kind = Stale
	}
	s.mu.Unlock()

	e.Hits.Add(1)
	if kind == Fresh {
		l.hits.Add(1)
	} else {
		l.staleHits.Add(1)
	}
	return e, kind
}

// Put inserts or replaces an entry, evicting the shard's least-recently-used
// entries if the per-shard budget is exceeded.
func (l *LRU) Put(k Key, e *Entry) {
	if l == nil {
		return
	}
	s := l.shardFor(k)
	s.mu.Lock()
	if el, ok := s.items[k]; ok {
		it := el.Value.(*lruItem)
		s.bytes -= it.entry.Bytes
		l.bytes.Add(-int64(it.entry.Bytes))
		it.entry = e
		s.ll.MoveToFront(el)
	} else {
		el := s.ll.PushFront(&lruItem{key: k, entry: e})
		s.items[k] = el
		l.entries.Add(1)
	}
	s.bytes += e.Bytes
	l.bytes.Add(int64(e.Bytes))
	var dropped int
	budget := int(l.perShard.Load())
	byteBudget := l.maxBytes.Load() / int64(len(l.shards))
	for s.ll.Len() > budget || (byteBudget > 0 && s.bytes > int(byteBudget)) {
		oldest := s.ll.Back()
		if oldest == nil {
			break
		}
		it := oldest.Value.(*lruItem)
		s.ll.Remove(oldest)
		delete(s.items, it.key)
		s.bytes -= it.entry.Bytes
		l.bytes.Add(-int64(it.entry.Bytes))
		l.entries.Add(-1)
		dropped++
	}
	// Do not hold the lock while the global counters move: they are
	// independent atomics, so batching is a pure win.
	s.mu.Unlock()
	l.evictions.Add(uint64(dropped))
	l.inserts.Add(1)
}

// Delete removes one entry.
func (l *LRU) Delete(k Key) {
	s := l.shardFor(k)
	s.mu.Lock()
	if el, ok := s.items[k]; ok {
		it := el.Value.(*lruItem)
		s.ll.Remove(el)
		delete(s.items, k)
		s.bytes -= it.entry.Bytes
		s.mu.Unlock()
		l.bytes.Add(-int64(it.entry.Bytes))
		l.entries.Add(-1)
		return
	}
	s.mu.Unlock()
}

// Flush drops everything.
func (l *LRU) Flush() {
	for _, s := range l.shards {
		s.mu.Lock()
		n := s.ll.Len()
		b := s.bytes
		s.items = make(map[Key]*list.Element, 1024)
		s.ll.Init()
		s.bytes = 0
		s.mu.Unlock()
		l.entries.Add(-int64(n))
		l.bytes.Add(-int64(b))
	}
}

// Walk visits every live entry. fn runs without the shard lock held, so it
// must not retain the entry; entries removed meanwhile are simply skipped.
// Returning false from fn stops the walk.
func (l *LRU) Walk(fn func(Key, *Entry) bool) {
	for _, s := range l.shards {
		s.mu.Lock()
		items := make([]lruItem, 0, s.ll.Len())
		for el := s.ll.Front(); el != nil; el = el.Next() {
			it := el.Value.(*lruItem)
			items = append(items, *it)
		}
		s.mu.Unlock()
		for i := range items {
			if !fn(items[i].key, items[i].entry) {
				return
			}
		}
	}
}

// DeleteWhere visits every entry and removes those matching pred.
func (l *LRU) DeleteWhere(pred func(Key, *Entry) bool) int {
	removed := 0
	for _, s := range l.shards {
		s.mu.Lock()
		for el := s.ll.Front(); el != nil; {
			next := el.Next()
			it := el.Value.(*lruItem)
			if pred(it.key, it.entry) {
				s.ll.Remove(el)
				delete(s.items, it.key)
				s.bytes -= it.entry.Bytes
				l.bytes.Add(-int64(it.entry.Bytes))
				l.entries.Add(-1)
				removed++
			}
			el = next
		}
		s.mu.Unlock()
	}
	return removed
}

// TryPrefetch acquires a refresh slot without blocking. Callers use it to
// make sure a burst of stale hits triggers exactly one background refresh.
func (l *LRU) TryPrefetch() bool {
	select {
	case l.prefetchGate <- struct{}{}:
		return true
	default:
		return false
	}
}

// ReleasePrefetch returns a slot taken by TryPrefetch.
func (l *LRU) ReleasePrefetch() {
	select {
	case <-l.prefetchGate:
	default:
	}
}

// HitKind describes how a lookup was satisfied.
type HitKind uint8

const (
	// Miss means the answer must be resolved.
	Miss HitKind = iota
	// Fresh means the entry is inside its TTL.
	Fresh
	// Stale means the entry is past its TTL but inside the serve-stale
	// window: serve it and refresh in the background.
	Stale
)

func (h HitKind) String() string {
	switch h {
	case Fresh:
		return "hit"
	case Stale:
		return "stale"
	default:
		return "miss"
	}
}

// LRUStats is a snapshot of the LRU counters.
type LRUStats struct {
	Entries   int     `json:"entries"`
	Bytes     int64   `json:"bytes"`
	Hits      uint64  `json:"hits"`
	Misses    uint64  `json:"misses"`
	StaleHits uint64  `json:"stale_hits"`
	Inserts   uint64  `json:"inserts"`
	Evictions uint64  `json:"evictions"`
	Expired   uint64  `json:"expired"`
	HitRate   float64 `json:"hit_rate"`
	Shards    int     `json:"shards"`
	Capacity  int     `json:"capacity"`
}

// Stats snapshots the counters.
func (l *LRU) Stats() LRUStats {
	if l == nil {
		return LRUStats{}
	}
	hits := l.hits.Load()
	misses := l.misses.Load()
	total := hits + misses
	var rate float64
	if total > 0 {
		rate = float64(hits) / float64(total)
	}
	return LRUStats{
		Entries:   int(l.entries.Load()),
		Bytes:     l.bytes.Load(),
		Hits:      hits,
		Misses:    misses,
		StaleHits: l.staleHits.Load(),
		Inserts:   l.inserts.Load(),
		Evictions: l.evictions.Load(),
		Expired:   l.expired.Load(),
		HitRate:   rate,
		Shards:    len(l.shards),
		Capacity:  int(l.perShard.Load()) * len(l.shards),
	}
}

// SetCapacity resizes the per-shard budget. Existing entries are kept; the
// next insert past the new budget evicts.
func (l *LRU) SetCapacity(size, shards int) {
	if shards <= 0 {
		shards = len(l.shards)
	}
	shards = roundPow2(shards)
	budget := int64(max(1, size/shards))
	l.perShard.Store(budget)
	// A shrink has to reclaim memory now, not on the next insert: after a
	// config reload lowers the cache size, the old entries would otherwise
	// sit there until they happened to be overwritten.
	var evicted int
	for _, s := range l.shards[:min(shards, len(l.shards))] {
		s.mu.Lock()
		for s.ll.Len() > int(budget) {
			oldest := s.ll.Back()
			if oldest == nil {
				break
			}
			it, ok := oldest.Value.(*lruItem)
			if !ok {
				break
			}
			s.ll.Remove(oldest)
			delete(s.items, it.key)
			s.bytes -= it.entry.Bytes
			l.bytes.Add(-int64(it.entry.Bytes))
			l.entries.Add(-1)
			evicted++
		}
		s.mu.Unlock()
	}
	l.evictions.Add(uint64(evicted))
}