kvcoord: fix CPut handling for intra-batch conflicts

Previously, for condition evaluation part of CPuts, we would ignore
writes that were performed within the same KV batch. That was the case
since we performed the read from the buffer when applying the
transformation but would actually buffer writes after the KV server
response comes back, so previous writes within the same KV batch would
be invisible. This is now fixed by buffering the write when applying
the transformation.

Also avoid flushing any buffered writes when rolling back the txn.

Release note: None

Author: yuzefovich

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -131,10 +131,12 @@ func (twb *txnWriteBuffer) SendLocked(
 		return twb.wrapped.SendLocked(ctx, ba)
 	}
 
-	if _, ok := ba.GetArg(kvpb.EndTxn); ok {
-		// TODO(arul): should we only flush if the transaction is being committed?
-		// If the transaction is being rolled back, we shouldn't needlessly flush
-		// writes.
+	if etArg, ok := ba.GetArg(kvpb.EndTxn); ok {
+		if !etArg.(*kvpb.EndTxnRequest).Commit {
+			// We're performing a rollback, so there is no point in flushing
+			// anything.
+			return twb.wrapped.SendLocked(ctx, ba)
+		}
 		return twb.flushWithEndTxn(ctx, ba)
 	}
 
@@ -375,6 +377,9 @@ func (twb *txnWriteBuffer) applyTransformations(
 			// Send a locking Get request to the KV layer; we'll evaluate the
 			// condition locally based on the response.
 			baRemote.Requests = append(baRemote.Requests, getReqU)
+			// Buffer a Put under the optimistic assumption that the condition
+			// will be satisfied.
+			twb.addToBuffer(t.Key, t.Value, t.Sequence)
 
 		case *kvpb.PutRequest:
 			// If the MustAcquireExclusiveLock flag is set on the Put, then we need to
@@ -838,13 +843,14 @@ func (t transformation) toResp(
 			req.AllowIfDoesNotExist,
 		)
 		if condFailedErr != nil {
+			// TODO(yuzefovich): consider "poisoning" the txnWriteBuffer when we
+			// hit a condition failed error to avoid mistaken usages (e.g. an
+			// attempt to flush with the EndTxn request with Commit=true).
 			pErr := kvpb.NewErrorWithTxn(condFailedErr, txn)
 			pErr.SetErrorIndex(int32(t.index))
 			return kvpb.ResponseUnion{}, pErr
 		}
-		// The condition was satisfied; buffer a Put, and return a synthesized
-		// response.
-		twb.addToBuffer(req.Key, req.Value, req.Sequence)
+		// The condition was satisfied - return a synthesized response.
 		ru.MustSetInner(&kvpb.ConditionalPutResponse{})
 
 	case *kvpb.PutRequest:
@@ -897,12 +903,6 @@ func (t transformation) toResp(
 		panic("unimplemented")
 	}
 
-	// TODO(arul): in the future, when we'll evaluate CPuts locally, we'll have
-	// this function take in the result of the KVGet, save the CPut function
-	// locally on the transformation, and use these two things to evaluate the
-	// condition here, on the client. We'll then construct and return the
-	// appropriate response.
-
 	return ru, nil
 }
 

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer_test.go

@@ -1109,11 +1109,11 @@ func TestTxnWriteBufferDecomposesConditionalPuts(t *testing.T) {
 			require.IsType(t, &kvpb.ConditionFailedError{}, pErr.GoError())
 		}
 
-		// Lastly, commit the transaction. A put should only be flushed if the condition
-		// evaluated successfully.
+		// Lastly, commit or rollback the transaction. A put should only be
+		// flushed if the condition evaluated successfully.
 		ba = &kvpb.BatchRequest{}
 		ba.Header = kvpb.Header{Txn: &txn}
-		ba.Add(&kvpb.EndTxnRequest{Commit: true})
+		ba.Add(&kvpb.EndTxnRequest{Commit: condEvalSuccessful})
 
 		mockSender.MockSend(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
 			if condEvalSuccessful {

pkg/kv/kvclient/kvcoord/txn_test.go

@@ -2011,6 +2011,7 @@ func TestTxnBufferedWritesConditionalPuts(t *testing.T) {
 				if expErr {
 					require.Error(t, err)
 					require.IsType(t, &kvpb.ConditionFailedError{}, err)
+					return err
 				} else {
 					require.NoError(t, err)
 				}
@@ -2022,11 +2023,10 @@ func TestTxnBufferedWritesConditionalPuts(t *testing.T) {
 				}
 			})
 
-			if commit {
+			if commit && !expErr {
 				require.NoError(t, err)
 			} else {
 				require.Error(t, err)
-				testutils.IsError(err, "abort")
 			}
 
 			// Verify the values are visible only if the transaction commited

pkg/sql/logictest/testdata/logic_test/buffered_writes

@@ -82,3 +82,47 @@ query II
 SELECT * FROM t1
 ----
 1 1
+
+statement ok
+CREATE TABLE t2 (k INT PRIMARY KEY);
+
+statement ok
+BEGIN;
+
+statement error pgcode 23505 duplicate key value violates unique constraint "t2_pkey"
+INSERT INTO t2 VALUES (1), (1);
+
+statement ok
+ROLLBACK;
+
+statement ok
+BEGIN;
+
+statement ok
+INSERT INTO t2 VALUES (1);
+
+statement error pgcode 23505 duplicate key value violates unique constraint "t2_pkey"
+INSERT INTO t2 VALUES (1);
+
+statement ok
+ROLLBACK;
+
+statement ok
+BEGIN;
+
+statement ok
+INSERT INTO t2 VALUES (1);
+
+statement ok
+DELETE FROM t2 WHERE k = 1;
+
+statement ok
+INSERT INTO t2 VALUES (1);
+
+statement ok
+COMMIT;
+
+query I
+SELECT * FROM t2
+----
+1