Merge #143763

143763: kvclient: correctly handle DeleteRange requests in the txnWriteBuffer  r=arulajmani a=arulajmani

DeleteRange requests can delete an arbitrary number of keys over a
keyspan. We can't know the number or size of these until we perform
a scan, which must happen on the server. As a result, we've got two
options:
1. We could decompose the DeleteRange request into a (possibly locking)
scan, followed by buffered point Deletes from the returned result.
2. We could flush the buffer and send the DeleteRange request to KV.

We choose the latter, for its simplicity. DeleteRange requests typically
delete a large number of keys (which could cause us to go over budget)
and are only used by SQL in very specific circumstances (typically, for
tables without secondary indexes). For now, we choose simplicity even
though it pesimizes the case where a DeleteRange request only affects
a small number of keys. We could revisit this in the future.

Closes #143423

Release note: None

Co-authored-by: Arul Ajmani <[email protected]>

Author: craig[bot]

pkg/kv/kvclient/kvcoord/dist_sender.go

@@ -1088,7 +1088,7 @@ func (ds *DistSender) initAndVerifyBatch(ctx context.Context, ba *kvpb.BatchRequ
 				foundReverse = true
 
 			case *kvpb.QueryIntentRequest, *kvpb.EndTxnRequest,
-				*kvpb.GetRequest, *kvpb.ResolveIntentRequest, *kvpb.DeleteRequest:
+				*kvpb.GetRequest, *kvpb.ResolveIntentRequest, *kvpb.DeleteRequest, *kvpb.PutRequest:
 				// Accepted point requests that can be in batches with limit.
 
 			default:

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -181,6 +181,33 @@ func (twb *txnWriteBuffer) SendLocked(
 		return twb.flushBufferAndSendBatch(ctx, ba)
 	}
 
+	if _, ok := ba.GetArg(kvpb.DeleteRange); ok {
+		// 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.
+		return twb.flushBufferAndSendBatch(ctx, ba)
+	}
+
 	// Check if buffering writes from the supplied batch will run us over
 	// budget. If it will, we shouldn't buffer writes from the current batch,
 	// and flush the buffer.

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer_test.go

@@ -56,6 +56,12 @@ func delArgs(key roachpb.Key, seq enginepb.TxnSeq) *kvpb.DeleteRequest {
 	}
 }
 
+func delRangeArgs(key, endKey roachpb.Key, seq enginepb.TxnSeq) *kvpb.DeleteRangeRequest {
+	return &kvpb.DeleteRangeRequest{
+		RequestHeader: kvpb.RequestHeader{Key: key, EndKey: endKey, Sequence: seq},
+	}
+}
+
 func makeBufferedWrite(key roachpb.Key, vals ...bufferedValue) bufferedWrite {
 	return bufferedWrite{key: key, vals: vals}
 }
@@ -1495,3 +1501,116 @@ func TestTxnWriteBufferFlushesWhenOverBudget(t *testing.T) {
 	require.Len(t, br.Responses, 1)
 	require.IsType(t, &kvpb.EndTxnResponse{}, br.Responses[0].GetInner())
 }
+
+// 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) {
+	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)),
+	}
+	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)
+
+		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
+	})
+
+	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())
+
+	// Ensure the buffer is empty at this point.
+	require.Equal(t, 0, len(twb.testingBufferedWritesAsSlice()))
+
+	// 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)
+
+		require.IsType(t, &kvpb.PutRequest{}, ba.Requests[0].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.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})
+
+	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, 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

@@ -126,3 +126,32 @@ query I
 SELECT * FROM t2
 ----
 1
+
+# Ensure that DeleteRange requests work correctly with buffered writes. In
+# particular, a DeleteRange request results in a buffer flush.
+
+statement ok
+CREATE TABLE t3 (k INT PRIMARY KEY)
+
+statement ok
+INSERT INTO t3 VALUES (1)
+
+statement ok
+BEGIN
+
+statement ok
+INSERT INTO t3 VALUES (2)
+
+statement count 0
+DELETE FROM t3 WHERE k = 3
+
+statement count 2
+DELETE FROM t3 WHERE k < 10 AND k > 0
+
+statement ok
+COMMIT
+
+query I
+SELECT count(*) from t3
+----
+0