diff --git a/cmd/sower/main.go b/cmd/sower/main.go index 84c8a7c..91ce7ce 100644 --- a/cmd/sower/main.go +++ b/cmd/sower/main.go @@ -8,11 +8,13 @@ import ( "fmt" "io" "log/slog" + "math" "net" "net/http" "os" "os/signal" "runtime" + "runtime/debug" "strconv" "strings" "syscall" @@ -34,13 +36,24 @@ import ( var ( version, date string conf config.SowerConfig + // logLevel backs the tint handler so the admin console can adjust the + // log level at runtime; it defaults to INFO until the config is loaded. + logLevel slog.LevelVar ) +// defaultMemoryLimitMiB bounds the Go heap unless GOMEMLIMIT or +// SOWER_MEMORY_LIMIT_MB overrides it; see init for the rationale. +const defaultMemoryLimitMiB = 128 + +// maxMemoryLimitMiB caps the SOWER_MEMORY_LIMIT_MB override so the MiB +// shift cannot overflow int64. +const maxMemoryLimitMiB int64 = 1 << 20 // 1 TiB + func init() { fi, _ := os.Stdout.Stat() noColor := (fi.Mode() & os.ModeCharDevice) == 0 deferlog.SetDefault(slog.New(tint.NewHandler(os.Stdout, - &tint.Options{AddSource: true, NoColor: noColor}))) + &tint.Options{AddSource: true, NoColor: noColor, Level: &logLevel}))) if strings.HasSuffix(os.Args[0], ".test") { return } @@ -70,6 +83,35 @@ func init() { slog.Error("validate config", "error", err) os.Exit(1) } + + // Bound the Go heap with a soft memory limit. The default rule sets + // (adlist/chinalist/gfwlist) build ~40MB of suffix trees, and GOGC's + // 2x target on top of that pushes resident memory toward 250MB on a + // gateway that mostly idles. The soft limit makes GC reclaim eagerly + // and return memory to the OS. + // + // Precedence: an explicit standard GOMEMLIMIT wins; otherwise the + // default 128MiB applies unless SOWER_MEMORY_LIMIT_MB overrides it + // (a positive MiB value, or 0 to disable the soft limit). Malformed or + // negative values leave the default in place. + if os.Getenv("GOMEMLIMIT") == "" { + var limit int64 = defaultMemoryLimitMiB + if v := os.Getenv("SOWER_MEMORY_LIMIT_MB"); v != "" { + if n, err := strconv.ParseInt(v, 10, 64); err == nil { + switch { + case n == 0: + debug.SetMemoryLimit(math.MaxInt64) // disable + case n > 0 && n <= maxMemoryLimitMiB: + limit = n + } + } + } + if limit > 0 { + debug.SetMemoryLimit(limit << 20) + } + } + + logLevel.Set(conf.LogLevel) slog.Info("starting sower", "version", version, "date", date, @@ -94,6 +136,10 @@ func main() { } func run(ctx context.Context, stop context.CancelFunc, cfg config.SowerConfig) error { + stateStore := loadAdminState(cfg) + baseCfg := cfg + applyConfigOverrides(&cfg, stateStore.ConfigOverrides()) + upstreamDNS := effectiveUpstreamDNS(cfg) proxyDial, err := GenProxyDial(cfg.Remote.Type, cfg.Remote.Addr, cfg.Remote.Password, upstreamDNS, upstreamtls.Options{ ServerName: cfg.Remote.TLS.ServerName, @@ -126,6 +172,12 @@ func run(ctx context.Context, stop context.CancelFunc, cfg config.SowerConfig) e return err } + baseline := snapshotBaseline(r) + stateStore.SetBaseline(baseline) + applyRuleDeltas(r, stateStore) + rulesMgr := newAdminRules(r, stateStore, baseline) + configMgr := newAdminConfig(baseCfg, stateStore, r) + errCh := make(chan error, 8) if err := startDNSListeners(ctx, cfg, r, stats, errCh); err != nil { return err @@ -134,10 +186,10 @@ func run(ctx context.Context, stop context.CancelFunc, cfg config.SowerConfig) e return err } if _, shared := sharedAdminHTTPAddr(cfg); shared { - if err := startSharedHTTPListener(ctx, cfg, r, stats, errCh); err != nil { + if err := startSharedHTTPListener(ctx, cfg, r, rulesMgr, configMgr, stats, errCh); err != nil { return err } - } else if err := startAdminListener(ctx, cfg, r, stats, errCh); err != nil { + } else if err := startAdminListener(ctx, cfg, rulesMgr, configMgr, stats, errCh); err != nil { return err } @@ -162,6 +214,82 @@ func effectiveUpstreamDNS(cfg config.SowerConfig) string { return cfg.DNS.Fallback } +// loadAdminState opens the admin state store. An empty state file path or +// the --ignore-admin-state escape hatch yields an in-memory store, so the +// rest of the code never has to nil-check persistence. +func loadAdminState(cfg config.SowerConfig) *admin.StateStore { + switch { + case cfg.Admin.StateFile == "": + return admin.LoadStateStore("") + case cfg.IgnoreAdminState: + slog.Warn("ignoring admin state file", "file", cfg.Admin.StateFile) + return admin.LoadStateStore("") + default: + return admin.LoadStateStore(cfg.Admin.StateFile) + } +} + +// applyConfigOverrides applies the whitelisted admin-state config overrides +// to cfg before the router and dialer are built. Invalid values are dropped +// with a warning so a bad override never blocks startup; the +// --ignore-admin-state flag skips the file entirely. +func applyConfigOverrides(cfg *config.SowerConfig, o admin.ConfigOverrides) { + if o.LogLevel != "" { + var lv slog.Level + if err := lv.UnmarshalText([]byte(strings.ToUpper(o.LogLevel))); err != nil { + slog.Warn("ignore invalid admin state override", "field", "log_level", "error", err) + } else { + cfg.LogLevel = lv + } + } + if o.DNSUpstream != "" { + if net.ParseIP(o.DNSUpstream) == nil { + slog.Warn("ignore invalid admin state override", "field", "dns_upstream") + } else { + cfg.DNS.Upstream = o.DNSUpstream + } + } + if o.DNSFallback != "" { + if net.ParseIP(o.DNSFallback) == nil { + slog.Warn("ignore invalid admin state override", "field", "dns_fallback") + } else { + cfg.DNS.Fallback = o.DNSFallback + } + } + logLevel.Set(cfg.LogLevel) +} + +// applyRuleDeltas replays persisted admin rule changes onto the freshly +// loaded rule sets: tombstoned baseline rules leave, admin additions enter. +// The rule set is rebuilt via Replace because RuleSet.Remove rebuilds the +// suffix tree per call, which is expensive beyond a handful of tombstones. +func applyRuleDeltas(r *router.Router, state *admin.StateStore) { + sets := map[admin.Category]*router.RuleSet{ + admin.CategoryBlock: r.BlockRule, + admin.CategoryDirect: r.DirectRule, + admin.CategoryProxy: r.ProxyRule, + } + for cat, rs := range sets { + d := state.Delta(cat) + if len(d.Add) == 0 && len(d.Remove) == 0 { + continue + } + tombstoned := make(map[string]struct{}, len(d.Remove)) + for _, rule := range d.Remove { + tombstoned[rule] = struct{}{} + } + rules := rs.List() + effective := make([]string, 0, len(rules)+len(d.Add)) + for _, rule := range rules { + if _, ok := tombstoned[rule]; !ok { + effective = append(effective, rule) + } + } + rs.Replace(append(effective, d.Add...)...) + slog.Info("applied admin rule deltas", "category", cat, "added", len(d.Add), "removed", len(d.Remove)) + } +} + func newRouter(cfg config.SowerConfig, proxyDial router.ProxyDialFn) (*router.Router, error) { r, err := router.NewRouter([]string{cfg.DNS.Serve, cfg.DNS.Serve6}, cfg.DNS.Upstream, cfg.DNS.Fallback, cfg.Router.Country.MMDB, proxyDial) if err != nil { diff --git a/internal/admin/counting.go b/internal/admin/counting.go new file mode 100644 index 0000000..4bcb8e8 --- /dev/null +++ b/internal/admin/counting.go @@ -0,0 +1,141 @@ +package admin + +import ( + "net" + "sync/atomic" + "time" +) + +// countingConn wraps a client connection and feeds byte counts into Stats. +type countingConn struct { + net.Conn + stats *Stats + kind string + clientIP string + created time.Time + domain atomic.Value // string; empty until BindConn + closed atomic.Bool + + // Per-direction pending byte batches. Read and Write run on different + // relay goroutines, so each direction is owned by one goroutine and the + // atomics below are uncontended; Close swaps them safely under I/O. + pendingUp atomic.Uint64 + pendingDown atomic.Uint64 + nextFlushUp atomic.Int64 // unix nanos, upload flush due + nextFlushDown atomic.Int64 // unix nanos, download flush due +} + +// flushBytes and flushInterval bound how long per-domain byte attribution +// may lag: throughput connections flush every flushBytes, slow trickles at +// least every flushInterval, and Close always drains the remainder. The +// global BytesUp/BytesDown totals stay real-time via direct atomics; only +// per-domain attribution and the OTel byte counter are batched. +const ( + flushBytes = 32 << 10 + flushInterval = 500 * time.Millisecond +) + +func (c *countingConn) Read(p []byte) (int, error) { + n, err := c.Conn.Read(p) + if n > 0 { + c.stats.bytesUp.Add(uint64(n)) + c.addBytes(true, uint64(n)) + } + return n, err +} + +func (c *countingConn) Write(p []byte) (int, error) { + n, err := c.Conn.Write(p) + if n > 0 { + c.stats.bytesDown.Add(uint64(n)) + c.addBytes(false, uint64(n)) + } + return n, err +} + +// Close releases the active-connection slot exactly once and records the +// connection duration. The underlying conn is closed on every call, matching +// net.Conn semantics. +func (c *countingConn) Close() error { + if c.closed.CompareAndSwap(false, true) { + switch c.kind { + case "http": + c.stats.activeHTTP.Add(^uint64(0)) + case "https": + c.stats.activeHTTPS.Add(^uint64(0)) + case "socks5": + c.stats.activeSocks.Add(^uint64(0)) + } + c.stats.metrics.ConnClosed(c.kind) + c.stats.metrics.RecordConnDuration(time.Since(c.created)) + c.flush(true, c.pendingUp.Swap(0)) + c.flush(false, c.pendingDown.Swap(0)) + } + return c.Conn.Close() +} + +func (c *countingConn) bind(domain string) { + c.domain.Store(domain) +} + +// addBytes accumulates payload bytes per direction and flushes the batch +// when it reaches flushBytes or flushInterval has elapsed. The first I/O on +// a fresh connection flushes immediately (nextFlush starts zero), so short +// connections attribute their bytes right away. +func (c *countingConn) addBytes(up bool, n uint64) { + var pending *atomic.Uint64 + var due *atomic.Int64 + if up { + pending, due = &c.pendingUp, &c.nextFlushUp + } else { + pending, due = &c.pendingDown, &c.nextFlushDown + } + + acc := pending.Add(n) + now := time.Now().UnixNano() + if acc < flushBytes && now < due.Load() && !c.closed.Load() { + return + } + + due.Store(now + int64(flushInterval)) + c.flush(up, pending.Swap(0)) +} + +// flush attributes one direction's pending bytes to the OTel byte counter +// and, when the domain is known, to the bound domain. The OTel counter is +// recorded unconditionally: protocol bytes read before BindConn (e.g. TLS +// ClientHello) were counted there in the pre-batching implementation too. +// Bytes read before BindConn are dropped from per-domain attribution. +func (c *countingConn) flush(up bool, n uint64) { + if n == 0 { + return + } + c.stats.metrics.RecordBytes(up, int(n)) + + domain, _ := c.domain.Load().(string) + if domain == "" { + return + } + source := Source(c.kind) + var upBytes, downBytes uint64 + if up { + upBytes = n + } else { + downBytes = n + } + c.stats.record(domain, source, c.clientIP, 0, upBytes, downBytes, func(d *domainStat, src *sourceStat, dc *clientStat) { + if up { + d.bytesUp += n + src.bytesUp += n + if dc != nil { + dc.bytesUp += n + } + } else { + d.bytesDown += n + src.bytesDown += n + if dc != nil { + dc.bytesDown += n + } + } + }) +} diff --git a/internal/admin/eviction.go b/internal/admin/eviction.go new file mode 100644 index 0000000..73b3faa --- /dev/null +++ b/internal/admin/eviction.go @@ -0,0 +1,86 @@ +package admin + +import ( + "container/heap" + "time" +) + +// evictionCandidate is a potential least-recently-seen entry. The heap keeps +// its newest element at the root, allowing a single scan to retain only the +// oldest candidates required for a batch eviction. +type evictionCandidate struct { + key string + lastSeen time.Time +} + +type oldestCandidates []evictionCandidate + +func (h oldestCandidates) Len() int { return len(h) } + +func (h oldestCandidates) Less(i, j int) bool { + if h[i].lastSeen.Equal(h[j].lastSeen) { + return h[i].key > h[j].key + } + return h[i].lastSeen.After(h[j].lastSeen) +} + +func (h oldestCandidates) Swap(i, j int) { h[i], h[j] = h[j], h[i] } + +func (h *oldestCandidates) Push(x any) { *h = append(*h, x.(evictionCandidate)) } + +func (h *oldestCandidates) Pop() any { + old := *h + n := len(old) + item := old[n-1] + *h = old[:n-1] + return item +} + +func older(a, b evictionCandidate) bool { + if a.lastSeen.Equal(b.lastSeen) { + return a.key < b.key + } + return a.lastSeen.Before(b.lastSeen) +} + +func keepOldest(candidates *oldestCandidates, candidate evictionCandidate, limit int) { + if candidates.Len() < limit { + heap.Push(candidates, candidate) + return + } + if older(candidate, (*candidates)[0]) { + heap.Pop(candidates) + heap.Push(candidates, candidate) + } +} + +// evictOldestDomainsLocked frees a batch of least-recently-seen domains so +// a high-cardinality stream amortizes the O(N) scan across many insertions. +func (s *Stats) evictOldestDomainsLocked(count int) { + count = min(count, len(s.domains)) + if count == 0 { + return + } + candidates := make(oldestCandidates, 0, count) + for domain, d := range s.domains { + keepOldest(&candidates, evictionCandidate{key: domain, lastSeen: d.lastSeen}, count) + } + for _, candidate := range candidates { + delete(s.domains, candidate.key) + } +} + +// evictOldestClientsLocked frees a batch of least-recently-seen client totals. +func (s *Stats) evictOldestClientsLocked(count int) { + count = min(count, len(s.clients)) + if count == 0 { + return + } + candidates := make(oldestCandidates, 0, count) + for ip, c := range s.clients { + keepOldest(&candidates, evictionCandidate{key: ip, lastSeen: c.lastSeen}, count) + } + for _, candidate := range candidates { + delete(s.clients, candidate.key) + } +} diff --git a/internal/admin/memory_bench_test.go b/internal/admin/memory_bench_test.go new file mode 100644 index 0000000..df56f86 --- /dev/null +++ b/internal/admin/memory_bench_test.go @@ -0,0 +1,44 @@ +package admin + +import ( + "fmt" + "runtime" + "testing" +) + +// TestStatsMemoryFootprint measures the heap retained by the domain table at +// realistic and worst-case cardinalities. Run with -count=1. +func TestStatsMemoryFootprint(t *testing.T) { + if testing.Short() { + t.Skip("memory footprint test; skipped in -short mode") + } + for _, domains := range []int{10_000, 50_000, 100_000} { + s := newTestStats(t) + for i := 0; i < domains; i++ { + domain := fmt.Sprintf("host%d.example.com", i) + s.record(domain, SourceHTTPS, "192.168.1.10", 1, 1000, 2000, func(d *domainStat, src *sourceStat, dc *clientStat) { + d.conns++ + d.bytesUp += 1000 + d.bytesDown += 2000 + src.bytesUp += 1000 + src.bytesDown += 2000 + if dc != nil { + dc.bytesUp += 1000 + dc.bytesDown += 2000 + } + }) + } + + runtime.GC() + var m runtime.MemStats + runtime.ReadMemStats(&m) + s.mu.Lock() + retained := len(s.domains) + s.mu.Unlock() + t.Logf("%d records: retained=%d heap=%d KiB (%.0f B/domain)", + domains, retained, m.HeapAlloc/1024, float64(m.HeapAlloc)/float64(max(retained, 1))) + if retained != domains { + t.Errorf("expected %d retained domains, got %d", domains, retained) + } + } +} diff --git a/internal/admin/relay_bench_test.go b/internal/admin/relay_bench_test.go new file mode 100644 index 0000000..7f8a9d3 --- /dev/null +++ b/internal/admin/relay_bench_test.go @@ -0,0 +1,58 @@ +package admin + +import ( + "net" + "testing" +) + +// discardConn accepts and discards data without copying, isolating the +// accounting overhead from fake-conn buffer growth. +type discardConn struct { + net.Conn +} + +func (discardConn) Read(p []byte) (int, error) { return len(p), nil } +func (discardConn) Write(p []byte) (int, error) { return len(p), nil } +func (discardConn) Close() error { return nil } +func (discardConn) RemoteAddr() net.Addr { + return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 12345} +} +func (discardConn) LocalAddr() net.Addr { return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 8080} } + +// BenchmarkCountingConnWrite measures the per-write accounting cost of the +// relay path: atomic counters, OTel byte recording, and per-domain +// attribution under the stats mutex. +func BenchmarkCountingConnWrite(b *testing.B) { + s := newBenchmarkStats(b, 100) + conn := s.WrapConn(discardConn{}, "https") + s.BindConn(conn, "example.com") + chunk := make([]byte, 4096) + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if _, err := conn.Write(chunk); err != nil { + b.Fatal(err) + } + } +} + +// BenchmarkCountingConnWriteChained measures one Write plus one Read, the +// full duplex relay unit of work. +func BenchmarkCountingConnWriteChained(b *testing.B) { + s := newBenchmarkStats(b, 100) + conn := s.WrapConn(discardConn{}, "http") + s.BindConn(conn, "example.com") + chunk := make([]byte, 4096) + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if _, err := conn.Write(chunk); err != nil { + b.Fatal(err) + } + if _, err := conn.Read(chunk); err != nil { + b.Fatal(err) + } + } +} diff --git a/internal/admin/snapshot.go b/internal/admin/snapshot.go new file mode 100644 index 0000000..751e420 --- /dev/null +++ b/internal/admin/snapshot.go @@ -0,0 +1,403 @@ +package admin + +import ( + "container/heap" + "sort" + "time" +) + +// Snapshot returns the window view of the current stats, advancing the +// history ring along the way. Domains stale for over staleDomainAge are +// hidden from the window but kept in the map for the cumulative Totals view. +// sort selects the domain ordering, source restricts domains to one entry +// point, and client restricts them to one client IP; invalid values fall back +// to DomainSortBytes and SourceAll, and an empty client disables the filter. +// Clients aggregates per-client counters over the source-filtered view and +// is intentionally unaffected by the client filter, so the console's client +// chips stay stable while one is selected. +func (s *Stats) Snapshot(sort DomainSort, source Source, client string) TrafficSnapshot { + s.sample() + + now := time.Now() + snap := TrafficSnapshot{ + Uptime: int64(now.Sub(s.start).Seconds()), + DNSQueries: s.dnsQueries.Load(), + BytesUp: s.bytesUp.Load(), + BytesDown: s.bytesDown.Load(), + Domains: []DomainStat{}, + } + snap.Conns.HTTP = s.connsHTTP.Load() + snap.Conns.HTTPS = s.connsHTTPS.Load() + snap.Conns.Socks5 = s.connsSocks.Load() + snap.Active.HTTP = s.activeHTTP.Load() + snap.Active.HTTPS = s.activeHTTPS.Load() + snap.Active.Socks5 = s.activeSocks.Load() + snap.RuleHits.Block = s.ruleBlock.Load() + snap.RuleHits.Direct = s.ruleDirect.Load() + snap.RuleHits.Proxy = s.ruleProxy.Load() + snap.Rates.BytesUpPerSec, snap.Rates.BytesDownPerSec, snap.Rates.DNSPerSec, snap.Rates.ConnsPerSec = s.rates() + snap.System.Goroutines = uint64(goroutineCount()) + snap.System.HeapAlloc = heapAlloc() + + s.mu.Lock() + clientAgg := make(map[string]*ClientStat) + for domain, d := range s.domains { + if now.Sub(d.lastSeen) > staleDomainAge { + // Window view only: the entry stays in the map so the cumulative + // Totals view keeps its from-start counters. + continue + } + ds := DomainStat{ + Domain: domain, + Conns: d.conns, + BytesUp: d.bytesUp, + BytesDown: d.bytesDown, + LastSeen: d.lastSeen, + } + if source != SourceAll { + src, ok := d.bySource[source] // nil map read is safe + if !ok || (src.conns == 0 && src.bytesUp == 0 && src.bytesDown == 0) { + continue + } + ds.Conns = src.conns + ds.BytesUp = src.bytesUp + ds.BytesDown = src.bytesDown + ds.LastSeen = src.lastSeen + } + // Aggregate the client chips over the source-filtered view but before + // the client filter: the chip row is a navigation device and must stay + // stable (and globally counted) while a client filter is active. + for ip, c := range d.byClient { + agg := clientAgg[ip] + if agg == nil { + agg = &ClientStat{IP: ip} + clientAgg[ip] = agg + } + agg.Conns += c.conns + agg.BytesUp += c.bytesUp + agg.BytesDown += c.bytesDown + if c.lastSeen.After(agg.LastSeen) { + agg.LastSeen = c.lastSeen + } + } + + if client != "" { + dc, ok := d.byClient[client] // nil map read is safe + if !ok || (dc.conns == 0 && dc.bytesUp == 0 && dc.bytesDown == 0) { + continue + } + ds.Conns = dc.conns + ds.BytesUp = dc.bytesUp + ds.BytesDown = dc.bytesDown + ds.LastSeen = dc.lastSeen + ds.LastClientIP = client + } else { + // Most recent client across all sources for the domain; the + // per-client map has no source dimension, so this is the best + // approximation for a source-filtered view. + var latest time.Time + for ip, c := range d.byClient { + if c.lastSeen.After(latest) { + latest, ds.LastClientIP = c.lastSeen, ip + } + } + } + snap.Domains = append(snap.Domains, ds) + } + s.mu.Unlock() + + snap.Domains = topDomains(snap.Domains, snapshotTopN, sort) + snap.Clients = topClients(clientAgg) + return snap +} + +// domainLess reports whether a should appear before b in the requested order. +func domainLess(mode DomainSort) func(a, b DomainStat) bool { + switch mode { + case DomainSortRecent: + return func(a, b DomainStat) bool { + if !a.LastSeen.Equal(b.LastSeen) { + return a.LastSeen.After(b.LastSeen) + } + return a.Domain < b.Domain + } + case DomainSortConns: + return func(a, b DomainStat) bool { + if a.Conns != b.Conns { + return a.Conns > b.Conns + } + return a.Domain < b.Domain + } + default: + return func(a, b DomainStat) bool { + ta, tb := a.BytesUp+a.BytesDown, b.BytesUp+b.BytesDown + if ta != tb { + return ta > tb + } + return a.Domain < b.Domain + } + } +} + +type domainTopHeap struct { + values []DomainStat + less func(DomainStat, DomainStat) bool +} + +func (h domainTopHeap) Len() int { return len(h.values) } + +// Less puts the worst retained item at the root. +func (h domainTopHeap) Less(i, j int) bool { return h.less(h.values[j], h.values[i]) } + +func (h domainTopHeap) Swap(i, j int) { h.values[i], h.values[j] = h.values[j], h.values[i] } + +func (h *domainTopHeap) Push(x any) { h.values = append(h.values, x.(DomainStat)) } + +func (h *domainTopHeap) Pop() any { + old := h.values + n := len(old) + item := old[n-1] + h.values = old[:n-1] + return item +} + +// topDomains returns the top limit domains in their requested display order. +func topDomains(domains []DomainStat, limit int, mode DomainSort) []DomainStat { + if len(domains) <= limit { + sortDomains(domains, mode) + return domains + } + less := domainLess(mode) + top := domainTopHeap{values: make([]DomainStat, 0, limit), less: less} + for _, domain := range domains { + if top.Len() < limit { + heap.Push(&top, domain) + continue + } + if less(domain, top.values[0]) { + heap.Pop(&top) + heap.Push(&top, domain) + } + } + sort.Slice(top.values, func(i, j int) bool { return less(top.values[i], top.values[j]) }) + return top.values +} + +type clientTopHeap []ClientStat + +func (h clientTopHeap) Len() int { return len(h) } + +func (h clientTopHeap) Less(i, j int) bool { return clientLess(h[j], h[i]) } + +func (h clientTopHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } + +func (h *clientTopHeap) Push(x any) { *h = append(*h, x.(ClientStat)) } + +func (h *clientTopHeap) Pop() any { + old := *h + n := len(old) + item := old[n-1] + *h = old[:n-1] + return item +} + +func clientLess(a, b ClientStat) bool { + ta, tb := a.BytesUp+a.BytesDown, b.BytesUp+b.BytesDown + if ta == tb { + return a.IP < b.IP + } + return ta > tb +} + +// topClients orders per-IP aggregates by total bytes and caps the list. +func topClients(agg map[string]*ClientStat) []ClientStat { + top := make(clientTopHeap, 0, min(clientsTopN, len(agg))) + for _, c := range agg { + item := *c + if top.Len() < clientsTopN { + heap.Push(&top, item) + continue + } + if clientLess(item, top[0]) { + heap.Pop(&top) + heap.Push(&top, item) + } + } + sort.Slice(top, func(i, j int) bool { return clientLess(top[i], top[j]) }) + return top +} + +// Totals is the from-start cumulative traffic view, independent of the +// window's staleness filter. +type Totals struct { + Domains []DomainStat `json:"domains"` + Clients []ClientStat `json:"clients"` +} + +// Totals returns the cumulative top domains and clients since start, each +// ordered by total bytes. Unlike Snapshot it ignores staleness: counters +// accumulate until the bounded maps evict their least-recently-seen entries. +func (s *Stats) Totals() Totals { + s.mu.Lock() + + domains := make([]DomainStat, 0, len(s.domains)) + for domain, d := range s.domains { + ds := DomainStat{ + Domain: domain, + Conns: d.conns, + BytesUp: d.bytesUp, + BytesDown: d.bytesDown, + LastSeen: d.lastSeen, + } + var latest time.Time + for ip, c := range d.byClient { + if c.lastSeen.After(latest) { + latest, ds.LastClientIP = c.lastSeen, ip + } + } + domains = append(domains, ds) + } + domains = topDomains(domains, snapshotTopN, DomainSortBytes) + + clients := make([]ClientStat, 0, len(s.clients)) + for ip, c := range s.clients { + clients = append(clients, ClientStat{ + IP: ip, + Conns: c.conns, + BytesUp: c.bytesUp, + BytesDown: c.bytesDown, + LastSeen: c.lastSeen, + }) + } + s.mu.Unlock() + + sort.Slice(clients, func(i, j int) bool { + ti := clients[i].BytesUp + clients[i].BytesDown + tj := clients[j].BytesUp + clients[j].BytesDown + if ti == tj { + return clients[i].IP < clients[j].IP + } + return ti > tj + }) + if len(clients) > clientsTopN { + clients = clients[:clientsTopN] + } + + return Totals{Domains: domains, Clients: clients} +} + +// sortDomains orders domains by the selected mode: total bytes (default), +// most recently seen, or connection count. Ties break alphabetically. +func sortDomains(ds []DomainStat, mode DomainSort) { + less := func(a, b DomainStat) bool { return a.Domain < b.Domain } + switch mode { + case DomainSortRecent: + less = func(a, b DomainStat) bool { + if !a.LastSeen.Equal(b.LastSeen) { + return a.LastSeen.After(b.LastSeen) + } + return a.Domain < b.Domain + } + case DomainSortConns: + less = func(a, b DomainStat) bool { + if a.Conns != b.Conns { + return a.Conns > b.Conns + } + return a.Domain < b.Domain + } + default: + less = func(a, b DomainStat) bool { + ta, tb := a.BytesUp+a.BytesDown, b.BytesUp+b.BytesDown + if ta != tb { + return ta > tb + } + return a.Domain < b.Domain + } + } + sort.Slice(ds, func(i, j int) bool { return less(ds[i], ds[j]) }) +} + +// History returns a copy of the in-process history ring. +func (s *Stats) History() History { + s.histMu.Lock() + defer s.histMu.Unlock() + out := make([]HistorySample, len(s.history)) + copy(out, s.history) + return History{Samples: out} +} + +// sample advances the history ring with the deltas since the last sample. It +// is called from Snapshot, so the ring advances at the poll cadence. +func (s *Stats) sample() { + now := time.Now() + cur := sampleCounters{ + at: now, + bytesUp: s.bytesUp.Load(), + bytesDown: s.bytesDown.Load(), + dns: s.dnsQueries.Load(), + conns: s.connsHTTP.Load() + s.connsHTTPS.Load() + s.connsSocks.Load(), + block: s.ruleBlock.Load(), + direct: s.ruleDirect.Load(), + proxy: s.ruleProxy.Load(), + } + + s.histMu.Lock() + defer s.histMu.Unlock() + prev := s.lastSample + if !prev.at.IsZero() && now.Sub(prev.at) < historyMinGap { + return + } + s.lastSample = cur + if prev.at.IsZero() { + return // first sample establishes the baseline + } + + sample := HistorySample{ + At: now, + Active: s.activeHTTP.Load() + s.activeHTTPS.Load() + s.activeSocks.Load(), + ActiveHTTP: s.activeHTTP.Load(), + ActiveHTTPS: s.activeHTTPS.Load(), + ActiveSocks: s.activeSocks.Load(), + } + if now.Sub(prev.at) <= historyMaxGap { + sample.BytesUp = cur.bytesUp - prev.bytesUp + sample.BytesDown = cur.bytesDown - prev.bytesDown + sample.DNS = cur.dns - prev.dns + sample.Conns = cur.conns - prev.conns + sample.Block = cur.block - prev.block + sample.Direct = cur.direct - prev.direct + sample.Proxy = cur.proxy - prev.proxy + } + s.history = append(s.history, sample) + if len(s.history) > historyCapacity { + s.history = s.history[len(s.history)-historyCapacity:] + } +} + +// rates computes per-second rates over the sliding rateWindow from the +// history ring. +func (s *Stats) rates() (up, down, dns, conns float64) { + s.histMu.Lock() + defer s.histMu.Unlock() + now := time.Now() + cutoff := now.Add(-rateWindow) + var sumUp, sumDown, sumDNS, sumConns uint64 + var start time.Time + for _, h := range s.history { + if h.At.Before(cutoff) { + continue + } + if start.IsZero() { + start = h.At + } + sumUp += h.BytesUp + sumDown += h.BytesDown + sumDNS += h.DNS + sumConns += h.Conns + } + if start.IsZero() || now.Sub(start) < time.Second { + return 0, 0, 0, 0 + } + elapsed := now.Sub(start).Seconds() + return float64(sumUp) / elapsed, float64(sumDown) / elapsed, float64(sumDNS) / elapsed, float64(sumConns) / elapsed +} diff --git a/internal/admin/stats.go b/internal/admin/stats.go index 22a048a..662ae90 100644 --- a/internal/admin/stats.go +++ b/internal/admin/stats.go @@ -4,7 +4,6 @@ import ( "fmt" "net" rtmetrics "runtime/metrics" - "sort" "strings" "sync" "sync/atomic" @@ -27,6 +26,8 @@ const ( staleDomainAge = time.Hour snapshotTopN = 100 clientsTopN = 50 + domainEvictionBatch = 1024 + clientEvictionBatch = 64 // historyCapacity keeps one hour of samples at the 5s poll cadence. historyCapacity = 720 @@ -334,309 +335,6 @@ func (s *Stats) BindConn(conn net.Conn, domain string) { }) } -// Snapshot returns the window view of the current stats, advancing the -// history ring along the way. Domains stale for over staleDomainAge are -// hidden from the window but kept in the map for the cumulative Totals view. -// sort selects the domain ordering, source restricts domains to one entry -// point, and client restricts them to one client IP; invalid values fall back -// to DomainSortBytes and SourceAll, and an empty client disables the filter. -// Clients aggregates per-client counters over the source-filtered view and -// is intentionally unaffected by the client filter, so the console's client -// chips stay stable while one is selected. -func (s *Stats) Snapshot(sort DomainSort, source Source, client string) TrafficSnapshot { - s.sample() - - snap := TrafficSnapshot{ - Uptime: int64(time.Since(s.start).Seconds()), - DNSQueries: s.dnsQueries.Load(), - BytesUp: s.bytesUp.Load(), - BytesDown: s.bytesDown.Load(), - Domains: []DomainStat{}, - } - snap.Conns.HTTP = s.connsHTTP.Load() - snap.Conns.HTTPS = s.connsHTTPS.Load() - snap.Conns.Socks5 = s.connsSocks.Load() - snap.Active.HTTP = s.activeHTTP.Load() - snap.Active.HTTPS = s.activeHTTPS.Load() - snap.Active.Socks5 = s.activeSocks.Load() - snap.RuleHits.Block = s.ruleBlock.Load() - snap.RuleHits.Direct = s.ruleDirect.Load() - snap.RuleHits.Proxy = s.ruleProxy.Load() - snap.Rates.BytesUpPerSec, snap.Rates.BytesDownPerSec, snap.Rates.DNSPerSec, snap.Rates.ConnsPerSec = s.rates() - snap.System.Goroutines = uint64(goroutineCount()) - snap.System.HeapAlloc = heapAlloc() - - s.mu.Lock() - defer s.mu.Unlock() - clientAgg := make(map[string]*ClientStat) - for domain, d := range s.domains { - if time.Since(d.lastSeen) > staleDomainAge { - // Window view only: the entry stays in the map so the cumulative - // Totals view keeps its from-start counters. - continue - } - ds := DomainStat{ - Domain: domain, - Conns: d.conns, - BytesUp: d.bytesUp, - BytesDown: d.bytesDown, - LastSeen: d.lastSeen, - } - if source != SourceAll { - src, ok := d.bySource[source] // nil map read is safe - if !ok || (src.conns == 0 && src.bytesUp == 0 && src.bytesDown == 0) { - continue - } - ds.Conns = src.conns - ds.BytesUp = src.bytesUp - ds.BytesDown = src.bytesDown - ds.LastSeen = src.lastSeen - } - // Aggregate the client chips over the source-filtered view but before - // the client filter: the chip row is a navigation device and must stay - // stable (and globally counted) while a client filter is active. - for ip, c := range d.byClient { - agg := clientAgg[ip] - if agg == nil { - agg = &ClientStat{IP: ip} - clientAgg[ip] = agg - } - agg.Conns += c.conns - agg.BytesUp += c.bytesUp - agg.BytesDown += c.bytesDown - if c.lastSeen.After(agg.LastSeen) { - agg.LastSeen = c.lastSeen - } - } - - if client != "" { - dc, ok := d.byClient[client] // nil map read is safe - if !ok || (dc.conns == 0 && dc.bytesUp == 0 && dc.bytesDown == 0) { - continue - } - ds.Conns = dc.conns - ds.BytesUp = dc.bytesUp - ds.BytesDown = dc.bytesDown - ds.LastSeen = dc.lastSeen - ds.LastClientIP = client - } else { - // Most recent client across all sources for the domain; the - // per-client map has no source dimension, so this is the best - // approximation for a source-filtered view. - var latest time.Time - for ip, c := range d.byClient { - if c.lastSeen.After(latest) { - latest, ds.LastClientIP = c.lastSeen, ip - } - } - } - snap.Domains = append(snap.Domains, ds) - } - sortDomains(snap.Domains, sort) - if len(snap.Domains) > snapshotTopN { - snap.Domains = snap.Domains[:snapshotTopN] - } - snap.Clients = topClients(clientAgg) - return snap -} - -// topClients orders per-IP aggregates by total bytes and caps the list. -func topClients(agg map[string]*ClientStat) []ClientStat { - out := make([]ClientStat, 0, len(agg)) - for _, c := range agg { - out = append(out, *c) - } - sort.Slice(out, func(i, j int) bool { - ti := out[i].BytesUp + out[i].BytesDown - tj := out[j].BytesUp + out[j].BytesDown - if ti == tj { - return out[i].IP < out[j].IP - } - return ti > tj - }) - if len(out) > clientsTopN { - out = out[:clientsTopN] - } - return out -} - -// Totals is the from-start cumulative traffic view, independent of the -// window's staleness filter. -type Totals struct { - Domains []DomainStat `json:"domains"` - Clients []ClientStat `json:"clients"` -} - -// Totals returns the cumulative top domains and clients since start, each -// ordered by total bytes. Unlike Snapshot it ignores staleness: counters -// accumulate until the bounded maps evict their least-recently-seen entries. -func (s *Stats) Totals() Totals { - s.mu.Lock() - defer s.mu.Unlock() - - domains := make([]DomainStat, 0, len(s.domains)) - for domain, d := range s.domains { - ds := DomainStat{ - Domain: domain, - Conns: d.conns, - BytesUp: d.bytesUp, - BytesDown: d.bytesDown, - LastSeen: d.lastSeen, - } - var latest time.Time - for ip, c := range d.byClient { - if c.lastSeen.After(latest) { - latest, ds.LastClientIP = c.lastSeen, ip - } - } - domains = append(domains, ds) - } - sortDomains(domains, DomainSortBytes) - if len(domains) > snapshotTopN { - domains = domains[:snapshotTopN] - } - - clients := make([]ClientStat, 0, len(s.clients)) - for ip, c := range s.clients { - clients = append(clients, ClientStat{ - IP: ip, - Conns: c.conns, - BytesUp: c.bytesUp, - BytesDown: c.bytesDown, - LastSeen: c.lastSeen, - }) - } - sort.Slice(clients, func(i, j int) bool { - ti := clients[i].BytesUp + clients[i].BytesDown - tj := clients[j].BytesUp + clients[j].BytesDown - if ti == tj { - return clients[i].IP < clients[j].IP - } - return ti > tj - }) - if len(clients) > clientsTopN { - clients = clients[:clientsTopN] - } - - return Totals{Domains: domains, Clients: clients} -} - -// sortDomains orders domains by the selected mode: total bytes (default), -// most recently seen, or connection count. Ties break alphabetically. -func sortDomains(ds []DomainStat, mode DomainSort) { - less := func(a, b DomainStat) bool { return a.Domain < b.Domain } - switch mode { - case DomainSortRecent: - less = func(a, b DomainStat) bool { - if !a.LastSeen.Equal(b.LastSeen) { - return a.LastSeen.After(b.LastSeen) - } - return a.Domain < b.Domain - } - case DomainSortConns: - less = func(a, b DomainStat) bool { - if a.Conns != b.Conns { - return a.Conns > b.Conns - } - return a.Domain < b.Domain - } - default: - less = func(a, b DomainStat) bool { - ta, tb := a.BytesUp+a.BytesDown, b.BytesUp+b.BytesDown - if ta != tb { - return ta > tb - } - return a.Domain < b.Domain - } - } - sort.Slice(ds, func(i, j int) bool { return less(ds[i], ds[j]) }) -} - -// History returns a copy of the in-process history ring. -func (s *Stats) History() History { - s.histMu.Lock() - defer s.histMu.Unlock() - out := make([]HistorySample, len(s.history)) - copy(out, s.history) - return History{Samples: out} -} - -// sample advances the history ring with the deltas since the last sample. It -// is called from Snapshot, so the ring advances at the poll cadence. -func (s *Stats) sample() { - now := time.Now() - cur := sampleCounters{ - at: now, - bytesUp: s.bytesUp.Load(), - bytesDown: s.bytesDown.Load(), - dns: s.dnsQueries.Load(), - conns: s.connsHTTP.Load() + s.connsHTTPS.Load() + s.connsSocks.Load(), - block: s.ruleBlock.Load(), - direct: s.ruleDirect.Load(), - proxy: s.ruleProxy.Load(), - } - - s.histMu.Lock() - defer s.histMu.Unlock() - prev := s.lastSample - if !prev.at.IsZero() && now.Sub(prev.at) < historyMinGap { - return - } - s.lastSample = cur - if prev.at.IsZero() { - return // first sample establishes the baseline - } - - sample := HistorySample{ - At: now, - Active: s.activeHTTP.Load() + s.activeHTTPS.Load() + s.activeSocks.Load(), - ActiveHTTP: s.activeHTTP.Load(), - ActiveHTTPS: s.activeHTTPS.Load(), - ActiveSocks: s.activeSocks.Load(), - } - if now.Sub(prev.at) <= historyMaxGap { - sample.BytesUp = cur.bytesUp - prev.bytesUp - sample.BytesDown = cur.bytesDown - prev.bytesDown - sample.DNS = cur.dns - prev.dns - sample.Conns = cur.conns - prev.conns - sample.Block = cur.block - prev.block - sample.Direct = cur.direct - prev.direct - sample.Proxy = cur.proxy - prev.proxy - } - s.history = append(s.history, sample) - if len(s.history) > historyCapacity { - s.history = s.history[len(s.history)-historyCapacity:] - } -} - -// rates computes per-second rates over the sliding rateWindow from the -// history ring. -func (s *Stats) rates() (up, down, dns, conns float64) { - s.histMu.Lock() - defer s.histMu.Unlock() - now := time.Now() - cutoff := now.Add(-rateWindow) - var sumUp, sumDown, sumDNS, sumConns uint64 - var start time.Time - for _, h := range s.history { - if h.At.Before(cutoff) { - continue - } - if start.IsZero() { - start = h.At - } - sumUp += h.BytesUp - sumDown += h.BytesDown - sumDNS += h.DNS - sumConns += h.Conns - } - if start.IsZero() || now.Sub(start) < time.Second { - return 0, 0, 0, 0 - } - elapsed := now.Sub(start).Seconds() - return float64(sumUp) / elapsed, float64(sumDown) / elapsed, float64(sumDNS) / elapsed, float64(sumConns) / elapsed -} - // record attributes one event to a domain, its source, and optionally its // client IP, with the given counter deltas. The client IP may be empty when it // cannot be determined; the per-client stat is then nil and callers must @@ -648,7 +346,7 @@ func (s *Stats) record(domain string, source Source, clientIP string, conns, up, d := s.domains[domain] if d == nil { if len(s.domains) >= maxTotalDomains { - s.evictOldestLocked() + s.evictOldestDomainsLocked(domainEvictionBatch) } d = &domainStat{bySource: make(map[Source]*sourceStat)} s.domains[domain] = d @@ -681,7 +379,7 @@ func (s *Stats) record(domain string, source Source, clientIP string, conns, up, cc := s.clients[clientIP] if cc == nil { if len(s.clients) >= maxTotalClients { - s.evictOldestClientLocked() + s.evictOldestClientsLocked(clientEvictionBatch) } cc = &clientStat{} s.clients[clientIP] = cc @@ -693,121 +391,6 @@ func (s *Stats) record(domain string, source Source, clientIP string, conns, up, } } -// evictOldestLocked drops the single least-recently-seen entry to keep the -// domain map bounded. -func (s *Stats) evictOldestLocked() { - var oldestKey string - var oldest time.Time - for domain, d := range s.domains { - if oldestKey == "" || d.lastSeen.Before(oldest) { - oldestKey = domain - oldest = d.lastSeen - } - } - if oldestKey != "" { - delete(s.domains, oldestKey) - } -} - -// evictOldestClientLocked drops the single least-recently-seen client to keep -// the cumulative client table bounded. -func (s *Stats) evictOldestClientLocked() { - var oldestKey string - var oldest time.Time - for ip, c := range s.clients { - if oldestKey == "" || c.lastSeen.Before(oldest) { - oldestKey = ip - oldest = c.lastSeen - } - } - if oldestKey != "" { - delete(s.clients, oldestKey) - } -} - -// countingConn wraps a client connection and feeds byte counts into Stats. -type countingConn struct { - net.Conn - stats *Stats - kind string - clientIP string - created time.Time - domain atomic.Value // string; empty until BindConn - closed atomic.Bool -} - -func (c *countingConn) Read(p []byte) (int, error) { - n, err := c.Conn.Read(p) - if n > 0 { - c.stats.bytesUp.Add(uint64(n)) - c.stats.metrics.RecordBytes(true, n) - c.recordBytes(n, true) - } - return n, err -} - -func (c *countingConn) Write(p []byte) (int, error) { - n, err := c.Conn.Write(p) - if n > 0 { - c.stats.bytesDown.Add(uint64(n)) - c.stats.metrics.RecordBytes(false, n) - c.recordBytes(n, false) - } - return n, err -} - -// Close releases the active-connection slot exactly once and records the -// connection duration. The underlying conn is closed on every call, matching -// net.Conn semantics. -func (c *countingConn) Close() error { - if c.closed.CompareAndSwap(false, true) { - switch c.kind { - case "http": - c.stats.activeHTTP.Add(^uint64(0)) - case "https": - c.stats.activeHTTPS.Add(^uint64(0)) - case "socks5": - c.stats.activeSocks.Add(^uint64(0)) - } - c.stats.metrics.ConnClosed(c.kind) - c.stats.metrics.RecordConnDuration(time.Since(c.created)) - } - return c.Conn.Close() -} - -func (c *countingConn) bind(domain string) { - c.domain.Store(domain) -} - -func (c *countingConn) recordBytes(n int, up bool) { - domain, _ := c.domain.Load().(string) - if domain == "" { - return - } - source := Source(c.kind) - var upBytes, downBytes uint64 - if up { - upBytes = uint64(n) - } else { - downBytes = uint64(n) - } - c.stats.record(domain, source, c.clientIP, 0, upBytes, downBytes, func(d *domainStat, src *sourceStat, dc *clientStat) { - if up { - d.bytesUp += uint64(n) - src.bytesUp += uint64(n) - if dc != nil { - dc.bytesUp += uint64(n) - } - } else { - d.bytesDown += uint64(n) - src.bytesDown += uint64(n) - if dc != nil { - dc.bytesDown += uint64(n) - } - } - }) -} - // normalizeDomain lowercases and strips the trailing dot of a domain or qname. func normalizeDomain(domain string) string { domain = strings.TrimSpace(domain) diff --git a/internal/admin/stats_bench_test.go b/internal/admin/stats_bench_test.go new file mode 100644 index 0000000..68f98c0 --- /dev/null +++ b/internal/admin/stats_bench_test.go @@ -0,0 +1,130 @@ +package admin + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" +) + +var ( + benchmarkSnapshot TrafficSnapshot + benchmarkTotals Totals +) + +func newBenchmarkStats(b *testing.B, domains int) *Stats { + b.Helper() + s, err := NewStats() + if err != nil { + b.Fatalf("new stats: %v", err) + } + b.Cleanup(func() { + ctx, cancel := context.WithTimeout(b.Context(), time.Second) + defer cancel() + if err := s.Metrics().Shutdown(ctx); err != nil { + b.Errorf("shutdown metrics: %v", err) + } + }) + + now := time.Now() + func() { + s.mu.Lock() + defer s.mu.Unlock() + for i := 0; i < domains; i++ { + domain := fmt.Sprintf("host%d.example", i) + s.domains[domain] = &domainStat{ + conns: uint64(i + 1), + bytesUp: uint64(i + 1), + bytesDown: uint64(i + 1), + lastSeen: now, + bySource: map[Source]*sourceStat{ + SourceHTTP: {conns: 1, bytesUp: 1, bytesDown: 1, lastSeen: now}, + }, + } + } + }() + return s +} + +func BenchmarkStatsSnapshot(b *testing.B) { + for _, domains := range []int{1_000, 10_000, 100_000} { + b.Run(fmt.Sprintf("domains=%d", domains), func(b *testing.B) { + s := newBenchmarkStats(b, domains) + b.ReportAllocs() + b.ResetTimer() + for range b.N { + benchmarkSnapshot = s.Snapshot(DomainSortBytes, SourceAll, "") + } + }) + } +} + +func BenchmarkStatsTotals(b *testing.B) { + for _, domains := range []int{1_000, 10_000, 100_000} { + b.Run(fmt.Sprintf("domains=%d", domains), func(b *testing.B) { + s := newBenchmarkStats(b, domains) + b.ReportAllocs() + b.ResetTimer() + for range b.N { + benchmarkTotals = s.Totals() + } + }) + } +} + +func BenchmarkStatsRecordConcurrentWithSnapshots(b *testing.B) { + s := newBenchmarkStats(b, 10_000) + stop := make(chan struct{}) + var snapshots sync.WaitGroup + snapshots.Add(1) + go func() { + defer snapshots.Done() + ticker := time.NewTicker(time.Millisecond) + defer ticker.Stop() + for { + select { + case <-stop: + return + case <-ticker.C: + benchmarkSnapshot = s.Snapshot(DomainSortBytes, SourceAll, "") + } + } + }() + b.Cleanup(func() { + close(stop) + snapshots.Wait() + }) + + var next atomic.Uint64 + b.ReportAllocs() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + i := next.Add(1) % 10_000 + s.record(fmt.Sprintf("host%d.example", i), SourceHTTP, "", 0, 1, 1, + func(d *domainStat, src *sourceStat, dc *clientStat) { + d.bytesUp++ + d.bytesDown++ + src.bytesUp++ + src.bytesDown++ + }) + } + }) +} + +func BenchmarkStatsInsertAtDomainCapacity(b *testing.B) { + s := newBenchmarkStats(b, maxTotalDomains) + var next atomic.Uint64 + b.ReportAllocs() + b.ResetTimer() + for range b.N { + i := next.Add(1) + s.record(fmt.Sprintf("new%d.example", i), SourceHTTP, "", 1, 0, 0, + func(d *domainStat, src *sourceStat, dc *clientStat) { + d.conns++ + src.conns++ + }) + } +} diff --git a/internal/admin/stats_test.go b/internal/admin/stats_test.go index 7fd5b07..ecc64d4 100644 --- a/internal/admin/stats_test.go +++ b/internal/admin/stats_test.go @@ -1,8 +1,12 @@ package admin import ( + "bytes" "fmt" "net" + "slices" + "sort" + "sync" "testing" "time" ) @@ -77,6 +81,13 @@ func TestStatsBindConnAttributesBytesToDomain(t *testing.T) { t.Fatalf("write: %v", err) } + // Byte attribution is batched per connection: the pending batches are + // drained on Close (or when a threshold/interval is crossed), so close + // before asserting the per-domain totals. + if err := wrapped.Close(); err != nil { + t.Fatalf("close: %v", err) + } + snap := s.Snapshot(DomainSortBytes, SourceAll, "") if len(snap.Domains) != 1 { t.Fatalf("expected 1 domain, got %v", snap.Domains) @@ -107,6 +118,161 @@ func TestStatsBindConnIgnoresUnwrappedConn(t *testing.T) { } } +// TestCountingConnBatchWindow: bytes within the flush window stay pending; +// Close drains them so attribution is exact in the end. +func TestCountingConnBatchWindow(t *testing.T) { + s := newTestStats(t) + wrapped := s.WrapConn(&fakeConn{readBuf: bytes.Repeat([]byte("x"), 64)}, "http") + s.BindConn(wrapped, "example.com") + + buf := make([]byte, 1) + // First I/O flushes immediately (nextFlush starts zero): attributed now. + if _, err := wrapped.Read(buf); err != nil { + t.Fatalf("read: %v", err) + } + // Second I/O falls inside the flush window: kept pending. + if _, err := wrapped.Read(buf); err != nil { + t.Fatalf("read: %v", err) + } + + snap := s.Snapshot(DomainSortBytes, SourceAll, "") + if len(snap.Domains) != 1 || snap.Domains[0].BytesUp != 1 { + t.Fatalf("expected 1 pending byte unflushed, got %+v", snap.Domains) + } + + if err := wrapped.Close(); err != nil { + t.Fatalf("close: %v", err) + } + snap = s.Snapshot(DomainSortBytes, SourceAll, "") + if len(snap.Domains) != 1 || snap.Domains[0].BytesUp != 2 { + t.Fatalf("expected both bytes after close, got %+v", snap.Domains) + } +} + +// TestCountingConnBatchThreshold: sustained traffic flushes at the byte +// threshold without waiting for Close. +func TestCountingConnBatchThreshold(t *testing.T) { + s := newTestStats(t) + wrapped := s.WrapConn(&fakeConn{readBuf: bytes.Repeat([]byte("x"), 9*flushBytes)}, "https") + s.BindConn(wrapped, "example.com") + + chunk := make([]byte, 4096) + total := 0 + for total < 9*4096 { + n, err := wrapped.Read(chunk) + if err != nil { + t.Fatalf("read: %v", err) + } + total += n + } + if total != 9*4096 { + t.Fatalf("read %d bytes, want %d", total, 9*4096) + } + + // 9×4KiB crosses the 32KiB threshold (first chunk flushes immediately, + // the next eight accumulate to exactly the threshold), so the domain + // total is visible without Close. + snap := s.Snapshot(DomainSortBytes, SourceAll, "") + if len(snap.Domains) != 1 || snap.Domains[0].BytesUp != uint64(total) { + t.Fatalf("expected threshold flush of %d bytes, got %+v", total, snap.Domains) + } +} + +// TestCountingConnBatchConcurrent exercises the relay layout: Read and Write +// run on different goroutines, each direction's pending batch is owned by +// its goroutine, and Close drains both directions under concurrent I/O. +func TestCountingConnBatchConcurrent(t *testing.T) { + s := newTestStats(t) + wrapped := s.WrapConn(&fakeConn{readBuf: bytes.Repeat([]byte("y"), 64<<10)}, "socks5") + s.BindConn(wrapped, "example.com") + + const iters = 8 + readChunk, writeChunk := make([]byte, 4096), make([]byte, 4096) + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + for range iters { + if _, err := wrapped.Read(readChunk); err != nil { + t.Errorf("read: %v", err) + return + } + } + }() + go func() { + defer wg.Done() + for range iters { + if _, err := wrapped.Write(writeChunk); err != nil { + t.Errorf("write: %v", err) + return + } + } + }() + wg.Wait() + + if err := wrapped.Close(); err != nil { + t.Fatalf("close: %v", err) + } + + snap := s.Snapshot(DomainSortBytes, SourceAll, "") + if len(snap.Domains) != 1 { + t.Fatalf("expected 1 domain, got %v", snap.Domains) + } + d := snap.Domains[0] + want := uint64(iters * len(readChunk)) + if d.BytesUp != want || d.BytesDown != want { + t.Fatalf("expected %d up/%d down after concurrent relay, got %d/%d", + want, want, d.BytesUp, d.BytesDown) + } +} + +// TestCountingConnCloseConcurrent races Close against active relay I/O: the +// closed latch must force a final drain of bytes that land after Close's +// own swap, so no byte is left permanently pending. +func TestCountingConnCloseConcurrent(t *testing.T) { + s := newTestStats(t) + wrapped := s.WrapConn(&fakeConn{readBuf: bytes.Repeat([]byte("z"), 1<<20)}, "http") + s.BindConn(wrapped, "example.com") + + const iters = 500 + readChunk, writeChunk := make([]byte, 4096), make([]byte, 4096) + var wg sync.WaitGroup + wg.Add(3) + go func() { + defer wg.Done() + for i := 0; i < iters; i++ { + if _, err := wrapped.Read(readChunk); err != nil { + t.Errorf("read: %v", err) + return + } + } + }() + go func() { + defer wg.Done() + for i := 0; i < iters; i++ { + if _, err := wrapped.Write(writeChunk); err != nil { + t.Errorf("write: %v", err) + return + } + } + }() + go func() { + defer wg.Done() + if err := wrapped.Close(); err != nil { + t.Errorf("close: %v", err) + } + }() + wg.Wait() + + // After Close, any byte that raced in must still be drained by the + // closed-latch flush; totals therefore equal the writes that reported + // success, regardless of interleaving. + want := uint64(iters * len(readChunk)) + if got := s.bytesDown.Load(); got != want { + t.Fatalf("expected %d total bytes down, got %d", want, got) + } +} + func TestStatsRecordDNS(t *testing.T) { s := newTestStats(t) s.RecordDNS("Example.COM.", "") @@ -139,31 +305,45 @@ func TestStatsSnapshotEvictsStaleDomains(t *testing.T) { func TestStatsDomainCapIsEnforced(t *testing.T) { s := newTestStats(t) - for i := 0; i < maxTotalDomains+50; i++ { - s.record(fmt.Sprintf("host%d.example", i), SourceHTTP, "", 1, 0, 0, func(d *domainStat, src *sourceStat, dc *clientStat) {}) + now := time.Now() + s.mu.Lock() + for i := 0; i < maxTotalDomains; i++ { + s.domains[fmt.Sprintf("host%d.example", i)] = &domainStat{lastSeen: now.Add(time.Duration(i) * time.Nanosecond)} } + s.mu.Unlock() + + s.record("fresh.example", SourceHTTP, "", 1, 0, 0, func(d *domainStat, src *sourceStat, dc *clientStat) {}) + s.mu.Lock() n := len(s.domains) + _, oldestPresent := s.domains["host0.example"] + _, newestPresent := s.domains[fmt.Sprintf("host%d.example", maxTotalDomains-1)] + _, freshPresent := s.domains["fresh.example"] s.mu.Unlock() - if n > maxTotalDomains { - t.Fatalf("expected domain cap enforced, got %d entries", n) + want := maxTotalDomains - domainEvictionBatch + 1 + if n != want { + t.Fatalf("expected %d domains after batch eviction, got %d", want, n) } - if n == 0 { - t.Fatal("expected some domains to survive") + if oldestPresent || !newestPresent || !freshPresent { + t.Fatalf("unexpected retained domains: oldest=%v newest=%v fresh=%v", oldestPresent, newestPresent, freshPresent) } } -func TestEvictOldestLocked(t *testing.T) { +func TestEvictOldestDomainsLocked(t *testing.T) { s := newTestStats(t) + now := time.Now() s.mu.Lock() - s.domains["old.example"] = &domainStat{lastSeen: time.Now().Add(-time.Hour)} - s.domains["new.example"] = &domainStat{lastSeen: time.Now()} - s.evictOldestLocked() - _, oldOK := s.domains["old.example"] - _, newOK := s.domains["new.example"] + for i := 0; i < 4; i++ { + s.domains[fmt.Sprintf("host%d.example", i)] = &domainStat{lastSeen: now.Add(time.Duration(i) * time.Second)} + } + s.evictOldestDomainsLocked(2) + _, firstPresent := s.domains["host0.example"] + _, secondPresent := s.domains["host1.example"] + _, thirdPresent := s.domains["host2.example"] + _, fourthPresent := s.domains["host3.example"] s.mu.Unlock() - if oldOK || !newOK { - t.Fatalf("expected oldest evicted, old=%v new=%v", oldOK, newOK) + if firstPresent || secondPresent || !thirdPresent || !fourthPresent { + t.Fatalf("unexpected eviction: first=%v second=%v third=%v fourth=%v", firstPresent, secondPresent, thirdPresent, fourthPresent) } } @@ -577,17 +757,28 @@ func TestStatsTotalsClientsSurviveDomainEviction(t *testing.T) { func TestStatsClientCapIsEnforced(t *testing.T) { s := newTestStats(t) - for i := 0; i < maxTotalClients+10; i++ { - s.record(fmt.Sprintf("host%d.example", i), SourceHTTP, fmt.Sprintf("192.168.%d.%d", i/256, i%256), 1, 0, 0, func(d *domainStat, src *sourceStat, dc *clientStat) {}) + now := time.Now() + s.mu.Lock() + for i := 0; i < maxTotalClients; i++ { + s.clients[fmt.Sprintf("192.168.%d.%d", i/256, i%256)] = &clientStat{lastSeen: now.Add(time.Duration(i) * time.Nanosecond)} } + s.domains["existing.example"] = &domainStat{bySource: make(map[Source]*sourceStat), lastSeen: now} + s.mu.Unlock() + + s.record("existing.example", SourceHTTP, "203.0.113.1", 1, 0, 0, func(d *domainStat, src *sourceStat, dc *clientStat) {}) + s.mu.Lock() n := len(s.clients) + _, oldestPresent := s.clients["192.168.0.0"] + _, newestPresent := s.clients[fmt.Sprintf("192.168.%d.%d", (maxTotalClients-1)/256, (maxTotalClients-1)%256)] + _, freshPresent := s.clients["203.0.113.1"] s.mu.Unlock() - if n > maxTotalClients { - t.Fatalf("expected client cap enforced, got %d entries", n) + want := maxTotalClients - clientEvictionBatch + 1 + if n != want { + t.Fatalf("expected %d clients after batch eviction, got %d", want, n) } - if n == 0 { - t.Fatal("expected some clients to survive") + if oldestPresent || !newestPresent || !freshPresent { + t.Fatalf("unexpected retained clients: oldest=%v newest=%v fresh=%v", oldestPresent, newestPresent, freshPresent) } } @@ -603,3 +794,62 @@ func TestStatsClientsPerDomainCap(t *testing.T) { t.Fatalf("expected per-domain client cap, got %d entries", n) } } + +func TestTopDomainsMatchesFullSort(t *testing.T) { + now := time.Now() + input := make([]DomainStat, 0, snapshotTopN+37) + for i := 0; i < snapshotTopN+37; i++ { + input = append(input, DomainStat{ + Domain: fmt.Sprintf("host%03d.example", i), + Conns: uint64((i * 7) % 31), + BytesUp: uint64((i * 13) % 101), + BytesDown: uint64((i * 17) % 97), + LastSeen: now.Add(time.Duration((i*19)%71) * time.Second), + }) + } + + for _, mode := range []DomainSort{DomainSortBytes, DomainSortRecent, DomainSortConns} { + want := slices.Clone(input) + sortDomains(want, mode) + want = want[:snapshotTopN] + got := topDomains(slices.Clone(input), snapshotTopN, mode) + if !slices.Equal(got, want) { + t.Fatalf("top domains differ for %s:\n got %v\nwant %v", mode, got, want) + } + } + + tied := topDomains([]DomainStat{ + {Domain: "z.example", BytesUp: 1}, + {Domain: "a.example", BytesUp: 1}, + }, 2, DomainSortBytes) + if tied[0].Domain != "a.example" || tied[1].Domain != "z.example" { + t.Fatalf("expected domain tie break by name, got %v", tied) + } +} + +func TestTopClientsMatchesFullSort(t *testing.T) { + agg := make(map[string]*ClientStat, clientsTopN+17) + for i := 0; i < clientsTopN+17; i++ { + ip := fmt.Sprintf("192.0.2.%03d", i) + agg[ip] = &ClientStat{IP: ip, BytesUp: uint64((i * 11) % 113), BytesDown: uint64((i * 23) % 109)} + } + + want := make([]ClientStat, 0, len(agg)) + for _, client := range agg { + want = append(want, *client) + } + sort.Slice(want, func(i, j int) bool { return clientLess(want[i], want[j]) }) + want = want[:clientsTopN] + got := topClients(agg) + if !slices.Equal(got, want) { + t.Fatalf("top clients differ:\n got %v\nwant %v", got, want) + } + + tied := topClients(map[string]*ClientStat{ + "192.0.2.2": {IP: "192.0.2.2", BytesUp: 1}, + "192.0.2.1": {IP: "192.0.2.1", BytesUp: 1}, + }) + if tied[0].IP != "192.0.2.1" || tied[1].IP != "192.0.2.2" { + t.Fatalf("expected IP tie break, got %v", tied) + } +} diff --git a/internal/metrics/metrics.go b/internal/metrics/metrics.go index 5db439a..cbef258 100644 --- a/internal/metrics/metrics.go +++ b/internal/metrics/metrics.go @@ -92,15 +92,21 @@ func New() (*Metrics, error) { return m, nil } +// Precomputed AddOptions avoid an attribute-slice allocation on every +// RecordBytes call, which runs on the per-I/O relay hot path. +var ( + bytesAttrUp = metric.WithAttributes(attribute.String("direction", "up")) + bytesAttrDown = metric.WithAttributes(attribute.String("direction", "down")) +) + // RecordBytes counts proxied payload bytes. up is relative to the client: // reads from the client are uploads, writes to the client downloads. func (m *Metrics) RecordBytes(up bool, n int) { - dir := "down" if up { - dir = "up" + m.bytes.Add(context.Background(), int64(n), bytesAttrUp) + } else { + m.bytes.Add(context.Background(), int64(n), bytesAttrDown) } - m.bytes.Add(context.Background(), int64(n), - metric.WithAttributes(attribute.String("direction", dir))) } // RecordDNS counts one DNS query. diff --git a/pkg/suffixtree/suffix_tree.go b/pkg/suffixtree/suffix_tree.go index c05be31..3a3452f 100644 --- a/pkg/suffixtree/suffix_tree.go +++ b/pkg/suffixtree/suffix_tree.go @@ -6,7 +6,9 @@ import ( type Node struct { *node - sep string + sep string + // Deprecated: unused. Retained to avoid breaking the package API; new + // code should not depend on it. Count uint64 } type node struct { @@ -131,7 +133,13 @@ func (n *node) matchSecs(secs []string) bool { return n.index("**") >= 0 } -// index return the sec index in node, or -1 if not found +// indexMapThreshold: below this many children a node keeps a plain slice +// scan and no map, saving ~50+ bytes of map header and entries per small +// node; suffix trees are dominated by shallow fan-out nodes. +const indexMapThreshold = 8 + +// index return the sec index in node, or -1 if not found. Read-only: Match +// runs concurrently on a shared tree, so no map is built on this path. func (n *node) index(sec string) int { if n == nil { return -1 @@ -146,8 +154,6 @@ func (n *node) index(sec string) int { for i := range n.secs { if n.secs[i] == sec { - n.ensureIndexMap() - n.indexMap[sec] = i return i } } @@ -160,7 +166,9 @@ func (n *node) append(sec string, child *node) int { n.subNodes = append(n.subNodes, child) n.ensureIndexMap() idx := len(n.secs) - 1 - n.indexMap[sec] = idx + if n.indexMap != nil { + n.indexMap[sec] = idx + } return idx } @@ -173,7 +181,7 @@ func (n *node) prepend(sec string, child *node) int { } func (n *node) ensureIndexMap() { - if n.indexMap != nil { + if n.indexMap != nil || len(n.secs) < indexMapThreshold { return } diff --git a/pkg/suffixtree/suffix_tree_internal_test.go b/pkg/suffixtree/suffix_tree_internal_test.go new file mode 100644 index 0000000..4d39668 --- /dev/null +++ b/pkg/suffixtree/suffix_tree_internal_test.go @@ -0,0 +1,110 @@ +package suffixtree + +import ( + "fmt" + "sync" + "testing" +) + +// TestIndexMapThreshold: nodes with few children stay mapless (slice scan), +// nodes at or above the threshold carry a map. +func TestIndexMapThreshold(t *testing.T) { + rules := make([]string, 0, 20) + for i := 0; i < 20; i++ { + rules = append(rules, fmt.Sprintf("host%d.wweir.cc", i)) + } + n := NewNodeFromRules(rules...) + + // The "cc" node has one child ("wweir"), the "wweir" node has 20 children. + cc := n.node.subNodes[n.index("cc")] + if cc == nil || cc.indexMap != nil { + t.Fatalf("expected single-child node without indexMap, got %+v", cc) + } + wweir := cc.subNodes[cc.index("wweir")] + if wweir == nil || wweir.indexMap == nil { + t.Fatalf("expected 20-child node to carry indexMap, got %+v", wweir) + } + if len(wweir.indexMap) != 20 { + t.Fatalf("expected 20 indexMap entries, got %d", len(wweir.indexMap)) + } +} + +// TestConcurrentMatchNoRace: Match runs concurrently on a shared tree; the +// index lookup must be read-only (no lazy map building on the query path). +func TestConcurrentMatchNoRace(t *testing.T) { + rules := make([]string, 0, 10000) + for i := 0; i < 10000; i++ { + rules = append(rules, fmt.Sprintf("host%d.example.com", i)) + } + rules = append(rules, "**.example.com", "*.wweir.cc") + n := NewNodeFromRules(rules...) + + var wg sync.WaitGroup + for g := 0; g < 8; g++ { + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 2000; i++ { + n.Match(fmt.Sprintf("host%d.example.com", i%10000)) + n.Match("a.b.example.com") + n.Match("x.wweir.cc") + } + }() + } + wg.Wait() +} + +// TestIndexMapMemoryFootprint walks a realistic rule shape and reports how +// many nodes keep a map (>=8 children) versus the mapless small nodes that +// previously carried a ~48B+ map header each. +func TestIndexMapMemoryFootprint(t *testing.T) { + // Many distinct second-level domains create lots of 1-2 child nodes. + rules := make([]string, 0, 30000) + for i := 0; i < 10000; i++ { + rules = append(rules, fmt.Sprintf("ads%d.foo%d.com", i, i%50)) + rules = append(rules, fmt.Sprintf("track%d.bar%d.net", i, i%50)) + rules = append(rules, fmt.Sprintf("*.cdn%d.org", i%100)) + } + n := NewNodeFromRules(rules...) + + var nodes, mapped, small int + var walk func(*node) + walk = func(cur *node) { + if cur == nil { + return + } + nodes++ + if len(cur.secs) >= indexMapThreshold { + mapped++ + } else if len(cur.secs) > 0 { + small++ + } + for _, sub := range cur.subNodes { + walk(sub) + } + } + walk(n.node) + + // Each mapless small node saves a map header (~48B) plus per-entry + // overhead (~16B); the slice scan costs nothing extra in memory. + saved := small * 48 + t.Logf("nodes=%d mapped=%d small=%d estimated savings≈%d KiB", + nodes, mapped, small, saved/1024) + if small == 0 { + t.Fatal("expected mapless small nodes in this shape") + } +} + +// BenchmarkSuffixTreeBuildAlloc reports the retained per-rule cost of a +// large rule set (map header savings included). +func BenchmarkSuffixTreeBuildAlloc(b *testing.B) { + rules := make([]string, 0, 100000) + for i := 0; i < 100000; i++ { + rules = append(rules, fmt.Sprintf("host%d.sub.example.com", i)) + } + b.ReportAllocs() + b.ResetTimer() + for range b.N { + NewNodeFromRules(rules...) + } +} diff --git a/router/dns_bench_test.go b/router/dns_bench_test.go new file mode 100644 index 0000000..431fdde --- /dev/null +++ b/router/dns_bench_test.go @@ -0,0 +1,50 @@ +package router + +import ( + "fmt" + "testing" + "time" +) + +// BenchmarkDNSPrepareUpstreamState measures the per-query upstream state +// snapshot cost (mutex + address-slice copy). +func BenchmarkDNSPrepareUpstreamState(b *testing.B) { + r, err := NewRouter(nil, "", "223.5.5.5", "", nil) + if err != nil { + b.Fatal(err) + } + r.dns.upstreamAddrs = []string{"1.1.1.1:53", "8.8.8.8:53", "223.5.5.5:53"} + r.dns.upstreamIndex = 0 + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + r.prepareUpstreamState(time.Now()) + } +} + +// BenchmarkDNSRuleMatchPath measures the rule-classification cost of one +// DNS query: domain-list construction plus up to three RuleSet.Match calls. +func BenchmarkDNSRuleMatchPath(b *testing.B) { + r, err := NewRouter(nil, "", "223.5.5.5", "", nil) + if err != nil { + b.Fatal(err) + } + rules := make([]string, 0, 10000) + for i := 0; i < 10000; i++ { + rules = append(rules, fmt.Sprintf("host%d.example.com", i)) + } + r.BlockRule.Replace(rules...) + r.DirectRule.Replace("direct.example.com") + r.ProxyRule.Replace("proxy.example.com") + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + domains := dnsRouteDomains("www.example.com", 1) + _ = dnsRuleMatch(r.BlockRule, domains) + if !dnsRuleMatch(r.DirectRule, domains) { + _ = dnsRuleMatch(r.ProxyRule, domains) + } + } +} diff --git a/router/ruleset.go b/router/ruleset.go index a7b0d45..216d30e 100644 --- a/router/ruleset.go +++ b/router/ruleset.go @@ -10,6 +10,8 @@ import ( // RuleSet is a thread-safe rule container. It retains the raw rule list so // rules can be listed and removed at runtime; the suffixtree is derived // state and is rebuilt from the retained list whenever a rule is removed. +// Membership for Add/Remove uses the raw rule strings (set), so case and +// trailing-dot variants stay distinct entries in the configuration. type RuleSet struct { mu sync.RWMutex rules []string @@ -25,7 +27,8 @@ func NewRuleSet(rules ...string) *RuleSet { return rs } -// Add appends rules that are not already present. Duplicates are ignored. +// Add appends rules that are not already present, using the raw rule string +// as the membership key. Duplicates are ignored. func (rs *RuleSet) Add(rules ...string) { if len(rules) == 0 { return @@ -46,6 +49,32 @@ func (rs *RuleSet) Add(rules ...string) { } } +// Replace swaps the entire rule list atomically, rebuilding the suffix +// tree. Literal duplicates in the input are dropped. It is used to rebuild +// a rule set from the boot baseline after a delta reset. +func (rs *RuleSet) Replace(rules ...string) { + // Construct the full immutable candidate before taking the match lock. + // The field swap below is the operation's linearization point. + nextRules := make([]string, 0, len(rules)) + nextSet := make(map[string]struct{}, len(rules)) + nextTree := suffixtree.NewNodeFromRules() + for _, rule := range rules { + if _, ok := nextSet[rule]; ok { + continue + } + nextSet[rule] = struct{}{} + nextRules = append(nextRules, rule) + nextTree.Add(rule) + } + nextTree.GC() + + rs.mu.Lock() + rs.rules = nextRules + rs.set = nextSet + rs.tree = nextTree + rs.mu.Unlock() +} + // Compact runs garbage collection on the underlying tree. Call it after a // batch of adds to avoid per-add compaction overhead on large rule files. func (rs *RuleSet) Compact() { diff --git a/router/ruleset_bench_test.go b/router/ruleset_bench_test.go new file mode 100644 index 0000000..8a948be --- /dev/null +++ b/router/ruleset_bench_test.go @@ -0,0 +1,34 @@ +package router + +import ( + "fmt" + "testing" +) + +func BenchmarkRuleSetMatch(b *testing.B) { + rules := make([]string, 0, 10000) + for i := 0; i < 10000; i++ { + rules = append(rules, fmt.Sprintf("host%d.example.com", i)) + } + rs := NewRuleSet(rules...) + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + rs.Match("www.example.com") + } +} + +func BenchmarkRuleSetMatchHit(b *testing.B) { + rules := make([]string, 0, 10000) + for i := 0; i < 10000; i++ { + rules = append(rules, fmt.Sprintf("host%d.example.com", i)) + } + rs := NewRuleSet(rules...) + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + rs.Match("host5000.example.com") + } +} diff --git a/router/ruleset_memory_test.go b/router/ruleset_memory_test.go new file mode 100644 index 0000000..defb36c --- /dev/null +++ b/router/ruleset_memory_test.go @@ -0,0 +1,66 @@ +package router + +import ( + "fmt" + "runtime" + "testing" +) + +// keepGlobal pins the trees across GC; without it the compiler may treat +// them as dead after the Match checks. +var keepGlobal struct { + block *RuleSet + direct *RuleSet + proxy *RuleSet +} + +// TestRulesetMemoryFootprint builds rule shapes matching the production +// ad/china/gfw lists (synthetic, no external files) and reports the retained +// heap. Skipped under -short. +func TestRulesetMemoryFootprint(t *testing.T) { + if testing.Short() { + t.Skip("memory footprint test; skipped in -short mode") + } + t.Cleanup(func() { + keepGlobal = struct { + block *RuleSet + direct *RuleSet + proxy *RuleSet + }{} + }) + + // adlist-like: many exact domains under a few second-level domains. + block := make([]string, 0, 160000) + for i := 0; i < 160000; i++ { + block = append(block, fmt.Sprintf("ads%d.doubleclick.net", i)) + } + // chinalist-like: many exact domains under many second-level domains. + direct := make([]string, 0, 110000) + for i := 0; i < 110000; i++ { + direct = append(direct, fmt.Sprintf("host%d.example%d.com", i, i%2000)) + } + // gfwlist-like: wildcard-heavy. + proxy := make([]string, 0, 5000) + for i := 0; i < 5000; i++ { + proxy = append(proxy, fmt.Sprintf("**.site%d.com", i%100)) + } + t.Logf("rules: block=%d direct=%d proxy=%d", len(block), len(direct), len(proxy)) + + keepGlobal.block = NewRuleSet(block...) + keepGlobal.direct = NewRuleSet(direct...) + keepGlobal.proxy = NewRuleSet(proxy...) + + runtime.GC() + var m runtime.MemStats + runtime.ReadMemStats(&m) + before := m.HeapAlloc + + // Keep references live across GC. + if !keepGlobal.block.Match("ads1.doubleclick.net") || !keepGlobal.direct.Match("host1.example1.com") || !keepGlobal.proxy.Match("a.b.site1.com") { + t.Fatal("unexpected match result") + } + runtime.GC() + runtime.ReadMemStats(&m) + t.Logf("retained after GC: %d KiB (%.1f MB)", m.HeapAlloc/1024, float64(m.HeapAlloc)/1024/1024) + t.Logf("heap before/after: %d / %d", before, m.HeapAlloc) +} diff --git a/router/ruleset_test.go b/router/ruleset_test.go index b00f22f..9dc6ab0 100644 --- a/router/ruleset_test.go +++ b/router/ruleset_test.go @@ -3,6 +3,7 @@ package router import ( "fmt" "sync" + "sync/atomic" "testing" ) @@ -149,6 +150,32 @@ func TestRuleSetConcurrentAccess(t *testing.T) { } } +func TestRuleSetReplaceConcurrentWithMatch(t *testing.T) { + rs := NewRuleSet("seed.example.com") + var stop atomic.Bool + var wg sync.WaitGroup + for range 8 { + wg.Add(1) + go func() { + defer wg.Done() + for !stop.Load() { + _ = rs.Match("seed.example.com") + _ = rs.Match("replacement.example.com") + } + }() + } + + for i := 0; i < 100; i++ { + rs.Replace(fmt.Sprintf("replacement%d.example.com", i), "replacement.example.com") + } + stop.Store(true) + wg.Wait() + + if !rs.Match("replacement.example.com") || rs.Match("seed.example.com") { + t.Fatalf("unexpected final rules: %v", rs.List()) + } +} + func TestListFilteredPaginationAndSearch(t *testing.T) { rs := NewRuleSet("example.com", "**.cn", "GitHub.com", "mail.google.com", "a.cn")