From 562d8601747dab06a8b875ffbbe0b0a55c003fcb Mon Sep 17 00:00:00 2001
From: Jordan Lewis <jordanthelewis@gmail.com>
Date: Fri, 30 Apr 2021 20:41:44 -0300
Subject: [PATCH] row: track KV bytes read for GetRequests

After #61583, we emit GetRequests from SQL when possible. We forgot to
update the stats collection code that tracks the number of bytes read
from KV when GetRequests were emitted; this is now corrected.

Release note: None
---
 pkg/sql/row/BUILD.bazel   |  1 -
 pkg/sql/row/kv_fetcher.go | 19 ++++++++++---------
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/pkg/sql/row/BUILD.bazel b/pkg/sql/row/BUILD.bazel
index f36c3e526eab..6a8836d6309a 100644
--- a/pkg/sql/row/BUILD.bazel
+++ b/pkg/sql/row/BUILD.bazel
@@ -59,7 +59,6 @@ go_library(
         "//pkg/util/log/eventpb",
         "//pkg/util/metric",
         "//pkg/util/mon",
-        "//pkg/util/syncutil",
         "//pkg/util/timeutil",
         "//pkg/util/unique",
         "//pkg/util/uuid",
diff --git a/pkg/sql/row/kv_fetcher.go b/pkg/sql/row/kv_fetcher.go
index 57d99b9cb0fe..acb581239981 100644
--- a/pkg/sql/row/kv_fetcher.go
+++ b/pkg/sql/row/kv_fetcher.go
@@ -13,6 +13,7 @@ package row
 import (
 	"context"
 	"time"
+	"sync/atomic"
 
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -22,7 +23,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/mon"
-	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/errors"
 )
 
@@ -37,8 +37,8 @@ type KVFetcher struct {
 	newSpan       bool
 
 	// Observability fields.
-	mu struct {
-		syncutil.Mutex
+	// Note: these need to be read via an atomic op.
+	atomics struct {
 		bytesRead int64
 	}
 }
@@ -114,9 +114,7 @@ func (f *KVFetcher) GetBytesRead() int64 {
 	if f == nil {
 		return 0
 	}
-	f.mu.Lock()
-	defer f.mu.Unlock()
-	return f.mu.bytesRead
+	return atomic.LoadInt64(&f.atomics.bytesRead)
 }
 
 // MVCCDecodingStrategy controls if and how the fetcher should decode MVCC
@@ -192,9 +190,12 @@ func (f *KVFetcher) NextKV(
 			return false, kv, false, nil
 		}
 		f.newSpan = true
-		f.mu.Lock()
-		f.mu.bytesRead += int64(len(f.batchResponse))
-		f.mu.Unlock()
+		nBytes := len(f.batchResponse)
+		for i := range f.kvs {
+			nBytes += len(f.kvs[i].Key)
+			nBytes += len(f.kvs[i].Value.RawBytes)
+		}
+		atomic.AddInt64(&f.atomics.bytesRead, int64(nBytes))
 	}
 }