Listen for context end better in resolver

resolver/resolve.go

@@ -32,6 +32,10 @@ var (
 	ErrFailure = errors.New("query failed")
 	// ErrContinue is returned when the resolver has no answer, and the next resolver should be asked
 	ErrContinue = errors.New("resolver has no answer")
+	// ErrCancelled is returned when the request was cancelled.
+	ErrCancelled = errors.New("request cancelled")
+	// ErrShuttingDown is returned when the resolver is shutting down.
+	ErrShuttingDown = errors.New("resolver is shutting down")
 
 	// detailed errors
 
@@ -275,6 +279,8 @@ retry:
 			case <-time.After(maxRequestTimeout):
 				// something went wrong with the query, retry
 				goto retry
+			case <-ctx.Done():
+				return nil
 			}
 		} else {
 			// but that someone is taking too long
@@ -331,7 +337,7 @@ resolveLoop:
 	for i = 0; i < 2; i++ {
 		for _, resolver := range resolvers {
 			if module.IsStopping() {
-				return nil, errors.New("shutting down")
+				return nil, ErrShuttingDown
 			}
 
 			// check if resolver failed recently (on first run)
@@ -364,6 +370,10 @@ resolveLoop:
 					resolver.Conn.ReportFailure()
 					log.Tracer(ctx).Debugf("resolver: query to %s timed out", resolver.Info.ID())
 					continue
+				case errors.Is(err, ErrCancelled):
+					return nil, err
+				case errors.Is(err, ErrShuttingDown):
+					return nil, err
 				default:
 					resolver.Conn.ReportFailure()
 					log.Tracer(ctx).Debugf("resolver: query to %s failed: %s", resolver.Info.ID(), err)

resolver/resolver-mdns.go

@@ -418,6 +418,8 @@ func queryMulticastDNS(ctx context.Context, q *Query) (*RRCache, error) {
 		if err != nil {
 			return rrCache, nil
 		}
+	case <-ctx.Done():
+		return nil, ErrCancelled
 	}
 
 	// Respond with NXDomain.

resolver/resolver-tcp.go

@@ -105,7 +105,7 @@ func (tr *TCPResolver) UseTLS() *TCPResolver {
 	return tr
 }
 
-func (tr *TCPResolver) getOrCreateResolverConn() (*tcpResolverConn, error) {
+func (tr *TCPResolver) getOrCreateResolverConn(ctx context.Context) (*tcpResolverConn, error) {
 	tr.Lock()
 	defer tr.Unlock()
 
@@ -117,6 +117,10 @@ func (tr *TCPResolver) getOrCreateResolverConn() (*tcpResolverConn, error) {
 			return tr.resolverConn, nil
 		case <-time.After(heartbeatTimeout):
 			log.Warningf("resolver: heartbeat for dns client %s failed", tr.resolver.Info.DescriptiveName())
+		case <-ctx.Done():
+			return nil, ErrCancelled
+		case <-module.Stopping():
+			return nil, ErrShuttingDown
 		}
 	}
 
@@ -130,7 +134,6 @@ func (tr *TCPResolver) getOrCreateResolverConn() (*tcpResolverConn, error) {
 	}
 
 	// Connect to server.
-	var err error
 	conn, err := tr.dnsClient.Dial(tr.resolver.ServerAddress)
 	if err != nil {
 		log.Debugf("resolver: failed to connect to %s", tr.resolver.Info.DescriptiveName())
@@ -171,7 +174,7 @@ func (tr *TCPResolver) getOrCreateResolverConn() (*tcpResolverConn, error) {
 // Query executes the given query against the resolver.
 func (tr *TCPResolver) Query(ctx context.Context, q *Query) (*RRCache, error) {
 	// Get resolver connection.
-	resolverConn, err := tr.getOrCreateResolverConn()
+	resolverConn, err := tr.getOrCreateResolverConn(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -185,6 +188,10 @@ func (tr *TCPResolver) Query(ctx context.Context, q *Query) (*RRCache, error) {
 	// Submit query request to live connection.
 	select {
 	case resolverConn.queries <- tq:
+	case <-ctx.Done():
+		return nil, ErrCancelled
+	case <-module.Stopping():
+		return nil, ErrShuttingDown
 	case <-time.After(defaultRequestTimeout):
 		return nil, ErrTimeout
 	}
@@ -193,6 +200,10 @@ func (tr *TCPResolver) Query(ctx context.Context, q *Query) (*RRCache, error) {
 	var reply *dns.Msg
 	select {
 	case reply = <-tq.Response:
+	case <-ctx.Done():
+		return nil, ErrCancelled
+	case <-module.Stopping():
+		return nil, ErrShuttingDown
 	case <-time.After(defaultRequestTimeout):
 		return nil, ErrTimeout
 	}