kvserver: plumb nodeCapacityProvider

Previously, we added a NodeCapacityProvider field to StoreGossip but left it nil
in production code In addition, we had a TODO in
1b6a91f
to find a more appropriate home for node capacity provider, since the runtime
load monitor is a node-level concept and was being passed to NewStores. During
implementation, we identified a better home for it. This commit moves
NodeCapacityProvider to be a node level concept as part of StoreConfig and
passes it to each store for use when constructing the store descriptor. Note
that runtime load monitoring currently lacks caching and just computes
periodically. We may consider this as a future improvement.

Epic: none
Release note: none

Author: wenyihu6

pkg/base/testing_knobs.go

@@ -59,4 +59,5 @@ type TestingKnobs struct {
 	LicenseTestingKnobs            ModuleTestingKnobs
 	VecIndexTestingKnobs           ModuleTestingKnobs
 	PolicyRefresherTestingKnobs    ModuleTestingKnobs
+	NodeCapacityProviderKnobs      ModuleTestingKnobs
 }

pkg/kv/kvserver/asim/gossip/gossip.go

@@ -51,7 +51,6 @@ type storeGossiper struct {
 
 func newStoreGossiper(
 	descriptorGetter func(cached bool) roachpb.StoreDescriptor,
-	nodeCapacityProvider kvserver.NodeCapacityProvider,
 	clock timeutil.TimeSource,
 	st *cluster.Settings,
 ) *storeGossiper {
@@ -62,9 +61,8 @@ func newStoreGossiper(
 
 	desc := sg.descriptorGetter(false /* cached */)
 	knobs := kvserver.StoreGossipTestingKnobs{AsyncDisabled: true}
-	sg.local = kvserver.NewStoreGossip(sg, sg, knobs, &st.SV, clock, nodeCapacityProvider)
+	sg.local = kvserver.NewStoreGossip(sg, sg, knobs, &st.SV, clock)
 	sg.local.Ident = roachpb.StoreIdent{StoreID: desc.StoreID, NodeID: desc.Node.NodeID}
-
 	return sg
 }
 
@@ -119,26 +117,17 @@ func NewGossip(s state.State, settings *config.SimulationSettings) *gossip {
 	return g
 }
 
-var _ kvserver.NodeCapacityProvider = &simNodeCapacityProvider{}
-
-type simNodeCapacityProvider struct {
-	localNodeID state.NodeID
-	state       state.State
-}
-
-func (s simNodeCapacityProvider) GetNodeCapacity(_ bool) roachpb.NodeCapacity {
-	return s.state.NodeCapacity(s.localNodeID)
-}
-
 func (g *gossip) addStoreToGossip(s state.State, storeID state.StoreID, nodeID state.NodeID) {
 	// Add the store gossip in an "adding" state initially, this is to avoid
 	// recursive calls to get the store descriptor.
 	g.storeGossip[storeID] = &storeGossiper{addingStore: true}
 	g.storeGossip[storeID] = newStoreGossiper(
 		func(cached bool) roachpb.StoreDescriptor {
-			return s.StoreDescriptors(cached, storeID)[0]
+			nc := s.NodeCapacity(nodeID)
+			desc := s.StoreDescriptors(cached, storeID)[0]
+			desc.NodeCapacity = nc
+			return desc
 		},
-		simNodeCapacityProvider{localNodeID: nodeID, state: s},
 		s.Clock(), g.settings.ST)
 }
 

pkg/kv/kvserver/store.go

@@ -47,6 +47,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage/snaprecv"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness"
+	loadmonitor "github.com/cockroachdb/cockroach/pkg/kv/kvserver/load"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/logstore"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/multiqueue"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/raftentry"
@@ -897,27 +898,28 @@ type Store struct {
 	raftLogQueue         *raftLogQueue      // Raft log truncation queue
 	// Carries out truncations proposed by the raft log queue, and "replicated"
 	// via raft, when they are safe. Created in Store.Start.
-	raftTruncator       *raftLogTruncator
-	raftSnapshotQueue   *raftSnapshotQueue          // Raft repair queue
-	tsMaintenanceQueue  *timeSeriesMaintenanceQueue // Time series maintenance queue
-	scanner             *replicaScanner             // Replica scanner
-	consistencyQueue    *consistencyQueue           // Replica consistency check queue
-	consistencyLimiter  *quotapool.RateLimiter      // Rate limits consistency checks
-	metrics             *StoreMetrics
-	intentResolver      *intentresolver.IntentResolver
-	recoveryMgr         txnrecovery.Manager
-	storeLiveness       storeliveness.Fabric
-	syncWaiters         []*logstore.SyncWaiterLoop
-	raftEntryCache      *raftentry.Cache
-	limiters            batcheval.Limiters
-	txnWaitMetrics      *txnwait.Metrics
-	raftMetrics         *raft.Metrics
-	sstSnapshotStorage  snaprecv.SSTSnapshotStorage
-	protectedtsReader   spanconfig.ProtectedTSReader
-	ctSender            *sidetransport.Sender
-	policyRefresher     *policyrefresher.PolicyRefresher
-	storeGossip         *StoreGossip
-	rebalanceObjManager *RebalanceObjectiveManager
+	raftTruncator        *raftLogTruncator
+	raftSnapshotQueue    *raftSnapshotQueue          // Raft repair queue
+	tsMaintenanceQueue   *timeSeriesMaintenanceQueue // Time series maintenance queue
+	scanner              *replicaScanner             // Replica scanner
+	consistencyQueue     *consistencyQueue           // Replica consistency check queue
+	consistencyLimiter   *quotapool.RateLimiter      // Rate limits consistency checks
+	metrics              *StoreMetrics
+	intentResolver       *intentresolver.IntentResolver
+	recoveryMgr          txnrecovery.Manager
+	storeLiveness        storeliveness.Fabric
+	syncWaiters          []*logstore.SyncWaiterLoop
+	raftEntryCache       *raftentry.Cache
+	limiters             batcheval.Limiters
+	txnWaitMetrics       *txnwait.Metrics
+	raftMetrics          *raft.Metrics
+	sstSnapshotStorage   snaprecv.SSTSnapshotStorage
+	protectedtsReader    spanconfig.ProtectedTSReader
+	ctSender             *sidetransport.Sender
+	policyRefresher      *policyrefresher.PolicyRefresher
+	nodeCapacityProvider *loadmonitor.NodeCapacityProvider
+	storeGossip          *StoreGossip
+	rebalanceObjManager  *RebalanceObjectiveManager
 
 	// kvflowRangeControllerFactory is used for replication AC (flow control) V2
 	// to create new range controllers which mediate the flow of requests to
@@ -1189,6 +1191,10 @@ type StoreConfig struct {
 	// leaseholder replicas. One per node.
 	PolicyRefresher *policyrefresher.PolicyRefresher
 
+	// NodeCapacityProvider is used to provide node capacity information for
+	// store.Descriptor.
+	NodeCapacityProvider *loadmonitor.NodeCapacityProvider
+
 	// TimeSeriesDataStore is an interface used by the store's time series
 	// maintenance queue to dispatch individual maintenance tasks.
 	TimeSeriesDataStore TimeSeriesDataStore
@@ -1485,6 +1491,7 @@ func NewStore(
 		metrics:                           newStoreMetrics(cfg.HistogramWindowInterval),
 		ctSender:                          cfg.ClosedTimestampSender,
 		policyRefresher:                   cfg.PolicyRefresher,
+		nodeCapacityProvider:              cfg.NodeCapacityProvider,
 		ioThresholds:                      &iot,
 		rangeFeedSlowClosedTimestampNudge: singleflight.NewGroup("rangfeed-ct-nudge", "range"),
 	}
@@ -1726,9 +1733,13 @@ func NewStore(
 		updateSystemConfigUpdateQueueLimits)
 
 	if s.cfg.Gossip != nil {
-		// TODO(wenyihu6): pass nodeCapacityProvider properly in production code.
-		s.storeGossip = NewStoreGossip(cfg.Gossip,
-			s, cfg.TestingKnobs.GossipTestingKnobs, &cfg.Settings.SV, timeutil.DefaultTimeSource{}, nil)
+		s.storeGossip = NewStoreGossip(
+			cfg.Gossip,
+			s,
+			cfg.TestingKnobs.GossipTestingKnobs,
+			&cfg.Settings.SV,
+			timeutil.DefaultTimeSource{},
+		)
 
 		// Add range scanner and configure with queues.
 		s.scanner = newReplicaScanner(
@@ -3276,11 +3287,12 @@ func (s *Store) Descriptor(ctx context.Context, useCached bool) (*roachpb.StoreD
 
 	// Initialize the store descriptor.
 	return &roachpb.StoreDescriptor{
-		StoreID:    s.Ident.StoreID,
-		Attrs:      s.Attrs(),
-		Node:       *s.nodeDesc,
-		Capacity:   capacity,
-		Properties: s.Properties(),
+		StoreID:      s.Ident.StoreID,
+		Attrs:        s.Attrs(),
+		Node:         *s.nodeDesc,
+		Capacity:     capacity,
+		Properties:   s.Properties(),
+		NodeCapacity: s.nodeCapacityProvider.GetNodeCapacity(useCached),
 	}, nil
 }
 

pkg/kv/kvserver/store_gossip.go

@@ -249,7 +249,6 @@ type StoreGossip struct {
 	descriptorGetter StoreDescriptorProvider
 	sv               *settings.Values
 	clock            timeutil.TimeSource
-	ncProvider       NodeCapacityProvider
 }
 
 // StoreGossipTestingKnobs defines the testing knobs specific to StoreGossip.
@@ -281,7 +280,6 @@ func NewStoreGossip(
 	testingKnobs StoreGossipTestingKnobs,
 	sv *settings.Values,
 	clock timeutil.TimeSource,
-	nodeCapacityProvider NodeCapacityProvider,
 ) *StoreGossip {
 	return &StoreGossip{
 		cachedCapacity:   &cachedCapacity{},
@@ -290,7 +288,6 @@ func NewStoreGossip(
 		knobs:            testingKnobs,
 		sv:               sv,
 		clock:            clock,
-		ncProvider:       nodeCapacityProvider,
 	}
 }
 
@@ -382,12 +379,6 @@ func (s *StoreGossip) GossipStore(ctx context.Context, useCached bool) error {
 		return errors.Wrapf(err, "problem getting store descriptor for store %+v", s.Ident)
 	}
 
-	// TODO(wenyihu6): ncProvider is nil during production code. Only populated
-	// for asim.
-	if s.ncProvider != nil {
-		storeDesc.NodeCapacity = s.ncProvider.GetNodeCapacity(useCached)
-	}
-
 	// Set countdown target for re-gossiping capacity to be large enough that
 	// it would only occur when there has been significant changes. We
 	// currently gossip every 10 seconds, meaning that unless significant
@@ -582,7 +573,3 @@ func deltaExceedsThreshold(
 	exceeds = deltaAbsolute >= requiredMinDelta && deltaFraction >= requiredDeltaFraction
 	return exceeds, delta
 }
-
-type NodeCapacityProvider interface {
-	GetNodeCapacity(useCached bool) roachpb.NodeCapacity
-}

pkg/kv/kvserver/store_gossip_test.go

@@ -108,7 +108,6 @@ func TestStoreGossipDeltaTrigger(t *testing.T) {
 				cfg.TestingKnobs.GossipTestingKnobs,
 				&cluster.MakeTestingClusterSettings().SV,
 				timeutil.DefaultTimeSource{},
-				nil,
 			)
 			sg.cachedCapacity.cached = tc.cached
 			sg.cachedCapacity.lastGossiped = tc.lastGossiped

pkg/kv/kvserver/store_test.go

@@ -38,6 +38,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness/livenesspb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/load"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/rditer"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/stateloader"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/storeliveness"
@@ -259,6 +260,9 @@ func createTestStoreWithoutStart(
 
 	stores := NewStores(cfg.AmbientCtx, cfg.Clock)
 	nodeDesc := &roachpb.NodeDescriptor{NodeID: 1}
+	if cfg.NodeCapacityProvider == nil {
+		cfg.NodeCapacityProvider = load.NewNodeCapacityProvider(stopper, stores, nil)
+	}
 
 	rangeProv := &dummyFirstRangeProvider{}
 	var storeSender struct{ kv.Sender }
@@ -4215,3 +4219,52 @@ func BenchmarkStoreGetReplica(b *testing.B) {
 		}
 	})
 }
+
+// TestNewNodeCapacityProviderCluster tests the basic functionality of the
+// NodeCapacityProvider with a real cluster.
+func TestNewNodeCapacityProviderCluster(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	numNodes := 3
+	numStoresPerNode := 2
+	var storeSpecs []base.StoreSpec
+	for i := 0; i < numStoresPerNode; i++ {
+		storeSpecs = append(storeSpecs, base.StoreSpec{InMemory: true})
+	}
+	serverArgs := base.TestServerArgs{
+		Knobs: base.TestingKnobs{
+			NodeCapacityProviderKnobs: &load.NodeCapacityProviderTestingKnobs{
+				CpuUsageRefreshInterval:    1 * time.Millisecond,
+				CpuCapacityRefreshInterval: 1 * time.Millisecond,
+			},
+		}, StoreSpecs: storeSpecs}
+	tcArgs := base.TestClusterArgs{
+		ParallelStart:   true,
+		ReplicationMode: base.ReplicationManual, // saves time
+		ServerArgsPerNode: map[int]base.TestServerArgs{
+			0: serverArgs,
+			1: serverArgs,
+		},
+	}
+
+	ctx := context.Background()
+	tc := serverutils.StartCluster(t, numNodes, tcArgs)
+	defer tc.Stopper().Stop(ctx)
+	store, err := tc.Server(0).GetStores().(*Stores).GetStore(tc.Server(0).GetFirstStoreID())
+	require.NoError(t, err)
+	testutils.SucceedsSoon(t, func() error {
+		storeDesc, err := store.Descriptor(ctx, false /*useCached*/)
+		require.NoError(t, err)
+		nc := storeDesc.NodeCapacity
+		require.Equal(t, int32(numStoresPerNode), nc.NumStores)
+		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",
+				nc.NodeCPURateUsage, nc.NodeCPURateCapacity, nc.StoresCPURate)
+		}
+		// TODO(wenyihu6): NodeCPURateCapacity <= NodeCPURateUsage fails on CI and
+		// requires more investigation.
+		return nil
+	})
+}

pkg/kv/kvserver/stores.go

@@ -336,8 +336,17 @@ func (ls *Stores) GetStoreMetricRegistry(storeID roachpb.StoreID) *metric.Regist
 	return nil
 }
 
-// TODO(wenyihu6): Implment properly for production code. This now just returns
-// an empty struct and unused everywhere.
-func (ls *Stores) GetNodeCapacity(useCached bool) roachpb.NodeCapacity {
-	panic("unimplemented")
+// 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 {
+		c, err := s.Capacity(context.Background(), useCached)
+		if err != nil {
+			panic(err)
+		}
+		storesCPURate += int64(c.CPUPerSecond)
+		numStores++
+		return nil
+	})
+	return storesCPURate, numStores
 }

