bucketeer/internal/dns/pool.go

package dns

import (
	"context"
	"errors"
	mathrand "math/rand/v2"
	"net"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// CuratedServers are well-known, reliable public resolvers used by default.
func CuratedServers() []string {
	return []string{
		"1.1.1.1", "1.0.0.1", // Cloudflare
		"8.8.8.8", "8.8.4.4", // Google
		"9.9.9.9", "149.112.112.112", // Quad9
		"208.67.222.222", "208.67.220.220", // OpenDNS
		"94.140.14.14", "94.140.15.15", // AdGuard
		"76.76.2.0", "76.76.10.0", // ControlD
		"185.228.168.9", "185.228.169.9", // CleanBrowsing
		"156.154.70.1", "156.154.71.1", // Neustar
		"64.6.64.6", "64.6.65.6", // Verisign
		"74.82.42.42",                    // Hurricane Electric
		"216.146.35.35", "216.146.36.36", // Oracle/Dyn
	}
}

const (
	defaultAttemptTimeout  = 2 * time.Second
	maxQueryAttempts       = 3
	serverFailureThreshold = 3
	serverCooldown         = 5 * time.Minute
)

type server struct {
	r         Resolver
	failures  atomic.Int32
	downUntil atomic.Int64 // unix nanoseconds; 0 = healthy
}

func (s *server) available(now int64) bool { return s.downUntil.Load() <= now }

func (s *server) success() {
	s.failures.Store(0)
	s.downUntil.Store(0)
}

func (s *server) failure(now int64) {
	if s.failures.Add(1) >= serverFailureThreshold {
		s.downUntil.Store(now + int64(serverCooldown))
		s.failures.Store(0)
	}
}

// ServerPool load-balances DNS lookups across many resolvers. Each query picks
// a random healthy resolver and, on a transient error, fails over to a bounded
// number of others. Resolvers that repeatedly fail are put in cooldown, so a
// large list of mostly-dead public resolvers converges on the working subset.
//
// ponytail: random selection + cooldown, not latency-weighted or health-scored.
// Upgrade to EWMA latency weighting if query latency matters.
type ServerPool struct {
	servers        []*server
	attemptTimeout time.Duration
}

// NewServerPool builds a pool. Entries may be "host", "host:port" or "system".
// An empty list falls back to the OS resolver.
func NewServerPool(servers []string) *ServerPool {
	p := &ServerPool{}
	for _, s := range servers {
		s = strings.TrimSpace(s)
		switch {
		case s == "":
			continue
		case s == "system" || s == "local":
			p.servers = append(p.servers, &server{r: &GoResolver{R: net.DefaultResolver}})
		default:
			if !strings.Contains(s, ":") {
				s += ":53"
			}
			p.servers = append(p.servers, &server{r: &GoResolver{R: newServerResolver(s)}})
		}
	}
	if len(p.servers) == 0 {
		p.servers = append(p.servers, &server{r: &GoResolver{R: net.DefaultResolver}})
	}
	return p
}

// Len reports how many resolvers are in the pool.
func (p *ServerPool) Len() int { return len(p.servers) }

// Warm probes every resolver and puts non-responders into cooldown, so the
// pool starts on the working subset. Returns the healthy count.
func (p *ServerPool) Warm(ctx context.Context, probeHost string, workers int) int {
	if workers <= 0 {
		workers = 32
	}
	servers := p.servers
	probeTimeout := time.Second
	if p.attemptTimeout > 0 && p.attemptTimeout < probeTimeout {
		probeTimeout = p.attemptTimeout
	}
	jobs := make(chan *server)
	var (
		wg      sync.WaitGroup
		healthy atomic.Int32
	)
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for s := range jobs {
				rctx, cancel := context.WithTimeout(ctx, probeTimeout)
				_, err := s.r.LookupIP(rctx, probeHost)
				cancel()
				if err == nil || definitive(err) {
					healthy.Add(1) // answered, even if NXDOMAIN
					continue
				}
				s.downUntil.Store(time.Now().Add(serverCooldown).UnixNano())
			}
		}()
	}
	for _, s := range servers {
		select {
		case jobs <- s:
		case <-ctx.Done():
			close(jobs)
			wg.Wait()
			return int(healthy.Load())
		}
	}
	close(jobs)
	wg.Wait()
	return int(healthy.Load())
}

func newServerResolver(server string) *net.Resolver {
	return &net.Resolver{
		PreferGo: true,
		Dial: func(ctx context.Context, network, _ string) (net.Conn, error) {
			if network != "tcp" {
				network = "udp"
			}
			d := net.Dialer{Timeout: 5 * time.Second}
			return d.DialContext(ctx, network, server)
		},
	}
}

func (p *ServerPool) withAttemptTimeout(ctx context.Context) (context.Context, context.CancelFunc) {
	t := p.attemptTimeout
	if t <= 0 {
		t = defaultAttemptTimeout
	}
	return context.WithTimeout(ctx, t)
}

func triedContain(tried []int, n, idx int) bool {
	for i := 0; i < n; i++ {
		if tried[i] == idx {
			return true
		}
	}
	return false
}

// pickIndex returns an untried resolver, preferring healthy ones, by scanning
// cyclically from a random start. Scanning (rather than random sampling) means
// a healthy resolver is always found when one exists, which matters for small
// pools. With a warm pool of thousands the first candidate is almost always
// healthy, so it stays O(1) in practice. If all are in cooldown it still
// returns one, so a fully-cooled pool can recover.
func (p *ServerPool) pickIndex(now int64, tried []int, n int) int {
	total := len(p.servers)
	if total == 0 {
		return -1
	}
	start := mathrand.IntN(total)
	for i := 0; i < total; i++ {
		idx := (start + i) % total
		if !triedContain(tried, n, idx) && p.servers[idx].available(now) {
			return idx
		}
	}
	for i := 0; i < total; i++ {
		idx := (start + i) % total
		if !triedContain(tried, n, idx) {
			return idx
		}
	}
	return -1
}

// definitive reports whether err is a real DNS answer (NXDOMAIN) that should not
// be retried against other servers.
func definitive(err error) bool {
	var dnsErr *net.DNSError
	return errors.As(err, &dnsErr) && dnsErr.IsNotFound
}

func (p *ServerPool) lookup(ctx context.Context, fn func(context.Context, Resolver) error) error {
	var (
		tried [maxQueryAttempts]int
		last  error
	)
	for n := 0; n < maxQueryAttempts; n++ {
		idx := p.pickIndex(time.Now().UnixNano(), tried[:], n)
		if idx < 0 {
			break
		}
		tried[n] = idx
		actx, cancel := p.withAttemptTimeout(ctx)
		err := fn(actx, p.servers[idx].r)
		cancel()
		if err == nil {
			p.servers[idx].success()
			return nil
		}
		if definitive(err) {
			return err
		}
		p.servers[idx].failure(time.Now().UnixNano())
		last = err
		if ctx.Err() != nil {
			return err
		}
	}
	return last
}

func (p *ServerPool) LookupCNAME(ctx context.Context, host string) (string, error) {
	var result string
	err := p.lookup(ctx, func(c context.Context, r Resolver) error {
		v, e := r.LookupCNAME(c, host)
		if e == nil {
			result = v
		}
		return e
	})
	return result, err
}

func (p *ServerPool) LookupIP(ctx context.Context, host string) ([]net.IP, error) {
	var result []net.IP
	err := p.lookup(ctx, func(c context.Context, r Resolver) error {
		v, e := r.LookupIP(c, host)
		if e == nil {
			result = v
		}
		return e
	})
	return result, err
}