Merge #153757

153757: asim: record cpu utilization in StoreMetrics r=tbg a=wenyihu6

Resolves: #153742
Release note: none

---
**asim: plumb NodeCPURateCapacity everywhere**

Previously, node CPU rate capacity wasn’t consistently plumbed, and some tests
treated it as optional. This commit ensures it is passed through everywhere,
preventing misconfigurations from being silently missed.

---
**asim: record cpu utilization in StoreMetrics**

Previously, StoreMetrics didn't include CPU utilization, making it hard to
evaluate the simulation setup. This commit adds CPU utilization to the store
metrics. Note that due to the lack of node-level metrics, the same CPU capacity
will be repeated for stores on the same node.

---
**asim: move SetNodeLocality & s.SetNodeCPURateCapacity to AddNode**

Previously, s.SetNodeLocality and s.SetNodeCPURateCapacity were called after
s.AddNode, making it easy to miss populating these fields. This commit updates
s.AddNode to take CPU rate capacity and node capacity explicitly as arguments.



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

Author: craig[bot]

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

@@ -32,6 +32,11 @@ const (
 	defaultLBRebalancingObjective  = 0 // QPS
 )
 
+const DefaultNodeCPURateCapacityNanos = 8 * 1e9        // 8 vcpus
+const DefaultStoreDiskCapacityBytes = 1024 << 30       // 1024 GiB
+const DoubleDefaultNodeCPURateCapacityNanos = 16 * 1e9 // 16 vcpus
+const DoubleDefaultStoreDiskCapacityBytes = 2048 << 30 // 2048 GiB
+
 var (
 	// DefaultStartTime is used as the default beginning time for simulation
 	// runs. It isn't necessarily meaningful other than for logging and having

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

@@ -11,6 +11,7 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/kv/kvserver/asim/assertion",
+        "//pkg/kv/kvserver/asim/config",
         "//pkg/kv/kvserver/asim/history",
         "//pkg/kv/kvserver/asim/state",
         "//pkg/kv/kvserver/liveness/livenesspb",

pkg/kv/kvserver/asim/event/mutation_event.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/config"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/state"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness/livenesspb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -123,14 +124,14 @@ func (se SetSpanConfigEvent) String() string {
 
 func (ae AddNodeEvent) Func() EventFunc {
 	return MutationFunc(func(ctx context.Context, s state.State) {
-		node := s.AddNode()
+		// TDOO(wenyihu6): should we change AddNode to take in
+		var locality roachpb.Locality
 		if ae.LocalityString != "" {
-			var locality roachpb.Locality
 			if err := locality.Set(ae.LocalityString); err != nil {
 				panic(fmt.Sprintf("unable to set node locality %s", err.Error()))
 			}
-			s.SetNodeLocality(node.NodeID(), locality)
 		}
+		node := s.AddNode(config.DefaultNodeCPURateCapacityNanos, locality)
 		for i := 0; i < ae.NumStores; i++ {
 			if _, ok := s.AddStore(node.NodeID()); !ok {
 				panic(fmt.Sprintf("adding store to node=%d failed", node))

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

@@ -65,7 +65,7 @@ func (h *History) ShowRecordedValueAt(idx int, stat string) (string, bool) {
 		}
 		storeID := h.Recorded[idx][i].StoreID
 
-		if stat == "disk_fraction_used" {
+		if stat == "disk_fraction_used" || stat == "cpu_util" {
 			_, _ = fmt.Fprintf(&buf, "s%v=%.2f", storeID, v)
 		} else {
 			_, _ = fmt.Fprintf(&buf, "s%v=%.0f", storeID, v)

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

@@ -22,6 +22,7 @@ func MakeTS(metrics [][]StoreMetrics) map[string][][]float64 {
 	// are partially duplicated with the cluster tracker.
 	ret["qps"] = make([][]float64, stores)
 	ret["cpu"] = make([][]float64, stores)
+	ret["cpu_util"] = make([][]float64, stores)
 	ret["write"] = make([][]float64, stores)
 	ret["write_b"] = make([][]float64, stores)
 	ret["write_bytes_per_second"] = make([][]float64, stores)
@@ -40,6 +41,7 @@ func MakeTS(metrics [][]StoreMetrics) map[string][][]float64 {
 		for i, sm := range sms {
 			ret["qps"][i] = append(ret["qps"][i], float64(sm.QPS))
 			ret["cpu"][i] = append(ret["cpu"][i], float64(sm.CPU))
+			ret["cpu_util"][i] = append(ret["cpu_util"][i], sm.NodeCPUUtilization)
 			ret["write"][i] = append(ret["write"][i], float64(sm.WriteKeys))
 			ret["write_b"][i] = append(ret["write_b"][i], float64(sm.WriteBytes))
 			ret["write_bytes_per_second"][i] = append(ret["write_bytes_per_second"][i], float64(sm.WriteBytesPerSecond))

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

@@ -22,6 +22,7 @@ type StoreMetrics struct {
 	StoreID             int64
 	QPS                 int64
 	CPU                 int64
+	NodeCPUUtilization  float64
 	WriteKeys           int64
 	WriteBytes          int64
 	WriteBytesPerSecond int64
@@ -75,6 +76,8 @@ func (sm *StoreMetrics) GetMetricValue(stat string) float64 {
 		return float64(sm.RangeSplits)
 	case "disk_fraction_used":
 		return sm.DiskFractionUsed
+	case "cpu_util":
+		return sm.NodeCPUUtilization
 	default:
 		return 0
 	}
@@ -139,12 +142,21 @@ func (mt *Tracker) Tick(ctx context.Context, tick time.Time, s state.State) {
 		}
 
 		desc := store.Descriptor()
+		nodeCapacity := s.NodeCapacity(store.NodeID())
+
+		// NodeCPURateUsage is the same as StoresCPURate in asim.
+		if nodeCapacity.NodeCPURateCapacity == 0 {
+			panic(fmt.Sprintf("unexpected: node cpu rate capacity is 0 (node cpu rate usage = %d)",
+				nodeCapacity.NodeCPURateUsage))
+		}
+		cpuUtil := float64(nodeCapacity.NodeCPURateUsage) / float64(nodeCapacity.NodeCPURateCapacity)
 
 		sm := StoreMetrics{
 			Tick:                tick,
 			StoreID:             int64(storeID),
 			QPS:                 int64(desc.Capacity.QueriesPerSecond),
 			CPU:                 int64(desc.Capacity.CPUPerSecond),
+			NodeCPUUtilization:  cpuUtil,
 			WriteKeys:           u.WriteKeys,
 			WriteBytes:          u.WriteBytes,
 			WriteBytesPerSecond: int64(desc.Capacity.WriteBytesPerSecond),

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

@@ -26,7 +26,8 @@ var AllClusterOptions = [...]string{"single_region", "single_region_multi_store"
 // SingleRegionConfig is a simple cluster config with a single region and 3
 // zones, all have the same number of nodes.
 var SingleRegionConfig = ClusterInfo{
-	StoreDiskCapacityBytes: 1024 << 30, // 1024 GiB
+	NodeCPURateCapacityNanos: config.DefaultNodeCPURateCapacityNanos, // 8vpucs
+	StoreDiskCapacityBytes:   config.DefaultStoreDiskCapacityBytes,   // 1024 GiB
 	Regions: []Region{
 		{
 			Name: "US",
@@ -42,7 +43,8 @@ var SingleRegionConfig = ClusterInfo{
 // SingleRegionMultiStoreConfig is a simple cluster config with a single region
 // and 3 zones, all zones have 1 node and 5 stores per node.
 var SingleRegionMultiStoreConfig = ClusterInfo{
-	StoreDiskCapacityBytes: 1024 << 30, // 1024 GiB
+	NodeCPURateCapacityNanos: config.DefaultNodeCPURateCapacityNanos, // 8 vcpus
+	StoreDiskCapacityBytes:   config.DefaultStoreDiskCapacityBytes,   // 1024 GiB
 	Regions: []Region{
 		{
 			Name: "US",
@@ -57,7 +59,8 @@ var SingleRegionMultiStoreConfig = ClusterInfo{
 
 // MultiRegionConfig is a perfectly balanced cluster config with 3 regions.
 var MultiRegionConfig = ClusterInfo{
-	StoreDiskCapacityBytes: 2048 << 30, // 2048 GiB
+	NodeCPURateCapacityNanos: config.DoubleDefaultNodeCPURateCapacityNanos, // 16 vcpus
+	StoreDiskCapacityBytes:   config.DoubleDefaultStoreDiskCapacityBytes,   // 2048 GiB
 	Regions: []Region{
 		{
 			Name: "US_East",
@@ -88,7 +91,8 @@ var MultiRegionConfig = ClusterInfo{
 
 // ComplexConfig is an imbalanced multi-region cluster config.
 var ComplexConfig = ClusterInfo{
-	StoreDiskCapacityBytes: 2048 << 30, // 2048 GiB
+	NodeCPURateCapacityNanos: config.DoubleDefaultNodeCPURateCapacityNanos, // 16 vcpus
+	StoreDiskCapacityBytes:   config.DoubleDefaultStoreDiskCapacityBytes,   // 2048 GiB
 	Regions: []Region{
 		{
 			Name: "US_East",
@@ -371,9 +375,7 @@ func LoadClusterInfo(c ClusterInfo, settings *config.SimulationSettings) State {
 				Tiers: []roachpb.Tier{regionTier, zoneTier},
 			}
 			for i := 0; i < z.NodeCount; i++ {
-				node := s.AddNode()
-				s.SetNodeLocality(node.NodeID(), locality)
-				s.SetNodeCPURateCapacity(node.NodeID(), c.NodeCPURateCapacityNanos)
+				node := s.AddNode(c.NodeCPURateCapacityNanos, locality)
 				storesRequired := z.StoresPerNode
 				if storesRequired < 1 {
 					panic(fmt.Sprintf("storesPerNode cannot be less than one but found %v", storesRequired))

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

@@ -158,11 +158,11 @@ func TestLoadRangesInfo(t *testing.T) {
 		t.Run(tc.desc, func(t *testing.T) {
 			settings := config.DefaultSimulationSettings()
 			state := NewState(settings)
-			_, ok := state.AddStore(state.AddNode().NodeID())
+			_, ok := state.AddStore(state.AddNode(-1, roachpb.Locality{}).NodeID())
 			require.True(t, ok)
-			_, ok = state.AddStore(state.AddNode().NodeID())
+			_, ok = state.AddStore(state.AddNode(-1, roachpb.Locality{}).NodeID())
 			require.True(t, ok)
-			_, ok = state.AddStore(state.AddNode().NodeID())
+			_, ok = state.AddStore(state.AddNode(-1, roachpb.Locality{}).NodeID())
 			require.True(t, ok)
 
 			if tc.expectPanic {

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

@@ -419,9 +419,9 @@ func (s *state) Replicas(storeID StoreID) []Replica {
 	return repls
 }
 
-// AddNode modifies the state to include one additional node. This cannot
-// fail. The new Node is returned.
-func (s *state) AddNode() Node {
+// AddNode modifies the state to include one additional node. This cannot fail.
+// The new Node is returned.
+func (s *state) AddNode(nodeCPUCapacity int64, locality roachpb.Locality) Node {
 	s.nodeSeqGen++
 	nodeID := s.nodeSeqGen
 	mmAllocator := mmaprototype.NewAllocatorState(s.clock, rand.New(rand.NewSource(s.settings.Seed)))
@@ -437,6 +437,8 @@ func (s *state) AddNode() Node {
 	}
 	s.nodes[nodeID] = node
 	s.SetNodeLiveness(nodeID, livenesspb.NodeLivenessStatus_LIVE)
+	s.SetNodeLocality(nodeID, locality)
+	s.SetNodeCPURateCapacity(nodeID, nodeCPUCapacity)
 	return node
 }
 func (s *state) SetNodeLocality(nodeID NodeID, locality roachpb.Locality) {

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

@@ -27,6 +27,7 @@ func NewStateWithDistribution(
 	// Currently multi-store is not tested for correctness. Default to a single
 	// store per node.
 	clusterInfo := ClusterInfoWithStoreCount(numNodes, 1 /* storesPerNode */)
+	clusterInfo.NodeCPURateCapacityNanos = config.DefaultNodeCPURateCapacityNanos
 	s := LoadClusterInfo(clusterInfo, settings)
 
 	stores := make([]StoreID, numNodes)