asim: support region and nodes_per_region with gen_cluster

Previously, gen_cluster did not support specifying node localities, making
constraint testing difficult; users had to add nodes manually with add_node.
This commit adds region and nodes_per_region options to gen_cluster. Note taht
len(region) == len(nodes_per_region). For example, region=(a,b,c) and
nodes_per_region=(1,1,2) assigns 1 node to a, 1 to b, and 2 to c. No new tests
are added yet, but future commits will utilize this feature.

Epic: none
Release note: none

Author: wenyihu6

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

@@ -8,6 +8,7 @@ package gen
 import (
 	"fmt"
 	"math/rand"
+	"strings"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim"
@@ -190,26 +191,49 @@ type BasicCluster struct {
 	Nodes             int
 	StoresPerNode     int
 	StoreByteCapacity int64
+	Region            []string
+	NodesPerRegion    []int
 }
 
 func (bc BasicCluster) String() string {
-	return fmt.Sprintf("basic cluster with nodes=%d, stores_per_node=%d, store_byte_capacity=%d",
+	var b strings.Builder
+	_, _ = fmt.Fprintf(&b, "basic cluster with nodes=%d, stores_per_node=%d, store_byte_capacity=%d",
 		bc.Nodes, bc.StoresPerNode, bc.StoreByteCapacity)
+	if len(bc.Region) != 0 {
+		_, _ = fmt.Fprintf(&b, ", region=%v, nodes_per_region=%v", bc.Region, bc.NodesPerRegion)
+	}
+	return b.String()
 }
 
 // Generate returns a new simulator state, where the cluster is created with all
 // nodes having the same locality and with the specified number of stores/nodes
 // created. The cluster is created based on the stores and stores-per-node
 // values the basic cluster generator is created with.
 func (bc BasicCluster) Generate(seed int64, settings *config.SimulationSettings) state.State {
-	info := state.ClusterInfoWithStoreCount(bc.Nodes, bc.StoresPerNode)
+	info := bc.info()
 	info.StoreDiskCapacityBytes = bc.StoreByteCapacity
 	return state.LoadClusterInfo(info, settings)
 }
 
 func (bc BasicCluster) Regions() []state.Region {
-	info := state.ClusterInfoWithStoreCount(bc.Nodes, bc.StoresPerNode)
-	return info.Regions
+	return bc.info().Regions
+}
+
+func (bc BasicCluster) info() state.ClusterInfo {
+	if len(bc.Region) == 0 {
+		return state.ClusterInfoWithStoreCount(bc.Nodes, bc.StoresPerNode)
+	}
+
+	regionNodeWeights := make([]float64, len(bc.NodesPerRegion))
+	totalNodes := 0
+	for i, nodes := range bc.NodesPerRegion {
+		regionNodeWeights[i] = float64(nodes) / float64(bc.Nodes)
+		totalNodes += nodes
+	}
+	if totalNodes != bc.Nodes {
+		panic(fmt.Sprintf("total nodes %d does not match expected nodes %d", totalNodes, bc.Nodes))
+	}
+	return state.ClusterInfoWithDistribution(bc.Nodes, bc.StoresPerNode, bc.Region, regionNodeWeights)
 }
 
 // LoadedRanges implements the RangeGen interface.

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

@@ -258,13 +258,19 @@ func TestDataDriven(t *testing.T) {
 				var nodes = 3
 				var storesPerNode = 1
 				var storeByteCapacity int64 = 256 << 30 /* 256 GiB  */
+				var region []string
+				var nodesPerRegion []int
 				scanIfExists(t, d, "nodes", &nodes)
 				scanIfExists(t, d, "stores_per_node", &storesPerNode)
 				scanIfExists(t, d, "store_byte_capacity", &storeByteCapacity)
+				scanIfExists(t, d, "region", &region)
+				scanIfExists(t, d, "nodes_per_region", &nodesPerRegion)
 				clusterGen = gen.BasicCluster{
 					Nodes:             nodes,
 					StoresPerNode:     storesPerNode,
 					StoreByteCapacity: storeByteCapacity,
+					Region:            region,
+					NodesPerRegion:    nodesPerRegion,
 				}
 				return ""
 			case "load_cluster":

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

@@ -41,7 +41,7 @@ func getNodeLivenessStatus(s string) livenesspb.NodeLivenessStatus {
 func scanArg(t *testing.T, d *datadriven.TestData, key string, dest interface{}) {
 	var tmp string
 	switch dest := dest.(type) {
-	case *string, *int, *int64, *uint64, *bool, *time.Duration, *float64, *[]int, *[]float64:
+	case *string, *int, *int64, *uint64, *bool, *time.Duration, *float64, *[]int, *[]float64, *[]string:
 		d.ScanArgs(t, key, dest)
 	case *OutputFlags:
 		var flagsTmp []string

pkg/kv/kvserver/asim/tests/testdata/non_rand/decommission_conformance

@@ -4,19 +4,7 @@
 # constraint conjunction currently satisfied by the decommissioning replica
 # being replaced, it was possible that a valid replacement target would never
 # be selected.
-gen_cluster nodes=4
-----
-
-set_locality node=1 locality=region=a
-----
-
-set_locality node=2 locality=region=b
-----
-
-set_locality node=3 locality=region=c
-----
-
-set_locality node=4 locality=region=c
+gen_cluster nodes=4 region=(a,b,c) nodes_per_region=(1,1,2)
 ----
 
 # Generate 5 ranges, where initially there will be two replicas in region c and
@@ -45,10 +33,13 @@ OK
 topology
 ----
 a
-  └── [1]
+  a_1
+  │ └── [1]
 b
-  └── [2]
+  b_1
+  │ └── [2]
 c
-  └── [3 4]
+  c_1
+    └── [3 4]
 
 # vim:ft=sh