Merge #149582

149582: asim: port first batch of testing framework changes r=tbg a=tbg

This ports a few improvements from #149531 to master:

- ability to combine multiple load generators
- improved range creation (placement_type option, min/max key, ...)
- ability to disable split/lease/replicate queue
- plots show initial and final values

I also cleaned up the args handling to make sure that we don't silently skip args not supported by the test harness. This was helpful to find what still needed porting.

A first datadriven test file was moved over; the CPU/write tracking is still unsupported, so I commented it out with TODOs to re-enable once possible. Unfortunately, introducing this will be a bigger yak shave, since this data is communicated through the `roachpb` protos, and we may not want to change them without also first adding the code that populates these fields correctly in production builds.

Epic: CRDB-25222

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

Author: 3 people

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

@@ -110,6 +110,12 @@ type SimulationSettings struct {
 	// rebalancer would care to reconcile (via lease or replica rebalancing) between
 	// any two stores.
 	LBMinRequiredQPSDiff float64
+	// ReplicateQueueEnabled controls whether the replicate queue is enabled.
+	ReplicateQueueEnabled bool
+	// LeaseQueueEnabled controls whether the lease queue is enabled.
+	LeaseQueueEnabled bool
+	// SplitQueueEnabled controls whether the split queue is enabled.
+	SplitQueueEnabled bool
 }
 
 // DefaultSimulationSettings returns a set of default settings for simulation.
@@ -136,6 +142,9 @@ func DefaultSimulationSettings() *SimulationSettings {
 		LBRebalancingInterval:   defaultLBRebalancingInterval,
 		LBRebalanceQPSThreshold: defaultLBRebalanceQPSThreshold,
 		LBMinRequiredQPSDiff:    defaultLBMinRequiredQPSDiff,
+		ReplicateQueueEnabled:   true,
+		LeaseQueueEnabled:       true,
+		SplitQueueEnabled:       true,
 	}
 }
 

pkg/kv/kvserver/asim/gen/generator.go

@@ -93,6 +93,28 @@ func (ss StaticSettings) Generate(seed int64) config.SimulationSettings {
 	return ret
 }
 
