WIP shard distribution endpoint

api.go

@@ -1597,6 +1597,41 @@ func importExistenceColumns(qcx *Qcx, index *Index, columnIDs []uint64) error {
 	return ef.Import(qcx, existenceRowIDs, columnIDs, nil)
 }
 
+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
+		}
+		_, shards := api.cluster.shardDistributionByIndex(idx, calculatedMaxShard)
+		distByIndex[idx] = shards
+	}
+
+	return distByIndex
+}
+
+// ShardDistributionByIndex returns a slice of shards per node.
+func (api *API) ShardDistributionByIndex(ctx context.Context, index string, provideMaxShard bool, maxShard uint64) ([]Node, [][]uint64) {
+	span, _ := tracing.StartSpanFromContext(ctx, "API.ShardDistributionByIndex")
+	defer span.Finish()
+
+	calculatedMaxShard := uint64(0)
+	if provideMaxShard {
+		calculatedMaxShard = maxShard
+	} else {
+		// Get max shard from cluster.
+		maxShards := api.MaxShards(ctx)
+		if mx, ok := maxShards[index]; ok {
+			calculatedMaxShard = mx
+		}
+	}
+
+	return api.cluster.shardDistributionByIndex(index, calculatedMaxShard)
+}
+
 // MaxShards returns the maximum shard number for each index in a map.
 // TODO (2.0): This method has been deprecated. Instead, use
 // AvailableShardsByIndex.

cluster.go

@@ -978,12 +978,74 @@ func (c *cluster) translationNodes(to *cluster) (map[string][]*translationResize
 	return m, nil
 }
 
-// shardPartition returns the shard-partition that a shard belongs to.
-// NOTE: this is DIFFERENT from the key-partition
-func (c *cluster) shardToShardPartition(index string, shard uint64) int {
-	return shardToShardPartition(index, shard, c.partitionN)
+// shardDistributionByIndex returns a slice of shards per node for an index, up to maxShard.
+func (c *cluster) shardDistributionByIndex(index string, maxShard uint64) ([]Node, [][]uint64) {
+	m := make(map[Node][]uint64)
+
+	for i := range c.nodes {
+		m[*c.nodes[i]] = []uint64{}
+	}
+
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	for shard := uint64(0); shard <= maxShard; shard++ {
+		for _, node := range c.shardNodes(index, shard) {
+			m[*node] = append(m[*node], shard)
+		}
+	}
+
+	n := make([]Node, len(c.nodes))
+	s := make([][]uint64, len(c.nodes))
+
+	for i := range c.nodes {
+		n[i] = *c.nodes[i]
+		s[i] = m[*c.nodes[i]]
+	}
+
+	return n, s
 }
 
+// For specified index, return an object like
+/*
+{
+	"7aa98e81-0b53-43e4-9f98-c77d7fc2371a": {
+    	"primary-shards": [],
+		"replica-shards": []
+	},
+	"d4a7b7ff-529f-4d28-8d95-48d307751775": {
+    	"primary-shards": [],
+		"replica-shards": []
+	}
+	...
+}
+*/
+func (c *cluster) shardDistributionByIndex2(index string, maxShard uint64) map[string]interface{} {
+	dist := make(map[string]interface{})
+
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	for i := range c.nodes {
+		nodeDist := make(map[string][]uint64)
+		primaries := make([]uint64)
+		replicas := make([]uint64)
+		for shard := uint64(0); shard <= maxShard; shard++ {
+
+		}
+		dist[node.ID]["primary-shards"] = primaries
+		dist[node.ID]["replica-shards"] = replicas
+	}
+	return dist
+}
+
+// shardPartition returns the partition that a shard belongs to.
+func (c *cluster) shardPartition(index string, shard uint64) int {
+	return shardPartition(index, shard, c.partitionN)
+}
+
+// shardPartition returns the shard-partition that a shard belongs to.
+// NOTE: this is DIFFERENT from the key-partition
 func shardToShardPartition(index string, shard uint64, partitionN int) int {
 	var buf [8]byte
 	binary.BigEndian.PutUint64(buf[:], shard)

http/handler.go

@@ -393,6 +393,7 @@ func newRouter(handler *Handler) http.Handler {
 	router.HandleFunc("/ui/usage", handler.handleGetUsage).Methods("GET").Name("GetUsage")
 	router.HandleFunc("/ui/transaction", handler.handleGetTransactionList).Methods("GET").Name("GetTransactionList")
 	router.HandleFunc("/ui/transaction/", handler.handleGetTransactionList).Methods("GET").Name("GetTransactionList")
+	router.HandleFunc("/ui/shard-distribution", handler.handleGetShardDistribution).Methods("GET").Name("GetShardDistribution")
 
 	// /internal endpoints are for internal use only; they may change at any time.
 	// DO NOT rely on these for external applications!
@@ -684,6 +685,15 @@ func (h *Handler) handleGetUsage(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+// handleGetUsage handles GET /ui/shard-distribution requests.
+func (h *Handler) handleGetShardDistribution(w http.ResponseWriter, r *http.Request) {
+	dist := h.api.ShardDistribution(r.Context())
+	w.Header().Set("Content-Type", "application/json")
+	if err := json.NewEncoder(w).Encode(dist); err != nil {
+		h.logger.Printf("write status response error: %s", err)
+	}
+}
+
 // handleGetStatus handles GET /status requests.
 func (h *Handler) handleGetStatus(w http.ResponseWriter, r *http.Request) {
 	if !validHeaderAcceptJSON(r.Header) {