kvserver: requeue on priority inversion for replicate queue

Previously, a replica could enter the queue with high priority but, by the time
it was processed, the action planned for this replica may have a low priority,
causing us to perform low priority work. Specifically, we are mostly worried
about cases where the priority changes from any of the repair actions to
consider rebalance. Rebalancing could take a long time and block other ranges
enqueued with actual repair action needed. This commit ensures that such
replicas are requeued instead, avoiding priority inversions.

Author: wenyihu6

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

@@ -3245,3 +3245,62 @@ func replDescsToStoreIDs(descs []roachpb.ReplicaDescriptor) []roachpb.StoreID {
 	}
 	return ret
 }
+
+// roundToNearestPriorityCategory rounds a priority to the nearest 100. n should be non-negative.
+func roundToNearestPriorityCategory(n float64) float64 {
+	return math.Round(n/100.0) * 100
+}
+
+// WithinPriorityRange checks if a priority is within the range of possible
+// priorities for the allocator actions.
+func withinPriorityRange(priority float64) bool {
+	return AllocatorNoop.Priority() <= priority && priority <= AllocatorFinalizeAtomicReplicationChange.Priority()
+}
+
+// CheckPriorityInversion checks if the priority at enqueue time is higher than
+// the priority corresponding to the action computed at processing time. It
+// returns whether there was a priority inversion and whether the caller should
+// skip the processing of the range since the inversion is considered unfair.
+// Currently, we only consider the inversion as unfair if it has gone from a
+// repair action to lowest priority (AllocatorConsiderRebalance). We let
+// AllocatorRangeUnavailable, AllocatorNoop pass through since they are noop.
+//
+// NB: If shouldRequeue is true, isInversion must be true.
+func CheckPriorityInversion(
+	priorityAtEnqueue float64, actionAtProcessing AllocatorAction,
+) (isInversion bool, shouldRequeue bool) {
+	// priorityAtEnqueue of -1 is a special case reserved for processing logic to
+	// run even if there’s a priority inversion. If the priority is not -1, the
+	// range may be re-queued to be processed with the correct priority. It is
+	// used for things that call into baseQueue.process without going through the
+	// replicate priority queue. For example, s.ReplicateQueueDryRun or
+	// r.scatterRangeAndRandomizeLeases.
+
+	// NB: we need to check for when priorityAtEnqueue falls within the range
+	// of the allocator actions because store.Enqueue might enqueue things with
+	// a very high priority (1e5). In those cases, we do not want to requeue
+	// these actions or count it as an inversion.
+	if priorityAtEnqueue == -1 || !withinPriorityRange(priorityAtEnqueue) {
+		return false, false
+	}
+
+	if priorityAtEnqueue > AllocatorConsiderRebalance.Priority() && actionAtProcessing == AllocatorConsiderRebalance {
+		return true, true
+	}
+
+	// NB: Usually, the priority at enqueue time should correspond to
+	// action.Priority(). However, for AllocatorAddVoter,
+	// AllocatorRemoveDeadVoter, AllocatorRemoveVoter, the priority can be
+	// adjusted at enqueue time (See ComputeAction for more details). However, we
+	// expect the adjustment to be relatively small (<100). So we round the
+	// priority to the nearest 100 to compare against
+	// actionAtProcessing.Priority(). Without this rounding, we might treat going
+	// from 10000 to 999 as an inversion, but it was just due to the adjustment.
+	// Note that priorities at AllocatorFinalizeAtomicReplicationChange,
+	// AllocatorRemoveLearner, and AllocatorReplaceDeadVoter will be rounded to
+	// the same priority. They are so close to each other, so we don't really
+	// count it as an inversion among them.
+	normPriorityAtEnqueue := roundToNearestPriorityCategory(priorityAtEnqueue)
+	isInversion = normPriorityAtEnqueue > actionAtProcessing.Priority()
+	return isInversion, false
+}

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

@@ -9629,3 +9629,148 @@ func TestAllocatorRebalanceTargetVoterConstraintUnsatisfied(t *testing.T) {
 		})
 	}
 }
