Allow Vitess HTTP client timeout override

timvaillancourt · timvaillancourt · commit 4d14098e1f3c · 2020-04-24T21:48:55.000+02:00
diff --git a/go.mod b/go.mod
@@ -13,4 +13,5 @@ require (
 	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/rcrowley/go-metrics v0.0.0-20190826022208-cac0b30c2563
 	golang.org/x/sys v0.0.0-20190602015325-4c4f7f33c9ed // indirect
+	vitess.io/vitess v2.1.1+incompatible
 )
diff --git a/go.sum b/go.sum
@@ -11,6 +11,7 @@ github.com/circonus-labs/circonusllhist v0.1.3/go.mod h1:kMXHVDlOchFAehlya5ePtbp
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/go-sql-driver/mysql v1.5.0 h1:ozyZYNQW3x3HtqT1jira07DN2PArx2v7/mN66gGcHOs=
 github.com/go-sql-driver/mysql v1.5.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg=
+github.com/golang/protobuf v1.2.0 h1:P3YflyNX/ehuJFLhxviNdFxQPkGK5cDcApsge1SqnvM=
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/hashicorp/go-cleanhttp v0.5.0/go.mod h1:JpRdi6/HCYpAwUzNwuwqhbovhLtngrth3wmdIIUrZ80=
 github.com/hashicorp/go-immutable-radix v1.0.0 h1:AKDB1HM5PWEA7i4nhcpwOrO2byshxBjXVn/J/3+z5/0=
@@ -45,6 +46,9 @@ github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/tv42/httpunix v0.0.0-20150427012821-b75d8614f926/go.mod h1:9ESjWnEqriFuLhtthL60Sar/7RFoluCcXsuvEwTV5KM=
 golang.org/x/net v0.0.0-20181201002055-351d144fa1fc/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
+golang.org/x/sync v0.0.0-20181108010431-42b317875d0f h1:Bl/8QSvNqXvPGPGXa2z5xUTmV7VDcZyvRZ+QQXkXTZQ=
 golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20190602015325-4c4f7f33c9ed h1:uPxWBzB3+mlnjy9W58qY1j/cjyFjutgw/Vhan2zLy/A=
 golang.org/x/sys v0.0.0-20190602015325-4c4f7f33c9ed/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+vitess.io/vitess v2.1.1+incompatible h1:nuuGHiWYWpudD3gOCLeGzol2EJ25e/u5Wer2wV1O130=
