From 8fad2b92c2048ba2a336aaaa16e86901623ef93d Mon Sep 17 00:00:00 2001
From: Alan Bernstein <alan.aaron.bernstein@gmail.com>
Date: Fri, 7 Aug 2020 11:14:24 -0500
Subject: [PATCH] Add all node states to /status response

---
 api.go          | 6 ++++++
 cluster.go      | 6 ++++++
 http/handler.go | 8 +++++---
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/api.go b/api.go
index 501187512..4e11cc838 100644
--- a/api.go
+++ b/api.go
@@ -776,6 +776,12 @@ func (api *API) Hosts(ctx context.Context) []*Node {
 	return api.cluster.Nodes()
 }
 
+func (api *API) HostStates(ctx context.Context) map[string]string {
+	span, _ := tracing.StartSpanFromContext(ctx, "API.HostStates")
+	defer span.Finish()
+	return api.cluster.AllNodeStates()
+}
+
 // Node gets the ID, URI and coordinator status for this particular node.
 func (api *API) Node() *Node {
 	node := api.server.node()
diff --git a/cluster.go b/cluster.go
index 6e5450cf4..1762be18d 100644
--- a/cluster.go
+++ b/cluster.go
@@ -700,6 +700,12 @@ func (c *cluster) Nodes() []*Node {
 	return ret
 }
 
+func (c *cluster) AllNodeStates() map[string]string {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.Topology.nodeStates
+}
+
 // removeNodeBasicSorted removes a node from the cluster, maintaining the sort
 // order. Returns true if the node was removed. unprotected.
 func (c *cluster) removeNodeBasicSorted(nodeID string) bool {
diff --git a/http/handler.go b/http/handler.go
index f000a365c..186e857ea 100644
--- a/http/handler.go
+++ b/http/handler.go
@@ -702,6 +702,7 @@ func (h *Handler) handleGetStatus(w http.ResponseWriter, r *http.Request) {
 	status := getStatusResponse{
 		State:   h.api.State(),
 		Nodes:   h.api.Hosts(r.Context()),
+		NodeStates: h.api.HostStates(r.Context()),
 		LocalID: h.api.Node().ID,
 	}
 	w.Header().Set("Content-Type", "application/json")
@@ -758,9 +759,10 @@ type getSchemaResponse struct {
 }
 
 type getStatusResponse struct {
-	State   string         `json:"state"`
-	Nodes   []*pilosa.Node `json:"nodes"`
-	LocalID string         `json:"localID"`
+	State      string            `json:"state"`
+	Nodes      []*pilosa.Node    `json:"nodes"`
+	NodeStates map[string]string `json:"nodeStates"`
+	LocalID    string            `json:"localID"`
 }
 
 func hash(s string) string {