From b0b52b677429b2c11115a9eab812ad5b3f3116f9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zolt=C3=A1n=20Papp?= <zoltan.pmail@gmail.com>
Date: Fri, 24 Apr 2026 22:36:18 +0200
Subject: [PATCH] client/dns/mgmt: singleflight on-demand resolves
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

resolveOnDemand deduplicated the cache write but not the upstream
lookup — a burst of concurrent queries for a freshly-learned pool-root
subdomain (A + AAAA from a single getaddrinfo, or multiple peer
workers dialing the same foreign relay) each hit the bypass resolver
independently.

Wrap the lookupRecords call in refreshGroup.Do with an 'ondemand:'-
prefixed key so concurrent callers collapse into one request; the
prefix namespaces this key off scheduleRefresh's keyspace so the two
paths don't collide.

Error and empty-result paths preserved; the cache write still runs
under the existing TOCTOU-safe exists check + mutex so a winning
refresh isn't clobbered.
---
 client/internal/dns/mgmt/mgmt.go | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/client/internal/dns/mgmt/mgmt.go b/client/internal/dns/mgmt/mgmt.go
index 750c08466..34d158c51 100644
--- a/client/internal/dns/mgmt/mgmt.go
+++ b/client/internal/dns/mgmt/mgmt.go
@@ -8,6 +8,7 @@ import (
 	"net/url"
 	"os"
 	"slices"
+	"strconv"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -232,16 +233,26 @@ func (m *Resolver) resolveOnDemand(w dns.ResponseWriter, r *dns.Msg, question dn
 		return
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), dnsTimeout)
-	defer cancel()
-
-	records, err := m.lookupRecords(ctx, d, question)
+	// Collapse concurrent on-demand lookups for the same (name, qtype) into
+	// a single upstream query via singleflight. A burst of parallel queries
+	// for a freshly-learned pool-root subdomain (e.g. multiple peer workers
+	// dialing the same foreign relay, or A + AAAA racing each other) would
+	// otherwise each hit the bypass resolver independently. The prefix
+	// namespaces this key off scheduleRefresh's keyspace so the two paths
+	// can coexist without collisions.
+	key := "ondemand:" + question.Name + ":" + strconv.Itoa(int(question.Qtype))
+	result, err, _ := m.refreshGroup.Do(key, func() (any, error) {
+		ctx, cancel := context.WithTimeout(context.Background(), dnsTimeout)
+		defer cancel()
+		return m.lookupRecords(ctx, d, question)
+	})
 	if err != nil {
 		log.Debugf("on-demand resolve %s type=%s: %v",
 			d.SafeString(), dns.TypeToString[question.Qtype], err)
 		m.continueToNext(w, r)
 		return
 	}
+	records, _ := result.([]dns.RR)
 	if len(records) == 0 {
 		m.continueToNext(w, r)
 		return