diff --git a/cmd/fetch.go b/cmd/fetch.go index cfb2e40..5a6c754 100644 --- a/cmd/fetch.go +++ b/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)