kvclient: flush txnWriteBuffer if IncrementRequest is received

Previously we issued an error on IncrementRequest assuming this was
only used in (1) transactions that have already disabled the write
buffer because of DDL statements or (2) non-transactional requests.

However, it is also used by any test that sets
UseTransactionalDescIDGenerator to true. While we could call
txn.SetBufferedWritesEnabled from the transactional ID generator
itself, I opted to simply flush the buffer if we encounter an
IncrementRequest so that IncrementRequest does not have a hidden
requirement of disabling write buffering.

Epic: none

Release note: None

Author: stevendanna

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -217,31 +217,7 @@ func (twb *txnWriteBuffer) SendLocked(
 		return twb.flushBufferAndSendBatch(ctx, ba)
 	}
 
-	if _, ok := ba.GetArg(kvpb.DeleteRange); ok {
-		log.VEventf(ctx, 2, "DeleteRangeRequest forcing flush of write buffer")
-		// DeleteRange requests can delete an arbitrary number of keys over a
-		// given keyspan. We won't know the exact scope of the delete until
-		// we've scanned the keyspan, which must happen on the server. We've got
-		// a couple of options here:
-		// 1. We decompose the DeleteRange request into a (potentially locking)
-		// Scan followed by buffered point Deletes for each key in the scan's
-		// result.
-		// 2. We flush the buffer[1] and send the DeleteRange request to the KV
-		// layer.
-		//
-		// We choose option 2, as typically the number of keys deleted is large,
-		// and we may realize we're over budget after performing the initial
-		// scan of the keyspan. At that point, we'll have to flush the buffer
-		// anyway. Moreover, buffered writes are most impactful when a
-		// transaction is writing to a small number of keys. As such, it's fine
-		// to not optimize the DeleteRange case, as typically it results in a
-		// large writing transaction.
-		//
-		// [1] Technically, we only need to flush the overlapping portion of the
-		// buffer. However, for simplicity, the txnWriteBuffer doesn't support
-		// transactions with partially buffered writes and partially flushed
-		// writes. We could change this in the future if there's benefit to
-		// doing so.
+	if twb.batchRequiresFlush(ctx, ba) {
 		return twb.flushBufferAndSendBatch(ctx, ba)
 	}
 
@@ -295,6 +271,55 @@ func (twb *txnWriteBuffer) SendLocked(
 	return twb.mergeResponseWithTransformations(ctx, ts, br)
 }
 
+func (twb *txnWriteBuffer) batchRequiresFlush(ctx context.Context, ba *kvpb.BatchRequest) bool {
+	for _, ru := range ba.Requests {
+		req := ru.GetInner()
+		switch req.(type) {
+		case *kvpb.IncrementRequest:
+			// We don't typically see IncrementRequest in transactional batches that
+			// haven't already had write buffering disabled becuase of DDL statements.
+			//
+			// However, we do have at least a few users of the NewTransactionalGenerator
+			// in test code and builtins.
+			//
+			// We could handle this similar to how we handle ConditionalPut, but its
+			// not clear there is much value in that.
+			log.VEventf(ctx, 2, "%s forcing flush of write buffer", req.Method())
+			return true
+		case *kvpb.DeleteRangeRequest:
+			// DeleteRange requests can delete an arbitrary number of keys over a
+			// given keyspan. We won't know the exact scope of the delete until
+			// we've scanned the keyspan, which must happen on the server. We've got
+			// a couple of options here:
+			//
+			// 1. We decompose the DeleteRange request into a (potentially
+			//    locking) Scan followed by buffered point Deletes for each
+			//    key in the scan's result.
+			//
+			// 2. We flush the buffer[1] and send the DeleteRange request to
+			//    the KV layer.
+			//
+			// We choose option 2, as typically the number of keys deleted is large,
+			// and we may realize we're over budget after performing the initial
+			// scan of the keyspan. At that point, we'll have to flush the buffer
+			// anyway. Moreover, buffered writes are most impactful when a
+			// transaction is writing to a small number of keys. As such, it's fine
+			// to not optimize the DeleteRange case, as typically it results in a
+			// large writing transaction.
+			//
+			// [1] Technically, we only need to flush the overlapping portion of the
+			// buffer. However, for simplicity, the txnWriteBuffer doesn't support
+			// transactions with partially buffered writes and partially flushed
+			// writes. We could change this in the future if there's benefit to
+			// doing so.
+
+			log.VEventf(ctx, 2, "%s forcing flush of write buffer", req.Method())
+			return true
+		}
+	}
+	return false
+}
+
 // validateBatch returns an error if the batch is unsupported
 // by the txnWriteBuffer.
 func (twb *txnWriteBuffer) validateBatch(ba *kvpb.BatchRequest) error {
@@ -353,9 +378,9 @@ func (twb *txnWriteBuffer) validateRequests(ba *kvpb.BatchRequest) error {
 			}
 		case *kvpb.QueryLocksRequest, *kvpb.LeaseInfoRequest:
 		default:
-			// All other requests are unsupported. Note that we assume EndTxn and
-			// DeleteRange requests were handled explicitly before this method was
-			// called.
+			// All other requests are unsupported. Note that we assume that requests
+			// that should result in a buffer flush are handled explicitly before this
+			// method was called.
 			return unsupportedMethodError(t.Method())
 		}
 	}

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer_test.go

