Optionally fetch Marathon endpoint from ZK

configuration/configuration.go

@@ -45,6 +45,9 @@ func FromFile(filePath string) (Configuration, error) {
 	conf := &Configuration{}
 	err := conf.FromFile(filePath)
 	setValueFromEnv(&conf.Marathon.Endpoint, "MARATHON_ENDPOINT")
+	setValueFromEnv(&conf.Marathon.Zookeeper.Host, "MARATHON_ZK_HOST")
+	setValueFromEnv(&conf.Marathon.Zookeeper.Path, "MARATHON_ZK_PATH")
+	setBoolValueFromEnv(&conf.Marathon.UseZookeeper, "MARATHON_USE_ZK")
 	setValueFromEnv(&conf.Marathon.User, "MARATHON_USER")
 	setValueFromEnv(&conf.Marathon.Password, "MARATHON_PASSWORD")
 	setBoolValueFromEnv(&conf.Marathon.UseEventStream, "MARATHON_USE_EVENT_STREAM")

configuration/marathon.go

@@ -1,21 +1,66 @@
 package configuration
 
 import (
+	"github.com/QubitProducts/bamboo/Godeps/_workspace/src/github.com/samuel/go-zookeeper/zk"
 	"strings"
+	"time"
 )
 
 /*
 	Mesos Marathon configuration
 */
 type Marathon struct {
 	// comma separated marathon http endpoints including port number
-	Endpoint string
-	User     string
-	Password string
-
+	Endpoint       string
+	UseZookeeper   bool
+	Zookeeper      Zookeeper
+	User           string
+	Password       string
 	UseEventStream bool
 }
 
 func (m Marathon) Endpoints() []string {
+	if m.UseZookeeper {
+		endpoints, err := zkEndpoints(m.Zookeeper)
+		if err != nil {
+			return nil
+		}
+		return endpoints
+	}
 	return strings.Split(m.Endpoint, ",")
 }
+
+func zkEndpoints(zkConf Zookeeper) ([]string, error) {
+	// Only tested with marathon 0.11.1, assumes http:// for marathon
+	const scheme = "http://"
+	const leaderNode = "leader"
+
+	conn, _, err := zk.Connect(zkConf.ConnectionString(), time.Second*10)
+
+	if err != nil {
+		return nil, err
+	}
+
+	defer conn.Close()
+
+	var leaderPath = zkConf.Path + "/" + leaderNode
+
+	keys, _, err := conn.Children(leaderPath)
+
+	if err != nil {
+		return nil, err
+	}
+
+	endpoints := make([]string, 0, len(keys))
+
+	for _, childPath := range keys {
+		data, _, err := conn.Get(leaderPath + "/" + childPath)
+		if err != nil {
+			return nil, err
+		}
+		// TODO configurable http://??
+		endpoints = append(endpoints, scheme+string(data))
+	}
+
+	return endpoints, nil
+}

main/bamboo/bamboo.go

@@ -184,11 +184,11 @@ func listenToMarathonEventStream(conf *configuration.Configuration, sub api.Even
 	client := &http.Client{}
 	client.Timeout = 0 * time.Second
 
-	for _, marathon := range conf.Marathon.Endpoints() {
-		ticker := time.NewTicker(1 * time.Second)
-		eventsURL := marathon + "/v2/events"
-		go func() {
-			for _ = range ticker.C {
+	ticker := time.NewTicker(1 * time.Second)
+	go func() {
+		for _ = range ticker.C {
+			for _, marathon := range conf.Marathon.Endpoints() {
+				eventsURL := marathon + "/v2/events"
 				req, err := http.NewRequest("GET", eventsURL, nil)
 				req.Header.Set("Accept", "text/event-stream")
 				if len(conf.Marathon.User) > 0 && len(conf.Marathon.Password) > 0 {
@@ -236,8 +236,8 @@ func listenToMarathonEventStream(conf *configuration.Configuration, sub api.Even
 
 				log.Println("Event stream connection was closed. Re-opening...")
 			}
-		}()
-	}
+		}
+	}()
 }
 
 func configureLog() {