From 516fbad7f783b8a532a5b94645017cc7c81d221e Mon Sep 17 00:00:00 2001
From: secwall <secwall@yandex-team.ru>
Date: Thu, 25 Jul 2024 14:09:44 +0200
Subject: [PATCH] Add lookup timeout to host provider

---
 internal/dcs/config.go           |  6 ++++--
 internal/dcs/zk_host_provider.go | 29 ++++++++++++++++++-----------
 2 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/internal/dcs/config.go b/internal/dcs/config.go
index 74d15958..7884ea4d 100644
--- a/internal/dcs/config.go
+++ b/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,
 	}
 }
 
diff --git a/internal/dcs/zk_host_provider.go b/internal/dcs/zk_host_provider.go
index dbfd6072..5ead36fc 100644
--- a/internal/dcs/zk_host_provider.go
+++ b/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)
 		}