kvcoord: unset AsyncConsensus on split batches

When batches are split downstream of the transaction pipeline
interceptor, they must unset AsyncConsensus to avoid cases where a read
misses an un-validated, pipelined write that failed.

Note that we don't do this when splitting requests _by range_, because
those splits won't separate a scan from an overlapping write.

This should be very rare since it requires a read/write batch that SQL
doesn't typically produce and both reads must be be locking requests so
that CanPipeline returns true for all requests in the batch. It seems
unlikely that such a batch is ever produced outside in the standard
configuration outside of KVNemesis.

Fixes #150304

Release note: None

Author: stevendanna

pkg/kv/kvclient/kvcoord/dist_sender.go

@@ -1147,23 +1147,31 @@ func (ds *DistSender) initAndVerifyBatch(ctx context.Context, ba *kvpb.BatchRequ
 // request.
 var errNo1PCTxn = kvpb.NewErrorf("cannot send 1PC txn to multiple ranges")
 
-// splitBatchAndCheckForRefreshSpans splits the batch according to the
-// canSplitET parameter and checks whether the batch can forward its
-// read timestamp. If the batch has its CanForwardReadTimestamp flag
-// set but is being split across multiple sub-batches then the flag in
-// the batch header is unset.
-func splitBatchAndCheckForRefreshSpans(
+// splitBatchAndCheckForIncompatibilities splits the batch according to the
+// canSplitET parameter and checks whether the batch has settings incompatible
+// with being split.
+//
+// - If the batch has its CanForwardReadTimestamp flag set but is being split
+// across multiple sub-batches then the flag in the batch header is unset.
+//
+// - If the batch has AsyncConsensus set to true but is being split, then the
+// flag in the batch header is unset.
+func splitBatchAndCheckForIncompatibilities(
 	ba *kvpb.BatchRequest, canSplitET bool,
 ) [][]kvpb.RequestUnion {
 	parts := ba.Split(canSplitET)
 
-	// If the batch is split and the header has its CanForwardReadTimestamp flag
-	// set then we much check whether any request would need to be refreshed in
-	// the event that the one of the partial batches was to forward its read
-	// timestamp during a server-side refresh. If any such request exists then
-	// we unset the CanForwardReadTimestamp flag.
 	if len(parts) > 1 {
+		// If the batch is split and the header has its CanForwardReadTimestamp flag
+		// set then we much check whether any request would need to be refreshed in
+		// the event that the one of the partial batches was to forward its read
+		// timestamp during a server-side refresh. If any such request exists then
+		// we unset the CanForwardReadTimestamp flag.
 		unsetCanForwardReadTimestampFlag(ba)
+		// If the batch is split and the header has AsyncConsensus, we
+		// unconditionally unset it because it may be unsafe if the write needs to
+		// be observed by different reads that have been split.
+		unsetAsyncConsensus(ba)
 	}
 
 	return parts
@@ -1191,6 +1199,17 @@ func unsetCanForwardReadTimestampFlag(ba *kvpb.BatchRequest) {
 	}
 }
 
+// unsetAsyncConsensus ensures that if a batch that has been split because of
+// incompatible request flags, AsyncConsensus is not set.
+//
+// NOTE(ssd): We could be "smarter" here and only unset this flag if we have a
+// pipeline-able write that intersects requests in two different splits. But,
+// batches that require splitting are uncommon and thus aren't worth unnecessary
+// complexity.
+func unsetAsyncConsensus(ba *kvpb.BatchRequest) {
+	ba.AsyncConsensus = false
+}
+
 // Send implements the batch.Sender interface. It subdivides the Batch
 // into batches admissible for sending (preventing certain illegal
 // mixtures of requests), executes each individual part (which may
@@ -1231,7 +1250,7 @@ func (ds *DistSender) Send(
 	if ba.Txn != nil && ba.Txn.Epoch > 0 && !require1PC {
 		splitET = true
 	}
-	parts := splitBatchAndCheckForRefreshSpans(ba, splitET)
+	parts := splitBatchAndCheckForIncompatibilities(ba, splitET)
 	var singleRplChunk [1]*kvpb.BatchResponse
 	rplChunks := singleRplChunk[:0:1]
 
@@ -1277,7 +1296,7 @@ func (ds *DistSender) Send(
 			} else if require1PC {
 				log.Fatalf(ctx, "required 1PC transaction cannot be split: %s", ba)
 			}
-			parts = splitBatchAndCheckForRefreshSpans(ba, true /* split ET */)
+			parts = splitBatchAndCheckForIncompatibilities(ba, true /* split ET */)
 			onePart = false
 			// Restart transaction of the last chunk as multiple parts with
 			// EndTxn in the last part.

pkg/kv/kvclient/kvcoord/dist_sender_server_test.go

@@ -4816,3 +4816,97 @@ func TestUnexpectedCommitOnTxnRecovery(t *testing.T) {
 	close(blockCh)
 	wg.Wait()
 }
+
+// TestBatchPutScanReverseScanWithFailedPutReplication is a regression test for
+// #150304 in which a batch containing a Put, Scan, and ReverseScan, failed to
+// read its own write during one of the Scans. This was the result of the batch
+// being split downstream of the txnPipeliner and thus the Put not being
+// verified before the overlapping scan was issued.
+func TestBatchPutScanReverseScanWithFailedPutReplication(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	keyA := roachpb.Key("a")
+	keyB := roachpb.Key("b")
+	keyC := roachpb.Key("c")
+
+	var (
+		targetTxnIDString atomic.Value
+		cmdID             atomic.Value
+	)
+	cmdID.Store(kvserverbase.CmdIDKey(""))
+	targetTxnIDString.Store("")
+	ctx := context.Background()
+	st := cluster.MakeTestingClusterSettings()
+
+	tc := testcluster.StartTestCluster(t, 3, base.TestClusterArgs{
+		ServerArgs: base.TestServerArgs{
+			Settings: st,
+			Knobs: base.TestingKnobs{
+				Store: &kvserver.StoreTestingKnobs{
+					TestingProposalFilter: func(fArgs kvserverbase.ProposalFilterArgs) *kvpb.Error {
+						if fArgs.Req.Header.Txn == nil ||
+							fArgs.Req.Header.Txn.ID.String() != targetTxnIDString.Load().(string) {
+							return nil // not our txn
+						}
+						putReq, ok := fArgs.Req.Requests[0].GetInner().(*kvpb.PutRequest)
+						// Only fail replication on the Put to keyB.
+						if ok && putReq.Key.Equal(keyB) {
+							t.Logf("will fail application for txn %s; req: %+v; raft cmdID: %x",
+								fArgs.Req.Header.Txn.ID.String(), putReq, fArgs.CmdID)
+							cmdID.Store(fArgs.CmdID)
+						}
+						return nil
+					},
+					TestingApplyCalledTwiceFilter: func(fArgs kvserverbase.ApplyFilterArgs) (int, *kvpb.Error) {
+						if fArgs.CmdID == cmdID.Load().(kvserverbase.CmdIDKey) {
+							t.Logf("failing application for raft cmdID: %x", cmdID)
+							return 0, kvpb.NewErrorf("test injected error")
+						}
+						return 0, nil
+					},
+				},
+			},
+		},
+	})
+	defer tc.Stopper().Stop(ctx)
+
+	db := tc.Server(0).DB()
+
+	err := db.Put(ctx, keyA, "initial_a")
+	require.NoError(t, err)
+
+	err = db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+		if txnID := targetTxnIDString.Load(); txnID == "" {
+			// Store the txnID for the testing knobs.
+			targetTxnIDString.Store(txn.ID().String())
+			t.Logf("txn ID is: %s", txn.ID())
+		}
+
+		// Create a batch that will be split into two batches because of scan
+		// directions. We've arranged for the put to fail application.
+		b := txn.NewBatch()
+		b.Put(keyB, "value_b")
+		b.ScanForUpdate(keyA, keyC, kvpb.GuaranteedDurability)
+		b.ReverseScanForShare(keyA, keyC, kvpb.GuaranteedDurability)
+
+		// If the batch isn't marked for async consensus, the batch will fail.
+		if err := txn.Run(ctx, b); err != nil {
+			return err
+		}
+
+		// If the batch doesn't fail, then it _must_ have both rows in the second
+		// scan.
+		require.Equal(t, 2, len(b.Results[2].Rows))
+		return nil
+	})
+
+	// The transaction should fail due to the replication failure.
+	require.Error(t, err)
+	require.ErrorContains(t, err, "test injected error")
+
+	// Verify that the Put didn't succeed by checking that keyB doesn't exist.
+	res, err := db.Get(ctx, keyB)
+	require.NoError(t, err)
+	require.False(t, res.Exists(), "keyB should not exist due to failed replication")
+}