roachtest: add mma-and-count versions of allocbench

Fixes #154638.

Author: tbg

pkg/cmd/roachtest/tests/allocation_bench.go

@@ -44,6 +44,7 @@ const (
 type allocationBenchSpec struct {
 	nodes, cpus int
 	load        allocBenchLoad
+	lbrMode     string // see kvserverbase.LoadBasedRebalancingMode
 
 	startRecord time.Duration
 	samples     int
@@ -172,7 +173,7 @@ func (r kvAllocBenchEventRunner) run(ctx context.Context, c cluster.Cluster, t t
 	return c.RunE(ctx, option.WithNodes(c.WorkloadNode()), runCmd)
 }
 func registerAllocationBench(r registry.Registry) {
-	for _, spec := range []allocationBenchSpec{
+	specTemplates := []allocationBenchSpec{
 		// TODO(kvoli): Add a background event runner and implement events for
 		// import and index backfills.
 		{
@@ -242,8 +243,19 @@ func registerAllocationBench(r registry.Registry) {
 				},
 			},
 		},
-	} {
-		registerAllocationBenchSpec(r, spec)
+	}
+	for _, spec := range specTemplates {
+		{
+			spec := spec
+			spec.lbrMode = "leases and replicas"
+			registerAllocationBenchSpec(r, spec)
+		}
+		{
+			spec := spec
+			spec.lbrMode = "multi-metric and count"
+			spec.load.desc += "/lbr=mmc"
+			registerAllocationBenchSpec(r, spec)
+		}
 	}
 }
 
@@ -273,11 +285,18 @@ func setupAllocationBench(
 	t.Status("starting cluster")
 	for i := 1; i <= spec.nodes; i++ {
 		// Don't start a backup schedule as this test reports to roachperf.
+		settings := install.MakeClusterSettings()
+		settings.Env = append(settings.Env, "COCKROACH_ALLOW_MMA=true")
 		startOpts := option.NewStartOpts(option.NoBackupSchedule)
 		startOpts.RoachprodOpts.ExtraArgs = append(startOpts.RoachprodOpts.ExtraArgs,
 			"--vmodule=store_rebalancer=2,allocator=2,replicate_queue=2")
-		c.Start(ctx, t.L(), startOpts, install.MakeClusterSettings(), c.Node(i))
+		c.Start(ctx, t.L(), startOpts, settings, c.Node(i))
 	}
+	require.NotEmpty(t, spec.lbrMode, "lbrMode must be set")
+	_, err := c.Conn(ctx, t.L(), 1).ExecContext(ctx, fmt.Sprintf(
+		`SET CLUSTER SETTING kv.allocator.load_based_rebalancing = '%s';`,
+		spec.lbrMode))
+	require.NoError(t, err)
 
 	return setupStatCollector(ctx, t, c, spec)
 }