kvserver: absorb tenant metrics ref

This should be simpler overall: there is a single struct that represents the metrics for a tenant on the store.

Author: tbg

pkg/kv/kvserver/metrics.go

@@ -3207,29 +3207,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
@@ -3304,7 +3281,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,17 +3296,13 @@ func (sm *TenantsStorageMetrics) acquireTenant(tenantID roachpb.TenantID) *tenan
 	for {
 		if m, ok := sm.tenants.Load(tenantID); ok {
 			if alreadyDestroyed := incRef(m); !alreadyDestroyed {
-				return &m.ref
+				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{
-				ref: tenantMetricsRef{
-					_tenantID: tenantID,
-				},
-			}
+			m := &tenantStorageMetrics{tenantID: tenantID}
 			m.mu.Lock()
 			_, loaded := sm.tenants.LoadOrStore(tenantID, m)
 			if loaded {
@@ -3361,33 +3334,28 @@ func (sm *TenantsStorageMetrics) acquireTenant(tenantID roachpb.TenantID) *tenan
 			m.SysCount = sm.SysCount.AddChild(tenantIDStr)
 			m.AbortSpanBytes = sm.AbortSpanBytes.AddChild(tenantIDStr)
 			m.mu.Unlock()
-			return &m.ref
+			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)
-
-	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
 	}
 
-	if atomic.SwapInt32(&ref._state, 1) != 0 {
-		log.FatalfDepth(ctx, 1, "tenant metrics already released in:\n%s", ref._stack.SafeStack)
-		return // unreachable
-	}
+	m.mu.released.Store(true)
 
 	// The refCount is zero, delete this instance after destroying its metrics.
 	// Note that concurrent attempts to create an instance will detect the zero
@@ -3419,27 +3387,26 @@ 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 {
-	ref tenantMetricsRef
-	mu  struct {
+	tenantID roachpb.TenantID
+	mu       struct {
 		syncutil.Mutex
 		refCount int
+		released atomic.Bool // allowed to read without holding mu
+		stack    debugutil.SafeStack
 	}
 
 	LiveBytes      *aggmetric.Gauge
@@ -3464,6 +3431,14 @@ 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{
@@ -4023,10 +3998,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)
@@ -4050,17 +4024,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/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