149 lines
3.0 KiB
Go
149 lines
3.0 KiB
Go
package qlog
|
|
|
|
import (
|
|
"math/rand/v2"
|
|
"sort"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// randFloat is a 0..1 sample draw, isolated so tests can reason about
|
|
// sampling without seeding the global source.
|
|
func randFloat() float64 { return rand.Float64() }
|
|
|
|
// domainCount is a rolling per-domain counter.
|
|
type domainCount struct {
|
|
name string
|
|
qtype uint16
|
|
qtypeS string
|
|
total uint64
|
|
blocked uint64
|
|
last time.Time
|
|
}
|
|
|
|
// domainCounts is a bounded, TTL-evicting map of domain counters. It replaces
|
|
// the hot half of "what is everyone asking for" without the cost of a sorted
|
|
// set on every query.
|
|
type domainCounts struct {
|
|
mu sync.Mutex
|
|
counts map[string]*domainCount
|
|
limit int
|
|
ttl time.Duration
|
|
now func() time.Time
|
|
}
|
|
|
|
func newDomainCounts(limit int, ttl time.Duration) *domainCounts {
|
|
return &domainCounts{
|
|
counts: make(map[string]*domainCount, 1024),
|
|
limit: limit,
|
|
ttl: ttl,
|
|
now: time.Now,
|
|
}
|
|
}
|
|
|
|
func countKey(name string, qtype uint16) string {
|
|
return name + "|" + itoa(int(qtype))
|
|
}
|
|
|
|
func itoa(n int) string {
|
|
if n == 0 {
|
|
return "0"
|
|
}
|
|
var buf [8]byte
|
|
i := len(buf)
|
|
for n > 0 {
|
|
i--
|
|
buf[i] = byte('0' + n%10)
|
|
n /= 10
|
|
}
|
|
return string(buf[i:])
|
|
}
|
|
|
|
// observe records a query against its domain counter.
|
|
func (d *domainCounts) observe(e Entry) {
|
|
if e.Name == "" {
|
|
return
|
|
}
|
|
key := countKey(e.Name, e.QType)
|
|
now := d.now()
|
|
d.mu.Lock()
|
|
c, ok := d.counts[key]
|
|
if !ok {
|
|
if len(d.counts) >= d.limit {
|
|
d.evictLocked(now)
|
|
}
|
|
c = &domainCount{name: trimDot(e.Name), qtype: e.QType, qtypeS: e.QTypeStr}
|
|
d.counts[key] = c
|
|
}
|
|
c.total++
|
|
if e.Blocked {
|
|
c.blocked++
|
|
}
|
|
c.last = now
|
|
d.mu.Unlock()
|
|
}
|
|
|
|
// evictLocked removes expired counters, and if that is not enough, the oldest
|
|
// quarter. Callers hold d.mu.
|
|
func (d *domainCounts) evictLocked(now time.Time) {
|
|
cutoff := now.Add(-d.ttl)
|
|
removed := 0
|
|
for k, c := range d.counts {
|
|
if c.last.Before(cutoff) {
|
|
delete(d.counts, k)
|
|
removed++
|
|
}
|
|
}
|
|
if removed == 0 && len(d.counts) > 0 {
|
|
// Nothing has expired (a very busy window): drop the oldest quarter.
|
|
type kv struct {
|
|
k string
|
|
t time.Time
|
|
}
|
|
all := make([]kv, 0, len(d.counts))
|
|
for k, c := range d.counts {
|
|
all = append(all, kv{k, c.last})
|
|
}
|
|
sort.Slice(all, func(i, j int) bool { return all[i].t.Before(all[j].t) })
|
|
n := len(all) / 4
|
|
if n < 1 {
|
|
n = 1
|
|
}
|
|
for i := 0; i < n && i < len(all); i++ {
|
|
delete(d.counts, all[i].k)
|
|
}
|
|
}
|
|
}
|
|
|
|
// top returns the n hottest domains.
|
|
func (d *domainCounts) top(n int, blockedOnly bool) []DomainCount {
|
|
if n <= 0 {
|
|
n = 25
|
|
}
|
|
d.mu.Lock()
|
|
out := make([]DomainCount, 0, n)
|
|
for _, c := range d.counts {
|
|
if blockedOnly && c.blocked == 0 {
|
|
continue
|
|
}
|
|
dc := DomainCount{Name: c.name, Count: c.total, Blocked: c.blocked, QType: c.qtype, QTypeS: c.qtypeS}
|
|
if blockedOnly {
|
|
dc.Count = c.blocked
|
|
}
|
|
out = append(out, dc)
|
|
}
|
|
d.mu.Unlock()
|
|
sort.Slice(out, func(i, j int) bool { return out[i].Count > out[j].Count })
|
|
if len(out) > n {
|
|
out = out[:n]
|
|
}
|
|
return out
|
|
}
|
|
|
|
func trimDot(s string) string {
|
|
if n := len(s); n > 1 && s[n-1] == '.' {
|
|
return s[:n-1]
|
|
}
|
|
return s
|
|
}
|