kv, sql: Accumulate CPU time from KV BatchResponses

This commit implements a solution to recording kv cpu time on
every query execution - similar to other fields recorded in topLevelQueryStats.

This is done by including a new CPUTime field in the BatchResponse Header
from KV. The CPUTime field tracks the time reported in Replica.Send which
is representative of the "synchronous" work done by KV for the request, i.e.
it does not include other async replication related work.

GetKVCPUTime is added to the KVBatchFetcher interface.
GetKVResponseCPUTime is added to the KVReader interface.
kvCPUTime is added to the mutationPlanNode interface.

Epic: https://cockroachlabs.atlassian.net/browse/CRDB-55080
Part of: https://cockroachlabs.atlassian.net/browse/CRDB-55922
Release note: None

Author: alyshanjahani-crl

pkg/cli/debug_send_kv_batch_test.go

@@ -78,7 +78,8 @@ func TestSendKVBatch(t *testing.T) {
 		{
 			"header": {
 				"Timestamp": {},
-				"now": {}
+				"now": {},
+				"cpuTime": {}
 			},
 			"responses": [
 				{"put": {
@@ -112,9 +113,11 @@ func TestSendKVBatch(t *testing.T) {
 		require.NoError(t, err)
 
 		// Clean and check the BatchResponse output, by removing first line
-		// (contains input command) and emptying out all HLC timestamp objects.
+		// (contains input command), emptying out all HLC timestamp objects,
+		// and zeroing out the cpuTime value.
 		output = strings.SplitN(output, "\n", 2)[1]
 		output = regexp.MustCompile(`(?s)\{\s*"wallTime":.*?\}`).ReplaceAllString(output, "{}")
+		output = regexp.MustCompile(`"cpuTime"\s*:\s*("[^"]*"|\d+)`).ReplaceAllString(output, `"cpuTime": {}`)
 		require.JSONEq(t, jsonResponse, output)
 
 		// Check that a structured log event was emitted.

pkg/kv/kvclient/kvstreamer/streamer.go

@@ -307,6 +307,9 @@ type Streamer struct {
 type streamerStatistics struct {
 	atomics struct {
 		kvPairsRead *int64
+		// kvCPUTime tracks the cumulative CPU time (in nanoseconds) that KV
+		// reports in the BatchResponse headers.
+		kvCPUTime *int64
 		// batchRequestsIssued tracks the number of BatchRequests issued by the
 		// Streamer. Note that this number is only used for the logging done by
 		// the Streamer itself and is separate from any possible bookkeeping
@@ -381,6 +384,8 @@ type sendFn func(context.Context, *kvpb.BatchRequest) (*kvpb.BatchResponse, erro
 //
 // kvPairsRead should be incremented atomically with the sum of NumKeys
 // parameters of all received responses.
+// kvCPUTime should be incremented atomically with the CPUTime from BatchResponse
+// headers.
 func NewStreamer(
 	distSender *kvcoord.DistSender,
 	metrics *Metrics,
@@ -393,6 +398,7 @@ func NewStreamer(
 	limitBytes int64,
 	acc *mon.BoundAccount,
 	kvPairsRead *int64,
+	kvCPUTime *int64,
 	lockStrength lock.Strength,
 	lockDurability lock.Durability,
 	reverse bool,
@@ -430,6 +436,10 @@ func NewStreamer(
 		kvPairsRead = new(int64)
 	}
 	s.atomics.kvPairsRead = kvPairsRead
+	if kvCPUTime == nil {
+		kvCPUTime = new(int64)
+	}
+	s.atomics.kvCPUTime = kvCPUTime
 	s.coordinator = workerCoordinator{
 		s:                      s,
 		sendFn:                 sendFn,
@@ -1452,6 +1462,11 @@ func (w *workerCoordinator) performRequestAsync(
 		}
 		atomic.AddInt64(&w.s.atomics.batchRequestsIssued, 1)
 
+		// Accumulate CPU time from the BatchResponse header.
+		if br.CPUTime > 0 {
+			atomic.AddInt64(w.s.atomics.kvCPUTime, br.CPUTime)
+		}
+
 		// First, we have to reconcile the memory budget. We do it
 		// separately from processing the results because we want to know
 		// how many Gets and Scans need to be allocated for the ResumeSpans,

pkg/kv/kvclient/kvstreamer/streamer_accounting_test.go

@@ -100,6 +100,7 @@ func TestStreamerMemoryAccounting(t *testing.T) {
 			math.MaxInt64,
 			&acc,
 			nil, /* kvPairsRead */
+			nil, /* kvCPUTime */
 			lock.None,
 			lock.Unreplicated,
 			reverse,

pkg/kv/kvclient/kvstreamer/streamer_test.go

@@ -70,6 +70,7 @@ func getStreamer(
 		limitBytes,
 		acc,
 		nil, /* kvPairsRead */
+		nil, /* kvCPUTime */
 		lock.None,
 		lock.Unreplicated,
 		reverse,
@@ -131,6 +132,7 @@ func TestStreamerLimitations(t *testing.T) {
 				math.MaxInt64, /* limitBytes */
 				nil,           /* acc */
 				nil,           /* kvPairsRead */
+				nil,           /* kvCpuTime */
 				lock.None,
 				lock.Unreplicated,
 				false, /* reverse */

pkg/kv/kvpb/api.go

@@ -754,6 +754,7 @@ func (h *BatchResponse_Header) combine(o BatchResponse_Header) error {
 	h.Now.Forward(o.Now)
 	h.RangeInfos = append(h.RangeInfos, o.RangeInfos...)
 	h.CollectedSpans = append(h.CollectedSpans, o.CollectedSpans...)
+	h.CPUTime += o.CPUTime
 	return nil
 }
 

pkg/kv/kvpb/api.proto

@@ -3207,6 +3207,10 @@ message BatchResponse {
     // The field is cleared by the DistSender because it refers routing
     // information not exposed by the KV API.
     repeated RangeInfo range_infos = 7 [(gogoproto.nullable) = false];
+    // CPUTime is tracked in Replica.Send, meaning it only tracks the CPU time
+    // for the request evaluation goroutine, and not the other replication related
+    // work. The value is either zero (unset) or positive.
+    int64 cpu_time = 8 [(gogoproto.customname) = "CPUTime"];
     // NB: if you add a field here, don't forget to update combine().
   }
   Header header = 1 [(gogoproto.nullable) = false, (gogoproto.embed) = true];

pkg/kv/kvpb/batch_test.go

@@ -705,7 +705,8 @@ func TestBatchResponseCombine(t *testing.T) {
 		)
 		brTxn := &BatchResponse{
 			BatchResponse_Header: BatchResponse_Header{
-				Txn: &txn,
+				Txn:     &txn,
+				CPUTime: 123,
 			},
 		}
 		if err := br.Combine(context.Background(), brTxn, nil, &BatchRequest{}); err != nil {
@@ -714,6 +715,9 @@ func TestBatchResponseCombine(t *testing.T) {
 		if br.Txn.Name != "test" {
 			t.Fatal("Combine() did not update the header")
 		}
+		if br.CPUTime != 123 {
+			t.Fatalf("Combine() did not accumulate CPUTime: expected 123, got %d", br.CPUTime)
+		}
 	}
 
 	br.Responses = make([]ResponseUnion, 1)

pkg/kv/kvserver/replica_send.go

@@ -206,14 +206,18 @@ func (r *Replica) SendWithWriteBytes(
 		}
 	}
 
+	cpuTime := grunning.Difference(startCPU, grunning.Time())
+	if br != nil {
+		br.CPUTime = int64(cpuTime)
+	}
 	if pErr == nil {
 		// Return range information if it was requested. Note that we don't return it
 		// on errors because the code doesn't currently support returning both a br
 		// and a pErr here. Also, some errors (e.g. NotLeaseholderError) have custom
 		// ways of returning range info.
 		r.maybeAddRangeInfoToResponse(ctx, ba, br)
 		// Handle load-based splitting, if necessary.
-		r.recordBatchForLoadBasedSplitting(ctx, ba, br, int(grunning.Difference(startCPU, grunning.Time())))
+		r.recordBatchForLoadBasedSplitting(ctx, ba, br, int(cpuTime))
 	}
 
 	// Record summary throughput information about the batch request for

pkg/sql/colexecop/operator.go

@@ -92,6 +92,9 @@ type KVReader interface {
 	// KV requests. It must be safe for concurrent use. It is used to calculate
 	// the SQL CPU time.
 	GetKVCPUTime() time.Duration
+	// GetKVResponseCPUTime returns the CPU time as reported by KV BatchResponses
+	// processed by the KVReader throughout its lifetime so far.
+	GetKVResponseCPUTime() int64
 	// UsedStreamer returns whether the Streamer API was used by the KVReader.
 	UsedStreamer() bool
 }

pkg/sql/colfetcher/cfetcher.go

@@ -271,16 +271,17 @@ type cFetcher struct {
 	// stableKVs indicates whether the KVs returned by nextKVer are stable (i.e.
 	// are not invalidated) across NextKV() calls.
 	stableKVs bool
-	// bytesRead, kvPairsRead, and batchRequestsIssued store the total number of
-	// bytes read, key-values pairs read, and of BatchRequests issued,
-	// respectively, by this cFetcher throughout its lifetime in case when the
-	// underlying row.KVFetcher has already been closed and nil-ed out.
+	// bytesRead, kvPairsRead, kvCPUTime, and batchRequestsIssued store the total number of
+	// bytes read, key-values pairs read, CPU time reported by KV BatchResponses, and of
+	// BatchRequests issued, respectively, by this cFetcher throughout its lifetime in
+	// case when the underlying row.KVFetcher has already been closed and nil-ed out.
 	//
 	// The fields should not be accessed directly by the users of the cFetcher -
-	// getBytesRead(), getKVPairsRead(), and getBatchRequestsIssued() should be
+	// getBytesRead(), getKVPairsRead(), getKVCpuTime(), and getBatchRequestsIssued() should be
 	// used instead.
 	bytesRead           int64
 	kvPairsRead         int64
+	kvCPUTime           int64
 	batchRequestsIssued int64
 	// cpuStopWatch tracks the CPU time spent by this cFetcher while fulfilling KV
 	// requests *in the current goroutine*.
@@ -1484,6 +1485,15 @@ func (cf *cFetcher) getKVPairsRead() int64 {
 	return cf.kvPairsRead
 }
 
+// getKVCPUTime returns the CPU time as reported by KV BatchResponses processed
+// by the cFetcher throughout its lifetime so far.
+func (cf *cFetcher) getKVCPUTime() int64 {
+	if cf.fetcher != nil {
+		return cf.fetcher.GetKVCPUTime()
+	}
+	return cf.kvCPUTime
+}
+
 // getBatchRequestsIssued returns the number of BatchRequests issued by the
 // cFetcher throughout its lifetime so far.
 func (cf *cFetcher) getBatchRequestsIssued() int64 {
@@ -1522,6 +1532,7 @@ func (cf *cFetcher) Close(ctx context.Context) {
 		if cf.fetcher != nil {
 			cf.bytesRead = cf.fetcher.GetBytesRead()
 			cf.kvPairsRead = cf.fetcher.GetKVPairsRead()
+			cf.kvCPUTime = cf.fetcher.GetKVCPUTime()
 			cf.batchRequestsIssued = cf.fetcher.GetBatchRequestsIssued()
 			cf.fetcher.Close(ctx)
 			cf.fetcher = nil