mmaprototype: pass context properly

This commit ensures `IsInConflictWithMMA` receives a proper context so that
its logs appear in traces.

Author: wenyihu6

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

@@ -1911,7 +1911,7 @@ func (a Allocator) RebalanceTarget(
 			// would have been deleted from the candidates set in bestRebalanceTarget,
 			// so we will not select it again.
 			if advisor := results[bestIdx].advisor; advisor != nil {
-				if a.as.IsInConflictWithMMA(target.store.StoreID, advisor, false) {
+				if a.as.IsInConflictWithMMA(ctx, target.store.StoreID, advisor, false) {
 					continue
 				}
 			} else {

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

@@ -1343,19 +1343,18 @@ func (a *allocatorState) BuildMMARebalanceAdvisor(
 // for making sure the MMARebalanceAdvisor is for the correct existing store and
 // candidate set.
 func (a *allocatorState) IsInConflictWithMMA(
-	cand roachpb.StoreID, advisor *MMARebalanceAdvisor, cpuOnly bool,
+	ctx context.Context, cand roachpb.StoreID, advisor *MMARebalanceAdvisor, cpuOnly bool,
 ) bool {
 	if advisor.disabled {
 		return false
 	}
 	if advisor.existingStoreSLS == nil {
-		// TODO(wenyihu6): pass in the actual ctx here
-		summary := a.cs.computeLoadSummary(context.Background(), advisor.existingStoreID,
+		summary := a.cs.computeLoadSummary(ctx, advisor.existingStoreID,
 			&advisor.means.storeLoad, &advisor.means.nodeLoad)
 		advisor.existingStoreSLS = &summary
 	}
 	existingSLS := advisor.existingStoreSLS
-	candSLS := a.cs.computeLoadSummary(context.Background(), cand, &advisor.means.storeLoad, &advisor.means.nodeLoad)
+	candSLS := a.cs.computeLoadSummary(ctx, cand, &advisor.means.storeLoad, &advisor.means.nodeLoad)
 	if cpuOnly {
 		return candSLS.dimSummary[CPURate] > existingSLS.dimSummary[CPURate]
 	}

pkg/kv/kvserver/mmaintegration/allocator_sync.go

@@ -6,6 +6,8 @@
 package mmaintegration
 
 import (
+	"context"
+
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/mmaprototype"
@@ -52,7 +54,7 @@ type mmaState interface {
 	BuildMMARebalanceAdvisor(existing roachpb.StoreID, cands []roachpb.StoreID) *mmaprototype.MMARebalanceAdvisor
 	// IsInConflictWithMMA is called by the allocator sync to determine if the
 	// given candidate is in conflict with the existing store.
-	IsInConflictWithMMA(cand roachpb.StoreID, advisor *mmaprototype.MMARebalanceAdvisor, cpuOnly bool) bool
+	IsInConflictWithMMA(ctx context.Context, cand roachpb.StoreID, advisor *mmaprototype.MMARebalanceAdvisor, cpuOnly bool) bool
 }
 
 // TODO(wenyihu6): make sure allocator sync can tolerate cluster setting

pkg/kv/kvserver/mmaintegration/thrashing.go

@@ -6,6 +6,8 @@
 package mmaintegration
 
 import (
+	"context"
+
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/mmaprototype"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -101,10 +103,13 @@ func (as *AllocatorSync) BuildMMARebalanceAdvisor(
 
 // IsInConflictWithMMA determines if a candidate conflicts with MMA's goals.
 func (as *AllocatorSync) IsInConflictWithMMA(
-	cand roachpb.StoreID, advisor *mmaprototype.MMARebalanceAdvisor, cpuOnly bool,
+	ctx context.Context,
+	cand roachpb.StoreID,
+	advisor *mmaprototype.MMARebalanceAdvisor,
+	cpuOnly bool,
 ) bool {
 	if as.knobs != nil && as.knobs.OverrideIsInConflictWithMMA != nil {
 		return as.knobs.OverrideIsInConflictWithMMA(cand)
 	}
-    return as.mmaAllocator.IsInConflictWithMMA(cand, advisor, cpuOnly)
+	return as.mmaAllocator.IsInConflictWithMMA(ctx, cand, advisor, cpuOnly)
 }