@@ -1508,118 +1508,164 @@ func TestTxnWriteBufferFlushesWhenOverBudget(t *testing.T) {
 	require.Equal(t, int64(1), twb.txnMetrics.TxnWriteBufferMemoryLimitExceeded.Count())
 }
 
-// TestTxnWriteBufferDeleteRange ensures that the txnWriteBuffer correctly
-// handles DeleteRange requests. In particular, whenever we see a batch with a
-// DeleteRange request, the write buffer is flushed and write buffering is
-// turned off for subsequent requests.
-func TestTxnWriteBufferDeleteRange(t *testing.T) {
+// TestTxnWriteBufferFlushesIfBatchRequiresFlushing ensures that the
+// txnWriteBuffer correctly handles requests that aren't currently
+// supported by the txnWriteBuffer by flushing the buffer before
+// processing the request.
+func TestTxnWriteBufferFlushesIfBatchRequiresFlushing(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 := "valA"
 
-	ba := &kvpb.BatchRequest{}
-	ba.Header = kvpb.Header{Txn: &txn}
-	putA := putArgs(keyA, valA, 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())
-	// Verify the Put was buffered correctly.
-	expBufferedWrites := []bufferedWrite{
-		makeBufferedWrite(keyA, makeBufferedValue("valA", 10)),
-		makeBufferedWrite(keyC, makeBufferedValue("", 10)),
+	type batchSendMock func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error)
+	type testCase struct {
+		name         string
+		ba           func(*kvpb.BatchRequest)
+		baSender     func(*testing.T) batchSendMock
+		validateResp func(*testing.T, *kvpb.BatchResponse, *kvpb.Error)
 	}
-	require.Equal(t, expBufferedWrites, twb.testingBufferedWritesAsSlice())
-
-	// Send a DeleteRange request. This should result in the entire buffer
-	// being flushed. Note that we're flushing the delete to key C as well, even
-	// though it doesn't overlap with the DeleteRange request.
-	ba = &kvpb.BatchRequest{}
-	ba.Header = kvpb.Header{Txn: &txn}
-	delRange := delRangeArgs(keyA, keyB, txn.Sequence)
-	ba.Add(delRange)
 
-	mockSender.MockSend(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
-		require.Len(t, ba.Requests, 3)
+	testCases := []testCase{
+		{
+			name: "DeleteRange",
+			ba: func(b *kvpb.BatchRequest) {
+				b.Add(delRangeArgs(keyA, keyB, b.Txn.Sequence))
+			},
+			baSender: func(t *testing.T) batchSendMock {
+				return func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+					require.Len(t, ba.Requests, 3)
+					require.IsType(t, &kvpb.PutRequest{}, ba.Requests[0].GetInner())
+					require.IsType(t, &kvpb.DeleteRequest{}, ba.Requests[1].GetInner())
+					require.IsType(t, &kvpb.DeleteRangeRequest{}, ba.Requests[2].GetInner())
+
+					br := ba.CreateReply()
+					br.Txn = ba.Txn
+					return br, nil
+				}
+			},
+			validateResp: func(t *testing.T, br *kvpb.BatchResponse, pErr *kvpb.Error) {
+				require.Nil(t, pErr)
+				require.NotNil(t, br)
+				require.Len(t, br.Responses, 1)
+				require.IsType(t, &kvpb.DeleteRangeResponse{}, br.Responses[0].GetInner())
+			},
+		},
+		{
+			name: "Increment",
+			ba: func(b *kvpb.BatchRequest) {
+				b.Add(&kvpb.IncrementRequest{
+					RequestHeader: kvpb.RequestHeader{Key: keyA},
+					Increment:     1,
+				})
+			},
+			baSender: func(t *testing.T) batchSendMock {
+				return func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+					require.Len(t, ba.Requests, 3)
+					require.IsType(t, &kvpb.PutRequest{}, ba.Requests[0].GetInner())
+					require.IsType(t, &kvpb.DeleteRequest{}, ba.Requests[1].GetInner())
+					require.IsType(t, &kvpb.IncrementRequest{}, ba.Requests[2].GetInner())
+
+					br := ba.CreateReply()
+					br.Txn = ba.Txn
+					return br, nil
+				}
+			},
+			validateResp: func(t *testing.T, br *kvpb.BatchResponse, pErr *kvpb.Error) {
+				require.Nil(t, pErr)
+				require.NotNil(t, br)
+				require.Len(t, br.Responses, 1)
+				require.IsType(t, &kvpb.IncrementResponse{}, br.Responses[0].GetInner())
+			},
+		},
+	}
 
-		require.IsType(t, &kvpb.PutRequest{}, ba.Requests[0].GetInner())
-		require.IsType(t, &kvpb.DeleteRequest{}, ba.Requests[1].GetInner())
-		require.IsType(t, &kvpb.DeleteRangeRequest{}, ba.Requests[2].GetInner())
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			twb, mockSender := makeMockTxnWriteBuffer(cluster.MakeClusterSettings())
 
-		br = ba.CreateReply()
-		br.Txn = ba.Txn
-		return br, nil
-	})
+			txn := makeTxnProto()
+			txn.Sequence = 10
+			ba := &kvpb.BatchRequest{}
+			ba.Header = kvpb.Header{Txn: &txn}
+			putA := putArgs(keyA, valA, txn.Sequence)
+			delC := delArgs(keyC, txn.Sequence)
+			ba.Add(putA)
+			ba.Add(delC)
 
