Merge #146526

146526: kvserver: improve tenant_id CPU observability r=tbg a=tbg

Motivated by https://github.com/cockroachlabs/support/issues/3297.

This commit

- adds a new `replicas.cpunanospersecond` metric, which aggregates the Replica
ReqCPUNanosPerSecond at the tenant level.
- Adds a tenant_id tag to CPU profiles.

This should simplify investigations related to tenant-induced overload: the new
metric should often help pinpoint the set of hot tenants. CPU profiles can help
dig into the specific code paths this tenant is exercising. This can then be
rounded out with the existing metrics for request counts by tenant (both
received by KV and sent by the tenant Pod) for a comprehensive picture.

Release note: the `replicas.cpunanospersecond` metric was added. Notably, when
child labels are enabled, it exposes evaluation-related Replica CPU usage by
tenant.
Epic: none


Co-authored-by: Tobias Grieger <[email protected]>

Author: craig[bot]

docs/generated/metrics/metrics.yaml

@@ -15153,6 +15153,14 @@ layers:
       unit: COUNT
       aggregation: AVG
       derivative: NONE
+    - name: replicas.cpunanospersecond
+      exported_name: replicas_cpunanospersecond
+      description: Nanoseconds of CPU time in Replica request processing including evaluation but not replication
+      y_axis_label: Nanoseconds
+      type: COUNTER
+      unit: NANOSECONDS
+      aggregation: AVG
+      derivative: NON_NEGATIVE_DERIVATIVE
     - name: replicas.leaders
       exported_name: replicas_leaders
       description: Number of raft leaders

pkg/kv/kvserver/metrics.go

@@ -249,6 +249,13 @@ var (
 		Unit:        metric.Unit_COUNT,
 	}
 
