Files
StasonJatham 9710efec1a fix memory blowup and keep direct matches from being dropped
- dns.TTLCache: stop preallocating every shard map to maxShard; this
  allocated the full theoretical cache capacity up front (hundreds of MB
  of empty buckets) and pushed the container over its memory limit,
  causing an OOM restart loop under Docker's default 1g cap.
- pipeline: write direct provider matches immediately without a blocking
  DNS enrichment step, so they can no longer be starved by the optional
  CNAME queue. Add TestDirectNotStarvedByDNSBacklog.
- config: add DEDUP_SIZE, lower DNS_CACHE_SIZE/DEDUP_SIZE defaults to 256k.
- .env.example / README: document DEDUP_SIZE and cache sizing.

Live: stable at ~250-300 MiB with no restarts, direct queue draining to 0.
2026-09-22 13:33:05 +02:00

104 lines
2.1 KiB
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
}