sql/row: reduce allocations initializing txnKVFetcher and txnKVStreamer

The `kvBatchFetcherHelper` no longer implements the `NextBatch` method
for the `txnKVFetcher` and `txnKVStreamer`. The latter two implement
this method themselves, eliminating an allocation of the `nextBatch`
closure. The `kvBatchFetcherHelper` has been renamed to `kvBatchMetrics`
now that its only concern is recording and providing metrics.

Release note: None

pkg/sql/row/kv_batch_fetcher.go

@@ -132,7 +132,7 @@ func (s *identifiableSpans) reset() {
 
 // txnKVFetcher handles retrieval of key/values.
 type txnKVFetcher struct {
-	kvBatchFetcherHelper
+	kvBatchMetrics
 	// "Constant" fields, provided by the caller.
 	sendFn sendFunc
 	// spans is the list of Spans that will be read by this KV Fetcher. If an
@@ -390,7 +390,7 @@ func newTxnKVFetcherInternal(args newTxnKVFetcherArgs) *txnKVFetcher {
 		args.admission.pacerFactory,
 		args.admission.settingsValues,
 	)
-	f.kvBatchFetcherHelper.init(f.nextBatch, args.kvPairsRead, args.batchRequestsIssued)
+	f.kvBatchMetrics.init(args.kvPairsRead, args.batchRequestsIssued)
 	return f
 }
 
@@ -802,6 +802,16 @@ func popBatch(
 	return nil, nil, colBatch, remainingColBatches
 }
 
+// NextBatch implements the KVBatchFetcher interface.
+func (f *txnKVFetcher) NextBatch(ctx context.Context) (KVBatchFetcherResponse, error) {
+	resp, err := f.nextBatch(ctx)
+	if !resp.MoreKVs || err != nil {
+		return resp, err
+	}
+	f.kvBatchMetrics.Record(resp)
+	return resp, nil
+}
+
 func (f *txnKVFetcher) nextBatch(ctx context.Context) (resp KVBatchFetcherResponse, err error) {
 	// The purpose of this loop is to unpack the two-level batch structure that is
 	// returned from the KV layer.
@@ -1058,33 +1068,24 @@ func spansToRequests(
 	return reqs
 }
 
-// kvBatchFetcherHelper is a small helper that extracts common logic for
-// implementing some methods of the KVBatchFetcher interface related to
-// observability.
-type kvBatchFetcherHelper struct {
-	nextBatch func(context.Context) (KVBatchFetcherResponse, error)
-	atomics   struct {
+// kvBatchMetrics tracks metrics and implements some of the methods for
+// the KVBatchFetcher interface related to observability.
+type kvBatchMetrics struct {
+	atomics struct {
 		bytesRead           int64
 		kvPairsRead         *int64
 		batchRequestsIssued *int64
 	}
 }
 
-func (h *kvBatchFetcherHelper) init(
-	nextBatch func(context.Context) (KVBatchFetcherResponse, error),
-	kvPairsRead, batchRequestsIssued *int64,
-) {
-	h.nextBatch = nextBatch
+func (h *kvBatchMetrics) init(kvPairsRead, batchRequestsIssued *int64) {
 	h.atomics.kvPairsRead = kvPairsRead
 	h.atomics.batchRequestsIssued = batchRequestsIssued
 }
 
-// NextBatch implements the KVBatchFetcher interface.
-func (h *kvBatchFetcherHelper) NextBatch(ctx context.Context) (KVBatchFetcherResponse, error) {
-	resp, err := h.nextBatch(ctx)
-	if !resp.MoreKVs || err != nil {
-		return resp, err
-	}
+// Record records metrics for the given batch response. It should be called
+// after each batch is fetched.
+func (h *kvBatchMetrics) Record(resp KVBatchFetcherResponse) {
 	atomic.AddInt64(h.atomics.kvPairsRead, resp.kvPairsRead)
 	// Note that if resp.ColBatch is nil, then GetBatchMemSize will return 0.
 	// TODO(yuzefovich, 23.1): for resp.ColBatch this includes the decoded
@@ -1097,27 +1098,26 @@ func (h *kvBatchFetcherHelper) NextBatch(ctx context.Context) (KVBatchFetcherRes
 		nBytes += len(resp.KVs[i].Value.RawBytes)
 	}
 	atomic.AddInt64(&h.atomics.bytesRead, int64(nBytes))
-	return resp, nil
 }
 
 // GetBytesRead implements the KVBatchFetcher interface.
-func (h *kvBatchFetcherHelper) GetBytesRead() int64 {
+func (h *kvBatchMetrics) GetBytesRead() int64 {
 	if h == nil {
 		return 0
 	}
 	return atomic.LoadInt64(&h.atomics.bytesRead)
 }
 
 // GetKVPairsRead implements the KVBatchFetcher interface.
-func (h *kvBatchFetcherHelper) GetKVPairsRead() int64 {
+func (h *kvBatchMetrics) GetKVPairsRead() int64 {
 	if h == nil || h.atomics.kvPairsRead == nil {
 		return 0
 	}
 	return atomic.LoadInt64(h.atomics.kvPairsRead)
 }
 
 // GetBatchRequestsIssued implements the KVBatchFetcher interface.
-func (h *kvBatchFetcherHelper) GetBatchRequestsIssued() int64 {
+func (h *kvBatchMetrics) GetBatchRequestsIssued() int64 {
 	if h == nil || h.atomics.batchRequestsIssued == nil {
 		return 0
 	}

pkg/sql/row/kv_batch_streamer.go

@@ -23,7 +23,7 @@ import (
 
 // txnKVStreamer handles retrieval of key/values.
 type txnKVStreamer struct {
-	kvBatchFetcherHelper
+	kvBatchMetrics
 	streamer       *kvstreamer.Streamer
 	lockStrength   lock.Strength
 	lockDurability lock.Durability
@@ -69,7 +69,7 @@ func newTxnKVStreamer(
 		acc:            acc,
 		rawMVCCValues:  rawMVCCValues,
 	}
-	f.kvBatchFetcherHelper.init(f.nextBatch, kvPairsRead, batchRequestsIssued)
+	f.kvBatchMetrics.init(kvPairsRead, batchRequestsIssued)
 	return f
 }
 
@@ -184,6 +184,16 @@ func (f *txnKVStreamer) releaseLastResult(ctx context.Context) {
 	f.lastResultState.Result = kvstreamer.Result{}
 }
 
+// NextBatch implements the KVBatchFetcher interface.
+func (f *txnKVStreamer) NextBatch(ctx context.Context) (KVBatchFetcherResponse, error) {
+	resp, err := f.nextBatch(ctx)
+	if !resp.MoreKVs || err != nil {
+		return resp, err
+	}
+	f.kvBatchMetrics.Record(resp)
+	return resp, nil
+}
+
 func (f *txnKVStreamer) nextBatch(ctx context.Context) (resp KVBatchFetcherResponse, _ error) {
 	// Check whether there are more batches in the current ScanResponse.
 	if len(f.lastResultState.remainingBatches) > 0 {
@@ -241,5 +251,5 @@ func (f *txnKVStreamer) Close(ctx context.Context) {
 	f.streamer.Close(ctx)
 	f.acc.Clear(ctx)
 	// Preserve observability-related fields.
-	*f = txnKVStreamer{kvBatchFetcherHelper: f.kvBatchFetcherHelper}
+	*f = txnKVStreamer{kvBatchMetrics: f.kvBatchMetrics}
 }