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.

Author: wenyihu6

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

@@ -124,17 +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 always require node cpu capacity and locality
-		// string as part of the input
-		s.SetNodeCPURateCapacity(node.NodeID(), config.DefaultNodeCPURateCapacityNanos)
+		// 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/state/config_loader.go

@@ -375,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/state.go

@@ -85,7 +85,7 @@ type State interface {
 	LeaseholderStore(RangeID) (Store, bool)
 	// AddNode modifies the state to include one additional node. This cannot
 	// fail. The new Node is returned.
-	AddNode() Node
+	AddNode(nodeCPUCapacity int64, locality roachpb.Locality) Node
 	// SetNodeLocality sets the locality of the node with ID NodeID to be the
 	// locality given.
 	SetNodeLocality(NodeID, roachpb.Locality)

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

@@ -25,7 +25,7 @@ import (
 
 func TestStateUpdates(t *testing.T) {
 	s := NewState(config.DefaultSimulationSettings())
-	node := s.AddNode()
+	node := s.AddNode(-1, roachpb.Locality{})
 	s.AddStore(node.NodeID())
 	require.Equal(t, 1, len(s.Nodes()))
 	require.Equal(t, 1, len(s.Stores()))
@@ -39,7 +39,7 @@ func TestRangeSplit(t *testing.T) {
 	k1 := MinKey
 	r1 := s.rangeFor(k1)
 
-	n1 := s.AddNode()
+	n1 := s.AddNode(-1, roachpb.Locality{})
 	s1, _ := s.AddStore(n1.NodeID())
 
 	repl1, _ := s.AddReplica(r1.RangeID(), s1.StoreID(), roachpb.VOTER_FULL)
@@ -126,7 +126,7 @@ func TestValidTransfer(t *testing.T) {
 
 	_, r1, _ := s.SplitRange(1)
 
-	n1 := s.AddNode()
+	n1 := s.AddNode(-1, roachpb.Locality{})
 	s1, _ := s.AddStore(n1.NodeID())
 	s2, _ := s.AddStore(n1.NodeID())
 
@@ -159,7 +159,7 @@ func TestTransferLease(t *testing.T) {
 
 	_, r1, _ := s.SplitRange(1)
 
-	n1 := s.AddNode()
+	n1 := s.AddNode(-1, roachpb.Locality{})
 	s1, _ := s.AddStore(n1.NodeID())
 	s2, _ := s.AddStore(n1.NodeID())
 
@@ -187,7 +187,7 @@ func TestValidReplicaTarget(t *testing.T) {
 
 	_, r1, _ := s.SplitRange(1)
 
-	n1 := s.AddNode()
+	n1 := s.AddNode(-1, roachpb.Locality{})
 	s1, _ := s.AddStore(n1.NodeID())
 	s2, _ := s.AddStore(n1.NodeID())
 
@@ -227,7 +227,7 @@ func TestAddReplica(t *testing.T) {
 	_, r1, _ := s.SplitRange(1)
 	_, r2, _ := s.SplitRange(2)
 
-	n1 := s.AddNode()
+	n1 := s.AddNode(-1, roachpb.Locality{})
 	s1, _ := s.AddStore(n1.NodeID())
 	s2, _ := s.AddStore(n1.NodeID())
 
@@ -249,7 +249,7 @@ func TestAddReplica(t *testing.T) {
 func TestWorkloadApply(t *testing.T) {
 	s := NewState(config.DefaultSimulationSettings())
 
-	n1 := s.AddNode()
+	n1 := s.AddNode(-1, roachpb.Locality{})
 	s1, _ := s.AddStore(n1.NodeID())
 	s2, _ := s.AddStore(n1.NodeID())
 	s3, _ := s.AddStore(n1.NodeID())
@@ -299,9 +299,9 @@ func TestReplicaLoadRangeUsageInfo(t *testing.T) {
 	s := NewState(settings)
 	start := settings.StartTime
 
-	n1 := s.AddNode()
-	n2 := s.AddNode()
-	n3 := s.AddNode()
+	n1 := s.AddNode(-1, roachpb.Locality{})
+	n2 := s.AddNode(-1, roachpb.Locality{})
+	n3 := s.AddNode(-1, roachpb.Locality{})
 	k1 := Key(100)
 	qps := 1000
 	s1, _ := s.AddStore(n1.NodeID())
@@ -532,7 +532,7 @@ func TestSetSpanConfig(t *testing.T) {
 
 	setupState := func() State {
 		s := newState(settings)
-		node := s.AddNode()
+		node := s.AddNode(-1, roachpb.Locality{})
 		_, ok := s.AddStore(node.NodeID())
 		require.True(t, ok)