99 lines
1.9 KiB
Go
99 lines
1.9 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 {
|
||
|
|
c.shards[i].m = make(map[string]cacheEntry[V], maxShard)
|
||
|
|
}
|
||
|
|
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
|
||
|
|
}
|