+type MultiLoad []BasicLoad
+
+// MultiLoad implements the LoadGen interface. It is a collection of
+// BasicLoad.
+var _ LoadGen = MultiLoad{}
+
+func (ml MultiLoad) String() string {
+	var str string
+	for _, load := range ml {
+		str += fmt.Sprintf("%s\n", load.String())
+	}
+	return str
+}
+
+func (ml MultiLoad) Generate(seed int64, settings *config.SimulationSettings) []workload.Generator {
+	var generators []workload.Generator
+	for _, load := range ml {
+		generators = append(generators, load.Generate(seed, settings)...)
+	}
+	return generators
+}
+
 // BasicLoad implements the LoadGen interface.
 type BasicLoad struct {
 	RWRatio        float64
@@ -103,6 +125,8 @@ type BasicLoad struct {
 	MinKey, MaxKey int64
 }
 
+var _ LoadGen = BasicLoad{}
+
 func (bl BasicLoad) String() string {
 	return fmt.Sprintf(
 		"basic load with rw_ratio=%0.2f, rate=%0.2f, skewed_access=%t, min_block_size=%d, max_block_size=%d, "+
@@ -212,6 +236,8 @@ const (
 	Skewed
 	Random
 	WeightedRandom
+	Weighted
+	ReplicaPlacement
 )
 
 func (p PlacementType) String() string {
@@ -224,6 +250,10 @@ func (p PlacementType) String() string {
 		return "random"
 	case WeightedRandom:
 		return "weighted_rand"
+	case Weighted:
+		return "weighted"
+	case ReplicaPlacement:
+		return "replica_placement"
 	default:
 		panic("unknown placement type")
 	}
@@ -239,6 +269,10 @@ func GetRangePlacementType(s string) PlacementType {
 		return Random
 	case "weighted_rand":
 		return WeightedRandom
+	case "weighted":
+		return Weighted
+	case "replica_placement":
+		return ReplicaPlacement
 	default:
 		panic(fmt.Sprintf("unknown placement type %s", s))
 	}
@@ -251,13 +285,14 @@ func GetRangePlacementType(s string) PlacementType {
 // WeightedRandomizedBasicRanges.
 type BaseRanges struct {
 	Ranges            int
-	KeySpace          int
+	MinKey, MaxKey    int64
 	ReplicationFactor int
 	Bytes             int64
+	ReplicaPlacement  state.ReplicaPlacement
 }
 
 func (b BaseRanges) String() string {
-	return fmt.Sprintf("ranges=%d, key_space=%d, replication_factor=%d, bytes=%d", b.Ranges, b.KeySpace, b.ReplicationFactor, b.Bytes)
+	return fmt.Sprintf("ranges=%d, min_key=%d, max_key=%d, replication_factor=%d, bytes=%d", b.Ranges, b.MinKey, b.MaxKey, b.ReplicationFactor, b.Bytes)
 }
 
 // GetRangesInfo generates and distributes ranges across stores based on
@@ -267,13 +302,25 @@ func (b BaseRanges) GetRangesInfo(
 ) state.RangesInfo {
 	switch pType {
 	case Even:
-		return state.RangesInfoEvenDistribution(numOfStores, b.Ranges, b.KeySpace, b.ReplicationFactor, b.Bytes)
+		return state.RangesInfoEvenDistribution(numOfStores, b.Ranges, b.MinKey, b.MaxKey, b.ReplicationFactor, b.Bytes)
 	case Skewed:
-		return state.RangesInfoSkewedDistribution(numOfStores, b.Ranges, b.KeySpace, b.ReplicationFactor, b.Bytes)
+		return state.RangesInfoSkewedDistribution(numOfStores, b.Ranges, b.MinKey, b.MaxKey, b.ReplicationFactor, b.Bytes)
 	case Random:
-		return state.RangesInfoRandDistribution(randSource, numOfStores, b.Ranges, b.KeySpace, b.ReplicationFactor, b.Bytes)
+		return state.RangesInfoRandDistribution(randSource, numOfStores, b.Ranges, b.MinKey, b.MaxKey, b.ReplicationFactor, b.Bytes)
 	case WeightedRandom:
-		return state.RangesInfoWeightedRandDistribution(randSource, weightedRandom, b.Ranges, b.KeySpace, b.ReplicationFactor, b.Bytes)
+		return state.RangesInfoWeightedRandDistribution(
+			randSource, weightedRandom, b.Ranges, b.MinKey, b.MaxKey, b.ReplicationFactor, b.Bytes)
+	case ReplicaPlacement:
+		// TODO(tbg): port this over from the prototype.
+		/*
+			return state.RangesInfoWithReplicaPlacement(
+				b.ReplicaPlacement,
+				b.Ranges,
+				state.DefaultSpanConfigWithRF(b.ReplicationFactor),
+				b.MinKey, b.MaxKey, b.Bytes,
+			)
+		*/
+		panic("unimplemented")
 	default:
 		panic(fmt.Sprintf("unexpected range placement type %v", pType))
 	}
@@ -292,6 +339,9 @@ func (b BaseRanges) LoadRangeInfo(s state.State, rangesInfo state.RangesInfo) {
 type BasicRanges struct {
 	BaseRanges
 	PlacementType PlacementType
+	// ReplicaWeights and LeaseWeights are only non-nil when the placement type
+	// is Weighted.
+	ReplicaWeights, LeaseWeights []float64
 }
 
 func (br BasicRanges) String() string {
@@ -311,3 +361,45 @@ func (br BasicRanges) Generate(
 	br.LoadRangeInfo(s, rangesInfo)
 	return s
 }
+
+// MultiRanges implements the RangeGen interface, supporting multiple
+// BasicRanges generation.
+type MultiRanges []BasicRanges
+
+var _ RangeGen = MultiRanges{}
+
+func (mr MultiRanges) String() string {
+	var str string
+	for _, ranges := range mr {
+		str += fmt.Sprintf("%s\n", ranges.String())
+	}
+	return str
+}
+
+func (mr MultiRanges) Generate(
+	seed int64, settings *config.SimulationSettings, s state.State,
+) state.State {
+	var rangeInfos []state.RangeInfo
+	for _, ranges := range mr {
+		var nextInfos state.RangesInfo
+		if ranges.PlacementType == Weighted {
+			// TODO(tbg): instead refactoring GetRangesInfo to be more general.
+			var storeIDs []state.StoreID
+			for _, store := range s.Stores() {
+				storeIDs = append(storeIDs, store.StoreID())
+			}
+			nextInfos = state.RangesInfoWithDistribution(storeIDs,
+				ranges.ReplicaWeights, ranges.LeaseWeights,
+				ranges.Ranges, state.DefaultSpanConfigWithRF(ranges.ReplicationFactor),
+				ranges.MinKey, ranges.MaxKey, ranges.Bytes)
+		} else {
+			if ranges.LeaseWeights != nil || ranges.ReplicaWeights != nil {
+				panic("leaseWeights and replicaWeights should be nil for non-weighted placement types")
+			}
+			nextInfos = ranges.GetRangesInfo(ranges.PlacementType, len(s.Stores()), nil, []float64{})
+		}
+		rangeInfos = append(rangeInfos, nextInfos...)
+	}
+	state.LoadRangeInfo(s, rangeInfos...)
+	return s
+}

pkg/kv/kvserver/asim/history/BUILD.bazel

@@ -8,5 +8,6 @@ go_library(
     deps = [
         "//pkg/kv/kvserver/asim/metrics",
         "//pkg/kv/kvserver/asim/state",
+        "@com_github_montanaflynn_stats//:stats",
     ],
 )

pkg/kv/kvserver/asim/history/history.go

@@ -7,15 +7,22 @@ package history
 
 import (
 	"context"
+	"fmt"
+	"strings"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/metrics"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/state"
+	"github.com/montanaflynn/stats"
 )
 
 // History contains recorded information that summarizes a simulation run.
 // Currently it only contains the store metrics of the run.
 // TODO(kvoli): Add a range log like structure to the history.
 type History struct {
+	// Recorded contains per-store metrics snapshots at each tick. The outer slice
+	// grows over time at each tick, while each inner slice has one StoreMetrics per
+	// store. E.g. Recorded[0] is the first tick, Recorded[0][0] is the first store's
+	// metrics at the first tick.
 	Recorded [][]metrics.StoreMetrics
 	S        state.State
 }
@@ -24,3 +31,33 @@ type History struct {
 func (h *History) Listen(ctx context.Context, sms []metrics.StoreMetrics) {
 	h.Recorded = append(h.Recorded, sms)
 }
+
+func (h *History) ShowRecordedValueAt(idx int, stat string) string {
+	var buf strings.Builder
+
+	storeMetricsAtTick := h.Recorded[idx]
+	values := make([]float64, 0, len(storeMetricsAtTick))
+
+	_, _ = fmt.Fprintf(&buf, "[")
+
+	// Extract values for each store. Note that h.Recorded[idx] is already sorted
+	// by store ID when appending to h.Recorded.
+	for i, sm := range storeMetricsAtTick {
+		if i > 0 {
+			_, _ = fmt.Fprintf(&buf, ", ")
+		}
+		value := sm.GetMetricValue(stat)
+		if stat == "disk_fraction_used" {
+			_, _ = fmt.Fprintf(&buf, "s%v=%.2f", sm.StoreID, value)
+		} else {
+			_, _ = fmt.Fprintf(&buf, "s%v=%.0f", sm.StoreID, value)
+		}
+		values = append(values, value)
+	}
+	_, _ = fmt.Fprintf(&buf, "]")
+	stddev, _ := stats.StandardDeviation(values)
+	mean, _ := stats.Mean(values)
+	sum, _ := stats.Sum(values)
+	_, _ = fmt.Fprintf(&buf, " (stddev=%.2f, mean=%.2f, sum=%.0f)", stddev, mean, sum)
+	return buf.String()
+}

pkg/kv/kvserver/asim/metrics/tracker.go

@@ -40,6 +40,44 @@ type StoreMetrics struct {
 	DiskFractionUsed   float64
 }
 
+// GetMetricValue extracts the requested metric value from StoreMetrics.
+func (sm *StoreMetrics) GetMetricValue(stat string) float64 {
+	switch stat {
+	case "qps":
+		return float64(sm.QPS)
+	// case "cpu":
+	// 	value = float64(sm.CPU)
+	// case "write_bytes_per_second":
+	// 	value = float64(sm.WriteBytesPerSecond)
+	case "write":
+		return float64(sm.WriteKeys)
+	case "write_b":
+		return float64(sm.WriteBytes)
+	case "read":
+		return float64(sm.ReadKeys)
+	case "read_b":
+		return float64(sm.ReadBytes)
+	case "replicas":
+		return float64(sm.Replicas)
+	case "leases":
+		return float64(sm.Leases)
+	case "lease_moves":
+		return float64(sm.LeaseTransfers)
+	case "replica_moves":
+		return float64(sm.Rebalances)
+	case "replica_b_rcvd":
+		return float64(sm.RebalanceRcvdBytes)
+	case "replica_b_sent":
+		return float64(sm.RebalanceSentBytes)
+	case "range_splits":
+		return float64(sm.RangeSplits)
+	case "disk_fraction_used":
+		return sm.DiskFractionUsed
+	default:
+		return 0
+	}
+}
+
 // the MetricsTracker to report new store metrics for a tick.
 type StoreMetricsListener interface {
 	Listen(context.Context, []StoreMetrics)

pkg/kv/kvserver/asim/queue/lease_queue.go

@@ -59,6 +59,11 @@ func NewLeaseQueue(
 // meets the criteria it is enqueued. The criteria is currently if the
 // allocator returns a lease transfer.
 func (lq *leaseQueue) MaybeAdd(ctx context.Context, replica state.Replica, s state.State) bool {
+	if !lq.settings.LeaseQueueEnabled {
+		// Nothing to do, disabled.
+		return false
+	}
+
 	repl := NewSimulatorReplica(replica, s)
 	lq.AddLogTag("r", repl.repl.Descriptor())
 	lq.AnnotateCtx(ctx)

pkg/kv/kvserver/asim/queue/replicate_queue.go

@@ -58,6 +58,11 @@ func NewReplicateQueue(
 // meets the criteria it is enqueued. The criteria is currently if the
 // allocator returns a non-noop, then the replica is added.
 func (rq *replicateQueue) MaybeAdd(ctx context.Context, replica state.Replica, s state.State) bool {
+	if !rq.settings.ReplicateQueueEnabled {
+		// Nothing to do, disabled.
+		return false
+	}
+
 	repl := NewSimulatorReplica(replica, s)
 	rq.AddLogTag("r", repl.repl.Descriptor())
 	rq.AnnotateCtx(ctx)

pkg/kv/kvserver/asim/queue/split_queue.go

@@ -42,6 +42,11 @@ func NewSplitQueue(
 // MaybeAdd proposes a range for being split. If it meets the criteria it is
 // enqueued.
 func (sq *splitQueue) MaybeAdd(ctx context.Context, replica state.Replica, state state.State) bool {
+	if !sq.settings.SplitQueueEnabled {
+		// Nothing to do, disabled.
+		return false
+	}
+
 	priority := sq.shouldSplit(sq.lastTick, replica.Range(), state)
 	if priority < 1 {
 		return false

pkg/kv/kvserver/asim/state/BUILD.bazel

@@ -10,6 +10,8 @@ go_library(
         "liveness.go",
         "load.go",
         "new_state.go",
+        "new_state_test_helper.go",
+        "parser_replica_placement.go",
         "split_decider.go",
         "state.go",
         "state_listener.go",