dnsforward: collect optimistic refreshes through the proxy

Replace the upstream wrapper with the callback added in
AdguardTeam/dnsproxy#520.

The wrapper could not work.  An exchange has to be attributed to the
request it belongs to, both so that the statistics ignore lists apply and
so that a background refresh can be told from a foreground query, and an
upstream.Upstream is given no way to do it: Exchange takes only a *dns.Msg
and no context, and ExchangeParallel and ExchangeAll copy the request once
per upstream before calling it, so no identity of the original survives.
The request-pointer key therefore missed in parallel and fastest-address
mode with more than one upstream, and queries from clients with ignored
statistics were counted after all.  Only the single-upstream path skips
the copy, which is why the first regression passed.

Foreground exchanges go back to being collected from
proxy.DNSContext.QueryStatistics, which is client-attributed and already
gated by ShouldCount, so the ignore lists behave exactly as they did
before this branch.  Background refreshes arrive through
proxy.Config.OnOptimisticRefresh, which fires once the optimistic cache
has refreshed an expired entry; they belong to no client by construction,
so there is nothing to attribute and nothing to double count.

A retried exchange is still left out, but the check moved to where the
statistics entry is built, and compares against the timeout the upstreams
of that request were constructed with: defaultLocalTimeout for private
rDNS, the configured one otherwise.

Both are covered end to end, and both regressions fail against the
unfixed implementation: the retry filter in load-balance, parallel, and
fastest-address mode, and the refresh collection against an entry that
expires while the test runs.

Author: vortexilation

CHANGELOG.md

