Getting hosts from vitess API

doc/mysql.md

@@ -79,6 +79,12 @@ Looking at clusters configuration:
       "PoolName": "my_prod4_pool"
     }
   },
+  "sharded": {
+    "VitessSettings": {
+      "API": "https://vtctld.example.com/api/",
+      "Keyspace": "my_sharded_ks"
+    }
+  },
   "local": {
     "User": "msandbox",
     "Password": "msandbox",

go/config/mysql_config.go

@@ -22,6 +22,7 @@ type MySQLClusterConfigurationSettings struct {
 	HttpCheckPath     string  //  Specify if different than specified by MySQLConfigurationSettings
 
 	HAProxySettings     HAProxyConfigurationSettings // If list of servers is to be acquired via HAProxy, provide this field
+	VitessSettings      VitessConfigurationSettings  // If list of servers is to be acquired via Vitess, provide this field
 	StaticHostsSettings StaticHostsConfigurationSettings
 }
 

go/config/vitess_config.go

@@ -0,0 +1,21 @@
+package config
+
+//
+// HAProxy-specific configuration
+//
+
+type VitessConfigurationSettings struct {
+	API      string
+	Keyspace string
+	Shard    string
+}
+
+func (settings *VitessConfigurationSettings) IsEmpty() bool {
+	if settings.API == "" {
+		return true
+	}
+	if settings.Keyspace == "" {
+		return true
+	}
+	return false
+}

go/throttle/throttler.go

@@ -12,6 +12,7 @@ import (
 	"github.com/github/freno/go/config"
 	"github.com/github/freno/go/haproxy"
 	"github.com/github/freno/go/mysql"
+	"github.com/github/freno/go/vitess"
 
 	"github.com/outbrain/golib/log"
 	"github.com/patrickmn/go-cache"
@@ -266,6 +267,29 @@ func (throttler *Throttler) refreshMySQLInventory() error {
 				throttler.mysqlClusterProbesChan <- clusterProbes
 				return nil
 			}
+
+			if !clusterSettings.VitessSettings.IsEmpty() {
+				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)
+				if err != nil {
+					return log.Errorf("Unable to get vitess hosts from %s, %s/%s: %+v", clusterSettings.VitessSettings.API, keyspace, shard, err)
+				}
+				log.Debugf("Read %+v hosts from vitess %s, %s/%s", len(tablets), clusterSettings.VitessSettings.API, keyspace, shard)
+				clusterProbes := &mysql.ClusterProbes{
+					ClusterName:      clusterName,
+					IgnoreHostsCount: clusterSettings.IgnoreHostsCount,
+					InstanceProbes:   mysql.NewProbes(),
+				}
+				for _, tablet := range tablets {
+					key := mysql.InstanceKey{Hostname: tablet.MysqlHostname, Port: int(tablet.MysqlPort)}
+					addInstanceKey(&key, clusterSettings, clusterProbes.InstanceProbes)
+				}
+				throttler.mysqlClusterProbesChan <- clusterProbes
+				return nil
+			}
+
 			if !clusterSettings.StaticHostsSettings.IsEmpty() {
 				clusterProbes := &mysql.ClusterProbes{
 					ClusterName:    clusterName,

go/vitess/api_client.go

@@ -0,0 +1,49 @@
+package vitess
+
+import (
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
+	"net/http"
+	"strings"
+	"time"
+)
+
+// Tablet represents information about a running instance of vttablet.
+type Tablet struct {
+	MysqlHostname string `json:"mysql_hostname,omitempty"`
+	MysqlPort     int32  `json:"mysql_port,omitempty"`
+}
+
+var httpClient = http.Client{
+	Timeout: 1 * time.Second,
+}
+
+func constructAPIURL(api string, keyspace string, shard string) (url string) {
+	api = strings.TrimRight(api, "/")
+	if !strings.HasSuffix(api, "/api") {
+		api = fmt.Sprintf("%s/api", api)
+	}
+	url = fmt.Sprintf("%s/ks_tablets/%s/%s", api, keyspace, shard)
+
+	return url
+}
+
+// ParseTablets reads from vitess /api/ks_tablets/<keyspace>/[shard] and returns a
+// tblet (mysql_hostname, mysql_port) listing
+func ParseTablets(api string, keyspace string, shard string) (tablets []Tablet, err error) {
+	url := constructAPIURL(api, keyspace, shard)
+	resp, err := httpClient.Get(url)
+	if err != nil {
+		return tablets, err
+	}
+
+	defer resp.Body.Close()
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return tablets, err
+	}
+
+	err = json.Unmarshal(body, &tablets)
+	return tablets, err
+}