Merge #155033 #156876 #156942

155033: asim: use kvserver.LoadBasedRebalancingObjective r=wenyihu6 a=wenyihu6

Previously, LBRebalancingObjective was part of the simulation settings, using
QPS as the load rebalance objective. However, the cluster setting
kvserver.LoadBasedRebalancingObjective defaults to CPU rebalancing instead. This
commit updates asim to rely on the cluster setting, adopting the default cluster
setting configuration.

Epic: CRDB-55052
Release note: none 
Resolves: #154512

156876: rttanalysis: deflake BenchmarkJobs r=rafiss a=rafiss

The benchmark started failing since the code doesn't handle a nil table descriptor. The fix is to create a dummy descriptor in the job payload details.

fixes #151302
Release note: None

156942: server: make mma Allocator calls more visible r=wenyihu6 a=tbg

They were kind of hidden before, and it has been
tripping me up every now and then for a while.
The problem was that NewAllocatorState returns
a private struct directly, so a bunch of the
calls later in the method do not show up when
you look at the call hierarchy for the Allocator
interface. This is now fixed.

Epic: CRDB-55052.

Co-authored-by: wenyihu6 <[email protected]>
Co-authored-by: Rafi Shamim <[email protected]>
Co-authored-by: Tobias Grieger <[email protected]>

Author: 4 people

pkg/bench/rttanalysis/BUILD.bazel

