Merge #149892

149892: asim: port over more asim changes from prototype r=tbg a=wenyihu6


**asim: correct data driven input command for example_fullfisk**

Previously, we started tracking logical bytes from follower replicas, which
increased the used store capacity. Since gen_cluster defaults to a 256 GiB
store_byte_capacity, this was no longer sufficient and resulted in a fraction
used over 100%. This commit updates the data driven input to match the current
prototype input.

Epic: none
Release note: none

---

**kvserver: track write bytes per sec properly**

Previously, we introduced write bytes per second as part of the store
descriptor's capacity. This commit updates the corresponding capacity tracking
logic in the production code properly, along with the gossip logic that handles
capacity change events. Note that the actual capacity info is not yet used by any
production logic in a meaningful way.

Epic: none
Release note: none

---

**asim: fix inconsistent range state**

Previously, capacity change events could be published while a range was still in
an inconsistent state (such as, after the leaseholder was removed but before a
new one was added). Without this fix, future commits that populate the node
capacity provider would trigger a panic: node capacity calls RangeUsageInfo to
aggregate CPU rate across replicas, which asserts on the range having a valid
leaseholder. In this inconsistent state, the leaseholder can't be found, leading
to a panic.

Epic: none
Release note: none

---

**asim: introduce NodeCapacityProvider**

Previously, we added state.NodeCapacity but didn’t implement the logic to
retrieve actual node capacity. This commit introduces a NodeCapacityProvider
interface, implemented by the asim’s state and Stores in production code.
Node capacity provider is saved as a field in StoreGossip. It will later be used
to populate node capacity of store descriptor before gossiping. Note that node
capacity provider remains nil in production to avoid changing existing behavior.

Epic: none
Release note: none

---

**asim: move store pool to be at the node level**

Previously, asim created a separate store pool per store, unlike production
code, which shares a single store pool among stores on the same node. This
commit updates asim to align with the behaviour.

Epic: none
Release note: none

---

**asim: introduce MMAllocator at node level**

This commit introduces construction of MMAllocator at the node level. Note that
the old allocator is still kept around at the store level since only MMA store
rebalancer would use MMAAllocator.

Epic: none
Release note: none

---

**asim: remove unused storepool field from store**

Previously, we moved the store pool from the store level to the node level and
removed the usage of store pool from the store. This commit cleans up the
now unused storePool field.

Epic: none
Release note: none

---

**asim: use cluster setting for some non-simulation settings**

Previously, we began populating cluster settings in the simulator to allow
direct use of non-simulation settings. This commit updates LoadThreshold,
kvserver.LBRebalancingObjective, and allocatorimpl.LoadRebalanceRequiredMinDiff
to use values from the populated cluster settings directly.

Epic: none
Release note: none

---

**asim: panic if ticket not found in store rebalancer**

This commit adds a panic assertion to ensure pending ticket exists in store
rebalancer asim since asim should not delete a ticket from controller.tickets.

Epic: none
Relese note: none



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

Author: craig[bot]

pkg/kv/kvserver/allocator/storepool/store_pool.go

