Merge #151916

151916: kvcoord: avoid piggybacking on any early-returning batch r=miraradeva a=stevendanna

To reduce the number of batches, we'd _like_ to piggy-back our buffer
flush on the user-provided batch. But, in some cases this is problematic
as the user may have set options on the batch that can result in the
batch being incompletely processed. Here, we check for those options and
flush the buffer via a sanitized batch before processing the user's
batch.

In a follow-up, we may also _always_ use a separate batch for midTxn
flushes to avoid nearly all read-write batches.

Epic: None
Release note: None

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

Author: craig[bot]

pkg/kv/kvclient/kvcoord/dist_sender.go

@@ -1897,8 +1897,7 @@ func (ds *DistSender) divideAndSendBatchToRanges(
 		}
 	}()
 
-	canParallelize := ba.Header.MaxSpanRequestKeys == 0 && ba.Header.TargetBytes == 0 &&
-		!ba.Header.ReturnElasticCPUResumeSpans
+	canParallelize := !ba.MightStopEarly()
 	if ba.IsSingleCheckConsistencyRequest() {
 		// Don't parallelize full checksum requests as they have to touch the
 		// entirety of each replica of each range they touch.
@@ -2002,9 +2001,8 @@ func (ds *DistSender) divideAndSendBatchToRanges(
 				ba.UpdateTxn(resp.reply.Txn)
 			}
 
-			mightStopEarly := ba.MaxSpanRequestKeys > 0 || ba.TargetBytes > 0 || ba.ReturnElasticCPUResumeSpans
 			// Check whether we've received enough responses to exit query loop.
-			if mightStopEarly {
+			if ba.MightStopEarly() {
 				var replyKeys int64
 				var replyBytes int64
 				for _, r := range resp.reply.Responses {

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -1515,10 +1515,10 @@ func (rr requestRecord) toResp(
 		// a lock, we add it to the buffer since we may need to flush it as
 		// replicated lock.
 		if rr.transformed {
-
 			transformedGetResponse := br.GetInner().(*kvpb.GetResponse)
 			valueWasPresent := transformedGetResponse.Value.IsPresent()
-			lockShouldHaveBeenAcquired := valueWasPresent || req.LockNonExisting
+			lockShouldHaveBeenAcquired := (valueWasPresent || req.LockNonExisting) &&
+				transformedGetResponse.ResumeSpan == nil
 
 			if lockShouldHaveBeenAcquired {
 				dla := &bufferedDurableLockAcquisition{
@@ -1771,10 +1771,7 @@ func (twb *txnWriteBuffer) flushBufferAndSendBatch(
 	}
 
 	midTxnFlush := !hasEndTxn
-
-	// SkipLocked reads cannot be in a batch with basically anything else. If we
-	// encounter one, we need to flush our buffer in its own batch.
-	splitBatchRequired := ba.WaitPolicy == lock.WaitPolicy_SkipLocked
+	splitBatchRequired := separateBatchIsNeeded(ba)
 
 	// Flush all buffered writes by pre-pending them to the requests being sent
 	// in the batch.
@@ -1816,17 +1813,21 @@ func (twb *txnWriteBuffer) flushBufferAndSendBatch(
 	})
 
 	if splitBatchRequired {
-		log.VEventf(ctx, 2, "flushing buffer via separate batch")
 		flushBatch := ba.ShallowCopy()
-		flushBatch.WaitPolicy = 0
+		clearBatchRequestOptions(flushBatch)
 		flushBatch.Requests = reqs
+		log.VEventf(ctx, 2, "flushing %d buffered requests via separate batch", len(reqs))
+
 		br, pErr := twb.wrapped.SendLocked(ctx, flushBatch)
 		if pErr != nil {
 			pErr.Index = nil
 			return nil, pErr
 		}
-
+		if err := requireAllFlushedRequestsProcessed(br.Responses); err != nil {
+			return nil, kvpb.NewError(err)
+		}
 		ba.UpdateTxn(br.Txn)
+
 		return twb.wrapped.SendLocked(ctx, ba)
 	} else {
 		ba = ba.ShallowCopy()
@@ -1835,13 +1836,56 @@ func (twb *txnWriteBuffer) flushBufferAndSendBatch(
 		if pErr != nil {
 			return nil, twb.adjustErrorUponFlush(ctx, numRevisionsBuffered, pErr)
 		}
-
+		if err := requireAllFlushedRequestsProcessed(br.Responses[0:numRevisionsBuffered]); err != nil {
+			return nil, kvpb.NewError(err)
+		}
 		// Strip out responses for all the flushed buffered writes.
 		br.Responses = br.Responses[numRevisionsBuffered:]
 		return br, nil
 	}
 }
 
+func requireAllFlushedRequestsProcessed(responses []kvpb.ResponseUnion) error {
+	for _, resp := range responses {
+		if resp.GetInner().Header().ResumeSpan != nil {
+			return errors.AssertionFailedf("response from buffered request has non-nil resume span")
+		}
+	}
+	return nil
+}
+
+// separateBatchIsNeeded returns true if BatchRequest contains any options that
+// require us to flush buffered requests using a separate batch.
+//
+// NB: If you are updating this function, you need to update
+// clearBatchRequestOptions as well.
+func separateBatchIsNeeded(ba *kvpb.BatchRequest) bool {
+	return ba.MightStopEarly() ||
+		ba.ReadConsistency != 0 ||
+		ba.WaitPolicy != 0 ||
+		ba.WriteOptions != nil && (*ba.WriteOptions != kvpb.WriteOptions{}) ||
+		ba.IsReverse
+}
+
+// clearBatchRequestOptions clears any options that should not be present on a
+// batch used to send previously buffered requests.
+//
+// NB: If you are updating this function, you need to update
+// separateBatchIsNeeded as well.
+func clearBatchRequestOptions(ba *kvpb.BatchRequest) {
+	// If read consistency is set to anything but CONSISTENT, our flush will fail
+	// because we only allow inconsistent reads for read only requests.
+	ba.ReadConsistency = 0
+	// If WaitPolicy is set to SkipLocked, our request may fail validation.
+	ba.WaitPolicy = 0
+	// Reset options that could result in an early batch return.
+	ba.MaxSpanRequestKeys = 0
+	ba.TargetBytes = 0
+	ba.ReturnElasticCPUResumeSpans = false
+	ba.IsReverse = false
+	ba.WriteOptions = nil
+}
+
 // hasBufferedWrites returns whether the interceptor has buffered any writes
 // locally.
 func (twb *txnWriteBuffer) hasBufferedWrites() bool {