kvclient: fix handling of ResumeSpans in txnWriteBuffer

Previously, we incorrectly handled Scans and ReverseScans when they had
ResumeSpan set. In particular, we would include buffered writes that
overlapped with the key span until the "full end", ignoring where the
ResumeSpan started. This meant that we could include the same buffered
write multiple times into the merged response. This is now fixed by
overlapping the btree only with the key span that was scanned on the
current "page".

Release note: None

Author: yuzefovich

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -1477,6 +1477,11 @@ type respIter struct {
 	// BATCH_RESPONSE are supported right now.
 	scanFormat kvpb.ScanFormat
 
+	// resumeSpan, if set, is the ResumeSpan of the response. When non-nil, it
+	// means that the response is being paginated, so we need to overlap the
+	// buffer only with the part of the original span that was actually scanned.
+	resumeSpan *roachpb.Span
+
 	// rows is the Rows field of the corresponding response.
 	//
 	// Only set with KEY_VALUES scan format.
@@ -1514,6 +1519,7 @@ func newScanRespIter(req *kvpb.ScanRequest, resp *kvpb.ScanResponse) *respIter {
 	return &respIter{
 		scanReq:        req,
 		scanFormat:     req.ScanFormat,
+		resumeSpan:     resp.ResumeSpan,
 		rows:           resp.Rows,
 		batchResponses: resp.BatchResponses,
 	}
@@ -1530,6 +1536,7 @@ func newReverseScanRespIter(
 	return &respIter{
 		reverseScanReq: req,
 		scanFormat:     req.ScanFormat,
+		resumeSpan:     resp.ResumeSpan,
 		rows:           resp.Rows,
 		batchResponses: resp.BatchResponses,
 	}
@@ -1583,13 +1590,27 @@ func (s *respIter) startKey() roachpb.Key {
 	if s.scanReq != nil {
 		return s.scanReq.Key
 	}
+	// For ReverseScans, the EndKey of the ResumeSpan is updated to indicate the
+	// start key for the "next" page, which is exactly the last key that was
+	// reverse-scanned for the current response.
+	// TODO(yuzefovich): we should have some unit tests that exercise the
+	// ResumeSpan case.
+	if s.resumeSpan != nil {
+		return s.resumeSpan.EndKey
+	}
 	return s.reverseScanReq.Key
 }
 
 // endKey returns the end key of the request in response to which the iterator
 // was created.
 func (s *respIter) endKey() roachpb.Key {
 	if s.scanReq != nil {
+		// For Scans, the Key of the ResumeSpan is updated to indicate the start
+		// key for the "next" page, which is exactly the last key that was
+		// scanned for the current response.
+		if s.resumeSpan != nil {
+			return s.resumeSpan.Key
+		}
 		return s.scanReq.EndKey
 	}
 	return s.reverseScanReq.EndKey

pkg/sql/logictest/testdata/logic_test/buffered_writes

@@ -439,3 +439,43 @@ query I
 SELECT count(*) FROM uvw
 ----
 64
+
+# Create some large blobs that will require the scans to be paginated (in 10MiB
+# chunks).
+statement ok
+CREATE TABLE large (k PRIMARY KEY, blob) AS SELECT i * 100, repeat('a', 11 * 1024 * 1024) FROM generate_series(1, 4) AS g(i);
+
+statement ok
+BEGIN;
+
+statement ok
+INSERT INTO large SELECT 11 + i * 100, 'b' FROM generate_series(1, 4) AS g(i);
+
+# Forward scan.
+query I
+SELECT k FROM large ORDER BY k;
+----
+100
+111
+200
+211
+300
+311
+400
+411
+
+# Reverse scan.
+query I
+SELECT k FROM large ORDER BY k DESC;
+----
+411
+400
+311
+300
+211
+200
+111
+100
+
+statement ok
+COMMIT;