Merge #149948

149948: asim: port over mma integration and asim store rebalancer  r=tbg a=wenyihu6

**Stacked on top of #149930

---

**asim: pass node id to queues for logging**

This commit passes in the node id to different queues for logging purposes.

Epic: none
Release note: none

---

**asim: add AllocatorSync to asim node**

This commit adds the AllocatorSync struct to the asim node. Each
MMAllocator will have a corresponding AllocatorSync responsible for coordinating
changes between the various queues and the MMA logic. This struct is currently
unused and will be integrated in future commits.

Epic: none
Release note: none

---

**asim: integrate mma store rebalancer**

This commit integrates mma store rebalancers to asim.

---

**asim: add AllocatorSync to asim queues**

This commit populates the AllocatorSync field in asim queues. It will
be used in future commits to coordinate changes with mma.

Epic: none
Release note: none


Co-authored-by: wenyihu6 <wenyi@cockroachlabs.com>

Author: craig[bot]

pkg/kv/kvserver/asim/BUILD.bazel

@@ -10,6 +10,7 @@ go_library(
         "//pkg/kv/kvserver/asim/gossip",
         "//pkg/kv/kvserver/asim/history",
         "//pkg/kv/kvserver/asim/metrics",
+        "//pkg/kv/kvserver/asim/mmaintegration",
         "//pkg/kv/kvserver/asim/op",
         "//pkg/kv/kvserver/asim/queue",
         "//pkg/kv/kvserver/asim/scheduled",

pkg/kv/kvserver/asim/asim.go

@@ -13,6 +13,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/gossip"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/history"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/metrics"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/mmaintegration"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/op"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/queue"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/scheduled"
@@ -29,8 +30,8 @@ type Simulator struct {
 	curr time.Time
 	end  time.Time
 	// interval is the step between ticks for active simulaton components, such
-	// as the queues, store rebalancer and state changers. It should be set
-	// lower than the bgInterval, as updated occur more frequently.
+	// as the queues, store rebalancer and state changers. It should be set lower
+	// than the bgInterval, as updates occur more frequently.
 	interval time.Duration
 
 	// The simulator can run multiple workload Generators in parallel.
@@ -47,6 +48,8 @@ type Simulator struct {
 	sqs map[state.StoreID]queue.RangeQueue
 	// Store rebalancers.
 	srs map[state.StoreID]storerebalancer.StoreRebalancer
+	// Multi-metric store rebalancers.
+	mmSRs map[state.StoreID]*mmaintegration.MMAStoreRebalancer
 	// Store operation controllers.
 	controllers map[state.StoreID]op.Controller
 
@@ -87,6 +90,7 @@ func NewSimulator(
 	lqs := make(map[state.StoreID]queue.RangeQueue)
 	sqs := make(map[state.StoreID]queue.RangeQueue)
 	srs := make(map[state.StoreID]storerebalancer.StoreRebalancer)
+	mmSRs := make(map[state.StoreID]*mmaintegration.MMAStoreRebalancer)
 	changer := state.NewReplicaChanger()
 	controllers := make(map[state.StoreID]op.Controller)
 
@@ -103,6 +107,7 @@ func NewSimulator(
 		sqs:            sqs,
 		controllers:    controllers,
 		srs:            srs,
+		mmSRs:          mmSRs,
 		pacers:         pacers,
 		gossip:         gossip.NewGossip(initialState, settings),
 		metrics:        m,
@@ -133,19 +138,24 @@ func (s *Simulator) StoreAddNotify(storeID state.StoreID, _ state.State) {
 func (s *Simulator) addStore(storeID state.StoreID, tick time.Time) {
 	allocator := s.state.Allocator(storeID)
 	storePool := s.state.StorePool(storeID)
+	store, _ := s.state.Store(storeID)
 	s.rqs[storeID] = queue.NewReplicateQueue(
 		storeID,
+		store.NodeID(),
 		s.changer,
 		s.settings,
 		allocator,
+		s.state.Node(store.NodeID()).AllocatorSync(),
 		storePool,
 		tick,
 	)
 	s.lqs[storeID] = queue.NewLeaseQueue(
 		storeID,
+		store.NodeID(),
 		s.changer,
 		s.settings,
 		allocator,
+		s.state.Node(store.NodeID()).AllocatorSync(),
 		storePool,
 		tick,
 	)
@@ -175,6 +185,25 @@ func (s *Simulator) addStore(storeID state.StoreID, tick time.Time) {
 		s.settings,
 		storerebalancer.GetStateRaftStatusFn(s.state),
 	)
+	// TODO: We add the store to every node's allocator in the cluster
+	// immediately. This is also updated via gossip, however there is a delay
+	// after startup. When calling `mma.ProcessStoreLeaseholderMsg` via
+	// tickMMStoreRebalancers, there may be not be a store state for some
+	// replicas. Setting it here ensures that the store is always present and
+	// initiated in each node's allocator. We should instead handle this in mma,
+	// or integration component.
+	for _, node := range s.state.Nodes() {
+		node.MMAllocator().SetStore(state.StoreAttrAndLocFromDesc(
+			s.state.StoreDescriptors(false, storeID)[0]))
+	}
+	s.mmSRs[storeID] = mmaintegration.NewMMAStoreRebalancer(
+		storeID,
+		store.NodeID(),
+		s.state.Node(store.NodeID()).MMAllocator(),
+		s.state.Node(store.NodeID()).AllocatorSync(),
+		s.controllers[storeID],
+		s.settings,
+	)
 }
 
 // GetNextTickTime returns a simulated tick time, or an indication that the
@@ -242,6 +271,9 @@ func (s *Simulator) RunSim(ctx context.Context) {
 		// Simulate the store rebalancer logic.
 		s.tickStoreRebalancers(ctx, tick, stateForAlloc)
 
+		// Simulate the multi-metric store rebalancer logic.
+		s.tickMMStoreRebalancers(ctx, tick, stateForAlloc)
+
 		// Print tick metrics.
 		s.tickMetrics(ctx, tick)
 	}
@@ -342,6 +374,16 @@ func (s *Simulator) tickStoreRebalancers(ctx context.Context, tick time.Time, st
 	}
 }
 
+// tickStoreRebalancers iterates over the multi-metric store rebalancers in the
+// cluster and ticks their control loop.
+func (s *Simulator) tickMMStoreRebalancers(ctx context.Context, tick time.Time, state state.State) {
+	stores := s.state.Stores()
+	s.shuffler(len(stores), func(i, j int) { stores[i], stores[j] = stores[j], stores[i] })
+	for _, store := range stores {
+		s.mmSRs[store.StoreID()].Tick(ctx, tick, state)
+	}
+}
+
 // tickMetrics prints the metrics up to the given tick.
 func (s *Simulator) tickMetrics(ctx context.Context, tick time.Time) {
 	s.metrics.Tick(ctx, tick, s.state)

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

@@ -16,6 +16,7 @@ go_library(
         "//pkg/kv/kvpb",
         "//pkg/kv/kvserver/allocator",
         "//pkg/kv/kvserver/allocator/allocatorimpl",
+        "//pkg/kv/kvserver/allocator/mmaprototypehelpers",
         "//pkg/kv/kvserver/allocator/plan",
         "//pkg/kv/kvserver/allocator/storepool",
         "//pkg/kv/kvserver/asim/config",

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

@@ -8,9 +8,11 @@ package queue
 import (
 	"container/heap"
 	"context"
+	"fmt"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/allocatorimpl"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/mmaprototypehelpers"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/plan"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/storepool"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/config"
@@ -27,14 +29,17 @@ type leaseQueue struct {
 	planner   plan.ReplicationPlanner
 	clock     *hlc.Clock
 	settings  *config.SimulationSettings
+	as        *mmaprototypehelpers.AllocatorSync
 }
 
 // NewLeaseQueue returns a new lease queue.
 func NewLeaseQueue(
 	storeID state.StoreID,
+	nodeID state.NodeID,
 	stateChanger state.Changer,
 	settings *config.SimulationSettings,
 	allocator allocatorimpl.Allocator,
+	allocatorSync *mmaprototypehelpers.AllocatorSync,
 	storePool storepool.AllocatorStorePool,
 	start time.Time,
 ) RangeQueue {
@@ -50,8 +55,10 @@ func NewLeaseQueue(
 		planner:   plan.NewLeasePlanner(allocator, storePool),
 		storePool: storePool,
 		clock:     storePool.Clock(),
+		as:        allocatorSync,
 	}
 	lq.AddLogTag("lease", nil)
+	lq.AddLogTag(fmt.Sprintf("n%ds%d", nodeID, storeID), "")
 	return &lq
 }
 

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

@@ -140,9 +140,11 @@ func TestLeaseQueue(t *testing.T) {
 			store, _ := s.Store(testingStore)
 			lq := NewLeaseQueue(
 				store.StoreID(),
+				store.NodeID(),
 				changer,
 				testSettings,
 				s.Allocator(store.StoreID()),
+				s.Node(store.NodeID()).AllocatorSync(),
 				s.StorePool(store.StoreID()),
 				start,
 			)

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

@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/allocatorimpl"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/mmaprototypehelpers"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/plan"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/storepool"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/config"
@@ -26,14 +27,17 @@ type replicateQueue struct {
 	planner  plan.ReplicationPlanner
 	clock    *hlc.Clock
 	settings *config.SimulationSettings
+	as       *mmaprototypehelpers.AllocatorSync
 }
 
 // NewReplicateQueue returns a new replicate queue.
 func NewReplicateQueue(
 	storeID state.StoreID,
+	nodeID state.NodeID,
 	stateChanger state.Changer,
 	settings *config.SimulationSettings,
 	allocator allocatorimpl.Allocator,
+	allocatorSync *mmaprototypehelpers.AllocatorSync,
 	storePool storepool.AllocatorStorePool,
 	start time.Time,
 ) RangeQueue {
@@ -49,8 +53,10 @@ func NewReplicateQueue(
 		planner: plan.NewReplicaPlanner(
 			allocator, storePool, plan.ReplicaPlannerTestingKnobs{}),
 		clock: storePool.Clock(),
+		as:    allocatorSync,
 	}
 	rq.AddLogTag("replica", nil)
+	rq.AddLogTag(fmt.Sprintf("n%ds%d", nodeID, storeID), "")
 	return &rq
 }
 
@@ -64,7 +70,7 @@ func (rq *replicateQueue) MaybeAdd(ctx context.Context, replica state.Replica, s
 	}
 
 	repl := NewSimulatorReplica(replica, s)
-	rq.AddLogTag("r", repl.repl.Descriptor())
+	rq.AddLogTag("r", repl.Desc().RangeID)
 	rq.AnnotateCtx(ctx)
 
 	desc := repl.Desc()

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

@@ -289,9 +289,11 @@ func TestReplicateQueue(t *testing.T) {
 			store, _ := s.Store(testingStore)
 			rq := NewReplicateQueue(
 				store.StoreID(),
+				store.NodeID(),
 				changer,
 				testSettings,
 				s.Allocator(store.StoreID()),
+				s.Node(store.NodeID()).AllocatorSync(),
 				s.StorePool(store.StoreID()),
 				start,
 			)

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

@@ -26,6 +26,7 @@ go_library(
         "//pkg/kv/kvserver/allocator",
         "//pkg/kv/kvserver/allocator/allocatorimpl",
         "//pkg/kv/kvserver/allocator/mmaprototype",
+        "//pkg/kv/kvserver/allocator/mmaprototypehelpers",
         "//pkg/kv/kvserver/allocator/storepool",
         "//pkg/kv/kvserver/asim/config",
         "//pkg/kv/kvserver/asim/workload",
@@ -45,6 +46,7 @@ go_library(
         "//pkg/util/metric",
         "//pkg/util/stop",
         "//pkg/util/timeutil",
+        "@com_github_cockroachdb_logtags//:logtags",
         "@com_github_google_btree//:btree",
         "@org_golang_google_protobuf//proto",
     ],

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

@@ -23,6 +23,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/allocatorimpl"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/mmaprototype"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/mmaprototypehelpers"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/storepool"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/config"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/workload"
@@ -36,6 +37,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/spanconfig/spanconfigreporter"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
+	"github.com/cockroachdb/logtags"
 	"github.com/google/btree"
 )
 
@@ -324,6 +326,11 @@ func (s *state) Nodes() []Node {
 	return nodes
 }
 
+func (s *state) Node(nodeID NodeID) Node {
+	node := s.nodes[nodeID]
+	return node
+}
+
 // RangeFor returns the range containing Key in [StartKey, EndKey). This
 // cannot fail.
 func (s *state) RangeFor(key Key) Range {
@@ -426,6 +433,7 @@ func (s *state) AddNode() Node {
 		stores:      []StoreID{},
 		storepool:   sp,
 		mmAllocator: mmAllocator,
+		as:          mmaprototypehelpers.NewAllocatorSync(sp, mmAllocator),
 	}
 	s.nodes[nodeID] = node
 	s.SetNodeLiveness(nodeID, livenesspb.NodeLivenessStatus_LIVE)
@@ -1136,7 +1144,24 @@ func (s *state) UpdateStorePool(
 		detail := storeDescriptors[gossipStoreID]
 		copiedDetail := detail.Copy()
 		node.storepool.Details.StoreDetails.Store(gossipStoreID, copiedDetail)
+		copiedDesc := *copiedDetail.Desc
 		// TODO(mma): Support origin timestamps.
+		ts := s.clock.Now()
+		storeLoadMsg := mmaprototypehelpers.MakeStoreLoadMsg(copiedDesc, ts.UnixNano())
+		node.mmAllocator.SetStore(StoreAttrAndLocFromDesc(copiedDesc))
+		ctx := logtags.AddTag(context.Background(), fmt.Sprintf("n%d", nodeID), "")
+		ctx = logtags.AddTag(ctx, "t", ts.Sub(s.settings.StartTime))
+		node.mmAllocator.ProcessStoreLoadMsg(ctx, &storeLoadMsg)
+	}
+}
+
+func StoreAttrAndLocFromDesc(desc roachpb.StoreDescriptor) mmaprototype.StoreAttributesAndLocality {
+	return mmaprototype.StoreAttributesAndLocality{
+		StoreID:      desc.StoreID,
+		NodeID:       desc.Node.NodeID,
+		NodeAttrs:    desc.Node.Attrs,
+		NodeLocality: desc.Node.Locality,
+		StoreAttrs:   desc.Attrs,
 	}
 }
 
@@ -1400,6 +1425,7 @@ type node struct {
 	stores      []StoreID
 	storepool   *storepool.StorePool
 	mmAllocator mmaprototype.Allocator
+	as          *mmaprototypehelpers.AllocatorSync
 }
 
 // NodeID returns the ID of this node.
@@ -1421,6 +1447,10 @@ func (n *node) MMAllocator() mmaprototype.Allocator {
 	return n.mmAllocator
 }
 
+func (n *node) AllocatorSync() *mmaprototypehelpers.AllocatorSync {
+	return n.as
+}
+
 // store is an implementation of the Store interface.
 type store struct {
 	storeID StoreID

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

@@ -13,6 +13,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/allocatorimpl"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/mmaprototype"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/mmaprototypehelpers"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/storepool"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/workload"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness/livenesspb"
@@ -59,6 +60,7 @@ type State interface {
 	// first flag is false, then the capacity is generated from scratch,
 	// otherwise the last calculated capacity values are used for each store.
 	StoreDescriptors(bool, ...StoreID) []roachpb.StoreDescriptor
+	Node(NodeID) Node
 	// Nodes returns all nodes that exist in this state.
 	Nodes() []Node
 	// RangeFor returns the range containing Key in [StartKey, EndKey). This
@@ -217,6 +219,8 @@ type Node interface {
 	Descriptor() roachpb.NodeDescriptor
 	// TODO(wenyihu6): use this in mma store rebalancer
 	MMAllocator() mmaprototype.Allocator
+	// AllocatorSync returns the AllocatorSync for this node.
+	AllocatorSync() *mmaprototypehelpers.AllocatorSync
 }
 
 // Store is a container for replicas.