kvclient: enable savepoint rollbacks with buffered writes

This patch correctly handles savepoint rollbacks when buffered writes
are enabled. Whenever a savepoint is rolled back, any writes that belong
to the savepoint are removed from the buffer.

Closes #139059

Release note: None

Author: arulajmani

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"encoding/binary"
 	"slices"
+	"sort"
 	"unsafe"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
@@ -423,7 +424,36 @@ func (twb *txnWriteBuffer) epochBumpedLocked() {}
 func (twb *txnWriteBuffer) createSavepointLocked(context.Context, *savepoint) {}
 
 // rollbackToSavepointLocked is part of the txnInterceptor interface.
-func (twb *txnWriteBuffer) rollbackToSavepointLocked(ctx context.Context, s savepoint) {}
+func (twb *txnWriteBuffer) rollbackToSavepointLocked(ctx context.Context, s savepoint) {
+	toDelete := make([]*bufferedWrite, 0)
+	it := twb.buffer.MakeIter()
+	for it.First(); it.Valid(); it.Next() {
+		bufferedVals := it.Cur().vals
+		// NB: the savepoint is being rolled back to s.seqNum (inclusive). So,
+		// idx is the index of the first value that is considered rolled back.
+		idx := sort.Search(len(bufferedVals), func(i int) bool {
+			return bufferedVals[i].seq >= s.seqNum
+		})
+		if idx == len(bufferedVals) {
+			// No writes are being rolled back.
+			continue
+		}
+		// Update size bookkeeping for the values we're rolling back.
+		for i := idx; i < len(bufferedVals); i++ {
+			twb.bufferSize -= bufferedVals[i].size()
+		}
+		// Rollback writes by truncating the buffered values.
+		it.Cur().vals = bufferedVals[:idx]
+		if len(it.Cur().vals) == 0 {
+			// All writes have been rolled back; we should remove this key from
+			// the buffer entirely.
+			toDelete = append(toDelete, it.Cur())
+		}
+	}
+	for _, bw := range toDelete {
+		twb.removeFromBuffer(bw)
+	}
+}
 
 // closeLocked implements the txnInterceptor interface.
 func (twb *txnWriteBuffer) closeLocked() {}
@@ -1057,6 +1087,12 @@ func (twb *txnWriteBuffer) addToBuffer(key roachpb.Key, val roachpb.Value, seq e
 	}
 }
 
+// removeFromBuffer removes all buffered writes on a given key from the buffer.
+func (twb *txnWriteBuffer) removeFromBuffer(bw *bufferedWrite) {
+	twb.buffer.Delete(bw)
+	twb.bufferSize -= bw.size()
+}
+
 // flushBufferAndSendBatch flushes all buffered writes when sending the supplied
 // batch request to the KV layer. This is done by pre-pending the buffered
 // writes to the requests in the batch.
