bucketeer/internal/dns/cache.go

package dns

import (
	"sync"
	"time"
)

type cacheEntry[V any] struct {
	val V
	exp time.Time
}

type shard[V any] struct {
	mu sync.Mutex
	m  map[string]cacheEntry[V]
}

// TTLCache is a sharded, memory-bounded, in-memory TTL cache. It is used for
// positive DNS results, negative DNS results, CNAME targets and hostname dedup.
//
// ponytail: once a shard is full, an arbitrary entry is evicted rather than the
// least-recently-used one. Swap in an LRU only if the hit rate actually matters.
type TTLCache[V any] struct {
	shards   []shard[V]
	mask     uint64
	maxShard int
}

func NewTTLCache[V any](maxEntries, shards int) *TTLCache[V] {
	if shards < 1 {
		shards = 1
	}
	n := 1
	for n < shards {
		n <<= 1
	}
	maxShard := maxEntries / n
	if maxShard < 1 {
		maxShard = 1
	}
	c := &TTLCache[V]{mask: uint64(n - 1), maxShard: maxShard}
	c.shards = make([]shard[V], n)
	for i := range c.shards {
		// Deliberately no capacity hint: pre-sizing every shard to maxShard
		// allocates the full theoretical capacity (hundreds of MB across the
		// caches) up front. Let the maps grow on demand; maxShard still bounds
		// them via eviction.
		c.shards[i].m = make(map[string]cacheEntry[V])
	}
	return c
}

func hashString(s string) uint64 {
	const prime = 1099511628211
	h := uint64(14695981039346656037)
	for i := 0; i < len(s); i++ {
		h ^= uint64(s[i])
		h *= prime
	}
	return h
}

func (c *TTLCache[V]) Get(key string) (V, bool) {
	s := &c.shards[hashString(key)&c.mask]
	s.mu.Lock()
	e, ok := s.m[key]
	if ok && time.Now().Before(e.exp) {
		s.mu.Unlock()
		return e.val, true
	}
	if ok {
		delete(s.m, key)
	}
	s.mu.Unlock()
	var zero V
	return zero, false
}

func (c *TTLCache[V]) Set(key string, val V, ttl time.Duration) {
	if ttl <= 0 {
		return
	}
	s := &c.shards[hashString(key)&c.mask]
	s.mu.Lock()
	if len(s.m) >= c.maxShard {
		for k := range s.m {
			delete(s.m, k)
			break
		}
	}
	s.m[key] = cacheEntry[V]{val: val, exp: time.Now().Add(ttl)}
	s.mu.Unlock()
}

func (c *TTLCache[V]) Len() int {
	total := 0
	for i := range c.shards {
		c.shards[i].mu.Lock()
		total += len(c.shards[i].m)
		c.shards[i].mu.Unlock()
	}
	return total
}