@@ -65,6 +65,7 @@ go_test(
         "//pkg/security/username",
         "//pkg/server",
         "//pkg/server/serverpb",
+        "//pkg/sql/catalog/descpb",
         "//pkg/testutils/pgurlutils",
         "//pkg/testutils/serverutils",
         "//pkg/testutils/skip",

pkg/bench/rttanalysis/jobs_test.go

@@ -14,13 +14,27 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/security/username"
 	"github.com/cockroachdb/cockroach/pkg/server"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 )
 
 func BenchmarkJobs(b *testing.B) { reg.Run(b) }
 func init() {
+	// Create a minimal table descriptor for the import job.
+	tableDesc := &descpb.TableDescriptor{
+		ID:            100,
+		ParentID:      1,
+		Name:          "benchmark_table",
+		FormatVersion: descpb.InterleavedFormatVersion,
+		Version:       1,
+	}
+
 	payloadBytes, err := protoutil.Marshal(&jobspb.Payload{
-		Details:       jobspb.WrapPayloadDetails(jobspb.ImportDetails{}),
+		Details: jobspb.WrapPayloadDetails(jobspb.ImportDetails{
+			Table: jobspb.ImportDetails_Table{
+				Desc: tableDesc,
+			},
+		}),
 		UsernameProto: username.RootUserName().EncodeProto(),
 	})
 	if err != nil {

pkg/kv/kvserver/asim/config/settings.go

@@ -29,8 +29,6 @@ const (
 	defaultSplitStatRetention      = 10 * time.Minute
 	defaultSeed                    = 42
 	defaultLBRebalancingInterval   = time.Minute
-	defaultLBRebalanceQPSThreshold = 0.1
-	defaultLBRebalancingObjective  = 0 // QPS
 )
 
 const DefaultNodeCPURateCapacityNanos = 8 * 1e9        // 8 vcpus
@@ -101,8 +99,6 @@ type SimulationSettings struct {
 	// SplitStatRetention is the duration which recorded load will be retained
 	// and factored into load based splitting decisions.
 	SplitStatRetention time.Duration
-	// LBRebalancingObjective is the load objective to balance.
-	LBRebalancingObjective int64
 	// LBRebalancingInterval controls how often the store rebalancer will
 	// consider opportunities for rebalancing.
 	LBRebalancingInterval time.Duration
@@ -141,7 +137,6 @@ func DefaultSimulationSettings() *SimulationSettings {
 		StateExchangeDelay:      defaultStateExchangeDelay,
 		SplitQPSThreshold:       defaultSplitQPSThreshold,
 		SplitStatRetention:      defaultSplitStatRetention,
-		LBRebalancingObjective:  defaultLBRebalancingObjective,
 		LBRebalancingInterval:   defaultLBRebalancingInterval,
 		ReplicateQueueEnabled:   true,
 		LeaseQueueEnabled:       true,

pkg/kv/kvserver/asim/state/impl.go

@@ -1402,6 +1402,8 @@ func (s *state) SetClusterSetting(Key string, Value interface{}) {
 	switch Key {
 	case "LBRebalancingMode":
 		kvserverbase.LoadBasedRebalancingMode.Override(context.Background(), &s.settings.ST.SV, kvserverbase.LBRebalancingMode(Value.(int64)))
+	case "LBRebalancingObjective":
+		kvserver.LoadBasedRebalancingObjective.Override(context.Background(), &s.settings.ST.SV, kvserver.LBRebalancingObjective(Value.(int64)))
 	default:
 		panic("other cluster settings not supported")
 	}

pkg/kv/kvserver/asim/storerebalancer/store_rebalancer.go

@@ -126,11 +126,11 @@ type simRebalanceObjectiveProvider struct {
 
 // Objective returns the current rebalance objective.
 func (s simRebalanceObjectiveProvider) Objective() kvserver.LBRebalancingObjective {
-	return kvserver.LBRebalancingObjective(s.settings.LBRebalancingObjective)
+	return kvserver.LoadBasedRebalancingObjective.Get(&s.settings.ST.SV)
 }
 
 func (src *storeRebalancerControl) scorerOptions() *allocatorimpl.LoadScorerOptions {
-	dim := kvserver.LBRebalancingObjective(src.settings.LBRebalancingObjective).ToDimension()
+	dim := kvserver.LoadBasedRebalancingObjective.Get(&src.settings.ST.SV).ToDimension()
 	return &allocatorimpl.LoadScorerOptions{
 		BaseScorerOptions: allocatorimpl.BaseScorerOptions{
 			IOOverload:    src.allocator.IOOverloadOptions(),
@@ -191,7 +191,7 @@ func (src *storeRebalancerControl) phasePrologue(
 		ctx, src.scorerOptions(),
 		hottestRanges(
 			s, src.storeID,
-			kvserver.LBRebalancingObjective(src.settings.LBRebalancingObjective).ToDimension(),
+			kvserver.LoadBasedRebalancingObjective.Get(&src.settings.ST.SV).ToDimension(),
 		),
 		kvserverbase.LoadBasedRebalancingMode.Get(&src.settings.ST.SV),
 	)

pkg/kv/kvserver/asim/storerebalancer/store_rebalancer_test.go

@@ -10,6 +10,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/config"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/gossip"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/op"
@@ -36,6 +37,8 @@ func TestStoreRebalancer(t *testing.T) {
 	start := testSettings.StartTime
 	testSettings.ReplicaChangeBaseDelay = 5 * time.Second
 	testSettings.StateExchangeDelay = 0
+	ctx := context.Background()
+	kvserver.LoadBasedRebalancingObjective.Override(ctx, &testSettings.ST.SV, kvserver.LBRebalancingQueries)
 
 	clusterInfo := state.ClusterInfoWithStoreCount(6, 1 /* storesPerNode */)
 
@@ -195,6 +198,8 @@ func TestStoreRebalancerBalances(t *testing.T) {
 	testSettings.ReplicaChangeBaseDelay = 1 * time.Second
 	testSettings.StateExchangeInterval = 1 * time.Second
 	testSettings.StateExchangeDelay = 0
+	ctx := context.Background()
+	kvserver.LoadBasedRebalancingObjective.Override(ctx, &testSettings.ST.SV, kvserver.LBRebalancingQueries)
 
 	distributeQPS := func(s state.State, qpsCounts map[state.StoreID]float64) {
 		dist := make([]float64, len(qpsCounts))

pkg/kv/kvserver/asim/tests/datadriven_simulation_test.go

@@ -724,7 +724,6 @@ func TestDataDriven(t *testing.T) {
 					dns = scanIfExists(t, d, "rebalance_range_threshold", &settingsGen.Settings.RangeRebalanceThreshold) || dns
 					dns = scanIfExists(t, d, "gossip_delay", &settingsGen.Settings.StateExchangeDelay) || dns
 					dns = scanIfExists(t, d, "range_size_split_threshold", &settingsGen.Settings.RangeSizeSplitThreshold) || dns
-					dns = scanIfExists(t, d, "rebalance_objective", &settingsGen.Settings.LBRebalancingObjective) || dns
 					var snapshotRateMiB int
 					dns = scanIfExists(t, d, "rebalancing_snapshot_rate_mib", &snapshotRateMiB) || dns
 					if snapshotRateMiB != 0 {
@@ -746,6 +745,17 @@ func TestDataDriven(t *testing.T) {
 								Value:            rebalanceMode,
 							}})
 					}
+
+					var rebalanceObjective int64
+					if scanIfExists(t, d, "rebalance_objective", &rebalanceObjective) {
+						events = append(events, scheduled.ScheduledEvent{
+							At: settingsGen.Settings.StartTime.Add(delay),
+							TargetEvent: event.SetSimulationSettingsEvent{
+								IsClusterSetting: true,
+								Key:              "LBRebalancingObjective",
+								Value:            rebalanceObjective,
+							}})
+					}
 					return ""
 				default:
 					return fmt.Sprintf("unknown command: %s", d.Cmd)

pkg/kv/kvserver/asim/tests/testdata/generated/.gitignore

@@ -6,7 +6,7 @@
 # Except this specific directory (relative to .gitignore)
 # Note that this intentionally does not include subdirectories (such as traces)
 # which are not deterministic.
-!sma/rebalancing/**
+!sma/rebalancing_qps/**
 **/traces
 # And the readme file.
 !README.md