Fetch Mesos leader via Agent master flag

cmd/fetch.go

@@ -16,8 +16,11 @@ package cmd
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
+	"io/ioutil"
 	"net/http"
+	"strings"
 
 	"github.com/aurora-scheduler/australis/internal"
 	realis "github.com/aurora-scheduler/gorealis/v2"
@@ -31,7 +34,12 @@ const (
 )
 
 type mesosAgentState struct {
-	MasterHostname string `json:"master_hostname,omitempty""`
+	Flags mesosAgentFlags `json:"flags,omitempty"`
+}
+
+type mesosAgentFlags struct {
+	Master    string `json:"master,omitempty"`
+	hasMaster bool   // indicates if the master flag contains direct Master's address
 }
 
 func init() {
@@ -134,9 +142,9 @@ Pass Zookeeper nodes separated by a space as an argument to this command.`,
 }
 
 var mesosCmd = &cobra.Command{
-	Use:               "mesos",
-	PreRun:            setConfig,
-	Short:             "Fetch information from Mesos.",
+	Use:    "mesos",
+	PreRun: setConfig,
+	Short:  "Fetch information from Mesos.",
 }
 
 var mesosLeaderCmd = &cobra.Command{
@@ -264,18 +272,19 @@ func fetchLeader(cmd *cobra.Command, args []string) {
 }
 
 func fetchMesosLeader(cmd *cobra.Command, args []string) {
-	var url string
 	if len(args) < 1 {
-		url, err := fetchMesosLeaderFromAgent(localAgentStateURL)
-		if err != nil || url == "" {
-			log.Debugf("unable to fetch Mesos leader from local Mesos agent: %v", err)
+		mesosAgentFlags, err := fetchMasterFromAgent(localAgentStateURL)
+		if err != nil || mesosAgentFlags.Master == "" {
+			log.Debugf("unable to fetch Mesos leader via local Mesos agent: %v", err)
 			args = append(args, "localhost")
-		} else {
-			fmt.Println(url)
+		} else if mesosAgentFlags.hasMaster {
+			fmt.Println(mesosAgentFlags.Master)
 			return
+		} else {
+			args = append(args, strings.Split(mesosAgentFlags.Master, ",")...)
 		}
 	}
-	log.Infof("Fetching Mesos-master leader from %v \n", args)
+	log.Infof("Fetching Mesos-master leader from Zookeeper node(s): %v \n", args)
 
 	url, err := realis.MesosFromZKOpts(realis.ZKEndpoints(args...), realis.ZKPath(cmd.Flag("zkPath").Value.String()))
 
@@ -286,7 +295,7 @@ func fetchMesosLeader(cmd *cobra.Command, args []string) {
 	fmt.Println(url)
 }
 
-func fetchMesosLeaderFromAgent(url string) (mesosLeaderHostName string, err error) {
+func fetchMasterFromAgent(url string) (mesosAgentFlags mesosAgentFlags, err error) {
 	resp, err := http.Get(url)
 	if err != nil {
 		return
@@ -301,10 +310,44 @@ func fetchMesosLeaderFromAgent(url string) (mesosLeaderHostName string, err erro
 	if err != nil {
 		return
 	}
-	mesosLeaderHostName = state.MasterHostname
+	mesosAgentFlags = state.Flags
+	err = updateMasterFlag(&mesosAgentFlags)
 	return
 }
 
+/*
+ Master flag can be passed as one of :
+ host:port
+ zk://host1:port1,host2:port2,.../path
+ zk://username:password@host1:port1,host2:port2,.../path
+ file:///path/to/file
+ This function takes care of all the above cases and updates flags with parsed values
+*/
+func updateMasterFlag(flags *mesosAgentFlags) error {
+	zkPathPrefix := "zk://"
+	filePathPrefix := "file://"
+	if strings.HasPrefix(flags.Master, zkPathPrefix) {
+		beginIndex := len(zkPathPrefix)
+		if strings.Contains(flags.Master, "@") {
+			beginIndex = strings.Index(flags.Master, "@") + 1
+		}
+		flags.Master = flags.Master[beginIndex:strings.LastIndex(flags.Master, "/")]
+	} else if strings.HasPrefix(flags.Master, filePathPrefix) {
+		content, err := ioutil.ReadFile(flags.Master)
+		if err != nil {
+			return err
+		}
+		if strings.Contains(string(content), filePathPrefix) {
+			return errors.New("invalid master file content")
+		}
+		flags.Master = string(content)
+		return updateMasterFlag(flags)
+	} else {
+		flags.hasMaster = true
+	}
+	return nil
+}
+
 // TODO: Expand this to be able to filter by job name and environment.
 func fetchJobs(cmd *cobra.Command, args []string) {
 	log.Infof("Fetching tasks under role: %s \n", *role)