refactor: simplify cluster shards

redis · Sep 20, 2023 · 8361761 · 8361761
1 parent 8503e0d
commit 8361761
Showing 1 changed file with 36 additions and 72 deletions.
diff --git a/cluster.go b/cluster.go
@@ -262,35 +262,33 @@ type group struct {
 	slots [][2]int64
 }
 
+func parseEndpoint(fallback, endpoint string, port int64) string {
+	switch endpoint {
+	case "":
+		endpoint, _, _ = net.SplitHostPort(fallback)
+	case "?":
+		return ""
+	}
+	return net.JoinHostPort(endpoint, strconv.FormatInt(port, 10))
+}
+
 // parseSlots - map redis slots for each redis nodes/addresses
 // defaultAddr is needed in case the node does not know its own IP
 func parseSlots(slots RedisMessage, defaultAddr string) map[string]group {
 	groups := make(map[string]group, len(slots.values))
 	for _, v := range slots.values {
-		var master string
-		switch v.values[2].values[0].string {
-		case "":
-			master = defaultAddr
-		case "?":
+		master := parseEndpoint(defaultAddr, v.values[2].values[0].string, v.values[2].values[1].integer)
+		if master == "" {
 			continue
-		default:
-			master = net.JoinHostPort(v.values[2].values[0].string, strconv.FormatInt(v.values[2].values[1].integer, 10))
 		}
 		g, ok := groups[master]
 		if !ok {
 			g.slots = make([][2]int64, 0)
 			g.nodes = make([]string, 0, len(v.values)-2)
 			for i := 2; i < len(v.values); i++ {
-				var dst string
-				switch v.values[i].values[0].string {
-				case "":
-					dst = defaultAddr
-				case "?":
-					continue
-				default:
-					dst = net.JoinHostPort(v.values[i].values[0].string, strconv.FormatInt(v.values[i].values[1].integer, 10))
+				if dst := parseEndpoint(defaultAddr, v.values[i].values[0].string, v.values[i].values[1].integer); dst != "" {
+					g.nodes = append(g.nodes, dst)
 				}
-				g.nodes = append(g.nodes, dst)
 			}
 		}
 		g.slots = append(g.slots, [2]int64{v.values[0].integer, v.values[1].integer})
@@ -302,70 +300,36 @@ func parseSlots(slots RedisMessage, defaultAddr string) map[string]group {
 // parseShards - map redis shards for each redis nodes/addresses
 // defaultAddr is needed in case the node does not know its own IP
 func parseShards(shards RedisMessage, defaultAddr string, tls bool) map[string]group {
-	parseNodeEndpoint := func(msg map[string]RedisMessage) string {
-		endpoint := msg["endpoint"].string
-		switch endpoint {
-		case "":
-			return defaultAddr
-		case "?":
-			return ""
-		}
-
-		port := msg["port"].integer
-		if tls && msg["tls-port"].integer > 0 {
-			port = msg["tls-port"].integer
-		}
-
-		return net.JoinHostPort(endpoint, strconv.FormatInt(port, 10))
-	}
-
 	groups := make(map[string]group, len(shards.values))
 	for _, v := range shards.values {
-		slotsAndNodes, _ := v.ToMap()
-		var (
-			master    string
-			masterPos int
-		)
-		nodes := slotsAndNodes["nodes"].values
-		for i := 0; i < len(nodes); i++ {
-			dict, _ := nodes[i].ToMap()
-			if dict["role"].string != "master" {
-				continue
-			}
-			master = parseNodeEndpoint(dict)
-			masterPos = i
-			break
+		m := -1
+		shard, _ := v.ToMap()
+		slots := shard["slots"].values
+		nodes := shard["nodes"].values
+		g := group{
+			nodes: make([]string, 0, len(nodes)),
+			slots: make([][2]int64, len(slots)/2),
 		}
-
-		if master == "" {
-			continue
+		for i := range g.slots {
+			g.slots[i] = [2]int64{slots[i*2].integer, slots[i*2+1].integer}
 		}
-
-		g, ok := groups[master]
-		if !ok {
-			g.slots = make([][2]int64, 0)
-			g.nodes = make([]string, 0, len(nodes))
-			g.nodes = append(g.nodes, master)
-			for i := 0; i < len(nodes); i++ {
-				if i == masterPos {
-					continue
-				}
-				dict, _ := nodes[i].ToMap()
-				dst := parseNodeEndpoint(dict)
-				if dst == "" {
-					continue
-				}
+		for _, n := range nodes {
+			dict, _ := n.ToMap()
+			port := dict["port"].integer
+			if tls && dict["tls-port"].integer > 0 {
+				port = dict["tls-port"].integer
+			}
+			if dst := parseEndpoint(defaultAddr, dict["endpoint"].string, port); dst != "" {
 				g.nodes = append(g.nodes, dst)
+				if dict["role"].string == "master" {
+					m = len(g.nodes) - 1
+				}
 			}
 		}
-		slots := slotsAndNodes["slots"]
-		arr, _ := slots.ToArray()
-		for i := 0; i+1 < len(arr); i += 2 {
-			start, _ := arr[i].AsInt64()
-			end, _ := arr[i+1].AsInt64()
-			g.slots = append(g.slots, [2]int64{start, end})
+		if m >= 0 {
+			g.nodes[0], g.nodes[m] = g.nodes[m], g.nodes[0]
+			groups[g.nodes[0]] = g
 		}
-		groups[master] = g
 	}
 	return groups
 }