kvserver: return error properly for GetAggregatedStoreStats

Previously, GetAggregatedStoreStats would panic on error for store.Capacity
instead of returning it. This commit fixes the issue by properly returning the
error back to store.Descriptor.

Author: wenyihu6

pkg/kv/kvserver/load/node_capacity_provider.go

@@ -24,7 +24,7 @@ type StoresStatsAggregator interface {
 	// GetAggregatedStoreStats returns the total cpu usage across all stores and
 	// the count of stores. If useCached is true, it uses the cached store
 	// descriptor instead of computing new ones. Implemented by Stores.
-	GetAggregatedStoreStats(useCached bool) (aggregatedCPUUsage int64, totalStoreCount int32)
+	GetAggregatedStoreStats(useCached bool) (aggregatedCPUUsage int64, totalStoreCount int32, err error)
 }
 
 // NodeCapacityProvider reports node-level cpu usage and capacity by sampling
@@ -82,8 +82,11 @@ func (n *NodeCapacityProvider) Run(ctx context.Context) {
 // capacity and aggregated store-level cpu usage. If useCached is true, it will
 // use cached store descriptors to aggregate the sum of store-level cpu
 // capacity.
-func (n *NodeCapacityProvider) GetNodeCapacity(useCached bool) roachpb.NodeCapacity {
-	storesCPURate, numStores := n.stores.GetAggregatedStoreStats(useCached)
+func (n *NodeCapacityProvider) GetNodeCapacity(useCached bool) (roachpb.NodeCapacity, error) {
+	storesCPURate, numStores, err := n.stores.GetAggregatedStoreStats(useCached)
+	if err != nil {
+		return roachpb.NodeCapacity{}, err
+	}
 	// TODO(wenyihu6): may be unexpected to caller that useCached only applies to
 	// the stores stats but not runtime load monitor. We can change
 	// runtimeLoadMonitor to also fetch updated stats.
@@ -95,7 +98,7 @@ func (n *NodeCapacityProvider) GetNodeCapacity(useCached bool) roachpb.NodeCapac
 		NumStores:           numStores,
 		NodeCPURateCapacity: cpuCapacityNanoPerSec,
 		NodeCPURateUsage:    cpuUsageNanoPerSec,
-	}
+	}, nil
 }
 
 // runtimeLoadMonitor polls cpu usage and capacity stats of the node

pkg/kv/kvserver/load/node_capacity_provider_test.go

@@ -25,8 +25,8 @@ type mockStoresStatsAggregator struct {
 
 func (m *mockStoresStatsAggregator) GetAggregatedStoreStats(
 	_ bool,
-) (totalCPUUsage int64, totalStoreCount int32) {
-	return m.cpuUsage, m.storeCount
+) (totalCPUUsage int64, totalStoreCount int32, _ error) {
+	return m.cpuUsage, m.storeCount, nil
 }
 
 // TestNodeCapacityProvider tests the basic functionality of the
@@ -51,7 +51,8 @@ func TestNodeCapacityProvider(t *testing.T) {
 
 	// Provider should have valid stats.
 	testutils.SucceedsSoon(t, func() error {
-		nc := provider.GetNodeCapacity(false)
+		nc, err := provider.GetNodeCapacity(false)
+		require.NoError(t, err)
 		if nc.NodeCPURateUsage == 0 || nc.NodeCPURateCapacity == 0 || nc.StoresCPURate == 0 {
 			return errors.Newf(
 				"CPU usage or capacity is 0: node cpu rate usage %v, node cpu rate capacity %v, stores cpu rate %v",
@@ -62,7 +63,8 @@ func TestNodeCapacityProvider(t *testing.T) {
 
 	cancel()
 	// GetNodeCapacity should still return valid stats after cancellation.
-	nc := provider.GetNodeCapacity(false)
+	nc, err := provider.GetNodeCapacity(false)
+	require.NoError(t, err)
 	require.Greater(t, nc.NodeCPURateCapacity, int64(0))
 	require.Greater(t, nc.NodeCPURateUsage, int64(0))
 }

pkg/kv/kvserver/store.go

@@ -3295,14 +3295,19 @@ func (s *Store) Descriptor(ctx context.Context, useCached bool) (*roachpb.StoreD
 		return nil, err
 	}
 
+	nc, err := s.nodeCapacityProvider.GetNodeCapacity(useCached)
+	if err != nil {
+		return nil, err
+	}
+
 	// Initialize the store descriptor.
 	return &roachpb.StoreDescriptor{
 		StoreID:      s.Ident.StoreID,
 		Attrs:        s.Attrs(),
 		Node:         *s.nodeDesc,
 		Capacity:     capacity,
 		Properties:   s.Properties(),
-		NodeCapacity: s.nodeCapacityProvider.GetNodeCapacity(useCached),
+		NodeCapacity: nc,
 	}, nil
 }
 

pkg/kv/kvserver/stores.go

@@ -338,15 +338,19 @@ func (ls *Stores) GetStoreMetricRegistry(storeID roachpb.StoreID) *metric.Regist
 
 // GetAggregatedStoreStats returns the aggregated cpu usage across all stores and
 // the count of stores.
-func (ls *Stores) GetAggregatedStoreStats(useCached bool) (storesCPURate int64, numStores int32) {
-	_ = ls.VisitStores(func(s *Store) error {
+func (ls *Stores) GetAggregatedStoreStats(
+	useCached bool,
+) (storesCPURate int64, numStores int32, _ error) {
+	if err := ls.VisitStores(func(s *Store) error {
 		c, err := s.Capacity(context.Background(), useCached)
 		if err != nil {
-			panic(err)
+			return err
 		}
 		storesCPURate += int64(c.CPUPerSecond)
 		numStores++
 		return nil
-	})
-	return storesCPURate, numStores
+	}); err != nil {
+		return 0, 0, err
+	}
+	return storesCPURate, numStores, nil
 }