pkg/server/BUILD.bazel

@@ -142,6 +142,7 @@ go_library(
         "//pkg/kv/kvserver/kvstorage",
         "//pkg/kv/kvserver/liveness",
         "//pkg/kv/kvserver/liveness/livenesspb",
+        "//pkg/kv/kvserver/load",
         "//pkg/kv/kvserver/logstore",
         "//pkg/kv/kvserver/loqrecovery",
         "//pkg/kv/kvserver/loqrecovery/loqrecoverypb",

pkg/server/server.go

@@ -48,6 +48,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness/livenesspb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/load"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/logstore"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/loqrecovery"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/protectedts"
@@ -161,13 +162,14 @@ type topLevelServer struct {
 	appRegistry  *metric.Registry
 	sysRegistry  *metric.Registry
 
-	recorder         *status.MetricsRecorder
-	runtime          *status.RuntimeStatSampler
-	ruleRegistry     *metric.RuleRegistry
-	promRuleExporter *metric.PrometheusRuleExporter
-	updates          *diagnostics.UpdateChecker
-	ctSender         *sidetransport.Sender
-	policyRefresher  *policyrefresher.PolicyRefresher
+	recorder             *status.MetricsRecorder
+	runtime              *status.RuntimeStatSampler
+	ruleRegistry         *metric.RuleRegistry
+	promRuleExporter     *metric.PrometheusRuleExporter
+	updates              *diagnostics.UpdateChecker
+	ctSender             *sidetransport.Sender
+	policyRefresher      *policyrefresher.PolicyRefresher
+	nodeCapacityProvider *load.NodeCapacityProvider
 
 	http            *httpServer
 	adminAuthzCheck privchecker.CheckerForRPCHandlers
@@ -688,6 +690,14 @@ func NewServer(cfg Config, stopper *stop.Stopper) (serverctl.ServerStartupInterf
 		policyRefresher = policyrefresher.NewPolicyRefresher(stopper, st, ctSender.GetLeaseholders,
 			rpcContext.RemoteClocks.AllLatencies, knobs)
 	}
+	var nodeCapacityProvider *load.NodeCapacityProvider
+	{
+		var knobs *load.NodeCapacityProviderTestingKnobs
+		if nodeCapacityProviderKnobs := cfg.TestingKnobs.NodeCapacityProviderKnobs; nodeCapacityProviderKnobs != nil {
+			knobs = nodeCapacityProviderKnobs.(*load.NodeCapacityProviderTestingKnobs)
+		}
+		nodeCapacityProvider = load.NewNodeCapacityProvider(stopper, stores, knobs)
+	}
 
 	// The Executor will be further initialized later, as we create more
 	// of the server's components. There's a circular dependency - many things
@@ -907,6 +917,7 @@ func NewServer(cfg Config, stopper *stop.Stopper) (serverctl.ServerStartupInterf
 		ClosedTimestampSender:        ctSender,
 		ClosedTimestampReceiver:      ctReceiver,
 		PolicyRefresher:              policyRefresher,
+		NodeCapacityProvider:         nodeCapacityProvider,
 		ProtectedTimestampReader:     protectedTSReader,
 		EagerLeaseAcquisitionLimiter: eagerLeaseAcquisitionLimiter,
 		KVMemoryMonitor:              kvMemoryMonitor,
@@ -1340,6 +1351,7 @@ func NewServer(cfg Config, stopper *stop.Stopper) (serverctl.ServerStartupInterf
 		updates:                   updates,
 		ctSender:                  ctSender,
 		policyRefresher:           policyRefresher,
+		nodeCapacityProvider:      nodeCapacityProvider,
 		runtime:                   runtimeSampler,
 		http:                      sHTTP,
 		adminAuthzCheck:           adminAuthzCheck,
@@ -2177,6 +2189,10 @@ func (s *topLevelServer) PreStart(ctx context.Context) error {
 	// closed timestamp policies for ranges periodically.
 	s.policyRefresher.Run(workersCtx)
 
+	// Start node capacity provider in the background. It refreshes node cpu usage
+	// and capacity for store descriptor.
+	s.nodeCapacityProvider.Run(workersCtx)
+
 	// Start dispatching extant flow tokens.
 	if err := s.raftTransport.Start(workersCtx); err != nil {
 		return err