@@ -50,8 +50,8 @@ NOTE: Add new changes BELOW THIS COMMENT.
   ([#8435]).  An optimistic cache hit is answered from the cache right away and the expired entry
   is refreshed by a background query, and those background queries used to be left out of the
   statistics.  The average was therefore based on cache misses alone, which are skewed towards
-  rare domain names that the upstream itself resolves slower.  Every exchange with an upstream
-  server is now counted, including the background ones.
+  rare domain names that the upstream itself resolves slower.  Those refreshes are now counted
+  as well, through the new `OnOptimisticRefresh` callback of the DNS proxy.
 
 - Upstream response times on the dashboard being far higher than the actual network latency to
   the upstream servers ([#8457]).  A plain DNS upstream retries once when an attempt times out,

internal/client/upstreammanager.go

@@ -16,27 +16,10 @@ import (
 	"github.com/AdguardTeam/golibs/timeutil"
 )
 
-// UpstreamConfigWrapper wraps parsed upstream configurations, e.g. to collect
-// statistics about the exchanges with them.
-type UpstreamConfigWrapper interface {
-	// WrapUpstreamConfig returns a wrapped copy of uc, or nil if uc is nil.  It
-	// must not modify uc.  timeout must be the timeout that the upstreams of uc
-	// have been constructed with.
-	WrapUpstreamConfig(
-		uc *proxy.UpstreamConfig,
-		timeout time.Duration,
-	) (wrapped *proxy.UpstreamConfig)
-}
-
 // CommonUpstreamConfig contains common settings for custom client upstream
 // configurations.
 type CommonUpstreamConfig struct {
-	Bootstrap upstream.Resolver
-
-	// UpstreamConfigWrapper wraps the parsed custom client upstream
-	// configurations.  If it's nil, they are used as is.
-	UpstreamConfigWrapper UpstreamConfigWrapper
-
+	Bootstrap               upstream.Resolver
 	UpstreamTimeout         time.Duration
 	BootstrapPreferIPv6     bool
 	EDNSClientSubnetEnabled bool
@@ -249,10 +232,6 @@ func newCustomUpstreamConfig(
 		panic(fmt.Errorf("creating custom upstream config: %w", err))
 	}
 
-	if w := conf.UpstreamConfigWrapper; w != nil {
-		upsConf = w.WrapUpstreamConfig(upsConf, conf.UpstreamTimeout)
-	}
-
 	return proxy.NewCustomUpstreamConfig(
 		upsConf,
 		cliConf.upstreamsCacheEnabled,

internal/dnsforward/config.go

@@ -354,6 +354,7 @@ func (s *Server) newProxyConfig(ctx context.Context) (conf *proxy.Config, err er
 		UpstreamConfig:            srvConf.UpstreamConfig,
 		PrivateRDNSUpstreamConfig: srvConf.PrivateRDNSUpstreamConfig,
 		RequestHandler:            ratelimitMw.Wrap(logMw.Wrap(s.Wrap(s))),
+		OnOptimisticRefresh:       s.handleOptimisticRefresh,
 		EnableEDNSClientSubnet:    srvConf.EDNSClientSubnet.Enabled,
 		MaxGoroutines:             srvConf.MaxGoroutines,
 		UseDNS64:                  srvConf.UseDNS64,

internal/dnsforward/dnsforward.go

@@ -121,24 +121,6 @@ type Server struct {
 	// stats is the statistics collector for client's DNS usage data.
 	stats stats.Interface
 
-	// ignoredReqs contains the request messages of the queries whose statistics
-	// must not be counted, keyed by *dns.Msg.  An upstream exchange carries no
-	// client identity, so [statsUpstream] cannot apply the ignore lists itself;
-	// the request handler records the decision here, while the client is still
-	// known, see [Server.markIgnoredReq].
-	//
-	// Only the ignored queries are stored, so the map is empty unless some
-	// client or domain is actually ignored.
-	ignoredReqs sync.Map
-
-	// upstreamStats is the same collector as stats, but it's set once and never
-	// reset, so that it can be used from [statsUpstream] without acquiring
-	// serverLock.  An upstream exchange may happen while serverLock is already
-	// held for reading, see [Server.Resolve], so acquiring it again there could
-	// deadlock with a concurrent writer.  Updating an already closed collector
-	// only loses the data, which is acceptable during the shutdown.
-	upstreamStats stats.Interface
-
 	// sysResolvers used to fetch system resolvers to use by default for private
 	// PTR resolving.
 	sysResolvers SystemResolvers
@@ -260,14 +242,13 @@ func NewServer(p DNSCreateParams) (s *Server, err error) {
 	}
 
 	s = &Server{
-		dnsFilter:     p.DNSFilter,
-		dhcpServer:    p.DHCPServer,
-		stats:         p.Stats,
-		upstreamStats: p.Stats,
-		queryLog:      p.QueryLog,
-		privateNets:   p.PrivateNets,
-		baseLogger:    p.Logger,
-		logger:        p.Logger.With(slogutil.KeyPrefix, "dnsforward"),
+		dnsFilter:   p.DNSFilter,
+		dhcpServer:  p.DHCPServer,
+		stats:       p.Stats,
+		queryLog:    p.QueryLog,
+		privateNets: p.PrivateNets,
+		baseLogger:  p.Logger,
+		logger:      p.Logger.With(slogutil.KeyPrefix, "dnsforward"),
 		// TODO(e.burkov):  Use some case-insensitive string comparison.
 		localDomainSuffix: strings.ToLower(localDomainSuffix),
 		etcHosts:          etcHosts,
@@ -586,10 +567,9 @@ func (s *Server) prepareUpstreamSettings(ctx context.Context, boot upstream.Reso
 		return fmt.Errorf("preparing upstream config: %w", err)
 	}
 
-	s.conf.UpstreamConfig = s.WrapUpstreamConfig(uc, s.conf.UpstreamTimeout)
+	s.conf.UpstreamConfig = uc
 	s.conf.ClientsContainer.UpdateCommonUpstreamConfig(&client.CommonUpstreamConfig{
 		Bootstrap:               boot,
-		UpstreamConfigWrapper:   s,
 		UpstreamTimeout:         s.conf.UpstreamTimeout,
 		BootstrapPreferIPv6:     s.conf.BootstrapPreferIPv6,
 		EDNSClientSubnetEnabled: s.conf.EDNSClientSubnet.Enabled,
@@ -685,16 +665,11 @@ func (s *Server) prepareInternalDNS(ctx context.Context) (err error) {
 		return err
 	}
 
-	privateUC, err := s.prepareLocalResolvers(ctx)
+	s.conf.PrivateRDNSUpstreamConfig, err = s.prepareLocalResolvers(ctx)
 	if err != nil {
 		return err
 	}
 
-	// NOTE:  The private rDNS upstreams are constructed with defaultLocalTimeout,
-	// see prepareLocalResolvers, so that is the timeout a retried exchange with
-	// them is measured against.
-	s.conf.PrivateRDNSUpstreamConfig = s.WrapUpstreamConfig(privateUC, defaultLocalTimeout)
-
 	err = s.prepareInternalProxy()
 	if err != nil {
 		return fmt.Errorf("preparing internal proxy: %w", err)
@@ -723,7 +698,7 @@ func (s *Server) setupFallbackDNS() (uc *proxy.UpstreamConfig, err error) {
 		return nil, err
 	}
 
-	return s.WrapUpstreamConfig(uc, s.conf.UpstreamTimeout), nil
+	return uc, nil
 }
 
 // setupAddrProc initializes the address processor.  It assumes s.serverLock is

internal/dnsforward/process.go

@@ -480,11 +480,6 @@ func (s *Server) processUpstream(
 
 	s.setCustomUpstream(ctx, l, pctx, dctx.clientID)
 
-	// Record whether the statistics ignore this query before resolving it,
-	// since the upstream wrappers that collect the response times have no way
-	// of knowing its client, see [Server.markIgnoredReq].
-	defer s.markIgnoredReq(dctx)()
-
 	// Process the request further since it wasn't filtered.
 	prx := s.proxy()
 	if prx == nil {

internal/dnsforward/stats.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"log/slog"
 	"net"
+	"net/netip"
 	"time"
 
 	"github.com/AdguardTeam/AdGuardHome/internal/aghnet"
@@ -29,10 +30,19 @@ func (s *Server) processQueryLogsAndStats(
 	host := aghnet.NormalizeDomain(q.Name)
 	processingTime := time.Since(dctx.startTime)
 
-	ip, ipStr, ids := s.clientIdentity(dctx)
+	ip := pctx.Addr.Addr().AsSlice()
+	s.anonymizer.Load()(ip)
+	ipStr := net.IP(ip).String()
 
 	l.DebugContext(ctx, "client ip for stats and querylog", "ip", ipStr)
 
+	ids := []string{ipStr}
+	if dctx.clientID != "" {
+		// Use the ClientID first because it has a higher priority.  Filters
+		// have the same priority, see applyAdditionalFiltering.
+		ids = []string{dctx.clientID, ipStr}
+	}
+
 	qt, cl := q.Qtype, q.Qclass
 
 	// Synchronize access to s.queryLog and s.stats so they won't be suddenly
@@ -66,25 +76,6 @@ func (s *Server) processQueryLogsAndStats(
 	return resultCodeSuccess
 }
 
-// clientIdentity returns the anonymized address of the client of dctx, its
-// string form, and the identifiers by which the query log and the statistics
-// know it.  dctx must not be nil.
-func (s *Server) clientIdentity(dctx *dnsContext) (ip net.IP, ipStr string, ids []string) {
-	addr := dctx.proxyCtx.Addr.Addr().AsSlice()
-	s.anonymizer.Load()(addr)
-
-	ip = net.IP(addr)
-	ipStr = ip.String()
-
-	if dctx.clientID != "" {
-		// Use the ClientID first because it has a higher priority.  Filters
-		// have the same priority, see applyAdditionalFiltering.
-		return ip, ipStr, []string{dctx.clientID, ipStr}
-	}
-
-	return ip, ipStr, []string{ipStr}
-}
-
 // shouldLog returns true if the query with the given data should be logged in
 // the query log.  s.serverLock is expected to be locked.
 func (s *Server) shouldLog(host string, qt, cl uint16, ids []string) (ok bool) {
@@ -149,14 +140,104 @@ func (s *Server) logQuery(dctx *dnsContext, ip net.IP, processingTime time.Durat
 	s.queryLog.Add(p)
 }
 
+// retryThreshold returns the duration at or above which a successful exchange
+// must have retried after an attempt that timed out, since a single attempt
+// cannot outlast the timeout it was made with.  pctx must not be nil.
+//
+// s.serverLock is expected to be locked.
+func (s *Server) retryThreshold(pctx *proxy.DNSContext) (d time.Duration) {
+	if pctx.RequestedPrivateRDNS != (netip.Prefix{}) {
+		// The private rDNS upstreams are constructed with a timeout of their
+		// own, see prepareLocalResolvers.
+		return defaultLocalTimeout
+	}
+
+	return s.conf.UpstreamTimeout
+}
+
+// appendCountedUpstreams appends those of us that should be counted in the
+// statistics to stats.
+//
+// A plain DNS upstream retries once when an attempt times out, for example
+// when a UDP datagram is lost, and reports the retried exchange as an ordinary
+// success.  Its duration is then at least the whole timeout, ten seconds by
+// default, even though the successful attempt itself took a millisecond.
+// Averaging such a sample in makes the reported response time an order of
+// magnitude higher than the actual one, so leave it out: it describes the retry
+// policy and the configured timeout rather than the speed of the upstream.
+//
+// See https://github.com/AdguardTeam/AdGuardHome/issues/8457.
+func appendCountedUpstreams(
+	stats []*proxy.UpstreamStatistics,
+	us []*proxy.UpstreamStatistics,
+	threshold time.Duration,
+) (appended []*proxy.UpstreamStatistics) {
+	for _, u := range us {
+		if threshold > 0 && u.Error == nil && u.QueryDuration >= threshold {
+			continue
+		}
+
+		stats = append(stats, u)
+	}
+
+	return stats
+}
+
+// handleOptimisticRefresh records the response times of an optimistic cache
+// refresh.  It implements [proxy.Config.OnOptimisticRefresh].  dctx must not be
+// nil.
+//
+// Such a refresh is performed in the background once an expired entry has
+// already been answered from the cache, so it reaches no request handler and
+// belongs to no client.  Without it the response times would only ever be
+// sampled from cache misses, and with the optimistic cache enabled the popular
+// names, which are exactly the ones kept warm, would never be sampled at all.
+//
+// See https://github.com/AdguardTeam/AdGuardHome/issues/8435.
+func (s *Server) handleOptimisticRefresh(dctx *proxy.DNSContext) {
+	qs := dctx.QueryStatistics()
+	if qs == nil || dctx.Req == nil || len(dctx.Req.Question) == 0 {
+		return
+	}
+
+	domain := aghnet.NormalizeDomain(dctx.Req.Question[0].Name)
+
+	// Synchronize access to s.stats so it won't be suddenly uninitialized while
+	// in use, the same way processQueryLogsAndStats does.
+	s.serverLock.RLock()
+	defer s.serverLock.RUnlock()
+
+	if s.stats == nil {
+		return
+	}
+
+	threshold := s.retryThreshold(dctx)
+	for _, u := range appendCountedUpstreams(nil, qs.Main(), threshold) {
+		if u.IsCached || u.Error != nil {
+			continue
+		}
+
+		s.stats.UpdateUpstream(&stats.UpstreamEntry{
+			Address:       u.Address,
+			Domain:        domain,
+			QueryDuration: u.QueryDuration,
+		})
+	}
+}
+
 // updateStats writes the request data into statistics.
 func (s *Server) updateStats(dctx *dnsContext, clientIP string, processingTime time.Duration) {
 	pctx := dctx.proxyCtx
 
-	// NOTE:  The upstream response times are not taken from
-	// [proxy.DNSContext.QueryStatistics] here, since they are collected for
-	// every exchange, including the background ones, see [statsUpstream].
+	var upstreamStats []*proxy.UpstreamStatistics
+	if qs := pctx.QueryStatistics(); qs != nil {
+		threshold := s.retryThreshold(pctx)
+		upstreamStats = appendCountedUpstreams(upstreamStats, qs.Main(), threshold)
+		upstreamStats = appendCountedUpstreams(upstreamStats, qs.Fallback(), threshold)
+	}
+
 	e := &stats.Entry{
+		UpstreamStats:  upstreamStats,
 		Domain:         aghnet.NormalizeDomain(pctx.Req.Question[0].Name),
 		Result:         stats.RNotFiltered,
 		ProcessingTime: processingTime,