From f2905ca4553d548b2a5ccd2d291693366c26fa4a Mon Sep 17 00:00:00 2001
From: Venkatraju V <venkatraju@slack-corp.com>
Date: Fri, 31 May 2024 15:14:18 -0700
Subject: [PATCH] debug page for viewing tabletbalancer state

Signed-off-by: Venkatraju V <venkatraju@slack-corp.com>
---
 go/vt/vtgate/balancer/balancer.go | 15 +++++++++++++++
 go/vt/vtgate/tabletgateway.go     | 10 ++++++++++
 go/vt/vtgate/vtgate.go            |  7 +++++++
 3 files changed, 32 insertions(+)

diff --git a/go/vt/vtgate/balancer/balancer.go b/go/vt/vtgate/balancer/balancer.go
index 80486695a5b..462ccfda901 100644
--- a/go/vt/vtgate/balancer/balancer.go
+++ b/go/vt/vtgate/balancer/balancer.go
@@ -20,6 +20,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"math/rand"
+	"net/http"
 	"sync"
 
 	"vitess.io/vitess/go/vt/discovery"
@@ -89,6 +90,9 @@ converge on the desired balanced query load.
 type TabletBalancer interface {
 	// Randomly shuffle the tablets into an order for routing queries
 	ShuffleTablets(target *querypb.Target, tablets []*discovery.TabletHealth)
+
+	// Balancer debug page request
+	DebugHandler(w http.ResponseWriter, r *http.Request)
 }
 
 func NewTabletBalancer(localCell string, vtGateCells []string) TabletBalancer {
@@ -146,6 +150,17 @@ func (b *tabletBalancer) print() string {
 		b.localCell, b.vtGateCells, string(allocations))
 }
 
+func (b *tabletBalancer) DebugHandler(w http.ResponseWriter, _ *http.Request) {
+	w.Header().Set("Content-Type", "text/plain")
+	fmt.Fprintf(w, "Local Cell: %v\r\n", b.localCell)
+	fmt.Fprintf(w, "Vtgate Cells: %v\r\n", b.vtGateCells)
+
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	allocations, _ := json.MarshalIndent(b.allocations, "", "  ")
+	fmt.Fprintf(w, "Allocations: %v\r\n", string(allocations))
+}
+
 // ShuffleTablets is the main entry point to the balancer.
 //
 // It shuffles the tablets into a preference list for routing a given query.
diff --git a/go/vt/vtgate/tabletgateway.go b/go/vt/vtgate/tabletgateway.go
index ce8c12019d5..eff7cdd02f4 100644
--- a/go/vt/vtgate/tabletgateway.go
+++ b/go/vt/vtgate/tabletgateway.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"fmt"
 	"math/rand"
+	"net/http"
 	"sort"
 	"sync"
 	"sync/atomic"
@@ -257,6 +258,15 @@ func (gw *TabletGateway) CacheStatus() TabletCacheStatusList {
 	return res
 }
 
+func (gw *TabletGateway) DebugBalancerHandler(w http.ResponseWriter, r *http.Request) {
+	if balancerEnabled {
+		gw.balancer.DebugHandler(w, r)
+	} else {
+		w.Header().Set("Content-Type", "text/plain")
+		w.Write([]byte("not enabled"))
+	}
+}
+
 // withRetry gets available connections and executes the action. If there are retryable errors,
 // it retries retryCount times before failing. It does not retry if the connection is in
 // the middle of a transaction. While returning the error check if it maybe a result of
diff --git a/go/vt/vtgate/vtgate.go b/go/vt/vtgate/vtgate.go
index 23bdf94d90d..6c60c1ff7db 100644
--- a/go/vt/vtgate/vtgate.go
+++ b/go/vt/vtgate/vtgate.go
@@ -300,6 +300,7 @@ func Init(
 	})
 	rpcVTGate.registerDebugHealthHandler()
 	rpcVTGate.registerDebugEnvHandler()
+	rpcVTGate.registerDebugBalancerHandler()
 	err := initQueryLogger(rpcVTGate)
 	if err != nil {
 		log.Fatalf("error initializing query logger: %v", err)
@@ -369,6 +370,12 @@ func (vtg *VTGate) registerDebugHealthHandler() {
 	})
 }
 
+func (vtg *VTGate) registerDebugBalancerHandler() {
+	http.HandleFunc("/debug/balancer", func(w http.ResponseWriter, r *http.Request) {
+		vtg.Gateway().DebugBalancerHandler(w, r)
+	})
+}
+
 // IsHealthy returns nil if server is healthy.
 // Otherwise, it returns an error indicating the reason.
 func (vtg *VTGate) IsHealthy() error {