refactor: extract helpers from main.go, name magic numbers, add doc comments

cmd/picodns/main.go

@@ -30,10 +30,32 @@ func main() {
 
 	cacheStore := cache.New(cfg.CacheSize, nil)
 
+	res, cached, rec, upstream := buildResolver(cfg, logger, cacheStore, ctx)
+
+	srv := server.New(cfg, logger, res)
+	if cached != nil && cfg.Stats {
+		srv.SetCacheCounters(func() (uint64, uint64) {
+			return cached.CacheHits.Load(), cached.CacheMiss.Load()
+		})
+	}
+
+	err := srv.Start(ctx)
+	if cfg.Stats {
+		printStats(logger, cached, rec, upstream)
+	}
+
+	if err != nil && !errors.Is(err, context.Canceled) {
+		logger.Error("server exited", "error", err)
+		os.Exit(1)
+	}
+}
+
+func buildResolver(cfg config.Config, logger *slog.Logger, cacheStore *cache.Cache, ctx context.Context) (types.Resolver, *resolver.Cached, *resolver.Recursive, *resolver.Upstream) {
 	var res types.Resolver
 	var cached *resolver.Cached
 	var rec *resolver.Recursive
 	var upstream *resolver.Upstream
+
 	if cfg.Recursive || len(cfg.Upstreams) == 0 {
 		logger.Info("using recursive resolver")
 		rec = resolver.NewRecursive()
@@ -62,70 +84,48 @@ func main() {
 		res = cached
 	}
 
-	srv := server.New(cfg, logger, res)
-	if cached != nil && cfg.Stats {
-		srv.SetCacheCounters(func() (uint64, uint64) {
-			return cached.CacheHits.Load(), cached.CacheMiss.Load()
-		})
-	}
+	return res, cached, rec, upstream
+}
 
-	err := srv.Start(ctx)
-	if cfg.Stats {
-		if rec != nil {
-			logger.Info("recursive internal cache stats",
-				"ns_cache", rec.NSCacheStatsSnapshot(),
-				"delegation_cache", rec.DelegationCacheStatsSnapshot(),
-			)
-			addr := rec.TransportAddrCacheStatsSnapshot()
-			addrHitRate := 0.0
-			if addr.Gets > 0 {
-				addrHitRate = float64(addr.Hits) / float64(addr.Gets)
-			}
-			logger.Info("transport addr cache stats",
-				"gets", addr.Gets,
-				"hits", addr.Hits,
-				"misses", addr.Misses,
-				"sets", addr.Sets,
-				"deletes", addr.Deletes,
-				"len", addr.Len,
-				"hit_rate", addrHitRate,
-			)
-		}
-		if upstream != nil {
-			addr := upstream.TransportAddrCacheStatsSnapshot()
-			addrHitRate := 0.0
-			if addr.Gets > 0 {
-				addrHitRate = float64(addr.Hits) / float64(addr.Gets)
-			}
-			logger.Info("transport addr cache stats",
-				"gets", addr.Gets,
-				"hits", addr.Hits,
-				"misses", addr.Misses,
-				"sets", addr.Sets,
-				"deletes", addr.Deletes,
-				"len", addr.Len,
-				"hit_rate", addrHitRate,
-			)
+func printStats(logger *slog.Logger, cached *resolver.Cached, rec *resolver.Recursive, upstream *resolver.Upstream) {
+	if rec != nil {
+		logger.Info("recursive internal cache stats",
+			"ns_cache", rec.NSCacheStatsSnapshot(),
+			"delegation_cache", rec.DelegationCacheStatsSnapshot(),
+		)
+		addr := rec.TransportAddrCacheStatsSnapshot()
+		addrHitRate := 0.0
+		if addr.Gets > 0 {
+			addrHitRate = float64(addr.Hits) / float64(addr.Gets)
 		}
-
-		if cached != nil {
-			snap := cached.StatsSnapshot()
-			hitRate := 0.0
-			if snap.Hits+snap.Miss > 0 {
-				hitRate = float64(snap.Hits) / float64(snap.Hits+snap.Miss)
-			}
-
-			logger.Info("resolver cache stats",
-				"cache_hits", snap.Hits,
-				"cache_miss", snap.Miss,
-				"cache_hit_rate", hitRate,
-			)
+		logger.Info("transport addr cache stats",
+			"gets", addr.Gets, "hits", addr.Hits, "misses", addr.Misses,
+			"sets", addr.Sets, "deletes", addr.Deletes, "len", addr.Len,
+			"hit_rate", addrHitRate,
+		)
+	}
+	if upstream != nil {
+		addr := upstream.TransportAddrCacheStatsSnapshot()
+		addrHitRate := 0.0
+		if addr.Gets > 0 {
+			addrHitRate = float64(addr.Hits) / float64(addr.Gets)
 		}
+		logger.Info("transport addr cache stats",
+			"gets", addr.Gets, "hits", addr.Hits, "misses", addr.Misses,
+			"sets", addr.Sets, "deletes", addr.Deletes, "len", addr.Len,
+			"hit_rate", addrHitRate,
+		)
 	}
-
-	if err != nil && !errors.Is(err, context.Canceled) {
-		logger.Error("server exited", "error", err)
-		os.Exit(1)
+	if cached != nil {
+		snap := cached.StatsSnapshot()
+		hitRate := 0.0
+		if snap.Hits+snap.Miss > 0 {
+			hitRate = float64(snap.Hits) / float64(snap.Hits+snap.Miss)
+		}
+		logger.Info("resolver cache stats",
+			"cache_hits", snap.Hits, "cache_miss", snap.Miss,
+			"cache_hit_rate", hitRate,
+		)
 	}
 }
 

internal/resolver/cache.go

@@ -15,6 +15,8 @@ import (
 	"picodns/internal/types"
 )
 
+// Cached wraps a Resolver with a read-through cache, inflight deduplication,
+// stale-while-revalidate, and optional prefetch.
 type Cached struct {
 	cache      *cache.Cache
 	upstream   types.Resolver
@@ -46,6 +48,7 @@ type Cached struct {
 	}
 }
 
+// NewCached creates a Cached resolver that fronts the given upstream with a cache layer.
 func NewCached(cacheStore *cache.Cache, upstream types.Resolver) *Cached {
 	c := &Cached{
 		cache:    cacheStore,
@@ -87,11 +90,13 @@ func NewCached(cacheStore *cache.Cache, upstream types.Resolver) *Cached {
 	return c
 }
 
+// CachedStatsSnapshot holds a point-in-time snapshot of cache hit/miss counters.
 type CachedStatsSnapshot struct {
 	Hits uint64
 	Miss uint64
 }
 
+// StatsSnapshot returns a point-in-time snapshot of cache hit/miss counters.
 func (c *Cached) StatsSnapshot() CachedStatsSnapshot {
 	return CachedStatsSnapshot{
 		Hits: c.CacheHits.Load(),
@@ -150,6 +155,8 @@ func readQuestions(req []byte, qdCount uint16) ([]dns.Question, bool) {
 	return qs, true
 }
 
+// ResolveFromCache attempts to serve a response from cache without forwarding upstream.
+// It returns the response, a cleanup function, and whether the cache contained an entry.
 func (c *Cached) ResolveFromCache(ctx context.Context, req []byte) ([]byte, func(), bool) {
 	sampled := c.tracers.resolveFromCache.ShouldSample()
 	done := c.tracers.resolveFromCache.TraceSampled(sampled)
@@ -173,6 +180,8 @@ func (c *Cached) ResolveFromCache(ctx context.Context, req []byte) ([]byte, func
 	return nil, nil, false
 }
 
+// Resolve looks up the answer in cache first, falling back to the upstream resolver on a miss.
+// It deduplicates concurrent requests for the same question and caches successful responses.
 func (c *Cached) Resolve(ctx context.Context, req []byte) ([]byte, func(), error) {
 	sampled := c.tracers.resolve.ShouldSample()
 	done := c.tracers.resolve.TraceSampled(sampled)

internal/resolver/conn_pool.go

@@ -57,7 +57,7 @@ func (p *connPool) get(ctx context.Context) (*net.UDPConn, func(bad bool), error
 
 	p.mu.Lock()
 	p.getCount++
-	prune := p.getCount%64 == 0 || (p.lastPrune.IsZero() || now.Sub(p.lastPrune) >= ConnPoolIdleTimeout)
+	prune := p.getCount%connPoolPruneEvery == 0 || (p.lastPrune.IsZero() || now.Sub(p.lastPrune) >= ConnPoolIdleTimeout)
 	if prune {
 		p.lastPrune = now
 		for i := 0; i < len(p.conns); {

internal/resolver/constants.go

@@ -44,6 +44,11 @@ const (
 	prefetchTimeout        = 10 * time.Second // Timeout for background prefetch operations
 	serveStaleFor          = 30 * time.Second // Serve expired cache entries for this long (stale-while-revalidate)
 
+	// EWMA and tracking constants
+	ewmaWeight         = 5  // EWMA denominator: new = (prev*(weight-1) + sample) / weight
+	maxTimeoutCount    = 6  // Maximum consecutive timeout count before capping
+	connPoolPruneEvery = 64 // Prune idle connections every N gets
+
 	// Parallel query settings
 	defaultMaxServers  = 3                      // Maximum concurrent servers for normal queries
 	glueMaxServers     = 2                      // Maximum concurrent servers for glue queries

internal/resolver/recursive.go

@@ -23,6 +23,8 @@ type inflightRecursive struct {
 	err  error
 }
 
+// Recursive performs full iterative DNS resolution starting from the root servers.
+// It maintains a connection pool, RTT tracker, NS cache, and delegation cache.
 type Recursive struct {
 	transport       types.Transport
 	bufPool         *pool.Bytes
@@ -99,6 +101,8 @@ func NewRecursive(opts ...Option) *Recursive {
 	return r
 }
 
+// SetObsEnabled enables or disables observability (stats collection) on the
+// resolver and its internal caches and transport.
 func (r *Recursive) SetObsEnabled(enabled bool) {
 	r.ObsEnabled = enabled
 	if r.nsCache != nil {
@@ -112,20 +116,23 @@ func (r *Recursive) SetObsEnabled(enabled bool) {
 	}
 }
 
+// NSCacheStatsSnapshot returns a point-in-time snapshot of the NS name cache statistics.
 func (r *Recursive) NSCacheStatsSnapshot() cache.TTLStatsSnapshot {
 	if r.nsCache == nil {
 		return cache.TTLStatsSnapshot{}
 	}
 	return r.nsCache.StatsSnapshot()
 }
 
+// DelegationCacheStatsSnapshot returns a point-in-time snapshot of the delegation cache statistics.
 func (r *Recursive) DelegationCacheStatsSnapshot() cache.TTLStatsSnapshot {
 	if r.delegationCache == nil || r.delegationCache.TTL == nil {
 		return cache.TTLStatsSnapshot{}
 	}
 	return r.delegationCache.StatsSnapshot()
 }
 
+// TransportAddrCacheStatsSnapshot returns a point-in-time snapshot of the transport address cache statistics.
 func (r *Recursive) TransportAddrCacheStatsSnapshot() cache.TTLStatsSnapshot {
 	if t, ok := r.transport.(*udpTransport); ok && t.addrCache != nil {
 		return t.addrCache.StatsSnapshot()
@@ -140,6 +147,8 @@ type resolutionStats struct {
 	glueLookups  int           // NS name resolution queries (when no glue records)
 }
 
+// Resolve resolves a DNS query by iteratively querying authoritative servers
+// from the root. It deduplicates concurrent identical queries.
 func (r *Recursive) Resolve(ctx context.Context, req []byte) ([]byte, func(), error) {
 	defer r.tracers.resolve.Trace()()
 

internal/resolver/rtt_tracker.go

@@ -58,7 +58,7 @@ func (t *rttTracker) Update(ctx context.Context, server string, d time.Duration)
 	if !ok {
 		t.rtts[server] = d
 	} else {
-		t.rtts[server] = (prev*4 + d) / 5
+		t.rtts[server] = (prev*(ewmaWeight-1) + d) / ewmaWeight
 	}
 	delete(t.timeouts, server)
 	delete(t.cooldown, server)
@@ -77,8 +77,8 @@ func (t *rttTracker) Failure(ctx context.Context, server string) {
 		}
 	}
 	count := t.timeouts[server] + 1
-	if count > 6 {
-		count = 6
+	if count > maxTimeoutCount {
+		count = maxTimeoutCount
 	}
 	backoff := baseTimeoutBackoff << (count - 1)
 	if backoff > maxTimeoutBackoff {

internal/resolver/upstream.go

@@ -13,24 +13,29 @@ var (
 	ErrNoUpstreams = errors.New("resolver: no upstreams configured")
 )
 
+// Upstream forwards DNS queries to one or more configured upstream resolvers,
+// trying each in order until one succeeds.
 type Upstream struct {
 	upstreams []string
 	transport types.Transport
 }
 
+// SetObsEnabled enables or disables observability on the upstream transport.
 func (u *Upstream) SetObsEnabled(enabled bool) {
 	if t, ok := u.transport.(*udpTransport); ok {
 		t.SetObsEnabled(enabled)
 	}
 }
 
+// TransportAddrCacheStatsSnapshot returns a point-in-time snapshot of the transport address cache statistics.
 func (u *Upstream) TransportAddrCacheStatsSnapshot() cache.TTLStatsSnapshot {
 	if t, ok := u.transport.(*udpTransport); ok && t.addrCache != nil {
 		return t.addrCache.StatsSnapshot()
 	}
 	return cache.TTLStatsSnapshot{}
 }
 
+// NewUpstream creates an Upstream resolver that forwards queries to the given addresses.
 func NewUpstream(upstreamAddrs []string) (*Upstream, error) {
 	if len(upstreamAddrs) == 0 {
 		return nil, ErrNoUpstreams
@@ -48,6 +53,7 @@ func NewUpstream(upstreamAddrs []string) (*Upstream, error) {
 	}, nil
 }
 
+// Resolve forwards the query to each upstream in order, returning the first successful response.
 func (u *Upstream) Resolve(ctx context.Context, req []byte) ([]byte, func(), error) {
 	if len(u.upstreams) == 0 {
 		return nil, nil, ErrNoUpstreams