+vitess.io/vitess v2.1.1+incompatible/go.mod h1:h4qvkyNYTOC0xI+vcidSWoka0gQAZc9ZPHbkHo48gP0=
diff --git a/pkg/throttle/throttler.go b/pkg/throttle/throttler.go
@@ -323,7 +323,8 @@ func (throttler *Throttler) refreshMySQLInventory() error {
 				log.Debugf("getting vitess data from %s", clusterSettings.VitessSettings.API)
 				keyspace := clusterSettings.VitessSettings.Keyspace
 				shard := clusterSettings.VitessSettings.Shard
-				vtManager := vitess.NewManager(config.Settings().VitessHTTPTimeoutSec)
+				timeout := time.Duration(config.Settings().VitessHTTPTimeoutSec) * time.Second
+				vtManager := vitess.NewManager(timeout)
 				tablets, err := vtManager.ParseTablets(clusterSettings.VitessSettings.API, keyspace, shard)
 				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
@@ -7,12 +7,20 @@ import (
 	"net/http"
 	"strings"
 	"time"
+
+	"vitess.io/vitess/go/vt/proto/topodata"
 )
 
 // Tablet represents information about a running instance of vttablet.
 type Tablet struct {
-	MysqlHostname string `json:"mysql_hostname,omitempty"`
-	MysqlPort     int32  `json:"mysql_port,omitempty"`
+	MysqlHostname string              `json:"mysql_hostname,omitempty"`
+	MysqlPort     int32               `json:"mysql_port,omitempty"`
+	Type          topodata.TabletType `json:"type,omitempty"`
+}
+
+// IsValidReplica returns a bool reflecting if a tablet type is REPLICA
+func (t Tablet) IsValidReplica() bool {
+	return t.Type == topodata.TabletType_REPLICA
 }
 
 // Manager gathers info from Vitess
@@ -21,10 +29,12 @@ type Manager struct {
 }
 
 // NewManager returns a new manager for Vitess
-func NewManager(timeoutSec int) *Manager {
-	return &Manager{httpClient: http.Client{
-		Timeout: time.Duration(timeoutSec) * time.Second,
-	}}
+func NewManager(timeout time.Duration) *Manager {
+	return &Manager{
+		httpClient: http.Client{
+			Timeout: timeout,
+		},
+	}
 }
 
 func constructAPIURL(api string, keyspace string, shard string) (url string) {
@@ -37,8 +47,18 @@ func constructAPIURL(api string, keyspace string, shard string) (url string) {
 	return url
 }
 
+// filterReplicaTablets parses a list of tablets, returning replica tablets only
+func filterReplicaTablets(tablets []Tablet) (replicas []Tablet) {
+	for _, tablet := range tablets {
+		if tablet.IsValidReplica() {
+			replicas = append(replicas, tablet)
+		}
+	}
+	return replicas
+}
+
 // ParseTablets reads from vitess /api/ks_tablets/<keyspace>/[shard] and returns a
-// tblet (mysql_hostname, mysql_port) listing
+// listing (mysql_hostname, mysql_port, type) of REPLICA tablets
 func (m *Manager) ParseTablets(api string, keyspace string, shard string) (tablets []Tablet, err error) {
 	url := constructAPIURL(api, keyspace, shard)
 	resp, err := m.httpClient.Get(url)
@@ -53,5 +73,5 @@ func (m *Manager) ParseTablets(api string, keyspace string, shard string) (table
 	}
 
 	err = json.Unmarshal(body, &tablets)
-	return tablets, err
+	return filterReplicaTablets(tablets), err
 }
diff --git a/pkg/vitess/api_client_test.go b/pkg/vitess/api_client_test.go
@@ -0,0 +1,88 @@
+package vitess
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"vitess.io/vitess/go/vt/proto/topodata"
+)
+
+func TestParseTablets(t *testing.T) {
+	vitessApi := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.String() {
+		case "/api/keyspace/test/tablets/00":
+			data, _ := json.Marshal([]Tablet{
+				{
+					MysqlHostname: "master",
+					Type:          topodata.TabletType_MASTER,
+				},
+				{
+					MysqlHostname: "replica",
+					Type:          topodata.TabletType_REPLICA,
+				},
+				{
+					MysqlHostname: "spare",
+					Type:          topodata.TabletType_SPARE,
+				},
+				{
+					MysqlHostname: "batch",
+					Type:          topodata.TabletType_BATCH,
+				},
+				{
+					MysqlHostname: "backup",
+					Type:          topodata.TabletType_BACKUP,
+				},
+				{
+
+					MysqlHostname: "restore",
+					Type:          topodata.TabletType_RESTORE,
+				},
+			})
+			fmt.Fprint(w, string(data))
+		default:
+			w.WriteHeader(http.StatusNotFound)
+			fmt.Fprint(w, "[]")
+		}
+	}))
+	defer vitessApi.Close()
+
+	m := NewManager(time.Duration(1) * time.Second)
+
+	t.Run("success", func(t *testing.T) {
+		tablets, err := m.ParseTablets(vitessApi.URL, "test", "00")
+		if err != nil {
+			t.Fatalf("Expected no error, got %q", err)
+		}
+
+		if len(tablets) != 1 {
+			t.Fatalf("Expected 1 tablet, got %d", len(tablets))
+		}
+
+		if tablets[0].MysqlHostname != "replica" {
+			t.Fatalf("Expected hostname %q, got %q", "replica", tablets[0].MysqlHostname)
+		}
+	})
+
+	t.Run("not-found", func(t *testing.T) {
+		tablets, err := m.ParseTablets(vitessApi.URL, "not-found", "00")
+		if err != nil {
+			t.Fatalf("Expected no error, got %q", err)
+		}
+
+		if len(tablets) > 0 {
+			t.Fatalf("Expected 0 tablets, got %d", len(tablets))
+		}
+	})
+
+	t.Run("failed", func(t *testing.T) {
+		vitessApi.Close() // kill the mock vitess API
+		_, err := m.ParseTablets(vitessApi.URL, "fail", "00")
+		if err == nil {
+			t.Fatal("Expected error, got nil")
+		}
+	})
+}
