From fbdca3c622e720a97565d607fe0c36c676d98bb8 Mon Sep 17 00:00:00 2001
From: Travis <travis@pilosa.com>
Date: Wed, 3 Feb 2021 20:58:25 -0600
Subject: [PATCH] refactor the PrimaryNodeID logic

---
 etcd/embed.go        | 34 +++++++++++++++-------------------
 topology/noder.go    | 19 +++++++++++++++++++
 topology/snapshot.go | 12 ++++++++++++
 3 files changed, 46 insertions(+), 19 deletions(-)

diff --git a/etcd/embed.go b/etcd/embed.go
index 924e4c507..7a0bc2b79 100644
--- a/etcd/embed.go
+++ b/etcd/embed.go
@@ -1044,25 +1044,18 @@ func (e *Etcd) RemoveShard(ctx context.Context, index, field string, shard uint6
 	return nil
 }
 
-// Nodes implements the Noder interface.
+// Nodes implements the Noder interface. It returns the sorted list of nodes
+// based on the etcd peers.
 func (e *Etcd) Nodes() []*topology.Node {
-	return e.nodes(true)
-}
-
-// nodes is a helper function used to get the sorted list of nodes based on the
-// etcd peers.
-func (e *Etcd) nodes(includeMeta bool) []*topology.Node {
 	peers := e.Peers()
 	nodes := make([]*topology.Node, len(peers))
 	for i, peer := range peers {
 		node := &topology.Node{}
 
-		if includeMeta {
-			if meta, err := e.Metadata(context.Background(), peer.ID); err != nil {
-				log.Println(err, "getting metadata") // TODO: handle this with a logger
-			} else if err := json.Unmarshal(meta, node); err != nil {
-				log.Println(err, "unmarshaling json metadata")
-			}
+		if meta, err := e.Metadata(context.Background(), peer.ID); err != nil {
+			log.Println(err, "getting metadata") // TODO: handle this with a logger
+		} else if err := json.Unmarshal(meta, node); err != nil {
+			log.Println(err, "unmarshaling json metadata")
 		}
 
 		node.ID = peer.ID
@@ -1078,14 +1071,17 @@ func (e *Etcd) nodes(includeMeta bool) []*topology.Node {
 
 // PrimaryNodeID implements the Noder interface.
 func (e *Etcd) PrimaryNodeID(hasher topology.Hasher) string {
-	nodes := e.nodes(false)
+	return topology.PrimaryNodeID(e.NodeIDs(), hasher)
+}
 
-	snap := topology.NewClusterSnapshot(topology.NewLocalNoder(nodes), hasher, 1)
-	primaryNode := snap.PrimaryFieldTranslationNode()
-	if primaryNode == nil {
-		return ""
+// NodeIDs returns the list of node IDs in the etcd cluster.
+func (e *Etcd) NodeIDs() []string {
+	peers := e.Peers()
+	ids := make([]string, len(peers))
+	for i, peer := range peers {
+		ids[i] = peer.ID
 	}
-	return primaryNode.ID
+	return ids
 }
 
 // SetNodes implements the Noder interface as NOP
diff --git a/topology/noder.go b/topology/noder.go
index c84523f7f..63067e199 100644
--- a/topology/noder.go
+++ b/topology/noder.go
@@ -47,6 +47,25 @@ func NewEmptyLocalNoder() *localNoder {
 	return &localNoder{}
 }
 
+// NewIDNoder is a helper function for wrapping an existing slice of Node IDs
+// with something which implements Noder.
+func NewIDNoder(ids []string) *localNoder {
+	nodes := make([]*Node, len(ids))
+	for i, id := range ids {
+		node := &Node{
+			ID: id,
+		}
+		nodes[i] = node
+	}
+
+	// Nodes must be sorted.
+	sort.Sort(ByID(nodes))
+
+	return &localNoder{
+		nodes: nodes,
+	}
+}
+
 // Nodes implements the Noder interface.
 func (n *localNoder) Nodes() []*Node {
 	return n.nodes
diff --git a/topology/snapshot.go b/topology/snapshot.go
index decccccfc..fc1a2d83f 100644
--- a/topology/snapshot.go
+++ b/topology/snapshot.go
@@ -280,3 +280,15 @@ func NodePositionByID(nodes []*Node, nodeID string) int {
 	}
 	return -1
 }
+
+// PrimaryNodeID returns the ID of the primary node, given a list of node IDs
+// and a hasher. The order of the node IDs provided does not matter because this
+// function will re-order them in a deterministic way.
+func PrimaryNodeID(nodeIDs []string, hasher Hasher) string {
+	snap := NewClusterSnapshot(NewIDNoder(nodeIDs), hasher, 1)
+	primaryNode := snap.PrimaryFieldTranslationNode()
+	if primaryNode == nil {
+		return ""
+	}
+	return primaryNode.ID
+}