From dd2f5968b198b1e636613399c01ae56e16f9eb85 Mon Sep 17 00:00:00 2001 From: Aditya Ujeniya Date: Wed, 8 Jul 2026 15:44:54 +0200 Subject: [PATCH] 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 --- cmd/cc-backend/main.go | 13 ++++--------- pkg/metricstore/metricstore.go | 13 ++++++++++--- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/cmd/cc-backend/main.go b/cmd/cc-backend/main.go index d515332f..c9eebeb5 100644 --- a/cmd/cc-backend/main.go +++ b/cmd/cc-backend/main.go @@ -367,15 +367,10 @@ func runServer(ctx context.Context) error { haveMetricstore := false mscfg := ccconf.GetPackageConfig("metric-store") if mscfg != nil { - metrics := metricstore.BuildMetricList() - // The NodeProvider MUST be set before Init: the checkpoint load inside - // Init consults it to load the full history for nodes with running - // jobs. InitMetrics is sync.Once-guarded, so the second call inside - // 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) + // The repository is injected as NodeProvider (breaking the import + // cycle) so the checkpoint load inside Init can fetch the full + // history for nodes with running jobs. + metricstore.Init(mscfg, metricstore.BuildMetricList(), repository.GetJobRepository(), &wg) metricstore.MetricStoreHandle = &metricstore.InternalMetricStore{} haveMetricstore = true diff --git a/pkg/metricstore/metricstore.go b/pkg/metricstore/metricstore.go index c763a903..2d9ba327 100644 --- a/pkg/metricstore/metricstore.go +++ b/pkg/metricstore/metricstore.go @@ -115,6 +115,9 @@ type MemoryStore struct { // Parameters: // - 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) +// - 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 // // 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 // 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() if rawConfig != nil { @@ -144,6 +147,9 @@ func Init(rawConfig json.RawMessage, metrics map[string]MetricConfig, wg *sync.W InitMetrics(metrics) ms := GetMemoryStore() + if provider != nil { + ms.SetNodeProvider(provider) + } d, err := time.ParseDuration(Keys.RetentionInMemory) 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 // consulted by Free (selective buffer retention), FromCheckpoint (full-history // loading for used hosts), and CleanupCheckpoints (skipping used hosts). -// It must be set before Init(): the checkpoint load inside Init reads it. -// If not set, all three fall back to their provider-less behavior. +// Server startup passes the provider to Init() directly; this setter serves +// 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) { ms.nodeProvider = provider }