-	br, pErr = twb.SendLocked(ctx, ba)
-	require.Nil(t, pErr)
-	require.NotNil(t, br)
-	// Even though we flushed some writes, it shouldn't make it back to the response.
-	require.Len(t, br.Responses, 1)
-	require.IsType(t, &kvpb.DeleteRangeResponse{}, br.Responses[0].GetInner())
-	require.Equal(t, int64(1), twb.txnMetrics.TxnWriteBufferDisabledAfterBuffering.Count())
+			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())
+			// Verify the Put was buffered correctly.
+			expBufferedWrites := []bufferedWrite{
+				makeBufferedWrite(keyA, makeBufferedValue("valA", 10)),
+				makeBufferedWrite(keyC, makeBufferedValue("", 10)),
+			}
+			require.Equal(t, expBufferedWrites, twb.testingBufferedWritesAsSlice())
 
-	// Ensure the buffer is empty at this point.
-	require.Equal(t, 0, len(twb.testingBufferedWritesAsSlice()))
+			// Send the batch that should require a flush
+			ba = &kvpb.BatchRequest{}
+			ba.Header = kvpb.Header{Txn: &txn}
+			tc.ba(ba)
+			mockSender.MockSend(tc.baSender(t))
+			br, pErr = twb.SendLocked(ctx, ba)
+			tc.validateResp(t, br, pErr)
+			// Ensure the buffer is empty at this point.
+			require.Equal(t, 0, len(twb.testingBufferedWritesAsSlice()))
+			require.Equal(t, int64(1), twb.txnMetrics.TxnWriteBufferDisabledAfterBuffering.Count())
 
