timvaillancourt · Apr 29, 2020 · Apr 24, 2020 · Apr 24, 2020 · Apr 24, 2020 · Apr 24, 2020
diff --git a/pkg/config/vitess_config.go b/pkg/config/vitess_config.go
@@ -5,9 +5,10 @@ package config
 //
 
 type VitessConfigurationSettings struct {
-	API      string
-	Keyspace string
-	Shard    string
+	API         string
+	Keyspace    string
+	Shard       string
+	TimeoutSecs uint
 }
 
 func (settings *VitessConfigurationSettings) IsEmpty() bool {

diff --git a/pkg/throttle/throttler.go b/pkg/throttle/throttler.go
@@ -323,7 +323,7 @@ func (throttler *Throttler) refreshMySQLInventory() error {
 				log.Debugf("getting vitess data from %s", clusterSettings.VitessSettings.API)
 				keyspace := clusterSettings.VitessSettings.Keyspace
 				shard := clusterSettings.VitessSettings.Shard
-				tablets, err := vitess.ParseTablets(clusterSettings.VitessSettings.API, keyspace, shard)
+				tablets, err := vitess.ParseTablets(clusterSettings.VitessSettings)
 				if err != nil {
 					return log.Errorf("Unable to get vitess hosts from %s, %s/%s: %+v", clusterSettings.VitessSettings.API, keyspace, shard, err)
 				}

diff --git a/pkg/vitess/api_client.go b/pkg/vitess/api_client.go
@@ -8,9 +8,12 @@ import (
 	"strings"
 	"time"
 
+	"github.com/github/freno/pkg/config"
 	"vitess.io/vitess/go/vt/proto/topodata"
 )
 
+const defaultTimeout = time.Duration(5) * time.Second
+
 // Tablet represents information about a running instance of vttablet.
 type Tablet struct {
 	MysqlHostname string              `json:"mysql_hostname,omitempty"`
@@ -24,15 +27,15 @@ func (t Tablet) IsValidReplica() bool {
 }
 
 var httpClient = http.Client{
-	Timeout: 1 * time.Second,
+	Timeout: defaultTimeout,
 }
 
-func constructAPIURL(api string, keyspace string, shard string) (url string) {
-	api = strings.TrimRight(api, "/")
+func constructAPIURL(settings config.VitessConfigurationSettings) (url string) {
+	api := strings.TrimRight(settings.API, "/")
 	if !strings.HasSuffix(api, "/api") {
 		api = fmt.Sprintf("%s/api", api)
 	}
-	url = fmt.Sprintf("%s/keyspace/%s/tablets/%s", api, keyspace, shard)
+	url = fmt.Sprintf("%s/keyspace/%s/tablets/%s", api, settings.Keyspace, settings.Shard)
 
 	return url
 }
@@ -49,8 +52,14 @@ func filterReplicaTablets(tablets []Tablet) (replicas []Tablet) {
 
 // ParseTablets reads from vitess /api/ks_tablets/<keyspace>/[shard] and returns a
 // listing (mysql_hostname, mysql_port, type) of REPLICA tablets
-func ParseTablets(api string, keyspace string, shard string) (tablets []Tablet, err error) {
-	url := constructAPIURL(api, keyspace, shard)
+func ParseTablets(settings config.VitessConfigurationSettings) (tablets []Tablet, err error) {
+	if settings.TimeoutSecs == 0 {
+		httpClient.Timeout = defaultTimeout
+	} else {
+		httpClient.Timeout = time.Duration(settings.TimeoutSecs) * time.Second
+	}
+
+	url := constructAPIURL(settings)
 	resp, err := httpClient.Get(url)
 	if err != nil {
 		return tablets, err

diff --git a/pkg/vitess/api_client_test.go b/pkg/vitess/api_client_test.go
@@ -6,7 +6,9 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
+	"github.com/github/freno/pkg/config"
 	"vitess.io/vitess/go/vt/proto/topodata"
 )
 
@@ -50,7 +52,12 @@ func TestParseTablets(t *testing.T) {
 	defer vitessApi.Close()
 
 	t.Run("success", func(t *testing.T) {
-		tablets, err := ParseTablets(vitessApi.URL, "test", "00")
+		tablets, err := ParseTablets(config.VitessConfigurationSettings{
+			API:         vitessApi.URL,
+			Keyspace:    "test",
+			Shard:       "00",
+			TimeoutSecs: 1,
+		})
 		if err != nil {
 			t.Fatalf("Expected no error, got %q", err)
 		}
@@ -62,22 +69,38 @@ func TestParseTablets(t *testing.T) {
 		if tablets[0].MysqlHostname != "replica" {
 			t.Fatalf("Expected hostname %q, got %q", "replica", tablets[0].MysqlHostname)
 		}
+
+		if httpClient.Timeout != time.Second {
+			t.Fatalf("Expected vitess client timeout of %v, got %v", time.Second, httpClient.Timeout)
+		}
 	})
 
 	t.Run("not-found", func(t *testing.T) {
-		tablets, err := ParseTablets(vitessApi.URL, "not-found", "00")
+		tablets, err := ParseTablets(config.VitessConfigurationSettings{
+			API:      vitessApi.URL,
+			Keyspace: "not-found",
+			Shard:    "40-80",
+		})
 		if err != nil {
 			t.Fatalf("Expected no error, got %q", err)
 		}
 
 		if len(tablets) > 0 {
 			t.Fatalf("Expected 0 tablets, got %d", len(tablets))
 		}
+
+		if httpClient.Timeout != defaultTimeout {
+			t.Fatalf("Expected vitess client timeout of %v, got %v", defaultTimeout, httpClient.Timeout)
+		}
 	})
 
 	t.Run("failed", func(t *testing.T) {
 		vitessApi.Close() // kill the mock vitess API
-		_, err := ParseTablets(vitessApi.URL, "fail", "00")
+		_, err := ParseTablets(config.VitessConfigurationSettings{
+			API:      vitessApi.URL,
+			Keyspace: "fail",
+			Shard:    "00",
+		})
 		if err == nil {
 			t.Fatal("Expected error, got nil")
 		}