mmaprototype: add tracing to rebalance_store test

Author: tbg

pkg/kv/kvserver/allocator/mmaprototype/BUILD.bazel

@@ -53,6 +53,7 @@ go_test(
         "//pkg/testutils/dd",
         "//pkg/util/randutil",
         "//pkg/util/timeutil",
+        "//pkg/util/tracing",
         "@com_github_cockroachdb_datadriven//:datadriven",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",

pkg/kv/kvserver/allocator/mmaprototype/cluster_state_test.go

@@ -21,7 +21,9 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/testutils/datapathutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/dd"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/datadriven"
+	"github.com/cockroachdb/redact"
 	"github.com/stretchr/testify/require"
 )
 
@@ -519,8 +521,13 @@ func TestClusterState(t *testing.T) {
 					storeID := dd.ScanArg[roachpb.StoreID](t, d, "store-id")
 					rng := rand.New(rand.NewSource(0))
 					dsm := newDiversityScoringMemo()
-					cs.rebalanceStores(context.Background(), storeID, rng, dsm)
-					return printPendingChangesTest(testingGetPendingChanges(t, cs))
+					tr := tracing.NewTracer()
+					ctx, finishAndGet := tracing.ContextWithRecordingSpan(context.Background(), tr, "rebalance-stores")
+					cs.rebalanceStores(ctx, storeID, rng, dsm)
+					rec := finishAndGet()
+					var sb redact.StringBuilder
+					rec.SafeFormatMinimal(&sb)
+					return sb.String() + printPendingChangesTest(testingGetPendingChanges(t, cs))
 
 				case "tick":
 					seconds := dd.ScanArg[int](t, d, "seconds")

pkg/kv/kvserver/allocator/mmaprototype/testdata/cluster_state/rebalance_stores_basic

@@ -55,10 +55,15 @@ store-id=2 node-id=2 status=ok accepting all reported=[cpu:10, write-bandwidth:0
 store-id=3 node-id=3 status=ok accepting all reported=[cpu:10, write-bandwidth:0, byte-size:0] adjusted=[cpu:10, write-bandwidth:0, byte-size:0] node-reported-cpu=10 node-adjusted-cpu=10 seq=1
   top-k-ranges (local-store-id=1) dim=WriteBandwidth: r1
 
-# TODO(tbg): maybe it'll be good to at least optionally be able to have the
-# trace printed in the output here.
+# s1 is overloaded and local, so rebalanceStores should try and succeed to shed
+# a lease from it.
 rebalance-stores store-id=1
 ----
+[mmaid=1] cluster means: (stores-load [cpu:33, write-bandwidth:0, byte-size:0]) (stores-capacity [cpu:100, write-bandwidth:100, byte-size:100]) (nodes-cpu-load 33) (nodes-cpu-capacity 100)
+[mmaid=1] overload-continued s1 ((store=overloadUrgent worst=CPURate cpu=overloadUrgent writes=loadNormal bytes=loadNormal node=overloadUrgent high_disk=false frac_pending=0.00,0.00(true))) - within grace period
+[mmaid=1] start processing shedding store s1: cpu node load overloadUrgent, store load overloadUrgent, worst dim CPURate
+[mmaid=1] top-K[CPURate] ranges for s1 with lease on local s1: r1:[cpu:60, write-bandwidth:0, byte-size:0]
+[mmaid=1] result(success): shedding r1 lease from s1 to s2 [change:r1=[transfer_to=2 cids=1,2]] with resulting loads source:[cpu:40, write-bandwidth:0, byte-size:0] target:[cpu:54, write-bandwidth:0, byte-size:0] (means: [cpu:33, write-bandwidth:0, byte-size:0]) (frac_pending: (src:0.00,target:0.50) (src:4.40,target:0.00))
 pending(2)
 change-id=1 store-id=1 node-id=1 range-id=1 load-delta=[cpu:-40, write-bandwidth:0, byte-size:0] start=0s
   prev=(replica-id=1 type=VOTER_FULL leaseholder=true)