+	metaReqCPUNanos = metric.Metadata{
+		Name:        "replicas.cpunanospersecond",
+		Help:        "Nanoseconds of CPU time in Replica request processing including evaluation but not replication",
+		Measurement: "Nanoseconds",
+		Unit:        metric.Unit_NANOSECONDS,
+	}
+
 	// Storage metrics.
 	metaLiveBytes = metric.Metadata{
 		Name:        "livebytes",
@@ -3207,29 +3214,6 @@ type StoreMetrics struct {
 	DiskWriteMaxBytesPerSecond *metric.Gauge
 }
 
-type tenantMetricsRef struct {
-	// All fields are internal. Don't access them.
-
-	_tenantID roachpb.TenantID
-	_state    int32 // atomic; 0=usable 1=poisoned
-
-	// _stack helps diagnose use-after-release when it occurs.
-	// This field is populated in releaseTenant and printed
-	// in assertions on failure.
-	_stack struct {
-		syncutil.Mutex
-		debugutil.SafeStack
-	}
-}
-
-func (ref *tenantMetricsRef) assert(ctx context.Context) {
-	if atomic.LoadInt32(&ref._state) != 0 {
-		ref._stack.Lock()
-		defer ref._stack.Unlock()
-		log.FatalfDepth(ctx, 1, "tenantMetricsRef already finalized in:\n%s", ref._stack.SafeStack)
-	}
-}
-
 // TenantsStorageMetrics are metrics which are aggregated over all tenants
 // present on the server. The struct maintains child metrics used by each
 // tenant to track their individual values. The struct expects that children
@@ -3238,6 +3222,7 @@ func (ref *tenantMetricsRef) assert(ctx context.Context) {
 type TenantsStorageMetrics struct {
 	// NB: If adding more metrics to this struct, be sure to
 	// also update tenantsStorageMetricsSet().
+	ReqCPUNanos    *aggmetric.AggCounterFloat64
 	LiveBytes      *aggmetric.AggGauge
 	KeyBytes       *aggmetric.AggGauge
 	ValBytes       *aggmetric.AggGauge
@@ -3274,6 +3259,7 @@ type TenantsStorageMetrics struct {
 // see kvbase.TenantsStorageMetricsSet for public access. Assigned in init().
 func tenantsStorageMetricsSet() map[string]struct{} {
 	return map[string]struct{}{
+		metaReqCPUNanos.Name:    {},
 		metaLiveBytes.Name:      {},
 		metaKeyBytes.Name:       {},
 		metaValBytes.Name:       {},
@@ -3304,7 +3290,7 @@ func (sm *TenantsStorageMetrics) MetricStruct() {}
 // method are reference counted with decrements occurring in the corresponding
 // releaseTenant call. This method must be called prior to adding or subtracting
 // MVCC stats.
-func (sm *TenantsStorageMetrics) acquireTenant(tenantID roachpb.TenantID) *tenantMetricsRef {
+func (sm *TenantsStorageMetrics) acquireTenant(tenantID roachpb.TenantID) *tenantStorageMetrics {
 	// incRef increments the reference count if it is not already zero indicating
 	// that the struct has already been destroyed.
 	incRef := func(m *tenantStorageMetrics) (alreadyDestroyed bool) {
@@ -3319,15 +3305,13 @@ func (sm *TenantsStorageMetrics) acquireTenant(tenantID roachpb.TenantID) *tenan
 	for {
 		if m, ok := sm.tenants.Load(tenantID); ok {
 			if alreadyDestroyed := incRef(m); !alreadyDestroyed {
-				return &tenantMetricsRef{
-					_tenantID: tenantID,
-				}
+				return m
 			}
 			// Somebody else concurrently took the reference count to zero, go back
 			// around. Because of the locking in releaseTenant, we know that we'll
 			// find a different value or no value at all on the next iteration.
 		} else {
-			m := &tenantStorageMetrics{}
+			m := &tenantStorageMetrics{tenantID: tenantID}
 			m.mu.Lock()
 			_, loaded := sm.tenants.LoadOrStore(tenantID, m)
 			if loaded {
@@ -3338,6 +3322,7 @@ func (sm *TenantsStorageMetrics) acquireTenant(tenantID roachpb.TenantID) *tenan
 			// Successfully stored a new instance, initialize it and then unlock it.
 			tenantIDStr := tenantID.String()
 			m.mu.refCount++
+			m.ReqCPUNanos = sm.ReqCPUNanos.AddChild(tenantIDStr)
 			m.LiveBytes = sm.LiveBytes.AddChild(tenantIDStr)
 			m.KeyBytes = sm.KeyBytes.AddChild(tenantIDStr)
 			m.ValBytes = sm.ValBytes.AddChild(tenantIDStr)
@@ -3359,37 +3344,35 @@ func (sm *TenantsStorageMetrics) acquireTenant(tenantID roachpb.TenantID) *tenan
 			m.SysCount = sm.SysCount.AddChild(tenantIDStr)
 			m.AbortSpanBytes = sm.AbortSpanBytes.AddChild(tenantIDStr)
 			m.mu.Unlock()
-			return &tenantMetricsRef{
-				_tenantID: tenantID,
-			}
+			return m
 		}
 	}
 }
 
 // releaseTenant releases the reference to the metrics for this tenant which was
 // acquired with acquireTenant. It will fatally log if no entry exists for this
 // tenant.
-func (sm *TenantsStorageMetrics) releaseTenant(ctx context.Context, ref *tenantMetricsRef) {
-	m := sm.getTenant(ctx, ref) // NB: asserts against use-after-release
-	if atomic.SwapInt32(&ref._state, 1) != 0 {
-		ref.assert(ctx) // this will fatal
-		return          // unreachable
-	}
-	ref._stack.Lock()
-	ref._stack.SafeStack = debugutil.Stack()
-	ref._stack.Unlock()
+func (sm *TenantsStorageMetrics) releaseTenant(ctx context.Context, m *tenantStorageMetrics) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
+	if m.mu.released.Load() {
+		log.FatalfDepth(ctx, 1, "tenant metrics already released in:\n%s", m.mu.stack)
+	}
 	m.mu.refCount--
-	if m.mu.refCount < 0 {
-		log.Fatalf(ctx, "invalid refCount on metrics for tenant %v: %d", ref._tenantID, m.mu.refCount)
-	} else if m.mu.refCount > 0 {
+	if n := m.mu.refCount; n < 0 {
+		log.Fatalf(ctx, "invalid refCount on metrics for tenant %v: %d", m.tenantID, n)
+	} else if n > 0 {
 		return
 	}
 
+	m.mu.released.Store(true)
+	m.mu.stack = debugutil.Stack()
+
 	// The refCount is zero, delete this instance after destroying its metrics.
 	// Note that concurrent attempts to create an instance will detect the zero
 	// refCount value and construct a new instance.
+	m.ReqCPUNanos.Unlink() // counter
+	m.ReqCPUNanos = nil
 	for _, gptr := range []**aggmetric.Gauge{
 		&m.LiveBytes,
 		&m.KeyBytes,
@@ -3417,28 +3400,30 @@ func (sm *TenantsStorageMetrics) releaseTenant(ctx context.Context, ref *tenantM
 		(*gptr).Unlink()
 		*gptr = nil
 	}
-	sm.tenants.Delete(ref._tenantID)
-}
-
-// getTenant is a helper method used to retrieve the metrics for a tenant. The
-// call will log fatally if no such tenant has been previously acquired.
-func (sm *TenantsStorageMetrics) getTenant(
-	ctx context.Context, ref *tenantMetricsRef,
-) *tenantStorageMetrics {
-	ref.assert(ctx)
-	m, ok := sm.tenants.Load(ref._tenantID)
-	if !ok {
-		log.Fatalf(ctx, "no metrics exist for tenant %v", ref._tenantID)
-	}
-	return m
+	sm.tenants.Delete(m.tenantID)
 }
 
+// tenantStorageMetrics is a struct that holds the metrics for all replicas (
+// within a Store) of a given tenant.
+//
+// Whenever it is guaranteed that the replica is not destroyed, the metrics
+// fields can be accessed directly (for example, when holding raftMu and having
+// previously checked the replica's destroyStatus).
+//
+// Whenever this is *not* guaranteed, use `TenantsStorageMetrics.acquireTenant`
+// ( followed by releaseTenant after completion of usage) instead to avoid
+// racing with a potential concurrent attempt to release the metrics.
 type tenantStorageMetrics struct {
-	mu struct {
+	tenantID roachpb.TenantID
+	mu       struct {
 		syncutil.Mutex
 		refCount int
+		released atomic.Bool // allowed to read without holding mu
+		stack    debugutil.SafeStack
 	}
 
+	ReqCPUNanos *aggmetric.CounterFloat64
+
 	LiveBytes      *aggmetric.Gauge
 	KeyBytes       *aggmetric.Gauge
 	ValBytes       *aggmetric.Gauge
@@ -3461,9 +3446,18 @@ type tenantStorageMetrics struct {
 	AbortSpanBytes *aggmetric.Gauge
 }
 
+func (tm *tenantStorageMetrics) assert(ctx context.Context) {
+	if tm.mu.released.Load() {
+		tm.mu.Lock()
+		defer tm.mu.Unlock()
+		log.Fatalf(ctx, "tenant metrics already released in:\n%s", tm.mu.stack)
+	}
+}
+
 func newTenantsStorageMetrics() *TenantsStorageMetrics {
 	b := aggmetric.MakeBuilder(multitenant.TenantIDLabel)
 	sm := &TenantsStorageMetrics{
+		ReqCPUNanos:    b.CounterFloat64(metaReqCPUNanos),
 		LiveBytes:      b.Gauge(metaLiveBytes),
 		KeyBytes:       b.Gauge(metaKeyBytes),
 		ValBytes:       b.Gauge(metaValBytes),
@@ -4020,10 +4014,9 @@ func newStoreMetrics(histogramWindow time.Duration) *StoreMetrics {
 // single snapshot of these gauges in the registry might mix the values of two
 // subsequent updates.
 func (sm *TenantsStorageMetrics) incMVCCGauges(
-	ctx context.Context, ref *tenantMetricsRef, delta enginepb.MVCCStats,
+	ctx context.Context, tm *tenantStorageMetrics, delta enginepb.MVCCStats,
 ) {
-	ref.assert(ctx)
-	tm := sm.getTenant(ctx, ref)
+	tm.assert(ctx)
 	tm.LiveBytes.Inc(delta.LiveBytes)
 	tm.KeyBytes.Inc(delta.KeyBytes)
 	tm.ValBytes.Inc(delta.ValBytes)
@@ -4047,17 +4040,17 @@ func (sm *TenantsStorageMetrics) incMVCCGauges(
 }
 
 func (sm *TenantsStorageMetrics) addMVCCStats(
-	ctx context.Context, ref *tenantMetricsRef, delta enginepb.MVCCStats,
+	ctx context.Context, tm *tenantStorageMetrics, delta enginepb.MVCCStats,
 ) {
-	sm.incMVCCGauges(ctx, ref, delta)
+	sm.incMVCCGauges(ctx, tm, delta)
 }
 
 func (sm *TenantsStorageMetrics) subtractMVCCStats(
-	ctx context.Context, ref *tenantMetricsRef, delta enginepb.MVCCStats,
+	ctx context.Context, tm *tenantStorageMetrics, delta enginepb.MVCCStats,
 ) {
 	var neg enginepb.MVCCStats
 	neg.Subtract(delta)
-	sm.incMVCCGauges(ctx, ref, neg)
+	sm.incMVCCGauges(ctx, tm, neg)
 }
 
 func (sm *StoreMetrics) updateEngineMetrics(m storage.Metrics) {

pkg/kv/kvserver/metrics_test.go

@@ -27,15 +27,14 @@ import (
 func TestTenantsStorageMetricsRelease(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	m := newTenantsStorageMetrics()
-	var refs []*tenantMetricsRef
+	var refs []*tenantStorageMetrics
 	const tenants = 7
 	for i := 0; i < tenants; i++ {
 		id := roachpb.MustMakeTenantID(roachpb.MinTenantID.InternalValue + uint64(i))
-		ref := m.acquireTenant(id)
-		tm := m.getTenant(context.Background(), ref)
+		tm := m.acquireTenant(id)
 		tm.SysBytes.Update(1023)
 		tm.KeyCount.Inc(123)
-		refs = append(refs, ref)
+		refs = append(refs, tm)
 	}
 	for i, ref := range refs {
 		require.Equal(t, int64(1023*(tenants-i)), m.SysBytes.Value(), i)

pkg/kv/kvserver/mvcc_gc_queue.go

@@ -664,7 +664,7 @@ func (mgcq *mvccGCQueue) process(
 ) (processed bool, err error) {
 	// Record the CPU time processing the request for this replica. This is
 	// recorded regardless of errors that are encountered.
-	defer repl.MeasureReqCPUNanos(grunning.Time())
+	defer repl.MeasureReqCPUNanos(ctx, grunning.Time())
 
 	// Lookup the descriptor and GC policy for the zone containing this key range.
 	desc, conf := repl.DescAndSpanConfig()

pkg/kv/kvserver/replica.go

@@ -548,13 +548,11 @@ type Replica struct {
 	// [^1]: TODO(pavelkalinnikov): we can but it'd be a larger refactor.
 	tenantLimiter tenantrate.Limiter
 
-	// tenantMetricsRef is a metrics reference indicating the tenant under
-	// which to track the range's contributions. This is determined by the
-	// start key of the Replica, once initialized.
-	// Its purpose is to help track down missing/extraneous release operations
-	// that would not be apparent or easy to resolve when refcounting at the store
-	// level only.
-	tenantMetricsRef *tenantMetricsRef
+	// tenantMetricsRef is a struct for per-tenant metrics contributed by this
+	// replica. It is determined by the start key of the Replica, once
+	// initialized. See tenantStorageMetrics for precautions that must be taken
+	// when accessing this.
+	tenantMetricsRef *tenantStorageMetrics
 
 	// sideTransportClosedTimestamp encapsulates state related to the closed
 	// timestamp's information about the range. Note that the
@@ -2770,9 +2768,23 @@ func init() {
 
 // MeasureReqCPUNanos measures the cpu time spent on this replica processing
 // requests.
-func (r *Replica) MeasureReqCPUNanos(start time.Duration) {
+func (r *Replica) MeasureReqCPUNanos(ctx context.Context, start time.Duration) {
 	r.measureNanosRunning(start, func(dur float64) {
 		r.loadStats.RecordReqCPUNanos(dur)
+		// NB: the caller also has a tenant ID, but we use the replica's here for
+		// simplicity. There is no established pattern for short-lived references
+		// to a specific tenant's metrics.
+		if r.tenantMetricsRef != nil {
+			// We can *not* use the tenant metrics directly because nothing in this
+			// current code path prevents the surrounding replica from getting
+			// destroyed, which could zero the refcount and release the metrics
+			// object. Instead, we go through acquireTenant, which gives us an object
+			// that is and remain valid. This is not an expensive operation in
+			// the common case (the replica still exists).
+			tm := r.store.metrics.acquireTenant(r.tenantMetricsRef.tenantID)
+			tm.ReqCPUNanos.Inc(dur)
+			r.store.metrics.releaseTenant(ctx, tm)
+		}
 	})
 }
 

pkg/kv/kvserver/replica_rate_limit.go

@@ -18,12 +18,13 @@ import (
 // maybeRateLimitBatch may block the batch waiting to be rate-limited. Note that
 // the replica must be initialized and thus there is no synchronization issue
 // on the tenantRateLimiter.
-func (r *Replica) maybeRateLimitBatch(ctx context.Context, ba *kvpb.BatchRequest) error {
+func (r *Replica) maybeRateLimitBatch(
+	ctx context.Context, ba *kvpb.BatchRequest, tenantIDOrZero roachpb.TenantID,
+) error {
 	if r.tenantLimiter == nil {
 		return nil
 	}
-	tenantID, ok := roachpb.ClientTenantFromContext(ctx)
-	if !ok || tenantID == roachpb.SystemTenantID {
+	if !tenantIDOrZero.IsSet() || tenantIDOrZero.IsSystem() {
 		return nil
 	}
 

pkg/kv/kvserver/replica_send.go

@@ -123,10 +123,23 @@ func (r *Replica) Send(
 func (r *Replica) SendWithWriteBytes(
 	ctx context.Context, ba *kvpb.BatchRequest,
 ) (*kvpb.BatchResponse, *kvadmission.StoreWriteBytes, *kvpb.Error) {
+	tenantIDOrZero, _ := roachpb.ClientTenantFromContext(ctx)
+
+	// Record the CPU time processing the request for this replica. This is
+	// recorded regardless of errors that are encountered.
+	startCPU := grunning.Time()
+	defer r.MeasureReqCPUNanos(ctx, startCPU)
+
 	if r.store.cfg.Settings.CPUProfileType() == cluster.CPUProfileWithLabels {
 		defer pprof.SetGoroutineLabels(ctx)
 		// Note: the defer statement captured the previous context.
-		ctx = pprof.WithLabels(ctx, pprof.Labels("range_str", r.rangeStr.ID()))
+		var lbls pprof.LabelSet
+		if tenantIDOrZero.IsSet() {
+			lbls = pprof.Labels("range_str", r.rangeStr.ID(), "tenant_id", tenantIDOrZero.String())
+		} else {
+			lbls = pprof.Labels("range_str", r.rangeStr.ID())
+		}
+		ctx = pprof.WithLabels(ctx, lbls)
 		pprof.SetGoroutineLabels(ctx)
 	}
 	if trace.IsEnabled() {
@@ -135,11 +148,6 @@ func (r *Replica) SendWithWriteBytes(
 	// Add the range log tag.
 	ctx = r.AnnotateCtx(ctx)
 
-	// Record the CPU time processing the request for this replica. This is
-	// recorded regardless of errors that are encountered.
-	startCPU := grunning.Time()
-	defer r.MeasureReqCPUNanos(startCPU)
-
 	isReadOnly := ba.IsReadOnly()
 	if err := r.checkBatchRequest(ba, isReadOnly); err != nil {
 		return nil, nil, kvpb.NewError(err)
@@ -148,7 +156,7 @@ func (r *Replica) SendWithWriteBytes(
 	if err := r.maybeBackpressureBatch(ctx, ba); err != nil {
 		return nil, nil, kvpb.NewError(err)
 	}
-	if err := r.maybeRateLimitBatch(ctx, ba); err != nil {
+	if err := r.maybeRateLimitBatch(ctx, ba, tenantIDOrZero); err != nil {
 		return nil, nil, kvpb.NewError(err)
 	}
 	if err := r.maybeCommitWaitBeforeCommitTrigger(ctx, ba); err != nil {