@@ -1092,11 +1128,11 @@ func (twb *txnWriteBuffer) flushBufferAndSendBatch(
 
 	reqs := make([]kvpb.RequestUnion, 0, numBuffered+len(ba.Requests))
 
-	// Next, remove the buffered writes from the buffer and collect them into requests.
+	// Next, remove the buffered writes from the buffer and collect them into
+	// requests.
 	for _, bw := range toFlushBufferedWrites {
 		reqs = append(reqs, bw.toRequest())
-		twb.buffer.Delete(&bw)
-		twb.bufferSize -= bw.size()
+		twb.removeFromBuffer(&bw)
 	}
 
 	// Layers below us expect that writes inside a batch are in sequence number

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer_test.go

@@ -1614,3 +1614,101 @@ func TestTxnWriteBufferDeleteRange(t *testing.T) {
 	require.Len(t, br.Responses, 1)
 	require.IsType(t, &kvpb.EndTxnResponse{}, br.Responses[0].GetInner())
 }
+
+// TestTxnWriteBufferRollbackToSavepoint tests the savepoint rollback logic.
+func TestTxnWriteBufferRollbackToSavepoint(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+	twb, mockSender := makeMockTxnWriteBuffer(cluster.MakeClusterSettings())
+
+	txn := makeTxnProto()
+	txn.Sequence = 10
+	keyA, keyB, keyC := roachpb.Key("a"), roachpb.Key("b"), roachpb.Key("c")
+	valA, valA2, valB := "valA", "valA2", "valB"
+
+	ba := &kvpb.BatchRequest{}
+	ba.Header = kvpb.Header{Txn: &txn}
+	putA := putArgs(keyA, valA, txn.Sequence)
+	txn.Sequence++
+	delC := delArgs(keyC, txn.Sequence)
+	ba.Add(putA)
+	ba.Add(delC)
+
+	numCalled := mockSender.NumCalled()
+	br, pErr := twb.SendLocked(ctx, ba)
+	require.Nil(t, pErr)
+	require.NotNil(t, br)
+	// All the requests should be buffered and not make it past the
+	// txnWriteBuffer. The response returned should be indistinguishable.
+	require.Equal(t, numCalled, mockSender.NumCalled())
+	require.Len(t, br.Responses, 2)
+	require.IsType(t, &kvpb.PutResponse{}, br.Responses[0].GetInner())
+	require.IsType(t, &kvpb.DeleteResponse{}, br.Responses[1].GetInner())
+	// Verify the
+	expBufferedWrites := []bufferedWrite{
+		makeBufferedWrite(keyA, makeBufferedValue("valA", 10)),
+		makeBufferedWrite(keyC, makeBufferedValue("", 11)),
+	}
+	require.Equal(t, expBufferedWrites, twb.testingBufferedWritesAsSlice())
+
+	// Create a savepoint. This is inclusive of the buffered Delete on keyC.
+	savepoint := &savepoint{seqNum: txn.Sequence}
+	twb.createSavepointLocked(ctx, savepoint)
+
+	// Add some new writes. A second write to keyA and a new one to keyB.
+	ba = &kvpb.BatchRequest{}
+	txn.Sequence++
+	putA2 := putArgs(keyA, valA2, txn.Sequence)
+	ba.Add(putA2)
+	txn.Sequence++
+	putB := putArgs(keyB, valB, txn.Sequence)
+	ba.Add(putB)
+
+	// All these writes should be buffered as well, with no KV requests sent yet.
+	numCalled = mockSender.NumCalled()
+	br, pErr = twb.SendLocked(ctx, ba)
+	require.Nil(t, pErr)
+	require.NotNil(t, br)
+
+	require.Equal(t, numCalled, mockSender.NumCalled())
+	require.Len(t, br.Responses, 2)
+	require.IsType(t, &kvpb.PutResponse{}, br.Responses[0].GetInner())
+	require.IsType(t, &kvpb.PutResponse{}, br.Responses[1].GetInner())
+	// Verify the state of the write buffer.
+	expBufferedWrites = []bufferedWrite{
+		makeBufferedWrite(keyA, makeBufferedValue("valA", 10), makeBufferedValue("valA2", 12)),
+		makeBufferedWrite(keyB, makeBufferedValue("valB", 13)),
+		makeBufferedWrite(keyC, makeBufferedValue("", 11)),
+	}
+	require.Equal(t, expBufferedWrites, twb.testingBufferedWritesAsSlice())
+
+	// Now, Rollback to the savepoint. This should leave just one write in the
+	// buffer, that on keyA at seqnum 10.
+	twb.rollbackToSavepointLocked(ctx, *savepoint)
+	expBufferedWrites = []bufferedWrite{
+		makeBufferedWrite(keyA, makeBufferedValue("valA", 10)),
+	}
+	require.Equal(t, expBufferedWrites, twb.testingBufferedWritesAsSlice())
+
+	// Commit the transaction.
+	ba = &kvpb.BatchRequest{}
+	ba.Header = kvpb.Header{Txn: &txn}
+	ba.Add(&kvpb.EndTxnRequest{Commit: true})
+
+	mockSender.MockSend(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+		require.Len(t, ba.Requests, 2)
+		require.IsType(t, &kvpb.PutRequest{}, ba.Requests[0].GetInner())
+		require.IsType(t, &kvpb.EndTxnRequest{}, ba.Requests[1].GetInner())
+
+		br = ba.CreateReply()
+		br.Txn = ba.Txn
+		return br, nil
+	})
+
+	br, pErr = twb.SendLocked(ctx, ba)
+	require.Nil(t, pErr)
+	require.NotNil(t, br)
+	require.Len(t, br.Responses, 1)
+	require.IsType(t, &kvpb.EndTxnResponse{}, br.Responses[0].GetInner())
+}

pkg/sql/logictest/testdata/logic_test/buffered_writes

@@ -78,7 +78,7 @@ INSERT INTO t1 VALUES (1,1)
 statement ok
 COMMIT
 
-query II
+query II rowsort
 SELECT * FROM t1
 ----
 1 1
@@ -122,7 +122,7 @@ INSERT INTO t2 VALUES (1);
 statement ok
 COMMIT;
 
-query I
+query I rowsort
 SELECT * FROM t2
 ----
 1
@@ -151,7 +151,103 @@ DELETE FROM t3 WHERE k < 10 AND k > 0
 statement ok
 COMMIT
 
-query I
+query I rowsort
 SELECT count(*) from t3
 ----
 0
+
+# Test savepoints, and in particular savepoint rollbacks, with buffered writes. 
+# We test both intermediate selects after rollbacks and the final state
+# the transaction has been committed.
+subtest savepoint_rollbacks
+
+# First, create a new table with a secondary index on it. That way, the DELETE
+# statements below will not use DeleteRange requets which cause the buffer to
+# be flushed.
+statement ok
+CREATE TABLE t4 (k INT PRIMARY KEY, v INT)
+
+statement ok
+CREATE INDEX idx ON t4 (v)
+
+statement ok
+BEGIN;
+INSERT INTO t4 VALUES(1, 100), (2, 200), (3, 300);
+SAVEPOINT s1;
+INSERT INTO t4 VALUES(4, 400), (5, 500), (6, 600)
+
+query II rowsort
+SELECT * FROM t4
+----
+1  100
+2  200
+3  300
+4  400
+5  500
+6  600
+
+statement ok
+SAVEPOINT s2; 
+INSERT INTO t4 VALUES(7, 700), (8, 800), (9, 900)
+
+query II rowsort
+SELECT * FROM t4
+----
+1  100
+2  200
+3  300
+4  400
+5  500
+6  600
+7  700
+8  800
+9  900
+
+# Throw in some Deletes.
+statement ok
+DELETE FROM t4 WHERE k = 1;
+DELETE FROM t4 WHERE k = 2;
+DELETE FROM t4 WHERE k = 3;
+
+query II rowsort
+SELECT * FROM t4
+----
+4  400
+5  500
+6  600
+7  700
+8  800
+9  900
+
+statement ok
+ROLLBACK TO SAVEPOINT s2
+
+query II rowsort
+SELECT * FROM t4
+----
+1  100
+2  200
+3  300
+4  400
+5  500
+6  600
+
+statement ok
+ROLLBACK TO SAVEPOINT s1;
+
+query II rowsort
+SELECT * FROM t4
+----
+1  100
+2  200
+3  300
+
+statement ok
+COMMIT
+
+query II rowsort
+SELECT * FROM t4
+----
+1  100
+2  200
+3  300