From 63cdb96c7ab6ce2fcfb49abe588368d3fb0f8639 Mon Sep 17 00:00:00 2001
From: Lawrence Wong <lawwong@paypal.com>
Date: Mon, 22 Aug 2022 09:55:41 -0700
Subject: [PATCH] Address code review comment from Tan

---
 zk.go | 45 +++++++++++++++++----------------------------
 1 file changed, 17 insertions(+), 28 deletions(-)

diff --git a/zk.go b/zk.go
index 8856357..f67c907 100644
--- a/zk.go
+++ b/zk.go
@@ -288,12 +288,11 @@ func MesosFromZKOpts(options ...ZKOpt) (string, error) {
 }
 
 // Retrieves current Mesos master nodes/leader from ZK with a custom configuration.
-func MasterNodesFromZKOpts(options ...ZKOpt) (map[string]string, error) {
-	result := make(map[string]string)
+func MasterNodesFromZKOpts(options ...ZKOpt) (map[string][]string, error) {
+	result := make(map[string][]string)
 
 	// Load the default configuration for Zookeeper followed by overriding values with those provided by the caller.
-	// 10s timeout is not enough. Use 100s
-	config := &zkConfig{backoff: defaultBackoff, timeout: time.Second * 100, logger: NoopLogger{}}
+	config := &zkConfig{backoff: defaultBackoff, timeout: time.Second * 10, logger: NoopLogger{}}
 	for _, opt := range options {
 		opt(config)
 	}
@@ -331,7 +330,7 @@ func MasterNodesFromZKOpts(options ...ZKOpt) (map[string]string, error) {
 
 		// Get all the master nodes through all the children in the given path
 		serviceInst := new(ServiceInstance)
-		var host string
+		var host []string
 		for _, child := range children {
 			childPath := config.path + "/" + child
 			data, _, err := c.Get(childPath)
@@ -359,23 +358,17 @@ func MasterNodesFromZKOpts(options ...ZKOpt) (map[string]string, error) {
 				}
 
 				for _, v := range serviceInst.AdditionalEndpoints {
-					result["leader"] = v.Host
+					result["leader"] = append(result["leader"], v.Host)
 				}
 			} else {
 				// data is not in a json format
-				hostname := string(data)
-				// Combine all master nodes into comma-separated
-				if len(host) > 0 {
-					host += ","
-				}
-				host += hostname
-				result["masterNodes"] = host
+				host = append(host, string(data))
 			}
-
 		}
+		result["masterNodes"] = host
 
 		// Master nodes data might not be available yet, try to fetch again.
-		if len(result["leader"]) == 0 || len(result["masterNodes"]) == 0 {
+		if len(result["masterNodes"]) == 0 {
 			return false, NewTemporaryError(errors.New("no master nodes found"))
 		}
 		return true, nil
@@ -390,12 +383,11 @@ func MasterNodesFromZKOpts(options ...ZKOpt) (map[string]string, error) {
 }
 
 // Retrieves current mesos master nodes/leader from ZK with a custom configuration.
-func MesosMasterNodesFromZKOpts(options ...ZKOpt) (map[string]string, error) {
-	result := make(map[string]string)
+func MesosMasterNodesFromZKOpts(options ...ZKOpt) (map[string][]string, error) {
+	result := make(map[string][]string)
 
-	// Load the default configuration for Zookeeper followed by overriding values with those provided by the caller.
-	// 10s timeout is not enough. Use 100s
-	config := &zkConfig{backoff: defaultBackoff, timeout: time.Second * 100, logger: NoopLogger{}}
+	// Load the default configuration for Zookeeper followed by overriding values with those provided by the caller.]
+	config := &zkConfig{backoff: defaultBackoff, timeout: time.Second * 10, logger: NoopLogger{}}
 	for _, opt := range options {
 		opt(config)
 	}
@@ -434,7 +426,7 @@ func MesosMasterNodesFromZKOpts(options ...ZKOpt) (map[string]string, error) {
 		// Get all the master nodes through all the children in the given path
 		minScore := math.MaxInt64
 		var mesosInstance MesosInstance
-		var host string
+		var host []string
 		for _, child := range children {
 			// Only the master nodes will start with json.info_
 			if strings.HasPrefix(child, "json.info_") {
@@ -466,20 +458,17 @@ func MesosMasterNodesFromZKOpts(options ...ZKOpt) (map[string]string, error) {
 				}
 				// Combine all master nodes into comma-separated
 				// Return hostname instead of ip to be consistent with aurora master nodes
-				if len(host) > 0 {
-					host += ","
-				}
-				host += mesosInstance.Address.Hostname
-				result["masterNodes"] = host
+				host = append(host, mesosInstance.Address.Hostname)
 				// get the leader from the child with the smallest score.
 				if score < minScore {
 					minScore = score
-					result["leader"] = mesosInstance.Address.Hostname
+					result["leader"] = append(result["leader"], mesosInstance.Address.Hostname)
 				}
 			}
 		}
+		result["masterNodes"] = host
 		// Master nodes data might not be available yet, try to fetch again.
-		if len(result["leader"]) == 0 || len(result["masterNodes"]) == 0 {
+		if len(result["masterNodes"]) == 0 {
 			return false, NewTemporaryError(errors.New("no mesos master nodes found"))
 		}
 		return true, nil