+
+// TestRoundToNearestPriorityCategory tests the RoundToNearestPriorityCategory
+// function.
+func TestRoundToNearestPriorityCategory(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	testCases := []struct {
+		name     string
+		input    float64
+		expected float64
+	}{
+		{
+			name:     "zero",
+			input:    0.0,
+			expected: 0.0,
+		},
+		{
+			name:     "exact multiple of 100",
+			input:    100.0,
+			expected: 100.0,
+		},
+		{
+			name:     "round down to nearest 100",
+			input:    149.0,
+			expected: 100.0,
+		},
+		{
+			name:     "round up to nearest 100",
+			input:    151.0,
+			expected: 200.0,
+		},
+		{
+			name:     "negative exact multiple of 100",
+			input:    -200.0,
+			expected: -200.0,
+		},
+		{
+			name:     "negative round down to nearest 100",
+			input:    -249.0,
+			expected: -200.0,
+		},
+		{
+			name:     "negative round up to nearest 100",
+			input:    -251.0,
+			expected: -300.0,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			require.Equal(t, tc.expected, roundToNearestPriorityCategory(tc.input))
+		})
+	}
+}
+
+// TestCheckPriorityInversion tests the CheckPriorityInversion function.
+func TestCheckPriorityInversion(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	for action := AllocatorNoop; action <= AllocatorFinalizeAtomicReplicationChange; action++ {
+		t.Run(action.String(), func(t *testing.T) {
+			if action == AllocatorConsiderRebalance || action == AllocatorNoop || action == AllocatorRangeUnavailable {
+				inversion, requeue := CheckPriorityInversion(action.Priority(), AllocatorConsiderRebalance)
+				require.False(t, inversion)
+				require.False(t, requeue)
+			} else {
+				inversion, requeue := CheckPriorityInversion(action.Priority(), AllocatorConsiderRebalance)
+				require.True(t, inversion)
+				require.True(t, requeue)
+			}
+		})
+	}
+
+	testCases := []struct {
+		name               string
+		priorityAtEnqueue  float64
+		actionAtProcessing AllocatorAction
+		expectedInversion  bool
+		expectedRequeue    bool
+	}{
+		{
+			name:               "AllocatorNoop at processing is noop",
+			priorityAtEnqueue:  AllocatorFinalizeAtomicReplicationChange.Priority(),
+			actionAtProcessing: AllocatorNoop,
+			expectedInversion:  true,
+			expectedRequeue:    false,
+		},
+		{
+			name:               "AllocatorRangeUnavailable at processing is noop",
+			priorityAtEnqueue:  AllocatorFinalizeAtomicReplicationChange.Priority(),
+			actionAtProcessing: AllocatorRangeUnavailable,
+			expectedInversion:  true,
+			expectedRequeue:    false,
+		},
+		{
+			name:               "priority -1 bypasses",
+			priorityAtEnqueue:  -1,
+			actionAtProcessing: AllocatorConsiderRebalance,
+			expectedInversion:  false,
+			expectedRequeue:    false,
+		},
+		{
+			name:               "above range priority(1e5)",
+			priorityAtEnqueue:  1e5,
+			actionAtProcessing: AllocatorConsiderRebalance,
+			expectedInversion:  false,
+			expectedRequeue:    false,
+		},
+		{
+			name:               "below range priority at -10",
+			priorityAtEnqueue:  -10,
+			actionAtProcessing: -100,
+			expectedInversion:  false,
+			expectedRequeue:    false,
+		},
+		{
+			name:               "inversion but small priority changes",
+			priorityAtEnqueue:  AllocatorFinalizeAtomicReplicationChange.Priority(),
+			actionAtProcessing: AllocatorReplaceDecommissioningNonVoter,
+			expectedInversion:  true,
+			expectedRequeue:    false,
+		},
+		{
+			name:               "inversion but small priority changes",
+			priorityAtEnqueue:  AllocatorRemoveDeadVoter.Priority(),
+			actionAtProcessing: AllocatorAddNonVoter,
+			expectedInversion:  true,
+			expectedRequeue:    false,
+		},
+		{
+			name:               "inversion but small priority changes",
+			priorityAtEnqueue:  AllocatorConsiderRebalance.Priority(),
+			actionAtProcessing: AllocatorNoop,
+			expectedInversion:  false,
+			expectedRequeue:    false,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			inversion, requeue := CheckPriorityInversion(tc.priorityAtEnqueue, tc.actionAtProcessing)
+			require.Equal(t, tc.expectedInversion, inversion)
+			require.Equal(t, tc.expectedRequeue, requeue)
+		})
+	}
+}

pkg/kv/kvserver/replicate_queue.go

@@ -722,6 +722,10 @@ func (rq *replicateQueue) process(
 		}
 
 		if err != nil {
+			if requeue {
+				log.KvDistribution.VEventf(ctx, 1, "re-queuing on errors: %v", err)
+				rq.maybeAdd(ctx, repl, rq.store.Clock().NowAsClockTimestamp())
+			}
 			return false, err
 		}
 
@@ -916,6 +920,22 @@ func (rq *replicateQueue) processOneChange(
 ) (requeue bool, _ error) {
 	change, err := rq.planner.PlanOneChange(
 		ctx, repl, desc, conf, plan.PlannerOptions{Scatter: scatter})
+
+	inversion, shouldRequeue := allocatorimpl.CheckPriorityInversion(priorityAtEnqueue, change.Action)
+	if inversion {
+		log.KvDistribution.VInfof(ctx, 2,
+			"priority inversion during process: shouldRequeue = %t action=%s, priority=%v, enqueuePriority=%v",
+			shouldRequeue, change.Action, change.Action.Priority(), priorityAtEnqueue)
+		if shouldRequeue && PriorityInversionRequeue.Get(&rq.store.cfg.Settings.SV) {
+			// Return true here to requeue the range. We can't return an error here
+			// because rq.process only requeue when error is nil. See
+			// replicateQueue.process for more details.
+			return true /*requeue*/, maybeAnnotateDecommissionErr(
+				errors.Errorf("requing due to priority inversion: action=%s, priority=%v, enqueuePriority=%v",
+					change.Action, change.Action.Priority(), priorityAtEnqueue), change.Action)
+		}
+	}
+
 	// When there is an error planning a change, return the error immediately
 	// and do not requeue. It is unlikely that the range or storepool state
 	// will change quickly enough in order to not get the same error and