asim: integrate mma store rebalancer

This commit integrates mma store rebalancers to asim.

Author: wenyihu6

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,
@@ -178,6 +183,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
@@ -245,6 +269,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)
 	}
@@ -345,6 +372,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/state/BUILD.bazel

@@ -46,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

@@ -37,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"
 )
 
@@ -325,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 {
@@ -1138,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,
 	}
 }
 

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

@@ -60,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

pkg/kv/kvserver/asim/tests/testdata/non_rand/example_zone_config.txt

@@ -53,22 +53,22 @@ US_West
 
 plot stat=replicas
 ----
- 54.70 ┤            ╭╮
- 51.05 ┤          ╭╮│╰╮ ╭╭─────────────────────────────────────────────────────────────
- 47.40 ┤          │╰╮╭╭──╯─────────────────────────────────────────────────────────────
- 43.76 ┤        ╭╮│╭──╯╯
- 40.11 ┤        │╰││╯╯
- 36.46 ┤        │ ╭╯
- 32.82 ┤        │╭│╯
- 29.17 ┤        ╭╭╯
- 25.52 ┤        ╭│╯
- 21.88 ┤        ││
- 18.23 ┼────────╮╯
- 14.59 ┼────────╰╮╮
- 10.94 ┤        ╰╰╮╮
-  7.29 ┤        ╰╰╰───╮╮
-  3.65 ┤         ╰╰╰╰─╰─────╮
-  0.00 ┤          ╰╰───╰╰──╰╰──────────────────────────────────────────────────────────
+ 56.70 ┤            ╭╮
+ 52.92 ┤          ╭╮│╰╮ ╭╮─────────────────────────────────────────────────────────────
+ 49.14 ┤          │││ ╭╮╭──────────────────────────────────────────────────────────────
+ 45.36 ┤        ╭╮│╰─╭│╰╯╰╯
+ 41.58 ┤        │╰│╭──╯
+ 37.80 ┤        │ ││╰╯
+ 34.02 ┤        │╭╭╯
+ 30.24 ┤        ╭╭│╯
+ 26.46 ┤        │╭╯
+ 22.68 ┤        ╭│
+ 18.90 ┤       ╭╭╯
+ 15.12 ┼─────────╮╮
+ 11.34 ┤        ╰╰╮╮
+  7.56 ┤        ╰╰╰─╮─╮
+  3.78 ┤         ╰╰╰╰─────╮╮
+  0.00 ┤          ╰╰──╰─╰─╰────────────────────────────────────────────────────────────
                                             replicas
 initial store values: [s1=17, s2=17, s3=17, s4=17, s5=17, s6=16, s7=16, s8=17, s9=17, s10=17, s11=17, s12=17, s13=16, s14=16, s15=16, s16=16, s17=17, s18=17, s19=17, s20=16, s21=17, s22=16, s23=17, s24=17, s25=17, s26=17, s27=17, s28=16, s29=17, s30=16, s31=17, s32=17, s33=16, s34=16, s35=17, s36=17] (stddev=0.47, mean=16.67, sum=600)
-last store values: [s1=50, s2=50, s3=49, s4=50, s5=50, s6=50, s7=50, s8=50, s9=50, s10=49, s11=51, s12=51, s13=0, s14=0, s15=0, s16=0, s17=0, s18=0, s19=0, s20=0, s21=0, s22=0, s23=0, s24=0, s25=0, s26=0, s27=0, s28=0, s29=0, s30=0, s31=0, s32=0, s33=0, s34=0, s35=0, s36=0] (stddev=23.57, mean=16.67, sum=600)
+last store values: [s1=49, s2=50, s3=50, s4=52, s5=49, s6=52, s7=49, s8=50, s9=50, s10=50, s11=49, s12=50, s13=0, s14=0, s15=0, s16=0, s17=0, s18=0, s19=0, s20=0, s21=0, s22=0, s23=0, s24=0, s25=0, s26=0, s27=0, s28=0, s29=0, s30=0, s31=0, s32=0, s33=0, s34=0, s35=0, s36=0] (stddev=23.58, mean=16.67, sum=600)