asim: support node_cpu_rate_capacity with gen_cluster

Previously, we added roachpb.NodeCapacity for MMA to compute resource
utilization across stores. This commit integrates it into the asim setup,
enabling gen_cluster to use the node_cpu_rate_capacity option. Note that no
functions currently access node capacity; future MMA integration commits will
utilize it.

Epic: none
Release note: none

Author: wenyihu6

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

@@ -193,17 +193,19 @@ func (lc LoadedCluster) Regions() []state.Region {
 
 // BasicCluster implements the ClusterGen interace.
 type BasicCluster struct {
-	Nodes             int
-	StoresPerNode     int
-	StoreByteCapacity int64
-	Region            []string
-	NodesPerRegion    []int
+	Nodes               int
+	StoresPerNode       int
+	StoreByteCapacity   int64
+	Region              []string
+	NodesPerRegion      []int
+	NodeCPURateCapacity int64
 }
 
 func (bc BasicCluster) String() string {
 	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)
+	_, _ = fmt.Fprintf(&b,
+		"basic cluster with nodes=%d, stores_per_node=%d, store_byte_capacity=%d, node_cpu_rate_capacity=%d",
+		bc.Nodes, bc.StoresPerNode, bc.StoreByteCapacity, bc.NodeCPURateCapacity)
 	if len(bc.Region) != 0 {
 		_, _ = fmt.Fprintf(&b, ", region=%v, nodes_per_region=%v", bc.Region, bc.NodesPerRegion)
 	}
@@ -217,6 +219,7 @@ func (bc BasicCluster) String() string {
 func (bc BasicCluster) Generate(seed int64, settings *config.SimulationSettings) state.State {
 	info := bc.info()
 	info.StoreDiskCapacityBytes = bc.StoreByteCapacity
+	info.NodeCPURateCapacityNanos = bc.NodeCPURateCapacity
 	return state.LoadClusterInfo(info, settings)
 }
 

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

@@ -276,8 +276,9 @@ type Region struct {
 // ClusterInfo contains cluster information needed for allocation decisions.
 // TODO(lidor): add cross region network latencies.
 type ClusterInfo struct {
-	Regions                []Region
-	StoreDiskCapacityBytes int64
+	Regions                  []Region
+	StoreDiskCapacityBytes   int64
+	NodeCPURateCapacityNanos int64
 }
 
 func (c ClusterInfo) String() (s string) {
@@ -373,6 +374,7 @@ func LoadClusterInfo(c ClusterInfo, settings *config.SimulationSettings) State {
 			for i := 0; i < z.NodeCount; i++ {
 				node := s.AddNode()
 				s.SetNodeLocality(node.NodeID(), locality)
+				s.SetNodeCPURateCapacity(node.NodeID(), c.NodeCPURateCapacityNanos)
 				storesRequired := z.StoresPerNode
 				if storesRequired < 1 {
 					panic(fmt.Sprintf("storesPerNode cannot be less than one but found %v", storesRequired))

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

@@ -437,6 +437,34 @@ func (s *state) SetNodeLocality(nodeID NodeID, locality roachpb.Locality) {
 	}
 }
 
+func (s *state) SetNodeCPURateCapacity(nodeID NodeID, cpuRateCapacity int64) {
+	node, ok := s.nodes[nodeID]
+	if !ok {
+		panic(fmt.Sprintf("programming error: node with ID %d doesn't exist", nodeID))
+	}
+	node.cpuRateCapacity = cpuRateCapacity
+}
+
+// NodeCapacity returns the capacity of the Node with ID NodeID. Note that it is
+// currently unused.
+// TODO(wenyihu6): MMA integration should later use it.
+func (s *state) NodeCapacity(nodeID NodeID) roachpb.NodeCapacity {
+	node := s.nodes[nodeID]
+	stores := node.Stores()
+	cpuRate := 0
+	for _, storeID := range stores {
+		capacity := s.capacity(storeID)
+		cpuRate += int(capacity.CPUPerSecond)
+	}
+
+	return roachpb.NodeCapacity{
+		StoresCPURate:       int64(cpuRate),
+		NumStores:           int32(len(stores)),
+		NodeCPURateUsage:    int64(cpuRate),
+		NodeCPURateCapacity: node.cpuRateCapacity,
+	}
+}
+
 // Topology represents the locality hierarchy information for a cluster.
 type Topology struct {
 	children map[string]*Topology
@@ -1351,8 +1379,9 @@ func (s *state) SetSimulationSettings(Key string, Value interface{}) {
 
 // node is an implementation of the Node interface.
 type node struct {
-	nodeID NodeID
-	desc   roachpb.NodeDescriptor
+	nodeID          NodeID
+	desc            roachpb.NodeDescriptor
+	cpuRateCapacity int64
 
 	stores []StoreID
 }

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

@@ -199,6 +199,11 @@ type State interface {
 	RegisterConfigChangeListener(ConfigChangeListener)
 	// SetSimulationSettings sets the simulation settings for the state.
 	SetSimulationSettings(Key string, Value interface{})
+	// NodeCapacity returns the capacity of the node with ID NodeID.
+	NodeCapacity(NodeID) roachpb.NodeCapacity
+	// SetNodeCPURateCapacity sets the CPU rate capacity for the node with ID
+	// NodeID to be equal to the value given.
+	SetNodeCPURateCapacity(NodeID, int64)
 }
 
 // Node is a container for stores and is part of a cluster.

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

@@ -263,19 +263,22 @@ func TestDataDriven(t *testing.T) {
 				var nodes = 3
 				var storesPerNode = 1
 				var storeByteCapacity int64 = 256 << 30 /* 256 GiB  */
+				var nodeCPURateCapacity int64
 				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)
+				scanIfExists(t, d, "node_cpu_rate_capacity", &nodeCPURateCapacity)
 				clusterGen = gen.BasicCluster{
-					Nodes:             nodes,
-					StoresPerNode:     storesPerNode,
-					StoreByteCapacity: storeByteCapacity,
-					Region:            region,
-					NodesPerRegion:    nodesPerRegion,
+					Nodes:               nodes,
+					StoresPerNode:       storesPerNode,
+					StoreByteCapacity:   storeByteCapacity,
+					Region:              region,
+					NodesPerRegion:      nodesPerRegion,
+					NodeCPURateCapacity: nodeCPURateCapacity,
 				}
 				return ""
 			case "load_cluster":

pkg/kv/kvserver/asim/tests/testdata/rand/default_settings

@@ -73,23 +73,23 @@ eval verbose=(config_gen)
 ----------------------------------
 sample1: start running
 configurations generated using seed 3440579354231278675
-	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944
+	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944, node_cpu_rate_capacity=0
 	basic ranges with placement_type=even, ranges=10, min_key=0, max_key=200000, replication_factor=3, bytes=0
 	basic load with rw_ratio=0.00, rate=0.00, skewed_access=false, min_block_size=1, max_block_size=1, min_key=0, max_key=200000, request_cpu_per_access=0, raft_cpu_per_write=0
 	number of mutation events=0, number of assertion events=0
 sample1: pass
 ----------------------------------
 sample2: start running
 configurations generated using seed 608747136543856411
-	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944
+	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944, node_cpu_rate_capacity=0
 	basic ranges with placement_type=even, ranges=10, min_key=0, max_key=200000, replication_factor=3, bytes=0
 	basic load with rw_ratio=0.00, rate=0.00, skewed_access=false, min_block_size=1, max_block_size=1, min_key=0, max_key=200000, request_cpu_per_access=0, raft_cpu_per_write=0
 	number of mutation events=0, number of assertion events=0
 sample2: pass
 ----------------------------------
 sample3: start running
 configurations generated using seed 5571782338101878760
-	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944
+	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944, node_cpu_rate_capacity=0
 	basic ranges with placement_type=even, ranges=10, min_key=0, max_key=200000, replication_factor=3, bytes=0
 	basic load with rw_ratio=0.00, rate=0.00, skewed_access=false, min_block_size=1, max_block_size=1, min_key=0, max_key=200000, request_cpu_per_access=0, raft_cpu_per_write=0
 	number of mutation events=0, number of assertion events=0
@@ -171,7 +171,7 @@ generating settings configurations using static option
 ----------------------------------
 sample1: start running
 configurations generated using seed 3440579354231278675
-	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944
+	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944, node_cpu_rate_capacity=0
 	basic ranges with placement_type=even, ranges=10, min_key=0, max_key=200000, replication_factor=3, bytes=0
 	basic load with rw_ratio=0.00, rate=0.00, skewed_access=false, min_block_size=1, max_block_size=1, min_key=0, max_key=200000, request_cpu_per_access=0, raft_cpu_per_write=0
 	number of mutation events=0, number of assertion events=0
@@ -181,7 +181,7 @@ sample1: pass
 ----------------------------------
 sample2: start running
 configurations generated using seed 608747136543856411
-	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944
+	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944, node_cpu_rate_capacity=0
 	basic ranges with placement_type=even, ranges=10, min_key=0, max_key=200000, replication_factor=3, bytes=0
 	basic load with rw_ratio=0.00, rate=0.00, skewed_access=false, min_block_size=1, max_block_size=1, min_key=0, max_key=200000, request_cpu_per_access=0, raft_cpu_per_write=0
 	number of mutation events=0, number of assertion events=0
@@ -191,7 +191,7 @@ sample2: pass
 ----------------------------------
 sample3: start running
 configurations generated using seed 5571782338101878760
-	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944
+	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944, node_cpu_rate_capacity=0
 	basic ranges with placement_type=even, ranges=10, min_key=0, max_key=200000, replication_factor=3, bytes=0
 	basic load with rw_ratio=0.00, rate=0.00, skewed_access=false, min_block_size=1, max_block_size=1, min_key=0, max_key=200000, request_cpu_per_access=0, raft_cpu_per_write=0
 	number of mutation events=0, number of assertion events=0
@@ -221,7 +221,7 @@ generating settings configurations using static option
 ----------------------------------
 sample1: start running
 configurations generated using seed 3440579354231278675
-	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944
+	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944, node_cpu_rate_capacity=0
 	basic ranges with placement_type=even, ranges=10, min_key=0, max_key=200000, replication_factor=3, bytes=0
 	basic load with rw_ratio=0.00, rate=0.00, skewed_access=false, min_block_size=1, max_block_size=1, min_key=0, max_key=200000, request_cpu_per_access=0, raft_cpu_per_write=0
 	number of mutation events=0, number of assertion events=0
@@ -236,7 +236,7 @@ sample1: pass
 ----------------------------------
 sample2: start running
 configurations generated using seed 608747136543856411
-	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944
+	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944, node_cpu_rate_capacity=0
 	basic ranges with placement_type=even, ranges=10, min_key=0, max_key=200000, replication_factor=3, bytes=0
 	basic load with rw_ratio=0.00, rate=0.00, skewed_access=false, min_block_size=1, max_block_size=1, min_key=0, max_key=200000, request_cpu_per_access=0, raft_cpu_per_write=0
 	number of mutation events=0, number of assertion events=0
@@ -251,7 +251,7 @@ sample2: pass
 ----------------------------------
 sample3: start running
 configurations generated using seed 5571782338101878760
-	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944
+	basic cluster with nodes=3, stores_per_node=1, store_byte_capacity=274877906944, node_cpu_rate_capacity=0
 	basic ranges with placement_type=even, ranges=10, min_key=0, max_key=200000, replication_factor=3, bytes=0
 	basic load with rw_ratio=0.00, rate=0.00, skewed_access=false, min_block_size=1, max_block_size=1, min_key=0, max_key=200000, request_cpu_per_access=0, raft_cpu_per_write=0
 	number of mutation events=0, number of assertion events=0

pkg/kv/kvserver/asim/tests/testdata/rand/weighted_rand

@@ -68,7 +68,7 @@ generating settings configurations using static option
 ----------------------------------
 sample1: start running
 configurations generated using seed 5571782338101878760
-	basic cluster with nodes=3, stores_per_node=2, store_byte_capacity=274877906944
+	basic cluster with nodes=3, stores_per_node=2, store_byte_capacity=274877906944, node_cpu_rate_capacity=0
 	weighted randomized ranges with placement_type=weighted_rand, weighted_rand=[0 0 0 0.3 0.3 0.4], ranges=563, min_key=0, max_key=160411, replication_factor=3, bytes=0
 	basic load with rw_ratio=0.00, rate=0.00, skewed_access=false, min_block_size=1, max_block_size=1, min_key=0, max_key=200000, request_cpu_per_access=0, raft_cpu_per_write=0
 	number of mutation events=0, number of assertion events=0
@@ -78,7 +78,7 @@ sample1: pass
 ----------------------------------
 sample2: start running
 configurations generated using seed 4299969443970870044
-	basic cluster with nodes=3, stores_per_node=2, store_byte_capacity=274877906944
+	basic cluster with nodes=3, stores_per_node=2, store_byte_capacity=274877906944, node_cpu_rate_capacity=0
 	weighted randomized ranges with placement_type=weighted_rand, weighted_rand=[0 0 0 0.3 0.3 0.4], ranges=299, min_key=0, max_key=9542, replication_factor=3, bytes=0
 	basic load with rw_ratio=0.00, rate=0.00, skewed_access=false, min_block_size=1, max_block_size=1, min_key=0, max_key=200000, request_cpu_per_access=0, raft_cpu_per_write=0
 	number of mutation events=0, number of assertion events=0
@@ -88,7 +88,7 @@ sample2: pass
 ----------------------------------
 sample3: start running
 configurations generated using seed 4157513341729910236
-	basic cluster with nodes=3, stores_per_node=2, store_byte_capacity=274877906944
+	basic cluster with nodes=3, stores_per_node=2, store_byte_capacity=274877906944, node_cpu_rate_capacity=0
 	weighted randomized ranges with placement_type=weighted_rand, weighted_rand=[0 0 0 0.3 0.3 0.4], ranges=521, min_key=0, max_key=82660, replication_factor=3, bytes=0
 	basic load with rw_ratio=0.00, rate=0.00, skewed_access=false, min_block_size=1, max_block_size=1, min_key=0, max_key=200000, request_cpu_per_access=0, raft_cpu_per_write=0
 	number of mutation events=0, number of assertion events=0