kv: add TestDecommissionPreCheckRetryThrottledStores

Previously, we made decommission prechecks retry on errors, since some transient
issues resolve quickly and shouldn’t cause the precheck to fail. This commit
adds a test that verifies the precheck retries when it encounters transient
throttled errors.

Author: wenyihu6

pkg/kv/kvserver/store.go

@@ -3965,6 +3965,13 @@ func (s *Store) AllocatorCheckRange(
 	collectTraces bool,
 	overrideStorePool storepool.AllocatorStorePool,
 ) (allocatorimpl.AllocatorAction, roachpb.ReplicationTarget, tracingpb.Recording, error) {
+	// Testing knob to inject errors.
+	if interceptor := s.cfg.TestingKnobs.AllocatorCheckRangeInterceptor; interceptor != nil {
+		if err := interceptor(); err != nil {
+			return allocatorimpl.AllocatorNoop, roachpb.ReplicationTarget{}, tracingpb.Recording{}, err
+		}
+	}
+
 	var spanOptions []tracing.SpanOption
 	if collectTraces {
 		spanOptions = append(spanOptions, tracing.WithRecording(tracingpb.RecordingVerbose))

pkg/kv/kvserver/testing_knobs.go

@@ -483,6 +483,12 @@ type StoreTestingKnobs struct {
 	// EnqueueReplicaInterceptor intercepts calls to `store.Enqueue()`.
 	EnqueueReplicaInterceptor func(queueName string, replica *Replica)
 
+	// AllocatorCheckRangeInterceptor intercepts calls to
+	// `Store.AllocatorCheckRange()` and can be used to inject errors for testing.
+	// If returns non-nil error, the error is returned instead of calling the
+	// actual AllocatorCheckRange logic.
+	AllocatorCheckRangeInterceptor func() error
+
 	// GlobalMVCCRangeTombstone will write a global MVCC range tombstone across
 	// the entire user keyspace during cluster bootstrapping. This will be written
 	// below all other data, and thus won't affect query results, but it does

pkg/server/storage_api/decommission_test.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -1050,3 +1051,51 @@ func checkRangeCheckResult(
 	}
 	passed = true
 }
+
+// TestDecommissionPreCheckRetryThrottledStores tests that the decommission
+// pre-check retries when it encounters transient throttled errors. Regression
+// test for #156849.
+func TestDecommissionPreCheckRetryThrottledStores(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+
+	var returnedError atomic.Bool
+	returnError := func() error {
+		if returnedError.CompareAndSwap(false, true) {
+			return errors.New("injected error")
+		}
+		return nil
+	}
+
+	tc := serverutils.StartCluster(t, 4, base.TestClusterArgs{
+		ServerArgs: base.TestServerArgs{
+			DefaultTestTenant: base.TestIsSpecificToStorageLayerAndNeedsASystemTenant,
+			Knobs: base.TestingKnobs{
+				Store: &kvserver.StoreTestingKnobs{
+					AllocatorCheckRangeInterceptor: returnError,
+				},
+			},
+		},
+		ReplicationMode: base.ReplicationManual,
+	})
+	defer tc.Stopper().Stop(ctx)
+
+	scratchKey := tc.ScratchRange(t)
+
+	// Add replicas to the scratch range on nodes 2 and 3.
+	scratchDesc := tc.AddVotersOrFatal(t, scratchKey, tc.Target(1), tc.Target(2))
+	require.Len(t, scratchDesc.InternalReplicas, 3)
+
+	// Perform decommission pre-check on node 3. AllocatorCheckRange will fail
+	// once due to the injected error, then succeed after retries.
+	decommissioningNodeIDs := []roachpb.NodeID{tc.Server(2).NodeID()}
+	result, err := tc.Server(0).DecommissionPreCheck(ctx, decommissioningNodeIDs,
+		true /* strictReadiness */, false /* collectTraces */, 0 /* maxErrors */)
+
+	require.NoError(t, err)
+	require.Greater(t, result.RangesChecked, 0)
+	require.Empty(t, result.RangesNotReady)
+	require.True(t, returnedError.Load())
+}