kvcoord: avoid unnecessary flushing in txnWriteBuffer

Previously, if the buffer was being flushed midway through a
transaction, we were flushing the entire buffer to ensure that future
savepoint rollbacks would result in the correct final value being
committed.

This was pessimistic, since many transactions don't use explicit
savepoints.

Here, we track whether an explicit savepoint has been created and only
flush those values required to rollback the earliest active savepoint.

Epic: none
Release note: None

Author: stevendanna

pkg/kv/kvclient/kvcoord/txn_coord_sender.go

@@ -217,6 +217,10 @@ type txnInterceptor interface {
 	// createSavepointLocked().
 	rollbackToSavepointLocked(context.Context, savepoint)
 
+	// releaseSavepointLocked is called when a savepoint is being
+	// released.
+	releaseSavepointLocked(context.Context, *savepoint)
+
 	// closeLocked closes the interceptor. It is called when the TxnCoordSender
 	// shuts down due to either a txn commit or a txn abort. The method will
 	// be called exactly once from cleanupTxnLocked.

pkg/kv/kvclient/kvcoord/txn_coord_sender_savepoints.go

@@ -187,7 +187,15 @@ func (tc *TxnCoordSender) ReleaseSavepoint(ctx context.Context, s kv.SavepointTo
 	}
 
 	sp := s.(*savepoint)
-	return tc.checkSavepointLocked(sp, "release")
+	if err := tc.checkSavepointLocked(sp, "release"); err != nil {
+		return err
+	}
+
+	for _, reqInt := range tc.interceptorStack {
+		reqInt.releaseSavepointLocked(ctx, sp)
+	}
+
+	return nil
 }
 
 // CanUseSavepoint is part of the kv.TxnSender interface.

pkg/kv/kvclient/kvcoord/txn_interceptor_committer.go

@@ -626,6 +626,9 @@ func (tc *txnCommitter) epochBumpedLocked() {}
 // createSavepointLocked is part of the txnInterceptor interface.
 func (*txnCommitter) createSavepointLocked(context.Context, *savepoint) {}
 
+// releaseSavepointLocked is part of the txnInterceptor interface.
+func (*txnCommitter) releaseSavepointLocked(context.Context, *savepoint) {}
+
 // rollbackToSavepointLocked is part of the txnInterceptor interface.
 func (*txnCommitter) rollbackToSavepointLocked(context.Context, savepoint) {}
 

pkg/kv/kvclient/kvcoord/txn_interceptor_heartbeater.go

@@ -302,6 +302,9 @@ func (h *txnHeartbeater) epochBumpedLocked() {}
 // createSavepointLocked is part of the txnInterceptor interface.
 func (*txnHeartbeater) createSavepointLocked(context.Context, *savepoint) {}
 
+// releaseSavepointLocked is part of the txnInterceptor interface.
+func (*txnHeartbeater) releaseSavepointLocked(context.Context, *savepoint) {}
+
 // rollbackToSavepointLocked is part of the txnInterceptor interface.
 func (*txnHeartbeater) rollbackToSavepointLocked(context.Context, savepoint) {}
 

pkg/kv/kvclient/kvcoord/txn_interceptor_metric_recorder.go

@@ -78,6 +78,9 @@ func (*txnMetricRecorder) epochBumpedLocked() {}
 // createSavepointLocked is part of the txnInterceptor interface.
 func (*txnMetricRecorder) createSavepointLocked(context.Context, *savepoint) {}
 
+// releaseSavepointLocked is part of the txnInterceptor interface.
+func (*txnMetricRecorder) releaseSavepointLocked(context.Context, *savepoint) {}
+
 // rollbackToSavepointLocked is part of the txnInterceptor interface.
 func (*txnMetricRecorder) rollbackToSavepointLocked(context.Context, savepoint) {}
 

pkg/kv/kvclient/kvcoord/txn_interceptor_pipeliner.go

