Add lookup timeout to host provider

internal/dcs/config.go

@@ -31,12 +31,14 @@ type ZookeeperConfig struct {
 }
 
 type RandomHostProviderConfig struct {
-	LookupTTL time.Duration `config:"lookup_ttl" yaml:"lookup_ttl"`
+	LookupTimeout time.Duration `config:"lookup_timeout" yaml:"lookup_timeout"`
+	LookupTTL     time.Duration `config:"lookup_ttl" yaml:"lookup_ttl"`
 }
 
 func DefaultRandomHostProviderConfig() RandomHostProviderConfig {
 	return RandomHostProviderConfig{
-		LookupTTL: 300 * time.Second,
+		LookupTimeout: 3 * time.Second,
+		LookupTTL:     300 * time.Second,
 	}
 }
 

internal/dcs/zk_host_provider.go

@@ -1,6 +1,7 @@
 package dcs
 
 import (
+	"context"
 	"fmt"
 	"math/rand"
 	"net"
@@ -11,20 +12,24 @@ import (
 )
 
 type RandomHostProvider struct {
-	lock       sync.Mutex
-	servers    []string
-	resolved   []string
-	tried      map[string]struct{}
-	logger     *log.Logger
-	lastLookup time.Time
-	lookupTTL  time.Duration
+	lock          sync.Mutex
+	servers       []string
+	resolved      []string
+	tried         map[string]struct{}
+	logger        *log.Logger
+	lastLookup    time.Time
+	lookupTTL     time.Duration
+	lookupTimeout time.Duration
+	resolver      *net.Resolver
 }
 
 func NewRandomHostProvider(config *RandomHostProviderConfig, logger *log.Logger) *RandomHostProvider {
 	return &RandomHostProvider{
-		lookupTTL: config.LookupTTL,
-		logger:    logger,
-		tried:     make(map[string]struct{}),
+		lookupTTL:     config.LookupTTL,
+		lookupTimeout: config.LookupTimeout,
+		logger:        logger,
+		tried:         make(map[string]struct{}),
+		resolver:      &net.Resolver{},
 	}
 }
 
@@ -50,7 +55,9 @@ func (rhp *RandomHostProvider) resolveHosts() error {
 		if err != nil {
 			return err
 		}
-		addrs, err := net.LookupHost(host)
+		ctx, cancel := context.WithTimeout(context.Background(), rhp.lookupTimeout)
+		defer cancel()
+		addrs, err := rhp.resolver.LookupHost(ctx, host)
 		if err != nil {
 			rhp.logger.Errorf("unable to resolve %s: %v", host, err)
 		}