From 518fd7e8f1d825d45565fa8e9f8d7ef9dc8a5ea1 Mon Sep 17 00:00:00 2001
From: Austen McClernon <austen@cockroachlabs.com>
Date: Fri, 5 May 2023 18:22:37 +0000
Subject: [PATCH] kvserver: annotate rebalance ctx with objective

Previously, the easiest method of determining the current rebalance
objective from logs was to view the cluster setting and check for
logging indicating a mixed version cluster - this was cumbersome.

This commit annotates the ctx in the store rebalancer loop with an
additional tag: `obj`. The `obj` tag indicates the current rebalance
objective, either `cpu` or `qps` currently.

resolves: #102812

Release note: None
---
 pkg/kv/kvserver/rebalance_objective.go | 16 ++++++++++++----
 pkg/kv/kvserver/store_rebalancer.go    |  3 +++
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/pkg/kv/kvserver/rebalance_objective.go b/pkg/kv/kvserver/rebalance_objective.go
index 3514f3c4fc99..268db3fb8c23 100644
--- a/pkg/kv/kvserver/rebalance_objective.go
+++ b/pkg/kv/kvserver/rebalance_objective.go
@@ -92,6 +92,17 @@ const (
 	LBRebalancingCPU
 )
 
+// LoadBasedRebalancingObjectiveMap maps the LoadBasedRebalancingObjective enum
+// value to a string.
+var LoadBasedRebalancingObjectiveMap map[int64]string = map[int64]string{
+	int64(LBRebalancingQueries): "qps",
+	int64(LBRebalancingCPU):     "cpu",
+}
+
+func (lbro LBRebalancingObjective) String() string {
+	return LoadBasedRebalancingObjectiveMap[int64(lbro)]
+}
+
 // LoadBasedRebalancingObjective is a cluster setting that defines the load
 // balancing objective of the cluster.
 var LoadBasedRebalancingObjective = settings.RegisterEnumSetting(
@@ -101,10 +112,7 @@ var LoadBasedRebalancingObjective = settings.RegisterEnumSetting(
 		"the cluster will attempt to balance qps among stores, if set to "+
 		"`cpu` the cluster will attempt to balance cpu usage among stores",
 	"cpu",
-	map[int64]string{
-		int64(LBRebalancingQueries): "qps",
-		int64(LBRebalancingCPU):     "cpu",
-	},
+	LoadBasedRebalancingObjectiveMap,
 ).WithPublic()
 
 // ToDimension returns the equivalent allocator load dimension of a rebalancing
diff --git a/pkg/kv/kvserver/store_rebalancer.go b/pkg/kv/kvserver/store_rebalancer.go
index 2f67ad2c0ddf..bc290abb6436 100644
--- a/pkg/kv/kvserver/store_rebalancer.go
+++ b/pkg/kv/kvserver/store_rebalancer.go
@@ -276,6 +276,9 @@ func (sr *StoreRebalancer) Start(ctx context.Context, stopper *stop.Stopper) {
 				continue
 			}
 			objective := sr.RebalanceObjective()
+			sr.AddLogTag("obj", objective)
+			ctx = sr.AnnotateCtx(ctx)
+
 			hottestRanges := sr.replicaRankings.TopLoad(objective.ToDimension())
 			options := sr.scorerOptions(ctx, objective.ToDimension())
 			rctx := sr.NewRebalanceContext(ctx, options, hottestRanges, mode)