@@ -984,6 +984,9 @@ func (tp *txnPipeliner) epochBumpedLocked() {
 // createSavepointLocked is part of the txnInterceptor interface.
 func (tp *txnPipeliner) createSavepointLocked(context.Context, *savepoint) {}
 
+// releaseSavepointLocked is part of the txnInterceptor interface.
+func (tp *txnPipeliner) releaseSavepointLocked(context.Context, *savepoint) {}
+
 // rollbackToSavepointLocked is part of the txnInterceptor interface.
 func (tp *txnPipeliner) rollbackToSavepointLocked(ctx context.Context, s savepoint) {
 	// Move all the writes in txnPipeliner that are not in the savepoint to the

pkg/kv/kvclient/kvcoord/txn_interceptor_seq_num_allocator.go

@@ -219,6 +219,9 @@ func (s *txnSeqNumAllocator) createSavepointLocked(ctx context.Context, sp *save
 	sp.seqNum = s.writeSeq
 }
 
+// releaseSavepointLocked is part of the txnInterceptor interface.
+func (*txnSeqNumAllocator) releaseSavepointLocked(context.Context, *savepoint) {}
+
 // rollbackToSavepointLocked is part of the txnInterceptor interface.
 func (s *txnSeqNumAllocator) rollbackToSavepointLocked(context.Context, savepoint) {
 	// Nothing to restore. The seq nums keep increasing. The TxnCoordSender has

pkg/kv/kvclient/kvcoord/txn_interceptor_span_refresher.go

@@ -764,6 +764,9 @@ func (sr *txnSpanRefresher) createSavepointLocked(ctx context.Context, s *savepo
 	s.refreshInvalid = sr.refreshInvalid
 }
 
+// releaseSavepointLocked is part of the txnInterceptor interface.
+func (sr *txnSpanRefresher) releaseSavepointLocked(context.Context, *savepoint) {}
+
 // rollbackToSavepointLocked is part of the txnInterceptor interface.
 func (sr *txnSpanRefresher) rollbackToSavepointLocked(ctx context.Context, s savepoint) {
 	if !KeepRefreshSpansOnSavepointRollback.Get(&sr.st.SV) {

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -177,6 +177,12 @@ type txnWriteBuffer struct {
 	// sees the next BatchRequest.
 	flushOnNextBatch bool
 
+	// firstExplicitSavepointSeq tracks the lowest explicit savepoint that hasn't
+	// been released or rolled back. If this savepoint is non-zero, then a
+	// mid-transaction flush must flush all revisions required to roll back this
+	// (or a later) savepoint.
+	firstExplicitSavepointSeq enginepb.TxnSeq
+
 	buffer        btree
 	bufferIDAlloc uint64
 	bufferSize    int64
@@ -605,6 +611,9 @@ func (twb *txnWriteBuffer) importLeafFinalState(context.Context, *roachpb.LeafTx
 
 // epochBumpedLocked implements the txnInterceptor interface.
 func (twb *txnWriteBuffer) epochBumpedLocked() {
+	// Sequence numbers are reset on epoch bumps so any retained savepoint is
+	// wrong.
+	twb.firstExplicitSavepointSeq = 0
 	twb.resetBuffer()
 }
 
@@ -613,8 +622,26 @@ func (twb *txnWriteBuffer) resetBuffer() {
 	twb.bufferSize = 0
 }
 
+func (twb *txnWriteBuffer) hasActiveSavepoint() bool {
+	return twb.firstExplicitSavepointSeq != enginepb.TxnSeq(0)
+}
+
 // createSavepointLocked is part of the txnInterceptor interface.
-func (twb *txnWriteBuffer) createSavepointLocked(context.Context, *savepoint) {}
+func (twb *txnWriteBuffer) createSavepointLocked(ctx context.Context, sp *savepoint) {
+	assertTrue(twb.firstExplicitSavepointSeq <= sp.seqNum,
+		"sequence number in created savepoint lower than retained savepoint")
+
+	if twb.firstExplicitSavepointSeq == enginepb.TxnSeq(0) {
+		twb.firstExplicitSavepointSeq = sp.seqNum
+	}
+}
+
+// rollbackToSavepointLocked is part of the txnInterceptor interface.
+func (twb *txnWriteBuffer) releaseSavepointLocked(ctx context.Context, sp *savepoint) {
+	if twb.firstExplicitSavepointSeq == sp.seqNum {
+		twb.firstExplicitSavepointSeq = 0
+	}
+}
 
 // rollbackToSavepointLocked is part of the txnInterceptor interface.
 func (twb *txnWriteBuffer) rollbackToSavepointLocked(ctx context.Context, s savepoint) {
@@ -647,6 +674,9 @@ func (twb *txnWriteBuffer) rollbackToSavepointLocked(ctx context.Context, s save
 	for _, bw := range toDelete {
 		twb.removeFromBuffer(bw)
 	}
+	if twb.firstExplicitSavepointSeq == s.seqNum {
+		twb.firstExplicitSavepointSeq = 0
+	}
 }
 
 // closeLocked implements the txnInterceptor interface.
@@ -1370,6 +1400,8 @@ func (twb *txnWriteBuffer) flushBufferAndSendBatch(
 		twb.txnMetrics.TxnWriteBufferDisabledAfterBuffering.Inc(1)
 	}
 
+	midTxnFlushWithExplicitSavepoint := !hasEndTxn && twb.hasActiveSavepoint()
+
 	// Flush all buffered writes by pre-pending them to the requests being sent
 	// in the batch.
 	//
@@ -1379,8 +1411,8 @@ func (twb *txnWriteBuffer) flushBufferAndSendBatch(
 	it := twb.buffer.MakeIter()
 	numRevisionsBuffered := 0
 	for it.First(); it.Valid(); it.Next() {
-		if !hasEndTxn {
-			revs := it.Cur().toAllRevisionRequests()
+		if midTxnFlushWithExplicitSavepoint {
+			revs := it.Cur().toAllRevisionRequests(twb.firstExplicitSavepointSeq)
 			numRevisionsBuffered += len(revs)
 			reqs = append(reqs, revs...)
 		} else {
@@ -1589,13 +1621,26 @@ func (bw *bufferedWrite) toRequest() kvpb.RequestUnion {
 }
 
 // toAllRevisionRequests returns requests for all revisions of the buffered
-// writes for the key. When the buffer is flushed before the end of a
-// transaction, all revisions must be written to storage to ensure that a future
-// savepoint rollback is properly handled.
-func (bw *bufferedWrite) toAllRevisionRequests() []kvpb.RequestUnion {
+// writes that need to be flushed given the minimum sequence number.
+//
+// When the buffer is flushed before the end of a transaction, previous
+// revisions must be written to storage to ensure that a future savepoint
+// rollback is properly handled.
+//
+// The given sequence number is the smallest sequence number associated with an
+// active savepoint.
+//
+// A write below the the minimum sequence number can be elided if there is a
+// subsequent write also below the minimum sequence number.
+func (bw *bufferedWrite) toAllRevisionRequests(minSeq enginepb.TxnSeq) []kvpb.RequestUnion {
 	rus := make([]kvpb.RequestUnion, 0, len(bw.vals))
-	for _, val := range bw.vals {
-		rus = append(rus, val.toRequestUnion(bw.key, bw.exclusionExpectedSinceTimestamp()))
+	maxIdx := len(bw.vals) - 1
+	for i, val := range bw.vals {
+		nextWriteLessThanSeq := (i+1 < maxIdx) && (bw.vals[i+1].seq < minSeq)
+		canElideRevision := val.seq < minSeq && nextWriteLessThanSeq
+		if !canElideRevision {
+			rus = append(rus, val.toRequestUnion(bw.key, bw.exclusionExpectedSinceTimestamp()))
+		}
 	}
 	return rus
 }

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer_test.go

@@ -1808,6 +1808,133 @@ func TestTxnWriteBufferRollbackToSavepoint(t *testing.T) {
 	require.IsType(t, &kvpb.EndTxnResponse{}, br.Responses[0].GetInner())
 }
 
+// TestTxnWriteBufferRollbackToSavepointMidTxn tests the savepoint rollback
+// logic in the presence of explicit savepoints.
+func TestTxnWriteBufferRollbackToSavepointMidTxn(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+
+	sendPut := func(t *testing.T, twb *txnWriteBuffer, mockSender *mockLockedSender, txn *roachpb.Transaction) {
+		txn.Sequence++
+		keyA := roachpb.Key("a")
+		valA := fmt.Sprintf("valA@%d", txn.Sequence)
+		putA := putArgs(keyA, valA, txn.Sequence)
+
+		ba := &kvpb.BatchRequest{}
+		ba.Header = kvpb.Header{Txn: txn}
+		ba.Add(putA)
+
+		mockSender.MockSend(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+			br := ba.CreateReply()
+			br.Txn = ba.Txn
+			return br, nil
+		})
+
+		numCalled := mockSender.NumCalled()
+		br, pErr := twb.SendLocked(ctx, ba)
+		require.Nil(t, pErr)
+		require.NotNil(t, br)
+		require.Equal(t, numCalled, mockSender.NumCalled())
+	}
+
+	delRangeBatch := func(txn *roachpb.Transaction) *kvpb.BatchRequest {
+		txn.Sequence++
+		keyB := roachpb.Key("b")
+		keyC := roachpb.Key("c")
+		delRangeReq := delRangeArgs(keyB, keyC, txn.Sequence)
+
+		ba := &kvpb.BatchRequest{}
+		ba.Header = kvpb.Header{Txn: txn}
+		ba.Add(delRangeReq)
+		return ba
+	}
+
+	savepoint := func(twb *txnWriteBuffer, txn *roachpb.Transaction) *savepoint {
+		txn.Sequence++
+		savepoint := &savepoint{seqNum: txn.Sequence}
+		twb.createSavepointLocked(ctx, savepoint)
+		return savepoint
+	}
+
+	t.Run("flush with no savepoint sends latest", func(t *testing.T) {
+		twb, mockSender := makeMockTxnWriteBuffer(cluster.MakeClusterSettings())
+		txn := makeTxnProto()
+		// Send 4 requests to the buffer
+		sendPut(t, &twb, mockSender, &txn)
+		sendPut(t, &twb, mockSender, &txn)
+		sendPut(t, &twb, mockSender, &txn)
+		sendPut(t, &twb, mockSender, &txn)
+		ba := delRangeBatch(&txn)
+
+		// Expect 1 Put and 1 DelRange.
+		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.DeleteRangeRequest{}, 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)
+	})
+
+	t.Run("flush with savepoint still elides unnecessary writes under savepoint", func(t *testing.T) {
+		twb, mockSender := makeMockTxnWriteBuffer(cluster.MakeClusterSettings())
+		txn := makeTxnProto()
+		sendPut(t, &twb, mockSender, &txn) // should be elided
+		sendPut(t, &twb, mockSender, &txn)
+		_ = savepoint(&twb, &txn)
+		sendPut(t, &twb, mockSender, &txn)
+		sendPut(t, &twb, mockSender, &txn)
+		ba := delRangeBatch(&txn)
+
+		// Expect 3 Put and 1 DelRange.
+		mockSender.MockSend(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+			require.Len(t, ba.Requests, 4)
+			require.IsType(t, &kvpb.PutRequest{}, ba.Requests[0].GetInner())
+			require.IsType(t, &kvpb.PutRequest{}, ba.Requests[1].GetInner())
+			require.IsType(t, &kvpb.PutRequest{}, ba.Requests[2].GetInner())
+			require.IsType(t, &kvpb.DeleteRangeRequest{}, ba.Requests[3].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)
+	})
+
+	t.Run("flush after release of earliest savepoint only sends latest", func(t *testing.T) {
+		twb, mockSender := makeMockTxnWriteBuffer(cluster.MakeClusterSettings())
+		txn := makeTxnProto()
+		sendPut(t, &twb, mockSender, &txn)
+		sendPut(t, &twb, mockSender, &txn)
+		sendPut(t, &twb, mockSender, &txn)
+		sp := savepoint(&twb, &txn)
+		sendPut(t, &twb, mockSender, &txn)
+		twb.releaseSavepointLocked(ctx, sp)
+
+		ba := delRangeBatch(&txn)
+		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.DeleteRangeRequest{}, 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)
+	})
+}
+
 // TestTxnWriteBufferFlushesAfterDisabling verifies that the txnWriteBuffer
 // flushes on the next batch after it is disabled if it buffered any writes.
 func TestTxnWriteBufferFlushesAfterDisabling(t *testing.T) {