Merge #159099

159099: mma: plumbing changes for existing MMA metrics r=wenyihu6 a=sumeerbhola

- Metrics are now per local store. This is similar to SMA metrics. For multi-store nodes, it is preferable to have counters reflecting what each local store was doing wrt changing leases and replicas.
- The metrics were not hooked up to a MetricRegistry -- this is fixed.
- MMAMetrics is moved to a different file and renamed counterMetrics. These are directly modified by the MMA code as it takes actions.

Many of the existing counter metrics are unused, and there will be cleanup in subsequent PRs to remove/add metrics and hook them up.

Epic: CRDB-55052

Release note: None

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

Author: craig[bot]

pkg/kv/kvserver/allocator/allocatorimpl/allocator_test.go

@@ -55,7 +55,7 @@ import (
 func MakeAllocatorSync(
 	sp *storepool.StorePool, st *cluster.Settings,
 ) *mmaintegration.AllocatorSync {
-	mmAllocator := mmaprototype.NewAllocatorState(timeutil.DefaultTimeSource{},
+	mmAllocator := mmaprototype.NewAllocatorState(timeutil.DefaultTimeSource{}, nil,
 		rand.New(rand.NewSource(timeutil.Now().UnixNano())))
 	return mmaintegration.NewAllocatorSync(sp, mmAllocator, st, nil)
 }

pkg/kv/kvserver/allocator/allocatorimpl/test_helpers.go

@@ -47,7 +47,8 @@ func CreateTestAllocatorWithKnobs(
 		liveness.TestTimeUntilNodeDeadOff, deterministic,
 		func() int { return numNodes },
 		livenesspb.NodeLivenessStatus_LIVE)
-	mmAllocator := mmaprototype.NewAllocatorState(timeutil.DefaultTimeSource{}, rand.New(rand.NewSource(timeutil.Now().UnixNano())))
+	mmAllocator := mmaprototype.NewAllocatorState(
+		timeutil.DefaultTimeSource{}, nil, rand.New(rand.NewSource(timeutil.Now().UnixNano())))
 	as := mmaintegration.NewAllocatorSync(storePool, mmAllocator, st, allocSyncKnobs)
 	a := MakeAllocator(st, as, deterministic, func(id roachpb.NodeID) (time.Duration, bool) {
 		return 0, true

pkg/kv/kvserver/allocator/mmaprototype/BUILD.bazel

@@ -13,6 +13,7 @@ go_library(
         "load.go",
         "memo_helper.go",
         "messages.go",
+        "mma_metrics.go",
         "range_change.go",
         "rebalance_advisor.go",
         "store_load_summary.go",

pkg/kv/kvserver/allocator/mmaprototype/allocator.go

@@ -30,7 +30,7 @@ type ChangeOptions struct {
 //     be less different than integration with the old allocator.
 type Allocator interface {
 	LoadSummaryForAllStores(context.Context) string
-	Metrics() *MMAMetrics
+
 	// Methods to update the state of the external world. The allocator starts
 	// with no knowledge.
 
@@ -59,13 +59,14 @@ type Allocator interface {
 	AdjustPendingChangeDisposition(change ExternalRangeChange, success bool)
 
 	// RegisterExternalChange informs this allocator about yet to complete
-	// changes to the cluster which were not initiated by this allocator. The
-	// ownership of all state inside change is handed off to the callee. If ok
-	// is true, the change was registered, and the caller is returned an
-	// ExternalRangeChange that it should subsequently use in a call to
-	// AdjustPendingChangeDisposition when the changes are completed, either
-	// successfully or not. If ok is false, the change was not registered.
-	RegisterExternalChange(change PendingRangeChange) (_ ExternalRangeChange, ok bool)
+	// changes to the cluster (on behalf of localStoreID) which were not
+	// initiated by this allocator. The ownership of all state inside change is
+	// handed off to the callee. If ok is true, the change was registered, and
+	// the caller is returned an ExternalRangeChange that it should subsequently
+	// use in a call to AdjustPendingChangeDisposition when the changes are
+	// completed, either successfully or not. If ok is false, the change was not
+	// registered.
+	RegisterExternalChange(localStoreID roachpb.StoreID, change PendingRangeChange) (_ ExternalRangeChange, ok bool)
 
 	// ComputeChanges is called periodically and frequently, say every 10s.
 	//

pkg/kv/kvserver/allocator/mmaprototype/allocator_state.go

@@ -26,128 +26,6 @@ import (
 	"github.com/cockroachdb/redact/interfaces"
 )
 
-type MMAMetrics struct {
-	DroppedDueToStateInconsistency  *metric.Counter
-	ExternalFailedToRegister        *metric.Counter
-	ExternaRegisterSuccess          *metric.Counter
-	ExternalReplicaRebalanceSuccess *metric.Counter
-	ExternalReplicaRebalanceFailure *metric.Counter
-	ExternalLeaseTransferSuccess    *metric.Counter
-	ExternalLeaseTransferFailure    *metric.Counter
-	MMAReplicaRebalanceSuccess      *metric.Counter
-	MMAReplicaRebalanceFailure      *metric.Counter
-	MMALeaseTransferSuccess         *metric.Counter
-	MMALeaseTransferFailure         *metric.Counter
-	MMARegisterLeaseSuccess         *metric.Counter
-	MMARegisterRebalanceSuccess     *metric.Counter
-}
-
-func makeMMAMetrics() *MMAMetrics {
-	return &MMAMetrics{
-		DroppedDueToStateInconsistency:  metric.NewCounter(metaDroppedDueToStateInconsistency),
-		ExternalFailedToRegister:        metric.NewCounter(metaExternalFailedToRegister),
-		ExternaRegisterSuccess:          metric.NewCounter(metaExternaRegisterSuccess),
-		MMARegisterLeaseSuccess:         metric.NewCounter(metaMMARegisterLeaseSuccess),
-		MMARegisterRebalanceSuccess:     metric.NewCounter(metaMMARegisterRebalanceSuccess),
-		ExternalReplicaRebalanceSuccess: metric.NewCounter(metaExternalReplicaRebalanceSuccess),
-		ExternalReplicaRebalanceFailure: metric.NewCounter(metaExternalReplicaRebalanceFailure),
-		ExternalLeaseTransferSuccess:    metric.NewCounter(metaExternalLeaseTransferSuccess),
-		ExternalLeaseTransferFailure:    metric.NewCounter(metaExternalLeaseTransferFailure),
-		MMAReplicaRebalanceSuccess:      metric.NewCounter(metaMMAReplicaRebalanceSuccess),
-		MMAReplicaRebalanceFailure:      metric.NewCounter(metaMMAReplicaRebalanceFailure),
-		MMALeaseTransferSuccess:         metric.NewCounter(metaMMALeaseTransferSuccess),
-		MMALeaseTransferFailure:         metric.NewCounter(metaMMALeaseTransferFailure),
-	}
-}
-
-var (
-	metaDroppedDueToStateInconsistency = metric.Metadata{
-		Name:        "mma.dropped",
-		Help:        "Number of operations dropped due to MMA state inconsistency",
-		Measurement: "Range Rebalances",
-		Unit:        metric.Unit_COUNT,
-	}
-	metaExternalFailedToRegister = metric.Metadata{
-		Name:        "mma.external.dropped",
-		Help:        "Number of external operations that failed to register with MMA",
-		Measurement: "Range Rebalances",
-		Unit:        metric.Unit_COUNT,
-	}
-	metaExternaRegisterSuccess = metric.Metadata{
-		Name:        "mma.external.success",
-		Help:        "Number of external operations successfully registered with MMA",
-		Measurement: "Range Rebalances",
-		Unit:        metric.Unit_COUNT,
-	}
-	metaMMARegisterLeaseSuccess = metric.Metadata{
-		Name:        "mma.lease.register.success",
-		Help:        "Number of lease transfers successfully registered with MMA",
-		Measurement: "Range Rebalances",
-		Unit:        metric.Unit_COUNT,
-	}
-	metaMMARegisterRebalanceSuccess = metric.Metadata{
-		Name:        "mma.rebalance.register.success",
-		Help:        "Number of rebalance operations successfully registered with MMA",
-		Measurement: "Range Rebalances",
-		Unit:        metric.Unit_COUNT,
-	}
-	metaExternalReplicaRebalanceSuccess = metric.Metadata{
-		Name:        "mma.rebalances.external.success",
-		Help:        "Number of successful external replica rebalance operations",
-		Measurement: "Range Rebalances",
-		Unit:        metric.Unit_COUNT,
-	}
-
-	metaExternalLeaseTransferSuccess = metric.Metadata{
-		Name:        "mma.lease.external.success",
-		Help:        "Number of successful external lease transfer operations",
-		Measurement: "Lease Transfers",
-		Unit:        metric.Unit_COUNT,
-	}
-
-	metaExternalReplicaRebalanceFailure = metric.Metadata{
-		Name:        "mma.rebalances.external.failure",
-		Help:        "Number of failed external replica rebalance operations",
-		Measurement: "Range Rebalances",
-		Unit:        metric.Unit_COUNT,
-	}
-
-	metaExternalLeaseTransferFailure = metric.Metadata{
-		Name:        "mma.lease.external.failure",
-		Help:        "Number of failed external lease transfer operations",
-		Measurement: "Lease Transfers",
-		Unit:        metric.Unit_COUNT,
-	}
-
-	metaMMAReplicaRebalanceSuccess = metric.Metadata{
-		Name:        "mma.rebalance.success",
-		Help:        "Number of successful MMA-initiated replica rebalance operations",
-		Measurement: "Range Rebalances",
-		Unit:        metric.Unit_COUNT,
-	}
-
-	metaMMAReplicaRebalanceFailure = metric.Metadata{
-		Name:        "mma.rebalance.failure",
-		Help:        "Number of failed MMA-initiated replica rebalance operations",
-		Measurement: "Range Rebalances",
-		Unit:        metric.Unit_COUNT,
-	}
-
-	metaMMALeaseTransferSuccess = metric.Metadata{
-		Name:        "mma.lease.success",
-		Help:        "Number of successful MMA-initiated lease transfer operations",
-		Measurement: "Lease Transfers",
-		Unit:        metric.Unit_COUNT,
-	}
-
-	metaMMALeaseTransferFailure = metric.Metadata{
-		Name:        "mma.lease.failure",
-		Help:        "Number of failed MMA-initiated lease transfer operations",
-		Measurement: "Lease Transfers",
-		Unit:        metric.Unit_COUNT,
-	}
-)
-
 type allocatorState struct {
 	// Locking.
 	//
@@ -195,9 +73,14 @@ type allocatorState struct {
 	// try-write-lock that could quickly return with failure then we could avoid
 	// this. We could of course build our own queueing mechanism instead of
 	// relying on the queueing in mutex.
-	mmaMetrics *MMAMetrics
+
+	// mrProvider can be nil in tests.
+	mrProvider MetricRegistryForStoreProvider
 	mu         syncutil.Mutex
-	cs         *clusterState
+	// TODO(sumeer): localStoreMetrics is also protected by mu. Nest in struct
+	// with mu, when locking story is cleaned up.
+	localStoreMetrics map[roachpb.StoreID]*counterMetrics
+	cs                *clusterState
 
 	// Ranges that are under-replicated, over-replicated, don't satisfy
 	// constraints, have low diversity etc. Avoids iterating through all ranges.
@@ -212,15 +95,29 @@ type allocatorState struct {
 
 var _ Allocator = &allocatorState{}
 
-func NewAllocatorState(ts timeutil.TimeSource, rand *rand.Rand) *allocatorState {
+type MetricRegistryForStoreProvider interface {
+	// GetStoreMetricRegistry returns the registry for the store, if it is
+	// known, else nil.
+	GetStoreMetricRegistry(storeID roachpb.StoreID) *metric.Registry
+}
+
+// NewAllocatorState constructs a new implementation of Allocator.
+//
+// The metricRegistryProvider allows the allocator to lazily initialize
+// per-local-store metrics once the StoreID is known. It can be nil in tests,
+// in which case no metrics will be collected.
+func NewAllocatorState(
+	ts timeutil.TimeSource, metricRegistryProvider MetricRegistryForStoreProvider, rand *rand.Rand,
+) *allocatorState {
 	interner := newStringInterner()
 	cs := newClusterState(ts, interner)
 	return &allocatorState{
+		mrProvider:             metricRegistryProvider,
+		localStoreMetrics:      map[roachpb.StoreID]*counterMetrics{},
 		cs:                     cs,
 		rangesNeedingAttention: map[roachpb.RangeID]struct{}{},
 		diversityScoringMemo:   newDiversityScoringMemo(),
 		rand:                   rand,
-		mmaMetrics:             makeMMAMetrics(),
 	}
 }
 
@@ -231,8 +128,23 @@ func NewAllocatorState(ts timeutil.TimeSource, rand *rand.Rand) *allocatorState
 const remoteStoreLeaseSheddingGraceDuration = 2 * time.Minute
 const overloadGracePeriod = time.Minute
 
-func (a *allocatorState) Metrics() *MMAMetrics {
-	return a.mmaMetrics
+func (a *allocatorState) ensureMetricsForLocalStoreLocked(
+	localStoreID roachpb.StoreID,
+) *counterMetrics {
+	m, ok := a.localStoreMetrics[localStoreID]
+	if ok {
+		return m
+	}
+	m = makeCounterMetrics()
+	if a.mrProvider != nil {
+		mr := a.mrProvider.GetStoreMetricRegistry(localStoreID)
+		if mr == nil {
+			panic(errors.AssertionFailedf("no MetricRegistry for store s%v", localStoreID))
+		}
+		mr.AddMetricStruct(*m)
+	}
+	a.localStoreMetrics[localStoreID] = m
+	return m
 }
 
 func (a *allocatorState) LoadSummaryForAllStores(ctx context.Context) string {
@@ -300,17 +212,18 @@ func (a *allocatorState) AdjustPendingChangeDisposition(change ExternalRangeChan
 
 // RegisterExternalChange implements the Allocator interface.
 func (a *allocatorState) RegisterExternalChange(
-	change PendingRangeChange,
+	localStoreID roachpb.StoreID, change PendingRangeChange,
 ) (_ ExternalRangeChange, ok bool) {
 	a.mu.Lock()
 	defer a.mu.Unlock()
+	counterMetrics := a.ensureMetricsForLocalStoreLocked(localStoreID)
 	if err := a.cs.preCheckOnApplyReplicaChanges(change); err != nil {
-		a.mmaMetrics.ExternalFailedToRegister.Inc(1)
+		counterMetrics.ExternalFailedToRegister.Inc(1)
 		log.KvDistribution.Infof(context.Background(),
 			"did not register external changes: due to %v", err)
 		return ExternalRangeChange{}, false
 	} else {
-		a.mmaMetrics.ExternaRegisterSuccess.Inc(1)
+		counterMetrics.ExternaRegisterSuccess.Inc(1)
 	}
 	a.cs.addPendingRangeChange(change)
 	return MakeExternalRangeChange(change), true
@@ -325,7 +238,8 @@ func (a *allocatorState) ComputeChanges(
 	if msg.StoreID != opts.LocalStoreID {
 		panic(fmt.Sprintf("ComputeChanges: expected StoreID %d, got %d", opts.LocalStoreID, msg.StoreID))
 	}
-	a.cs.processStoreLeaseholderMsg(ctx, msg, a.mmaMetrics)
+	counterMetrics := a.ensureMetricsForLocalStoreLocked(opts.LocalStoreID)
+	a.cs.processStoreLeaseholderMsg(ctx, msg, counterMetrics)
 	re := newRebalanceEnv(a.cs, a.rand, a.diversityScoringMemo, a.cs.ts.Now())
 	return re.rebalanceStores(ctx, opts.LocalStoreID)
 }

pkg/kv/kvserver/allocator/mmaprototype/cluster_state.go

@@ -1348,13 +1348,13 @@ func (cs *clusterState) processStoreLoadMsg(ctx context.Context, storeMsg *Store
 }
 
 func (cs *clusterState) processStoreLeaseholderMsg(
-	ctx context.Context, msg *StoreLeaseholderMsg, metrics *MMAMetrics,
+	ctx context.Context, msg *StoreLeaseholderMsg, metrics *counterMetrics,
 ) {
 	cs.processStoreLeaseholderMsgInternal(ctx, msg, numTopKReplicas, metrics)
 }
 
 func (cs *clusterState) processStoreLeaseholderMsgInternal(
-	ctx context.Context, msg *StoreLeaseholderMsg, numTopKReplicas int, metrics *MMAMetrics,
+	ctx context.Context, msg *StoreLeaseholderMsg, numTopKReplicas int, metrics *counterMetrics,
 ) {
 	now := cs.ts.Now()
 	cs.gcPendingChanges(now)