Merge #156895 #158175

156895: kv, sql: Accumulate CPU time from KV BatchResponses r=alyshanjahani-crl a=alyshanjahani-crl

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

158175: mmaprototype: extract analyzeFunc r=wenyihu6 a=wenyihu6

Epic: CRDB-55052 
Release note: none

---
**mmaprototype: extract analyzeFunc**

This commit moves analyzeFunc out of finishInit into its own helper improve
readability.

---
**mmaprototype: extract doneFunc**

This commit moves doneFunc out of analyzeFunc into its own helper function to
improve readability.

---
**mmaprototype: add more comments for finishInit**

This commit adds more comments for rangeAnalyzedConstraints.finishInit.

---
**mmaprototype: extract diversity score computation**

This commit moves diversity score computation out of analyzeFunc into its own
helper function to improve readability.

---
**mmaprototype: clean up diversityScore**

Previously, we were inefficiently invoking diversityFunc twice over voters X
voters diversity score calculation when refactoring the code. This change
reverts it back so diversity score is computed once and reused.

---
**mmaprototype: rename diversityFunc to diversityOfTwoStoreSets**


---
**mmaprototype: rename ac.analyzeFunc to ac.initialize**

This commit renames ac.analyzeFunc to ac.initialize and updates the
signature so constraints are passed directly to initialize, rather than
being populated by caller.

---
**mmaprototype: move isConstraintSatisfied to a closure**

This commit moves isConstraintSatisfied from a struct method to a helper closure
inside ac.initialize.

---
**mmaprototype: clean up comments**




Co-authored-by: alyshanjahani-crl <alyshan@cockroachlabs.com>
Co-authored-by: wenyihu6 <wenyi@cockroachlabs.com>

Author: 3 people

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/allocator/mmaprototype/allocator_state.go

@@ -703,6 +703,11 @@ func sortTargetCandidateSetAndPick(
 	return cands.candidates[j].StoreID
 }
 
+// ensureAnalyzedConstraints ensures that the constraints field of rangeState is
+// populated. It uses rangeState.{replicas,conf} as inputs to the computation.
+//
+// NB: Caller is responsible for calling clearAnalyzedConstraints when rstate or
+// the rstate.constraints is no longer needed.
 func (cs *clusterState) ensureAnalyzedConstraints(rstate *rangeState) {
 	if rstate.constraints != nil {
 		return