Merge #144160

144160: kvcoord: avoid erroneous rollback of buffered write r=stevendanna a=stevendanna

Previously, if a transaction had only done non-locking reads and bufffered writes, a rollback would erroneously rollback the last issued buffered write because savepoint creation only incremented the sequence number of a lock had been acquired.

This case is never encountered in SQL-generated transactions because we never completely buffer a Put or Delete without having previously issued at least 1 locking read that we expect to return a lock.

But, we do hit these kind of transactions in KVNemesis.

Epic: none

Release note: None

Co-authored-by: Steven Danna <[email protected]>

Author: craig[bot]

pkg/kv/kvclient/kvcoord/txn_coord_sender_savepoints.go

@@ -85,7 +85,16 @@ func (tc *TxnCoordSender) CreateSavepoint(ctx context.Context) (kv.SavepointToke
 	// happened before the savepoint and those that happened after.
 	// TODO(nvanbenschoten): once #113765 is resolved, we should make this
 	// unconditional and push it into txnSeqNumAllocator.createSavepointLocked.
-	if tc.interceptorAlloc.txnPipeliner.hasAcquiredLocks() {
+	//
+	// We also increment the write sequence if any writes are buffered since
+	// a buffered write will not appears as a lock acquisition in the
+	// pipeliner. For SQL-generated transactions, this additional check is
+	// unnecessary since we won't ever have buffered writes wihout also at
+	// least 1 locked row. However, for other types of KV transactions such
+	// as those generated by KVNemesis we may have blind Put or Delete
+	// requests that would be erroneously rolled back if we didn't increment
+	// the sequence here.
+	if tc.interceptorAlloc.txnPipeliner.hasAcquiredLocks() || tc.interceptorAlloc.txnWriteBuffer.hasBufferedWrites() {
 		if err := tc.interceptorAlloc.txnSeqNumAllocator.stepWriteSeqLocked(ctx); err != nil {
 			return nil, err
 		}

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -1163,7 +1163,7 @@ func (twb *txnWriteBuffer) flushBufferAndSendBatch(
 
 	// Layers below us expect that writes inside a batch are in sequence number
 	// order but the iterator above returns data in key order. Here we re-sort it
-	// which is unfortunate but required we make a change to the pipeliner.
+	// which is unfortunate but required unless we make a change to the pipeliner.
 	slices.SortFunc(reqs, func(a kvpb.RequestUnion, b kvpb.RequestUnion) int {
 		aHeader := a.GetInner().Header()
 		bHeader := b.GetInner().Header()

pkg/kv/kvclient/kvcoord/txn_test.go

@@ -2046,3 +2046,39 @@ func TestTxnBufferedWritesConditionalPuts(t *testing.T) {
 		}
 	})
 }
+
+// TestTxnBufferedWritesRollbackToSavepointAllBuffered is a regression
+// test for a bug encountered during development where the sequence
+// number of a savepoint was not correctly advanced when all writes in
+// a transaction had been buffered.
+func TestTxnBufferedWritesRollbackToSavepointAllBuffered(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	s := createTestDB(t)
+	defer s.Stop()
+
+	ctx := context.Background()
+	value1 := []byte("value1")
+	testutils.RunTrueAndFalse(t, "buffered_write", func(t *testing.T, bufferedWritesEnabled bool) {
+		keyA := []byte(fmt.Sprintf("keyA-%v", bufferedWritesEnabled))
+
+		err := s.DB.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+			txn.SetBufferedWritesEnabled(bufferedWritesEnabled)
+			// Should not be rolled back because the savepoint is
+			// created after it.
+			if err := txn.Put(ctx, keyA, value1); err != nil {
+				return err
+			}
+			sp, err := txn.CreateSavepoint(ctx)
+			if err != nil {
+				return err
+			}
+			return txn.RollbackToSavepoint(ctx, sp)
+		})
+		require.NoError(t, err)
+		actualA, err := s.DB.Get(ctx, keyA)
+		require.NoError(t, err)
+		require.Equal(t, actualA.ValueBytes(), value1)
+	})
+
+}