feat: pass NodeProvider to metricstore Init

Setting the provider after Init left the checkpoint restore inside Init
without job information and forced callers to pre-call InitMetrics and
SetNodeProvider in the right order. Init now receives the provider as a
parameter and injects it before the restore, so the ordering is enforced
by the signature. SetNodeProvider remains for callers that do not run
Init (tests, -cleanup-checkpoints).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Aditya Ujeniya
2026-07-08 15:44:54 +02:00
parent ea2bfbb6aa
commit dd2f5968b1
2 changed files with 14 additions and 12 deletions
+4 -9
View File
@@ -367,15 +367,10 @@ func runServer(ctx context.Context) error {
haveMetricstore := false haveMetricstore := false
mscfg := ccconf.GetPackageConfig("metric-store") mscfg := ccconf.GetPackageConfig("metric-store")
if mscfg != nil { if mscfg != nil {
metrics := metricstore.BuildMetricList() // The repository is injected as NodeProvider (breaking the import
// The NodeProvider MUST be set before Init: the checkpoint load inside // cycle) so the checkpoint load inside Init can fetch the full
// Init consults it to load the full history for nodes with running // history for nodes with running jobs.
// jobs. InitMetrics is sync.Once-guarded, so the second call inside metricstore.Init(mscfg, metricstore.BuildMetricList(), repository.GetJobRepository(), &wg)
// Init is a no-op. The repository is injected as NodeProvider to
// break the import cycle.
metricstore.InitMetrics(metrics)
metricstore.GetMemoryStore().SetNodeProvider(repository.GetJobRepository())
metricstore.Init(mscfg, metrics, &wg)
metricstore.MetricStoreHandle = &metricstore.InternalMetricStore{} metricstore.MetricStoreHandle = &metricstore.InternalMetricStore{}
haveMetricstore = true haveMetricstore = true
+10 -3
View File
@@ -115,6 +115,9 @@ type MemoryStore struct {
// Parameters: // Parameters:
// - rawConfig: JSON configuration for the metric store (see MetricStoreConfig); may be nil to use defaults // - rawConfig: JSON configuration for the metric store (see MetricStoreConfig); may be nil to use defaults
// - metrics: Map of metric names to their configurations (frequency and aggregation strategy) // - metrics: Map of metric names to their configurations (frequency and aggregation strategy)
// - provider: NodeProvider consulted during the checkpoint restore (and later by
// Free/CleanupCheckpoints) to preserve data for nodes with running jobs; may be
// nil, in which case all provider-aware paths fall back to their plain behavior
// - wg: WaitGroup that will be incremented for each background goroutine started // - wg: WaitGroup that will be incremented for each background goroutine started
// //
// The function will call cclog.Fatal on critical errors during initialization. // The function will call cclog.Fatal on critical errors during initialization.
@@ -122,7 +125,7 @@ type MemoryStore struct {
// //
// Note: Signal handling must be implemented by the caller. Call Shutdown() when // Note: Signal handling must be implemented by the caller. Call Shutdown() when
// receiving termination signals to ensure checkpoint data is persisted. // receiving termination signals to ensure checkpoint data is persisted.
func Init(rawConfig json.RawMessage, metrics map[string]MetricConfig, wg *sync.WaitGroup) { func Init(rawConfig json.RawMessage, metrics map[string]MetricConfig, provider NodeProvider, wg *sync.WaitGroup) {
startupTime := time.Now() startupTime := time.Now()
if rawConfig != nil { if rawConfig != nil {
@@ -144,6 +147,9 @@ func Init(rawConfig json.RawMessage, metrics map[string]MetricConfig, wg *sync.W
InitMetrics(metrics) InitMetrics(metrics)
ms := GetMemoryStore() ms := GetMemoryStore()
if provider != nil {
ms.SetNodeProvider(provider)
}
d, err := time.ParseDuration(Keys.RetentionInMemory) d, err := time.ParseDuration(Keys.RetentionInMemory)
if err != nil { if err != nil {
@@ -254,8 +260,9 @@ func (ms *MemoryStore) GetMetricFrequency(metricName string) (int64, error) {
// The provider supplies the set of nodes in use by running jobs, which is // The provider supplies the set of nodes in use by running jobs, which is
// consulted by Free (selective buffer retention), FromCheckpoint (full-history // consulted by Free (selective buffer retention), FromCheckpoint (full-history
// loading for used hosts), and CleanupCheckpoints (skipping used hosts). // loading for used hosts), and CleanupCheckpoints (skipping used hosts).
// It must be set before Init(): the checkpoint load inside Init reads it. // Server startup passes the provider to Init() directly; this setter serves
// If not set, all three fall back to their provider-less behavior. // callers that do not run Init (tests, the -cleanup-checkpoints CLI path).
// If not set, all provider-aware paths fall back to their plain behavior.
func (ms *MemoryStore) SetNodeProvider(provider NodeProvider) { func (ms *MemoryStore) SetNodeProvider(provider NodeProvider) {
ms.nodeProvider = provider ms.nodeProvider = provider
} }