From 94e82aea864d7493c75659683c26faefa37c8754 Mon Sep 17 00:00:00 2001
From: Alan Bernstein <alan.aaron.bernstein@gmail.com>
Date: Wed, 14 Oct 2020 18:45:53 -0500
Subject: [PATCH] Consider available shards

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

diff --git a/api.go b/api.go
index e0d7d6891..31813ae65 100644
--- a/api.go
+++ b/api.go
@@ -1599,18 +1599,13 @@ func importExistenceColumns(qcx *Qcx, index *Index, columnIDs []uint64) error {
 
 // ShardDistribution returns an object representing the distribution of shards
 // across nodes for each index, distinguishing between primary and replica.
-// The structure of this information is [indexName][nodeID][primaryOrReplica]uint64.
+// The structure of this information is [indexName][nodeID][primaryOrReplica][]uint64.
 // This function supports a view in the UI.
 func (api *API) ShardDistribution(ctx context.Context) map[string]interface{} {
 	distByIndex := make(map[string]interface{})
-	maxShards := api.MaxShards(ctx)
 
 	for idx := range api.holder.indexes {
-		calculatedMaxShard := uint64(0)
-		if mx, ok := maxShards[idx]; ok {
-			calculatedMaxShard = mx
-		}
-		dist := api.cluster.shardDistributionByIndex(idx, calculatedMaxShard)
+		dist := api.cluster.shardDistributionByIndex(idx)
 		distByIndex[idx] = dist
 	}
 
diff --git a/cluster.go b/cluster.go
index dd1ae1457..20e10cd1e 100644
--- a/cluster.go
+++ b/cluster.go
@@ -980,7 +980,7 @@ func (c *cluster) translationNodes(to *cluster) (map[string][]*translationResize
 
 // shardDistributionByIndex returns a map of [nodeID][primaryOrReplica][]uint64,
 // where the int slices are lists of shards.
-func (c *cluster) shardDistributionByIndex(index string, maxShard uint64) map[string]map[string][]uint64 {
+func (c *cluster) shardDistributionByIndex(indexName string) map[string]map[string][]uint64 {
 	dist := make(map[string]map[string][]uint64)
 
 	for _, node := range c.nodes {
@@ -990,11 +990,14 @@ func (c *cluster) shardDistributionByIndex(index string, maxShard uint64) map[st
 		dist[node.ID] = nodeDist
 	}
 
+	index := c.holder.Index(indexName)
+	available := index.AvailableShards(includeRemote).Slice()
+
 	c.mu.RLock()
 	defer c.mu.RUnlock()
 
-	for shard := uint64(0); shard <= maxShard; shard++ {
-		p := c.shardToShardPartition(index, shard)
+	for _, shard := range available {
+		p := c.shardToShardPartition(indexName, shard)
 		nodes := c.partitionNodes(p)
 		dist[nodes[0].ID]["primary-shards"] = append(dist[nodes[0].ID]["primary-shards"], shard)
 		for k := 1; k < len(nodes); k++ {
diff --git a/http/handler.go b/http/handler.go
index 1f30c1db2..3d916fbcc 100644
--- a/http/handler.go
+++ b/http/handler.go
@@ -685,14 +685,6 @@ func (h *Handler) handleGetUsage(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
-type getUsageResponse struct {
-	Disk diskUsage `json:"bytesOnDisk"`
-}
-type diskUsage struct {
-	Total   int64            `json:"total"`
-	Indexes map[string]int64 `json:"indexes"`
-}
-
 // handleGetUsage handles GET /ui/shard-distribution requests.
 func (h *Handler) handleGetShardDistribution(w http.ResponseWriter, r *http.Request) {
 	dist := h.api.ShardDistribution(r.Context())