kvserver: deflake TestProcessSplitAfterRightHandSideHasBeenRemoved

The test had a race condition between lease transfer verification and
partition activation. After transferring the lease from n1 to n3, there
was a window where n1 could re-acquire the lease before the partition
was activated, causing the subsequent increment to hang waiting for a
slow proposal on the now-partitioned n1.

Fix this by using PinnedLeasesKnob to prevent n1 from re-acquiring the
lease after it's transferred to n3. The lease is pinned before the
transfer to close the race window entirely.

Additionally:
- Simplify the test to only run with leader leases since the test logic
  is not lease-type specific.
- Remove the now-unnecessary conditional for leader lease handling.
- Add skip.UnderDuressWithIssue to skip under stress/race.
- Send requests directly to the leaseholder's store (via TestSender)
  after the partition is activated, avoiding DistSender routing issues.

Fixes (on 26.1) #159676
Fixes #158295.

Author: tbg

pkg/kv/kvserver/client_raft_test.go

@@ -5327,9 +5327,13 @@ func TestAckWriteBeforeApplication(t *testing.T) {
 func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
+	skip.UnderDuressWithIssue(t, 158295)
 	ctx := context.Background()
 
-	testutils.RunValues(t, "lease-type", roachpb.TestingAllLeaseTypes(), func(t *testing.T, leaseType roachpb.LeaseType) {
+	// Only run with leader leases for simplicity. The test logic is not
+	// lease-type specific.
+	leaseType := roachpb.LeaseLeader
+	{
 		ctx = logtags.AddTag(ctx, "gotest", t.Name())
 		noopProposalFilter := kvserverbase.ReplicaProposalFilter(func(args kvserverbase.ProposalFilterArgs) *kvpb.Error {
 			return nil
@@ -5340,11 +5344,13 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 			return proposalFilter.Load().(kvserverbase.ReplicaProposalFilter)(args)
 		}
 
-		increment := func(t *testing.T, db *kv.DB, key roachpb.Key, by int64) {
+		// increment sends an increment request directly via the provided sender.
+		// We use a direct store sender (via TestSender) rather than going through
+		// DistSender to avoid routing issues when n1 is partitioned. See #158295.
+		increment := func(t *testing.T, sender kv.Sender, key roachpb.Key, by int64) {
 			t.Helper()
-			b := &kv.Batch{}
-			b.AddRawRequest(incrementArgs(key, by))
-			require.NoError(t, db.Run(ctx, b))
+			_, pErr := kv.SendWrapped(ctx, sender, incrementArgs(key, by))
+			require.NoError(t, pErr.GoError())
 		}
 		ensureNoTombstone := func(t *testing.T, store *kvserver.Store, rangeID roachpb.RangeID) {
 			t.Helper()
@@ -5403,14 +5409,19 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 		// split to succeed and the RHS to eventually also be on all 3 nodes.
 		setup := func(t *testing.T) (
 			tc *testcluster.TestCluster,
-			db *kv.DB,
+			leaseholderStore *kvserver.Store,
 			keyA, keyB roachpb.Key,
 			lhsID roachpb.RangeID,
 			lhsPartition *testClusterPartitionedRange,
 		) {
 			lisReg := listenerutil.NewListenerRegistry()
 			const numServers int = 3
 			stickyServerArgs := make(map[int]base.TestServerArgs)
+			// Create a shared PinnedLeasesKnob so we can prevent the partitioned
+			// node (n1) from re-acquiring the lease after we transfer it away. This
+			// closes a race window where n1 could get the lease back between the
+			// transfer and partition activation. See #158295.
+			pinnedLeases := kvserver.NewPinnedLeases()
 			for i := 0; i < numServers; i++ {
 				st := cluster.MakeTestingClusterSettings()
 				kvserver.OverrideDefaultLeaseType(ctx, &st.SV, leaseType)
@@ -5437,6 +5448,7 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 							// trigger by being caught up by a later snapshot. See #154313.
 							DisableRaftLogQueue:   true,
 							TestingProposalFilter: testingProposalFilter,
+							PinnedLeases:          pinnedLeases,
 						},
 					},
 					RaftConfig: base.RaftConfig{
@@ -5458,7 +5470,12 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 				})
 
 			tc.Stopper().AddCloser(stop.CloserFn(lisReg.Close))
-			db = tc.GetFirstStoreFromServer(t, 1).DB()
+			// Use the leaseholder's store (n3, server index 2) to send requests
+			// directly after the partition is activated, avoiding DistSender
+			// routing issues when n1 is partitioned.
+			leaseholderStore = tc.GetFirstStoreFromServer(t, 2)
+			// Use DB for operations before the partition is activated.
+			db := leaseholderStore.DB()
 
 			// Split off a non-system range so we don't have to account for node liveness
 			// traffic.
@@ -5473,22 +5490,30 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 			lhsPartition, err := setupPartitionedRange(tc, desc.RangeID,
 				0 /* replicaID */, 0 /* partitionedNode */, false /* activated */, kvtestutils.UnreliableRaftHandlerFuncs{})
 			require.NoError(t, err)
-			// Wait for all nodes to catch up.
-			increment(t, db, keyA, 5)
+			// Wait for all nodes to catch up. Use DB since partition is not active.
+			_, pErr := kv.SendWrapped(ctx, db.NonTransactionalSender(), incrementArgs(keyA, 5))
+			require.NoError(t, pErr.GoError())
 			tc.WaitForValues(t, keyA, []int64{5, 5, 5})
 
+			// Pin the lease to n3 (tc.Target(2)) so that n1 cannot re-acquire it
+			// after we transfer it away. This prevents a race where n1 could get
+			// the lease back before the partition is activated. See #158295.
+			pinnedLeases.PinLease(desc.RangeID, tc.Target(2).StoreID)
 			// Transfer the lease off of node 0.
 			tc.TransferRangeLeaseOrFatal(t, desc, tc.Target(2))
 
-			// Make sure everybody knows about that transfer.
-			increment(t, db, keyA, 1)
+			// Make sure everybody knows about that transfer. Use DB since
+			// partition is not active yet.
+			_, pErr = kv.SendWrapped(ctx, db.NonTransactionalSender(), incrementArgs(keyA, 1))
+			require.NoError(t, pErr.GoError())
 			tc.WaitForValues(t, keyA, []int64{6, 6, 6})
 			log.KvExec.Infof(ctx, "activating LHS partition: %s", lhsPartition)
 			lhsPartition.activate()
 
-			increment(t, db, keyA, 1)
+			// After partition activation, use the leaseholder's store directly.
+			increment(t, leaseholderStore.TestSender(), keyA, 1)
 			tc.WaitForValues(t, keyA, []int64{6, 7, 7})
-			return tc, db, keyA, keyB, lhsID, lhsPartition
+			return tc, leaseholderStore, keyA, keyB, lhsID, lhsPartition
 		}
 
 		// In this case we only have the LHS partitioned. The RHS will learn about its
@@ -5497,21 +5522,21 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 		// partition the RHS and ensure that the split does not clobber the RHS's hard
 		// state.
 		t.Run("(1) no RHS partition", func(t *testing.T) {
-			tc, db, keyA, keyB, _, lhsPartition := setup(t)
+			tc, leaseholderStore, keyA, keyB, _, lhsPartition := setup(t)
 
 			defer tc.Stopper().Stop(ctx)
 			tc.SplitRangeOrFatal(t, keyB)
 
 			// Write a value which we can observe to know when the split has been
 			// applied by the LHS.
-			increment(t, db, keyA, 1)
+			increment(t, leaseholderStore.TestSender(), keyA, 1)
 			tc.WaitForValues(t, keyA, []int64{6, 8, 8})
 
-			increment(t, db, keyB, 6)
+			increment(t, leaseholderStore.TestSender(), keyB, 6)
 			// Wait for all non-partitioned nodes to catch up.
 			tc.WaitForValues(t, keyB, []int64{0, 6, 6})
 
-			rhsInfo, err := getRangeInfo(ctx, db, keyB)
+			rhsInfo, err := getRangeInfo(ctx, leaseholderStore.DB(), keyB)
 			require.NoError(t, err)
 			rhsID := rhsInfo.Desc.RangeID
 			_, store0Exists := rhsInfo.Desc.GetReplicaDescriptor(1)
@@ -5553,21 +5578,21 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 		// This case is like the previous case except the store crashes after
 		// laying down a tombstone.
 		t.Run("(2) no RHS partition, with restart", func(t *testing.T) {
-			tc, db, keyA, keyB, _, lhsPartition := setup(t)
+			tc, leaseholderStore, keyA, keyB, _, lhsPartition := setup(t)
 			defer tc.Stopper().Stop(ctx)
 
 			tc.SplitRangeOrFatal(t, keyB)
 
 			// Write a value which we can observe to know when the split has been
 			// applied by the LHS.
-			increment(t, db, keyA, 1)
+			increment(t, leaseholderStore.TestSender(), keyA, 1)
 			tc.WaitForValues(t, keyA, []int64{6, 8, 8})
 
-			increment(t, db, keyB, 6)
+			increment(t, leaseholderStore.TestSender(), keyB, 6)
 			// Wait for all non-partitioned nodes to catch up.
 			tc.WaitForValues(t, keyB, []int64{0, 6, 6})
 
-			rhsInfo, err := getRangeInfo(ctx, db, keyB)
+			rhsInfo, err := getRangeInfo(ctx, leaseholderStore.DB(), keyB)
 			require.NoError(t, err)
 			rhsID := rhsInfo.Desc.RangeID
 			_, store0Exists := rhsInfo.Desc.GetReplicaDescriptor(1)
@@ -5602,7 +5627,7 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 			curB := int64(6)
 			for curB < 100 {
 				curB++
-				increment(t, db, keyB, 1)
+				increment(t, leaseholderStore.TestSender(), keyB, 1)
 				tc.WaitForValues(t, keyB, []int64{0, curB, curB})
 			}
 
@@ -5631,22 +5656,22 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 		// the split is processed. We partition the RHS's new replica ID before
 		// processing the split to ensure that the RHS doesn't get initialized.
 		t.Run("(3) initial replica RHS partition, no restart", func(t *testing.T) {
-			tc, db, keyA, keyB, _, lhsPartition := setup(t)
+			tc, leaseholderStore, keyA, keyB, _, lhsPartition := setup(t)
 			defer tc.Stopper().Stop(ctx)
 			var rhsPartition *testClusterPartitionedRange
 			partitionReplicaOnSplit(t, tc, keyB, lhsPartition, &rhsPartition)
 			tc.SplitRangeOrFatal(t, keyB)
 
 			// Write a value which we can observe to know when the split has been
 			// applied by the LHS.
-			increment(t, db, keyA, 1)
+			increment(t, leaseholderStore.TestSender(), keyA, 1)
 			tc.WaitForValues(t, keyA, []int64{6, 8, 8})
 
-			increment(t, db, keyB, 6)
+			increment(t, leaseholderStore.TestSender(), keyB, 6)
 			// Wait for all non-partitioned nodes to catch up.
 			tc.WaitForValues(t, keyB, []int64{0, 6, 6})
 
-			rhsInfo, err := getRangeInfo(ctx, db, keyB)
+			rhsInfo, err := getRangeInfo(ctx, leaseholderStore.DB(), keyB)
 			require.NoError(t, err)
 			rhsID := rhsInfo.Desc.RangeID
 			_, store0Exists := rhsInfo.Desc.GetReplicaDescriptor(1)
@@ -5698,35 +5723,32 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 		// about its higher replica ID the store crashes. However, it doesn't forget
 		// this ID, so the RHS replica is still not initialized by the split.
 		t.Run("(4) initial replica RHS partition, with restart", func(t *testing.T) {
-			tc, db, keyA, keyB, _, lhsPartition := setup(t)
+			tc, leaseholderStore, keyA, keyB, _, lhsPartition := setup(t)
 			defer tc.Stopper().Stop(ctx)
 			var rhsPartition *testClusterPartitionedRange
 
 			partitionReplicaOnSplit(t, tc, keyB, lhsPartition, &rhsPartition)
 			tc.SplitRangeOrFatal(t, keyB)
 
-			if leaseType == roachpb.LeaseLeader {
-				// Since both LHS and RHS use the same store, let's remove the store
-				// partition from `rhsPartition` and keep it only in `lhsPartition`.
-				// This will help us control the store partition using one object.
-				// TODO(ibrahim): Make the test pass when both LHS and RHS ranges are
-				// recovered at the same time.
-				store, err := tc.Servers[0].GetStores().(*kvserver.Stores).
-					GetStore(tc.Servers[0].GetFirstStoreID())
-				require.NoError(t, err)
-				rhsPartition.removeStore(store.StoreID())
-			}
+			// Since both LHS and RHS use the same store (leader leases), remove
+			// the store partition from `rhsPartition` and keep it only in
+			// `lhsPartition`. This lets us control the store partition using
+			// one object.
+			store, err := tc.Servers[0].GetStores().(*kvserver.Stores).
+				GetStore(tc.Servers[0].GetFirstStoreID())
+			require.NoError(t, err)
+			rhsPartition.removeStore(store.StoreID())
 
 			// Write a value which we can observe to know when the split has been
 			// applied by the LHS.
-			increment(t, db, keyA, 1)
+			increment(t, leaseholderStore.TestSender(), keyA, 1)
 			tc.WaitForValues(t, keyA, []int64{6, 8, 8})
 
-			increment(t, db, keyB, 6)
+			increment(t, leaseholderStore.TestSender(), keyB, 6)
 			// Wait for all non-partitioned nodes to catch up.
 			tc.WaitForValues(t, keyB, []int64{0, 6, 6})
 
-			rhsInfo, err := getRangeInfo(ctx, db, keyB)
+			rhsInfo, err := getRangeInfo(ctx, leaseholderStore.DB(), keyB)
 			require.NoError(t, err)
 			rhsID := rhsInfo.Desc.RangeID
 			_, store0Exists := rhsInfo.Desc.GetReplicaDescriptor(1)
@@ -5765,7 +5787,7 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 			curB := int64(6)
 			for curB < 100 {
 				curB++
-				increment(t, db, keyB, 1)
+				increment(t, leaseholderStore.TestSender(), keyB, 1)
 				tc.WaitForValues(t, keyB, []int64{0, curB, curB})
 			}
 
@@ -5791,7 +5813,7 @@ func TestProcessSplitAfterRightHandSideHasBeenRemoved(t *testing.T) {
 			})
 			tc.WaitForValues(t, keyB, []int64{curB, curB, curB})
 		})
-	})
+	}
 }
 
 type noopRaftMessageResponseStream struct{}