From 78afbb2e6239dac56309173dc029da13d9aa47f3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= Date: Mon, 21 Sep 2026 12:26:53 +0800 Subject: [PATCH] Respect DNS query timeouts during deduplication --- dns/client.go | 106 +++++++++++++++++++++++++++++++------------------- 1 file changed, 67 insertions(+), 39 deletions(-) diff --git a/dns/client.go b/dns/client.go index d99aed0e..77fe1cbc 100644 --- a/dns/client.go +++ b/dns/client.go @@ -48,7 +48,7 @@ type Client struct { networkManager adapter.NetworkManager logger logger.ContextLogger cache *freelru.Cache[dnsCacheKey, *dns.Msg] - cacheLock compatible.Map[dnsCacheKey, chan struct{}] + cacheLock compatible.Map[dnsExchangeKey, chan struct{}] backgroundRefresh compatible.Map[dnsCacheKey, struct{}] } @@ -95,6 +95,11 @@ type dnsCacheKey struct { environment uint64 } +type dnsExchangeKey struct { + dnsCacheKey + timeout time.Duration +} + func (k dnsCacheKey) persistentName() string { name := k.transportTag if k.clientSubnet.IsValid() { @@ -251,6 +256,10 @@ func (o *exchangeOperation) release() { } func (c *Client) beginExchange(ctx context.Context, transport adapter.DNSTransport, message *dns.Msg, options adapter.DNSQueryOptions, responseChecker func(response *dns.Msg) bool, allowWait bool) (*exchangeOperation, *dns.Msg, exchangeStatus, error) { + err := ctx.Err() + if err != nil { + return nil, nil, exchangeDone, err + } if len(message.Question) == 0 { if c.logger != nil { c.logger.WarnContext(ctx, "bad question size: ", len(message.Question)) @@ -286,8 +295,33 @@ func (c *Client) beginExchange(ctx context.Context, transport adapter.DNSTranspo if !disableCache { cacheKey := c.newCacheKey(transport, question, message, options) operation.cacheKey = cacheKey - cond, loaded := c.cacheLock.LoadOrStore(cacheKey, make(chan struct{})) - if loaded { + exchangeKey := dnsExchangeKey{dnsCacheKey: cacheKey, timeout: options.Timeout} + cond, loaded := c.cacheLock.LoadOrStore(exchangeKey, make(chan struct{})) + if !loaded { + operation.releaseCond = func() { + c.cacheLock.Delete(exchangeKey) + close(cond) + } + } + for { + response, ttl, isStale := c.loadResponse(cacheKey) + if response != nil { + if isStale && !options.DisableOptimisticCache { + c.backgroundRefreshDNS(transport, cacheKey, message.Copy(), options, responseChecker) + logOptimisticResponse(c.logger, ctx, response) + response.Id = message.Id + operation.release() + return nil, response, exchangeDone, nil + } else if !isStale { + logCachedResponse(c.logger, ctx, response, ttl) + response.Id = message.Id + operation.release() + return nil, response, exchangeDone, nil + } + } + if !loaded { + break + } if !allowWait { return nil, nil, exchangeWait, nil } @@ -296,28 +330,13 @@ func (c *Client) beginExchange(ctx context.Context, transport adapter.DNSTranspo case <-ctx.Done(): return nil, nil, exchangeDone, ctx.Err() } + err = ctx.Err() + if err != nil { + return nil, nil, exchangeDone, err + } cacheKey = c.newCacheKey(transport, question, message, options) operation.cacheKey = cacheKey - } else { - operation.releaseCond = func() { - c.cacheLock.Delete(cacheKey) - close(cond) - } - } - response, ttl, isStale := c.loadResponse(cacheKey) - if response != nil { - if isStale && !options.DisableOptimisticCache { - c.backgroundRefreshDNS(transport, cacheKey, message.Copy(), options, responseChecker) - logOptimisticResponse(c.logger, ctx, response) - response.Id = message.Id - operation.release() - return nil, response, exchangeDone, nil - } else if !isStale { - logCachedResponse(c.logger, ctx, response, ttl) - response.Id = message.Id - operation.release() - return nil, response, exchangeDone, nil - } + loaded = false } } @@ -379,12 +398,17 @@ func (c *Client) finishExchange(transport adapter.DNSTransport, operation *excha } func (c *Client) Exchange(ctx context.Context, transport adapter.DNSTransport, message *dns.Msg, options adapter.DNSQueryOptions, responseChecker func(response *dns.Msg) bool) (*dns.Msg, error) { + if options.Timeout == 0 { + options.Timeout = c.timeout + } + ctx, cancel := context.WithTimeout(ctx, options.Timeout) + defer cancel() operation, earlyResponse, status, err := c.beginExchange(ctx, transport, message, options, responseChecker, true) if status != exchangeReady { return earlyResponse, err } defer operation.release() - response, err := c.exchangeToTransport(operation.ctx, transport, operation.message, options.Timeout) + response, err := c.exchangeToTransport(operation.ctx, transport, operation.message) if err != nil { return nil, err } @@ -392,18 +416,26 @@ func (c *Client) Exchange(ctx context.Context, transport adapter.DNSTransport, m } func (c *Client) ExchangeAsync(ctx context.Context, transport adapter.DNSTransport, message *dns.Msg, options adapter.DNSQueryOptions, responseChecker func(response *dns.Msg) bool, callback func(response *dns.Msg, err error)) { + if options.Timeout == 0 { + options.Timeout = c.timeout + } + ctx, cancel := context.WithTimeout(ctx, options.Timeout) operation, earlyResponse, status, err := c.beginExchange(ctx, transport, message, options, responseChecker, false) switch status { case exchangeDone: + cancel() callback(earlyResponse, err) return case exchangeWait: go func() { - callback(c.Exchange(ctx, transport, message, options, responseChecker)) + response, exchangeErr := c.Exchange(ctx, transport, message, options, responseChecker) + cancel() + callback(response, exchangeErr) }() return } finish := func(response *dns.Msg, exchangeErr error) { + cancel() if exchangeErr != nil { operation.release() callback(nil, exchangeErr) @@ -413,7 +445,7 @@ func (c *Client) ExchangeAsync(ctx context.Context, transport adapter.DNSTranspo operation.release() callback(finishedResponse, finishErr) } - c.exchangeToTransportAsync(operation.ctx, transport, operation.message, options.Timeout, finish) + c.exchangeToTransportAsync(operation.ctx, transport, operation.message, finish) } func (c *Client) Lookup(ctx context.Context, transport adapter.DNSTransport, domain string, options adapter.DNSQueryOptions, responseChecker func(response *dns.Msg) bool) ([]netip.Addr, error) { @@ -644,8 +676,14 @@ func (c *Client) backgroundRefreshDNS(transport adapter.DNSTransport, key dnsCac } go func() { defer c.backgroundRefresh.Delete(key) + timeout := options.Timeout + if timeout == 0 { + timeout = c.timeout + } ctx := adapter.ContextWithDNSTransportTag(c.ctx, transport.Tag()) - response, err := c.exchangeToTransport(ctx, transport, message, options.Timeout) + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + response, err := c.exchangeToTransport(ctx, transport, message) if err != nil { if c.logger != nil { c.logger.DebugContext(ctx, "optimistic refresh failed for ", FqdnToDomain(key.Name), ": ", err) @@ -707,12 +745,7 @@ func stripDNSPadding(response *dns.Msg) { } } -func (c *Client) exchangeToTransport(ctx context.Context, transport adapter.DNSTransport, message *dns.Msg, timeout time.Duration) (*dns.Msg, error) { - if timeout == 0 { - timeout = c.timeout - } - ctx, cancel := context.WithTimeout(ctx, timeout) - defer cancel() +func (c *Client) exchangeToTransport(ctx context.Context, transport adapter.DNSTransport, message *dns.Msg) (*dns.Msg, error) { response, err := transport.Exchange(ctx, message) if err == nil { stripDNSPadding(response) @@ -725,13 +758,8 @@ func (c *Client) exchangeToTransport(ctx context.Context, transport adapter.DNST return nil, err } -func (c *Client) exchangeToTransportAsync(ctx context.Context, transport adapter.DNSTransport, message *dns.Msg, timeout time.Duration, callback func(response *dns.Msg, err error)) { - if timeout == 0 { - timeout = c.timeout - } - ctx, cancel := context.WithTimeout(ctx, timeout) +func (c *Client) exchangeToTransportAsync(ctx context.Context, transport adapter.DNSTransport, message *dns.Msg, callback func(response *dns.Msg, err error)) { transport.ExchangeAsync(ctx, message, func(response *dns.Msg, err error) { - cancel() if err == nil { stripDNSPadding(response) callback(response, nil)