kvstreamer: remove some redundant function calls

This commit refactors the logic of processing the BatchResponse in order
to avoid getting the "inner" request for each response - it is
sufficient for us to do a type switch on the response itself.

I did run the microbenchmarks, and they showed no difference, but
I believe this change should be beneficial.

Release note: None

pkg/kv/kvclient/kvstreamer/streamer.go

@@ -221,6 +221,7 @@ type Streamer struct {
 	hints         Hints
 	maxKeysPerRow int32
 	budget        *budget
+	keyLocking    lock.Strength
 
 	streamerStatistics
 
@@ -354,6 +355,7 @@ func NewStreamer(
 	limitBytes int64,
 	acc *mon.BoundAccount,
 	batchRequestsIssued *int64,
+	keyLocking lock.Strength,
 ) *Streamer {
 	if txn.Type() != kv.LeafTxn {
 		panic(errors.AssertionFailedf("RootTxn is given to the Streamer"))
@@ -362,6 +364,7 @@ func NewStreamer(
 		distSender: distSender,
 		stopper:    stopper,
 		budget:     newBudget(acc, limitBytes),
+		keyLocking: keyLocking,
 	}
 	if batchRequestsIssued == nil {
 		batchRequestsIssued = new(int64)
@@ -1428,9 +1431,9 @@ func processSingleRangeResults(
 			subRequestIdx = req.subRequestIdx[i]
 		}
 		reply := resp.GetInner()
-		switch req.reqs[i].GetInner().(type) {
-		case *roachpb.GetRequest:
-			get := reply.(*roachpb.GetResponse)
+		switch response := reply.(type) {
+		case *roachpb.GetResponse:
+			get := response
 			if get.ResumeSpan != nil {
 				// This Get wasn't completed.
 				continue
@@ -1454,8 +1457,8 @@ func processSingleRangeResults(
 			}
 			s.results.addLocked(result)
 
-		case *roachpb.ScanRequest:
-			scan := reply.(*roachpb.ScanResponse)
+		case *roachpb.ScanResponse:
+			scan := response
 			if len(scan.BatchResponses) == 0 && scan.ResumeSpan != nil {
 				// Only the first part of the conditional is true whenever we
 				// received an empty response for the Scan request (i.e. there
@@ -1546,9 +1549,9 @@ func buildResumeSingleRangeBatch(
 	for i, resp := range br.Responses {
 		position := req.positions[i]
 		reply := resp.GetInner()
-		switch origRequest := req.reqs[i].GetInner().(type) {
-		case *roachpb.GetRequest:
-			get := reply.(*roachpb.GetResponse)
+		switch response := reply.(type) {
+		case *roachpb.GetResponse:
+			get := response
 			if get.ResumeSpan == nil {
 				emptyResponse = false
 				continue
@@ -1557,10 +1560,10 @@ func buildResumeSingleRangeBatch(
 			// can just reuse the original request since it hasn't been
 			// modified which is also asserted below).
 			if buildutil.CrdbTestBuild {
-				if !get.ResumeSpan.Equal(origRequest.Span()) {
+				if origSpan := req.reqs[i].GetInner().Header().Span(); !get.ResumeSpan.Equal(origSpan) {
 					panic(errors.AssertionFailedf(
 						"unexpectedly the ResumeSpan %s on the GetResponse is different from the original span %s",
-						get.ResumeSpan, origRequest.Span(),
+						get.ResumeSpan, origSpan,
 					))
 				}
 			}
@@ -1574,8 +1577,8 @@ func buildResumeSingleRangeBatch(
 			}
 			resumeReqIdx++
 
-		case *roachpb.ScanRequest:
-			scan := reply.(*roachpb.ScanResponse)
+		case *roachpb.ScanResponse:
+			scan := response
 			if scan.ResumeSpan == nil {
 				emptyResponse = false
 				continue
@@ -1586,7 +1589,7 @@ func buildResumeSingleRangeBatch(
 			scans = scans[1:]
 			newScan.req.SetSpan(*scan.ResumeSpan)
 			newScan.req.ScanFormat = roachpb.BATCH_RESPONSE
-			newScan.req.KeyLocking = origRequest.KeyLocking
+			newScan.req.KeyLocking = s.keyLocking
 			newScan.union.Scan = &newScan.req
 			resumeReq.reqs[resumeReqIdx].Value = &newScan.union
 			resumeReq.positions = append(resumeReq.positions, position)

pkg/kv/kvclient/kvstreamer/streamer_test.go

@@ -47,6 +47,7 @@ func getStreamer(
 		limitBytes,
 		acc,
 		nil, /* batchRequestsIssued */
+		lock.None,
 	)
 }
 
@@ -98,6 +99,7 @@ func TestStreamerLimitations(t *testing.T) {
 				math.MaxInt64, /* limitBytes */
 				nil,           /* acc */
 				nil,           /* batchRequestsIssued */
+				lock.None,
 			)
 		})
 	})

pkg/sql/row/kv_fetcher.go

@@ -149,6 +149,7 @@ func NewStreamingKVFetcher(
 		streamerBudgetLimit,
 		streamerBudgetAcc,
 		&batchRequestsIssued,
+		getKeyLockingStrength(lockStrength),
 	)
 	mode := kvstreamer.OutOfOrder
 	if maintainOrdering {