Merge pull request #951 from molecula/shard-dist-api

Add ui-specific endpoint to report shard distribution data

api.go

@@ -1597,6 +1597,21 @@ func importExistenceColumns(qcx *Qcx, index *Index, columnIDs []uint64) error {
 	return ef.Import(qcx, existenceRowIDs, columnIDs, nil)
 }
 
+// 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.
+// This function supports a view in the UI.
+func (api *API) ShardDistribution(ctx context.Context) map[string]interface{} {
+	distByIndex := make(map[string]interface{})
+
+	for idx := range api.holder.indexes {
+		dist := api.cluster.shardDistributionByIndex(idx)
+		distByIndex[idx] = dist
+	}
+
+	return distByIndex
+}
+
 // MaxShards returns the maximum shard number for each index in a map.
 // TODO (2.0): This method has been deprecated. Instead, use
 // AvailableShardsByIndex.
@@ -1772,6 +1787,8 @@ func (api *API) Info() serverInfo {
 		Memory:           mem,
 		TxSrc:            api.holder.txf.TxType(),
 		ReplicaN:         api.cluster.ReplicaN,
+		ShardHash:        api.cluster.Hasher.Name(),
+		KeyHash:          api.cluster.Topology.Hasher.Name(),
 	}
 }
 
@@ -2019,6 +2036,8 @@ func (api *API) TranslateFieldDB(ctx context.Context, indexName, fieldName strin
 type serverInfo struct {
 	ShardWidth       uint64 `json:"shardWidth"`
 	ReplicaN         int    `json:"replicaN"`
+	ShardHash        string `json:"shardHash"`
+	KeyHash          string `json:"keyHash"`
 	Memory           uint64 `json:"memory"`
 	CPUType          string `json:"cpuType"`
 	CPUPhysicalCores int    `json:"cpuPhysicalCores"`

api_test.go

@@ -490,6 +490,8 @@ func (*offsetModHasher) Hash(key uint64, n int) int {
 	return int(key+1) % n
 }
 
+func (*offsetModHasher) Name() string { return "mod" }
+
 func TestAPI_ClearFlagForImportAndImportValues(t *testing.T) {
 	c := test.MustRunCluster(t, 1,
 		[]server.CommandOption{

cluster.go

@@ -978,6 +978,36 @@ func (c *cluster) translationNodes(to *cluster) (map[string][]*translationResize
 	return m, nil
 }
 
+// shardDistributionByIndex returns a map of [nodeID][primaryOrReplica][]uint64,
+// where the int slices are lists of shards.
+func (c *cluster) shardDistributionByIndex(indexName string) map[string]map[string][]uint64 {
+	dist := make(map[string]map[string][]uint64)
+
+	for _, node := range c.nodes {
+		nodeDist := make(map[string][]uint64)
+		nodeDist["primary-shards"] = make([]uint64, 0)
+		nodeDist["replica-shards"] = make([]uint64, 0)
+		dist[node.ID] = nodeDist
+	}
+
+	index := c.holder.Index(indexName)
+	available := index.AvailableShards(includeRemote).Slice()
+
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	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++ {
+			dist[nodes[k].ID]["replica-shards"] = append(dist[nodes[k].ID]["replica-shards"], shard)
+		}
+	}
+
+	return dist
+}
+
 // 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 {
@@ -1193,9 +1223,10 @@ func (c *cluster) containsShards(index string, availableShards *roaring.Bitmap,
 type Hasher interface {
 	// Hashes the key into a number between [0,N).
 	Hash(key uint64, n int) int
+	Name() string
 }
 
-// jmphasher represents an implementation of jmphash. Implements Hasher.
+// Jmphasher represents an implementation of jmphash. Implements Hasher.
 type Jmphasher struct{}
 
 // Hash returns the integer hash for the given key.
@@ -1209,6 +1240,11 @@ func (h *Jmphasher) Hash(key uint64, n int) int {
 	return int(b)
 }
 
+// Name returns the name of this hash.
+func (h *Jmphasher) Name() string {
+	return "jump-hash"
+}
+
 func (c *cluster) setup() error {
 	// Cluster always comes up in state STARTING until cluster membership is determined.
 	c.state = ClusterStateStarting
@@ -1905,7 +1941,7 @@ func (n nodeIDs) Len() int           { return len(n) }
 func (n nodeIDs) Swap(i, j int)      { n[i], n[j] = n[j], n[i] }
 func (n nodeIDs) Less(i, j int) bool { return n[i] < n[j] }
 
-// ContainsID returns true if idi matches one of the nodesets's IDs.
+// ContainsID returns true if id matches one of the nodesets's IDs.
 func (n nodeIDs) ContainsID(id string) bool {
 	for _, nid := range n {
 		if nid == id {

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) {

server/handler_test.go

@@ -400,6 +400,32 @@ func TestHandler_Endpoints(t *testing.T) {
 		}
 	})
 
+	t.Run("UI/shard-distribution", func(t *testing.T) {
+		// This tests the response structure, not the shard distribution.
+		w := httptest.NewRecorder()
+		h.ServeHTTP(w, test.MustNewHTTPRequest("GET", "/ui/shard-distribution", nil))
+		if w.Code != gohttp.StatusOK {
+			t.Fatalf("unexpected status code: %d", w.Code)
+		}
+
+		ret := mustJSONDecode(t, w.Body)
+
+		for indexName := range ret {
+			indexData := ret[indexName].(map[string]interface{})
+			for nodeName := range indexData {
+				nodeData := indexData[nodeName].(map[string]interface{})
+				_, hasPrimary := nodeData["primary-shards"]
+				_, hasReplica := nodeData["replica-shards"]
+
+				responseOK := hasPrimary && hasReplica
+				if !responseOK {
+					t.Fatalf("unexpected response structure")
+				}
+			}
+
+		}
+	})
+
 	t.Run("Metrics", func(t *testing.T) {
 		w := httptest.NewRecorder()
 		h.ServeHTTP(w, test.MustNewHTTPRequest("GET", "/metrics", nil))

test/cluster.go

@@ -34,6 +34,8 @@ type ModHasher struct{}
 
 func (*ModHasher) Hash(key uint64, n int) int { return int(key) % n }
 
+func (*ModHasher) Name() string { return "mod" }
+
 // Cluster represents a Pilosa cluster (multiple Command instances)
 type Cluster struct {
 	Nodes []*Command

utils_internal_test.go

@@ -110,6 +110,8 @@ func NewTestModHasher() *TestModHasher { return &TestModHasher{} }
 
 func (*TestModHasher) Hash(key uint64, n int) int { return int(key) % n }
 
+func (*TestModHasher) Name() string { return "mod" }
+
 // ClusterCluster represents a cluster of test nodes, each of which
 // has a Cluster.
 // ClusterCluster implements Broadcaster interface.