-	// Subsequent batches should not be buffered.
-	ba = &kvpb.BatchRequest{}
-	ba.Header = kvpb.Header{Txn: &txn}
-	putC := putArgs(keyC, valA, txn.Sequence)
-	ba.Add(putC)
+			// Subsequent batches should not be buffered.
+			ba = &kvpb.BatchRequest{}
+			ba.Header = kvpb.Header{Txn: &txn}
+			putC := putArgs(keyC, valA, txn.Sequence)
+			ba.Add(putC)
 
-	mockSender.MockSend(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
-		require.Len(t, ba.Requests, 1)
+			mockSender.MockSend(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+				require.Len(t, ba.Requests, 1)
 
-		require.IsType(t, &kvpb.PutRequest{}, ba.Requests[0].GetInner())
+				require.IsType(t, &kvpb.PutRequest{}, ba.Requests[0].GetInner())
 
-		br = ba.CreateReply()
-		br.Txn = ba.Txn
-		return br, nil
-	})
+				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.PutResponse{}, br.Responses[0].GetInner())
+			br, pErr = twb.SendLocked(ctx, ba)
+			require.Nil(t, pErr)
+			require.NotNil(t, br)
+			require.Len(t, br.Responses, 1)
+			require.IsType(t, &kvpb.PutResponse{}, br.Responses[0].GetInner())
 
-	// Commit the transaction. We flushed the buffer already, and no subsequent
-	// writes were buffered, so the buffer should be empty. As such, no write
-	// requests should be added to the batch.
-	ba = &kvpb.BatchRequest{}
-	ba.Header = kvpb.Header{Txn: &txn}
-	ba.Add(&kvpb.EndTxnRequest{Commit: true})
+			// Commit the transaction. We flushed the buffer already, and no subsequent
+			// writes were buffered, so the buffer should be empty. As such, no write
+			// requests should be added to the batch.
+			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, 1)
-		require.IsType(t, &kvpb.EndTxnRequest{}, ba.Requests[0].GetInner())
+			mockSender.MockSend(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+				require.Len(t, ba.Requests, 1)
+				require.IsType(t, &kvpb.EndTxnRequest{}, ba.Requests[0].GetInner())
 
-		br = ba.CreateReply()
-		br.Txn = ba.Txn
-		return br, nil
-	})
+				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())
+			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())
+		})
+	}
 }
 
 // TestTxnWriteBufferRollbackToSavepoint tests the savepoint rollback logic.
@@ -1915,16 +1961,6 @@ func TestTxnWriteBufferBatchRequestValidation(t *testing.T) {
 				return b
 			},
 		},
-		{
-			name: "batch with Increment",
-			ba: func() *kvpb.BatchRequest {
-				b := &kvpb.BatchRequest{Header: kvpb.Header{Txn: &txn}}
-				b.Add(&kvpb.IncrementRequest{
-					RequestHeader: kvpb.RequestHeader{Key: keyA, Sequence: txn.Sequence},
-				})
-				return b
-			},
-		},
 		{
 			name: "batch with ReturnRawMVCCValues Scan",
 			ba: func() *kvpb.BatchRequest {