Merge #143304

143304: kvclient: fix small bug in txnWriteBuffer r=yuzefovich a=stevendanna

Previously, we were not resetting the endKey before starting our iteration. This could result in the iterator ending before reaching our desired key despite it being in the buffer.

Epic: none
Release note: None

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

Author: craig[bot]

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -503,6 +503,15 @@ func (twb *txnWriteBuffer) applyTransformations(
 	return baRemote, ts
 }
 
+// seekItemForSpan returns a bufferedWrite appropriate for use with a
+// write-buffer iterator. Point lookups should use a nil end key.
+func (twb *txnWriteBuffer) seekItemForSpan(key, endKey roachpb.Key) *bufferedWrite {
+	seek := &twb.bufferSeek
+	seek.key = key
+	seek.endKey = endKey
+	return seek
+}
+
 // maybeServeRead serves the supplied read request from the buffer if a write or
 // deletion tombstone on the key is present in the buffer. Additionally, a
 // boolean indicating whether the read request was served or not is also
@@ -511,9 +520,7 @@ func (twb *txnWriteBuffer) maybeServeRead(
 	key roachpb.Key, seq enginepb.TxnSeq,
 ) (*roachpb.Value, bool) {
 	it := twb.buffer.MakeIter()
-	seek := &twb.bufferSeek
-	seek.key = key
-
+	seek := twb.seekItemForSpan(key, nil)
 	it.FirstOverlap(seek)
 	if it.Valid() {
 		bufferedVals := it.Cur().vals
@@ -543,11 +550,7 @@ func (twb *txnWriteBuffer) maybeServeRead(
 // write.
 func (twb *txnWriteBuffer) scanOverlaps(key roachpb.Key, endKey roachpb.Key) bool {
 	it := twb.buffer.MakeIter()
-	seek := &twb.bufferSeek
-	seek.key = key
-	seek.endKey = endKey
-
-	it.FirstOverlap(seek)
+	it.FirstOverlap(twb.seekItemForSpan(key, endKey))
 	return it.Valid()
 }
 
@@ -617,9 +620,7 @@ func (twb *txnWriteBuffer) mergeBufferAndResp(
 	reverse bool,
 ) {
 	it := twb.buffer.MakeIter()
-	seek := &twb.bufferSeek
-	seek.key = respIter.startKey()
-	seek.endKey = respIter.endKey()
+	seek := twb.seekItemForSpan(respIter.startKey(), respIter.endKey())
 
 	if reverse {
 		it.LastOverlap(seek)
@@ -871,8 +872,7 @@ func (t transformations) Empty() bool {
 // addToBuffer adds a write to the given key to the buffer.
 func (twb *txnWriteBuffer) addToBuffer(key roachpb.Key, val roachpb.Value, seq enginepb.TxnSeq) {
 	it := twb.buffer.MakeIter()
-	seek := &twb.bufferSeek
-	seek.key = key
+	seek := twb.seekItemForSpan(key, nil)
 
 	it.FirstOverlap(seek)
 	if it.Valid() {

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer_test.go

@@ -671,6 +671,65 @@ func TestTxnWriteBufferServesPointReadsLocally(t *testing.T) {
 	require.IsType(t, &kvpb.EndTxnResponse{}, br.Responses[0].GetInner())
 }
 
+// TestTxnWriteBufferServesPointReadsAfterScan is a regression test
+// for a bug in which reused iterator state resulted in the end key of
+// a scan affecting subsequent GetRequests.
+func TestTxnWriteBufferServesPointReadsAfterScan(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+	twb, mockSender := makeMockTxnWriteBuffer()
+
+	txn := makeTxnProto()
+	txn.Sequence = 10
+	keyA, keyB, keyC := roachpb.Key("a"), roachpb.Key("b"), roachpb.Key("c")
+
+	ba := &kvpb.BatchRequest{}
+	ba.Header = kvpb.Header{Txn: &txn}
+	ba.Add(putArgs(keyA, "valA", txn.Sequence))
+	ba.Add(putArgs(keyB, "valB", txn.Sequence))
+	ba.Add(putArgs(keyC, "valC", txn.Sequence))
+	numCalled := mockSender.NumCalled()
+	br, pErr := twb.SendLocked(ctx, ba)
+	require.Nil(t, pErr)
+	require.NotNil(t, br)
+	// All writes should be buffered.
+	require.Equal(t, numCalled, mockSender.NumCalled())
+
+	// First, read [a, c) via ScanRequest.
+	txn.Sequence = 10
+	ba = &kvpb.BatchRequest{}
+	ba.Header = kvpb.Header{Txn: &txn}
+	ba.Add(&kvpb.ScanRequest{
+		RequestHeader: kvpb.RequestHeader{Key: keyA, EndKey: keyC, Sequence: txn.Sequence},
+	})
+	mockSender.MockSend(func(ba *kvpb.BatchRequest) (*kvpb.BatchResponse, *kvpb.Error) {
+		require.Len(t, ba.Requests, 1)
+		require.IsType(t, &kvpb.ScanRequest{}, 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.Equal(t, int64(2), br.Responses[0].GetScan().NumKeys)
+
+	// Perform a read on keyC.
+	ba = &kvpb.BatchRequest{}
+	getC := &kvpb.GetRequest{RequestHeader: kvpb.RequestHeader{Key: keyC, Sequence: txn.Sequence}}
+	ba.Add(getC)
+
+	numCalled = mockSender.NumCalled()
+	br, pErr = twb.SendLocked(ctx, ba)
+	require.Nil(t, pErr)
+	require.NotNil(t, br)
+	require.Len(t, br.Responses, 1)
+	require.True(t, br.Responses[0].GetGet().Value.IsPresent())
+	require.Equal(t, mockSender.NumCalled(), numCalled)
+}
+
 // TestTxnWriteBufferServesOverlappingReadsCorrectly ensures that Scan and
 // ReverseScan requests that overlap with buffered writes are correctly served
 // from the buffer.