@@ -547,6 +547,7 @@ func (sp *StorePool) UpdateLocalStoreAfterRebalance(
 		detail.Desc.Capacity.RangeCount++
 		detail.Desc.Capacity.LogicalBytes += rangeUsageInfo.LogicalBytes
 		detail.Desc.Capacity.WritesPerSecond += rangeUsageInfo.WritesPerSecond
+		detail.Desc.Capacity.WriteBytesPerSecond += rangeUsageInfo.WriteBytesPerSecond
 		if detail.Desc.Capacity.CPUPerSecond >= 0 {
 			detail.Desc.Capacity.CPUPerSecond += rangeUsageInfo.RaftCPUNanosPerSecond
 		}
@@ -562,6 +563,11 @@ func (sp *StorePool) UpdateLocalStoreAfterRebalance(
 		} else {
 			detail.Desc.Capacity.WritesPerSecond -= rangeUsageInfo.WritesPerSecond
 		}
+		if detail.Desc.Capacity.WriteBytesPerSecond <= rangeUsageInfo.WriteBytesPerSecond {
+			detail.Desc.Capacity.WriteBytesPerSecond = 0
+		} else {
+			detail.Desc.Capacity.WriteBytesPerSecond -= rangeUsageInfo.WriteBytesPerSecond
+		}
 		// When CPU attribution is unsupported, the store will set the
 		// CPUPerSecond of its store capacity to be -1.
 		if detail.Desc.Capacity.CPUPerSecond >= 0 {

pkg/kv/kvserver/asim/asim.go

@@ -131,7 +131,7 @@ func (s *Simulator) StoreAddNotify(storeID state.StoreID, _ state.State) {
 }
 
 func (s *Simulator) addStore(storeID state.StoreID, tick time.Time) {
-	allocator := s.state.MakeAllocator(storeID)
+	allocator := s.state.Allocator(storeID)
 	storePool := s.state.StorePool(storeID)
 	s.rqs[storeID] = queue.NewReplicateQueue(
 		storeID,

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

@@ -30,7 +30,6 @@ const (
 	defaultLBRebalancingMode       = 2 // Leases and replicas.
 	defaultLBRebalancingInterval   = time.Minute
 	defaultLBRebalanceQPSThreshold = 0.1
-	defaultLBMinRequiredQPSDiff    = 200
 	defaultLBRebalancingObjective  = 0 // QPS
 )
 
@@ -105,13 +104,6 @@ type SimulationSettings struct {
 	// LBRebalancingInterval controls how often the store rebalancer will
 	// consider opportunities for rebalancing.
 	LBRebalancingInterval time.Duration
-	// LBRebalanceQPSThreshold is the fraction above or below the mean store QPS,
-	// that a store is considered overfull or underfull.
-	LBRebalanceQPSThreshold float64
-	// LBMinQPSDifferenceForTransfers is the minimum QPS difference that the store
-	// rebalancer would care to reconcile (via lease or replica rebalancing) between
-	// any two stores.
-	LBMinRequiredQPSDiff float64
 	// ReplicateQueueEnabled controls whether the replicate queue is enabled.
 	ReplicateQueueEnabled bool
 	// LeaseQueueEnabled controls whether the lease queue is enabled.
@@ -147,8 +139,6 @@ func DefaultSimulationSettings() *SimulationSettings {
 		LBRebalancingMode:       defaultLBRebalancingMode,
 		LBRebalancingObjective:  defaultLBRebalancingObjective,
 		LBRebalancingInterval:   defaultLBRebalancingInterval,
-		LBRebalanceQPSThreshold: defaultLBRebalanceQPSThreshold,
-		LBMinRequiredQPSDiff:    defaultLBMinRequiredQPSDiff,
 		ReplicateQueueEnabled:   true,
 		LeaseQueueEnabled:       true,
 		SplitQueueEnabled:       true,

pkg/kv/kvserver/asim/gossip/gossip.go

@@ -51,6 +51,7 @@ type storeGossiper struct {
 
 func newStoreGossiper(
 	descriptorGetter func(cached bool) roachpb.StoreDescriptor,
+	nodeCapacityProvider kvserver.NodeCapacityProvider,
 	clock timeutil.TimeSource,
 	st *cluster.Settings,
 ) *storeGossiper {
@@ -61,7 +62,7 @@ func newStoreGossiper(
 
 	desc := sg.descriptorGetter(false /* cached */)
 	knobs := kvserver.StoreGossipTestingKnobs{AsyncDisabled: true}
-	sg.local = kvserver.NewStoreGossip(sg, sg, knobs, &st.SV, clock)
+	sg.local = kvserver.NewStoreGossip(sg, sg, knobs, &st.SV, clock, nodeCapacityProvider)
 	sg.local.Ident = roachpb.StoreIdent{StoreID: desc.StoreID, NodeID: desc.Node.NodeID}
 
 	return sg
@@ -110,21 +111,35 @@ func NewGossip(s state.State, settings *config.SimulationSettings) *gossip {
 		},
 	}
 	for _, store := range s.Stores() {
-		g.addStoreToGossip(s, store.StoreID())
+		g.addStoreToGossip(s, store.StoreID(), store.NodeID())
 	}
 	s.RegisterCapacityChangeListener(g)
 	s.RegisterCapacityListener(g)
 	s.RegisterConfigChangeListener(g)
 	return g
 }
 
-func (g *gossip) addStoreToGossip(s state.State, storeID state.StoreID) {
+var _ kvserver.NodeCapacityProvider = &simNodeCapacityProvider{}
+
+type simNodeCapacityProvider struct {
+	localNodeID state.NodeID
+	state       state.State
+}
+
+func (s simNodeCapacityProvider) GetNodeCapacity(_ bool) roachpb.NodeCapacity {
+	return s.state.NodeCapacity(s.localNodeID)
+}
+
+func (g *gossip) addStoreToGossip(s state.State, storeID state.StoreID, nodeID state.NodeID) {
 	// Add the store gossip in an "adding" state initially, this is to avoid
 	// recursive calls to get the store descriptor.
 	g.storeGossip[storeID] = &storeGossiper{addingStore: true}
-	g.storeGossip[storeID] = newStoreGossiper(func(cached bool) roachpb.StoreDescriptor {
-		return s.StoreDescriptors(cached, storeID)[0]
-	}, s.Clock(), g.settings.ST)
+	g.storeGossip[storeID] = newStoreGossiper(
+		func(cached bool) roachpb.StoreDescriptor {
+			return s.StoreDescriptors(cached, storeID)[0]
+		},
+		simNodeCapacityProvider{localNodeID: nodeID, state: s},
+		s.Clock(), g.settings.ST)
 }
 
 // Tick checks for completed gossip updates and triggers new gossip
@@ -137,7 +152,7 @@ func (g *gossip) Tick(ctx context.Context, tick time.Time, s state.State) {
 		// If the store gossip for this store doesn't yet exist, create it and
 		// add it to the map of store gossips.
 		if sg, ok = g.storeGossip[store.StoreID()]; !ok {
-			g.addStoreToGossip(s, store.StoreID())
+			g.addStoreToGossip(s, store.StoreID(), store.NodeID())
 		}
 
 		// If the interval between the last time this store was gossiped for
@@ -183,6 +198,9 @@ func (g *gossip) NewCapacityNotify(capacity roachpb.StoreCapacity, storeID state
 	if sg, ok := g.storeGossip[storeID]; ok {
 		if !sg.addingStore {
 			sg.local.UpdateCachedCapacity(capacity)
+			sg.local.RecordNewPerSecondStats(
+				capacity.QueriesPerSecond, capacity.WritesPerSecond,
+				capacity.WriteBytesPerSecond, capacity.CPUPerSecond)
 		}
 	} else {
 		panic(
@@ -194,7 +212,8 @@ func (g *gossip) NewCapacityNotify(capacity roachpb.StoreCapacity, storeID state
 
 // StoreAddNotify notifies that a new store has been added with ID storeID.
 func (g *gossip) StoreAddNotify(storeID state.StoreID, s state.State) {
-	g.addStoreToGossip(s, storeID)
+	store, _ := s.Store(storeID)
+	g.addStoreToGossip(s, storeID, store.NodeID())
 }
 
 func (g *gossip) maybeUpdateState(tick time.Time, s state.State) {
@@ -210,7 +229,7 @@ func (g *gossip) maybeUpdateState(tick time.Time, s state.State) {
 	for _, update := range updates {
 		updateMap[update.Desc.StoreID] = update
 	}
-	for _, store := range s.Stores() {
-		s.UpdateStorePool(store.StoreID(), updateMap)
+	for _, node := range s.Nodes() {
+		s.UpdateStorePool(node.NodeID(), updateMap)
 	}
 }

pkg/kv/kvserver/asim/op/controller_test.go

@@ -281,7 +281,7 @@ func TestRelocateRangeOp(t *testing.T) {
 				settings,
 			)
 			changer := state.NewReplicaChanger()
-			allocator := s.MakeAllocator(state.StoreID(1))
+			allocator := s.Allocator(state.StoreID(1))
 			storePool := s.StorePool(state.StoreID(1))
 			controller := NewController(changer, allocator, storePool, settings, 1 /* storeID */)
 

pkg/kv/kvserver/asim/queue/lease_queue_test.go

@@ -142,7 +142,7 @@ func TestLeaseQueue(t *testing.T) {
 				store.StoreID(),
 				changer,
 				testSettings,
-				s.MakeAllocator(store.StoreID()),
+				s.Allocator(store.StoreID()),
 				s.StorePool(store.StoreID()),
 				start,
 			)

pkg/kv/kvserver/asim/queue/replicate_queue_test.go

@@ -291,7 +291,7 @@ func TestReplicateQueue(t *testing.T) {
 				store.StoreID(),
 				changer,
 				testSettings,
-				s.MakeAllocator(store.StoreID()),
+				s.Allocator(store.StoreID()),
 				s.StorePool(store.StoreID()),
 				start,
 			)

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

@@ -25,6 +25,7 @@ go_library(
         "//pkg/kv/kvserver",
         "//pkg/kv/kvserver/allocator",
         "//pkg/kv/kvserver/allocator/allocatorimpl",
+        "//pkg/kv/kvserver/allocator/mmaprototype",
         "//pkg/kv/kvserver/allocator/storepool",
         "//pkg/kv/kvserver/asim/config",
         "//pkg/kv/kvserver/asim/workload",

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

@@ -63,7 +63,7 @@ func NewStorePool(
 	nodeLivenessFn storepool.NodeLivenessFunc,
 	hlc *hlc.Clock,
 	st *cluster.Settings,
-) (*storepool.StorePool, *cluster.Settings) {
+) *storepool.StorePool {
 	stopper := stop.NewStopper()
 	defer stopper.Stop(context.Background())
 
@@ -80,7 +80,7 @@ func NewStorePool(
 		nodeLivenessFn,
 		/* deterministic */ true,
 	)
-	return sp, st
+	return sp
 }
 
 // OffsetTick offsets start time by adding tick number of seconds to it.