kvcoord: set ResumeSpan on unprocessed point requests

Previously, if batch execution was stopped early in dist sender because
of MaxSpanRequestKeys or TargetBytes, the returned batch could have
non-Get point-requests requests that were not processed but which didn't
have a resume-span sent.

The resume span was not set since SQL does not create batches with a mix
of read and write requests with MaxSpanRequestKeys or TargetBytes set;
however, KVNemesis will soon create such batches, so we'd like to handle
them correctly. Further, setting the ResumeSpan allows the
txnWriteBuffer to double-check that all writes were processed just in
case a bug results in the write buffer producing a batch that is
partially processed.

Fixes #153357
Release note: None

Author: stevendanna

pkg/kv/kvclient/kvcoord/dist_sender.go

@@ -2411,8 +2411,8 @@ func fillSkippedResponses(
 func maybeSetResumeSpan(
 	req kvpb.Request, hdr *kvpb.ResponseHeader, nextKey roachpb.RKey, isReverse bool,
 ) {
-	if _, ok := req.(*kvpb.GetRequest); ok {
-		// This is a Get request. There are three possibilities:
+	if !kvpb.IsRange(req) {
+		// This is a point request. There are three possibilities:
 		//
 		//  1. The request was completed. In this case we don't want a ResumeSpan.
 		//
@@ -2443,10 +2443,6 @@ func maybeSetResumeSpan(
 		return
 	}
 
-	if !kvpb.IsRange(req) {
-		return
-	}
-
 	origHeader := req.Header()
 	if isReverse {
 		if hdr.ResumeSpan != nil {

pkg/kv/kvclient/kvcoord/dist_sender_server_test.go

@@ -4535,3 +4535,79 @@ func TestUnexpectedCommitOnTxnRecovery(t *testing.T) {
 	close(blockCh)
 	wg.Wait()
 }
+
+// TestUnprocessedWritesHaveResumeSpanSet tests that writes that weren't
+// processed because proceeding requests exhausted a TargetBytes or
+// MaxSpanRequestKeys limit have an appropriate ResumeSpan set.
+func TestUnprocessedWritesHaveResumeSpanSet(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+
+	s, db := startNoSplitMergeServer(t)
+	defer s.Stopper().Stop(ctx)
+
+	var (
+		existingKey = "a"
+		splitPoint  = "b"
+		putKey      = "c"
+	)
+
+	// Set up a split so that our read and write below are in different ranges.
+	_, _, err := s.SplitRange(roachpb.Key(splitPoint))
+	require.NoError(t, err)
+	require.NoError(t, db.Put(ctx, existingKey, "existing-value"))
+
+	err = db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+		b := txn.NewBatch()
+		b.Header.MaxSpanRequestKeys = 1
+		// We use a ScanForUpdate here rather than Scan to ensure that the
+		// quota-consuming request is also a write (see #153397)
+		b.ScanForUpdate(existingKey, splitPoint, kvpb.GuaranteedDurability)
+		b.Del(putKey)
+		if err := txn.Run(ctx, b); err != nil {
+			return err
+		}
+		require.Len(t, b.Results, 2)
+		putResult := b.Results[1]
+		require.NotNil(t, putResult.ResumeSpan, "Put response should have a non-nil ResumeSpan")
+		return nil
+	})
+	require.NoError(t, err)
+}
+
+// TestMaxSpanRequestKeysWithMixedReadWriteBatches tests whether
+// MaxSpanRequestKeys is respected for certain read-write batches. This test
+// currently asserts that it _isn't_ respected, demonstrating a bug discovered
+// when writing TestUnprocessedWritesHaveResumeSpanSet and filed as #153397.
+func TestMaxSpanRequestKeysWithMixedReadWriteBatches(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+
+	s, db := startNoSplitMergeServer(t)
+	defer s.Stopper().Stop(ctx)
+
+	var (
+		startKey = "a"
+		endKey   = "b"
+	)
+
+	require.NoError(t, db.Put(ctx, startKey, "existing-value"))
+
+	err := db.Txn(ctx, func(ctx context.Context, txn *kv.Txn) error {
+		b := txn.NewBatch()
+		b.Header.MaxSpanRequestKeys = 1
+		b.Scan(startKey, endKey)
+		b.ScanForUpdate(startKey, endKey, kvpb.GuaranteedDurability)
+		if err := txn.Run(ctx, b); err != nil {
+			return err
+		}
+		require.Len(t, b.Results, 2)
+		require.Nil(t, b.Results[0].ResumeSpan, "first scan should be fully processed")
+		// TODO(#153397): The second scan here should not have been fully processed.
+		require.Nil(t, b.Results[1].ResumeSpan, "second scan should not be fully processed")
+		return nil
+	})
+	require.NoError(t, err)
+}