Respect DNS query timeouts during deduplication

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)