add hardcoded

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

@@ -190,6 +190,29 @@ func (bc BasicCluster) Regions() []state.Region {
 	return info.Regions
 }
 
+type ClusterWithLocality struct {
+	StoresPerNode  int
+	Region         []string
+	NodesPerRegion []int
+}
+
+func (bc ClusterWithLocality) String() string {
+	return fmt.Sprintf("cluster with locality with region=%v, nodes_per_region=%v",
+		bc.Region, bc.NodesPerRegion)
+}
+
+func (bc ClusterWithLocality) Generate(
+	seed int64, settings *config.SimulationSettings,
+) state.State {
+	info := state.ClusterInfoWithLocality(bc.StoresPerNode, bc.Region, bc.NodesPerRegion)
+	return state.LoadClusterInfo(info, settings)
+}
+
+func (bc ClusterWithLocality) Regions() []state.Region {
+	info := state.ClusterInfoWithLocality(bc.StoresPerNode, bc.Region, bc.NodesPerRegion)
+	return info.Regions
+}
+
 // LoadedRanges implements the RangeGen interface.
 type LoadedRanges struct {
 	Info state.RangesInfo

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

@@ -224,6 +224,23 @@ func RangesInfoWithDistribution(
 	return ret
 }
 
+func ClusterInfoWithLocality(storesPerNode int, regions []string, nodes []int) ClusterInfo {
+	if len(regions) != len(nodes) {
+		panic("len(regions) != len(nodes)")
+	}
+	ret := ClusterInfo{}
+
+	ret.Regions = make([]Region, len(regions))
+	for i, name := range regions {
+		ret.Regions[i] = Region{
+			Name:  name,
+			Zones: []Zone{NewZone(name+"1", nodes[i], storesPerNode)},
+		}
+	}
+
+	return ret
+}
+
 // ClusterInfoWithDistribution returns a ClusterInfo. The ClusterInfo regions
 // have nodes added to them according to the regionNodeWeights and numNodes
 // given. In cases where the numNodes does not divide among the regions given
@@ -244,7 +261,7 @@ func ClusterInfoWithDistribution(
 		availableNodes -= allocatedNodes
 		ret.Regions[i] = Region{
 			Name:  name,
-			Zones: []Zone{NewZone(name+"_1", allocatedNodes, storesPerNode)},
+			Zones: []Zone{NewZone(name+"1", allocatedNodes, storesPerNode)},
 		}
 	}
 

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

@@ -232,6 +232,22 @@ func TestDataDriven(t *testing.T) {
 				scanIfExists(t, d, "sample", &sample)
 				top := runs[sample-1].S.Topology()
 				return (&top).String()
+			case "gen_cluster_with_locality":
+				var storesPerNode = 1
+				var regions []string
+				var nodesPerRegion []int
+				scanIfExists(t, d, "stores_per_node", &storesPerNode)
+				scanIfExists(t, d, "regions", &regions)
+				scanIfExists(t, d, "nodes_per_region", &nodesPerRegion)
+				clusterGen = gen.ClusterWithLocality{
+					StoresPerNode:  storesPerNode,
+					NodesPerRegion: nodesPerRegion,
+					Region:         regions,
+				}
+				if len(regions) != len(nodesPerRegion) {
+					return "regions and nodes_per_region must have the same length"
+				}
+				return ""
 			case "gen_cluster":
 				var nodes = 3
 				var storesPerNode = 1

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

@@ -46,7 +46,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_with_locality regions=(a, b, c) nodes_per_region=(1, 2, 2)
 ----
 
 # Generate 5 ranges, where initially there will be two replicas in region c and
@@ -45,10 +33,13 @@ OK
 topology
 ----
 a
-  └── [1]
+  a1
+  │ └── [1]
 b
-  └── [2]
+  b1
+  │ └── [2 3]
 c
-  └── [3 4]
+  c1
+    └── [4 5]
 
 # vim:ft=sh

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

@@ -4,20 +4,7 @@
 gen_cluster nodes=5
 ----
 
-# TODO(kvoli): Update gen_cluster to take in localities for this purpose.
-set_locality node=1 locality=region=a
-----
-
-set_locality node=2 locality=region=a
-----
-
-set_locality node=3 locality=region=b
-----
-
-set_locality node=4 locality=region=b
-----
-
-set_locality node=5 locality=region=b
+gen_cluster_with_locality regions=(a, b) nodes_per_region=(2,3)
 ----
 
 # Generate 10 ranges, one replica will be placed on each node initially.
@@ -49,8 +36,10 @@ OK
 topology
 ----
 a
-  └── [1 2]
+  a1
+  │ └── [1 2]
 b
-  └── [3 4 5]
+  b1
+    └── [3 4 5]
 
 # vim:ft=sh

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

@@ -3,19 +3,7 @@
 gen_cluster nodes=5
 ----
 
-set_locality node=1 locality=region=a
-----
-
-set_locality node=2 locality=region=a
-----
-
-set_locality node=3 locality=region=b
-----
-
-set_locality node=4 locality=region=b
-----
-
-set_locality node=5 locality=region=b
+gen_cluster_with_locality regions=(a, b) nodes_per_region=(2,3)
 ----
 
 # Setup three span configs, the first span config should produce no violations
@@ -45,8 +33,10 @@ OK
 topology
 ----
 a
-  └── [1 2]
+  a1
+  │ └── [1 2]
 b
-  └── [3 4 5]
+  b1
+    └── [3 4 5]
 
 # vim:ft=sh