initial commit

internal/qlog/counts.go

@@ -0,0 +1,148 @@
+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
+}