@@ -33,6 +33,13 @@ type Pool struct {
3333
3434 clusters * smap.Map [* Cluster ]
3535 synchronization * synchronization.Synchronize [queries.Cluster , * Cluster ]
36+ store clustersSyncStore
37+
38+ cancel context.CancelFunc
39+ lifecycleDone <- chan struct {}
40+ startupReady chan struct {}
41+ startupReadyOnce sync.Once
42+ lifecycleMu * sync.Mutex
3643}
3744
3845func localClusterConfig () * queries.Cluster {
@@ -55,33 +62,43 @@ func NewPool(
5562 config cfg.Config ,
5663) (* Pool , error ) {
5764 clusters := smap .New [* Cluster ]()
58-
59- localCluster := localClusterConfig ()
60-
65+ poolCtx , cancel := context .WithCancel (ctx )
6166 p := & Pool {
62- db : db ,
63- tel : tel ,
64- clusters : clusters ,
65- synchronization : synchronization .NewSynchronize (
66- "clusters-pool" ,
67- "Clusters pool" ,
68- clustersSyncStore {
69- config : config ,
70- db : db ,
71- tel : tel ,
72- clusters : clusters ,
73- local : localCluster ,
74- localDiscovery : localDiscovery ,
75- queryLogsProvider : queryLogsProvider ,
76- queryMetricsProvider : queryMetricsProvider ,
77- sandboxLogsReader : sandboxLogsReader ,
78- featureFlags : featureFlags ,
79- },
80- ),
81- }
82-
83- // Periodically sync clusters with the database
84- go p .synchronization .Start (ctx , clustersSyncInterval , clusterSyncTimeout , true )
67+ db : db ,
68+ tel : tel ,
69+ clusters : clusters ,
70+ cancel : cancel ,
71+ lifecycleDone : poolCtx .Done (),
72+ startupReady : make (chan struct {}),
73+ lifecycleMu : & sync.Mutex {},
74+ }
75+ p .store = clustersSyncStore {
76+ config : config ,
77+ db : db ,
78+ tel : tel ,
79+ clusters : clusters ,
80+ local : localClusterConfig (),
81+ localDiscovery : localDiscovery ,
82+ queryLogsProvider : queryLogsProvider ,
83+ queryMetricsProvider : queryMetricsProvider ,
84+ sandboxLogsReader : sandboxLogsReader ,
85+ featureFlags : featureFlags ,
86+ activateCluster : func (cluster * Cluster , mode initialSyncMode ) error {
87+ return p .activateCluster (poolCtx , cluster , mode )
88+ },
89+ }
90+ p .synchronization = synchronization .NewSynchronize ("clusters-pool" , "Clusters pool" , p .store )
91+
92+ coordinator := newStartupCoordinator (
93+ func (inventoryCtx context.Context ) ([]queries.Cluster , error ) {
94+ return p .initialInventory (poolCtx , inventoryCtx )
95+ },
96+ func (attemptCtx context.Context , cluster queries.Cluster ) error {
97+ return p .startInitialRemoteCluster (poolCtx , attemptCtx , cluster )
98+ },
99+ p .finishStartup ,
100+ )
101+ go coordinator .Start (poolCtx )
85102
86103 return p , nil
87104}
@@ -94,7 +111,16 @@ func (p *Pool) GetClusters() map[string]*Cluster {
94111 return p .clusters .Items ()
95112}
96113
114+ // StartupReady closes after the first successful inventory and every remote
115+ // target in that inventory has applied a snapshot or exhausted startup retries.
116+ func (p * Pool ) StartupReady () <- chan struct {} {
117+ return p .startupReady
118+ }
119+
97120func (p * Pool ) Close (ctx context.Context ) {
121+ p .lifecycleMu .Lock ()
122+ p .cancel ()
123+ p .lifecycleMu .Unlock ()
98124 p .synchronization .Close ()
99125
100126 wg := & sync.WaitGroup {}
@@ -122,6 +148,7 @@ type clustersSyncStore struct {
122148 sandboxLogsReader ClickhouseLogsReader
123149 featureFlags * featureflags.Client
124150 config cfg.Config
151+ activateCluster func (* Cluster , initialSyncMode ) error
125152}
126153
127154func (d clustersSyncStore ) SourceList (ctx context.Context ) ([]queries.Cluster , error ) {
@@ -169,30 +196,35 @@ func (d clustersSyncStore) PoolExists(_ context.Context, cluster queries.Cluster
169196}
170197
171198func (d clustersSyncStore ) PoolInsert (ctx context.Context , cluster queries.Cluster ) {
172- clusterID := cluster .ID .String ()
173-
174199 logger .L ().Info (ctx , "Initializing newly discovered cluster" , logger .WithClusterID (cluster .ID ))
175200
176- var c * Cluster
177- var err error
201+ c , err := d .newCluster (cluster )
202+ if err != nil {
203+ logger .L ().Error (ctx , "Initializing cluster failed" , zap .Error (err ), logger .WithClusterID (cluster .ID ))
178204
179- // Local cluster
180- if cluster . ID == consts . LocalClusterID {
181- c = newLocalCluster ( context . WithoutCancel ( ctx ), d . tel , d . localDiscovery , d . queryMetricsProvider , d . queryLogsProvider , d . sandboxLogsReader , d . featureFlags , d . config )
182- d . clusters . Insert ( clusterID , c )
183- logger .L ().Info (ctx , "Local cluster initialized successfully" , logger .WithClusterID (cluster .ID ))
205+ return
206+ }
207+
208+ if err := d . activateCluster ( c , syncImmediately ); err != nil {
209+ logger .L ().Info (ctx , "Skipped cluster initialization during shutdown" , zap . Error ( err ) , logger .WithClusterID (cluster .ID ))
184210
185211 return
186212 }
187213
188- // Remote cluster
214+ logger .L ().Info (ctx , "Cluster initialized successfully" , logger .WithClusterID (cluster .ID ))
215+ }
216+
217+ func (d clustersSyncStore ) newCluster (cluster queries.Cluster ) (* Cluster , error ) {
218+ if cluster .ID == consts .LocalClusterID {
219+ return newLocalCluster (d .tel , d .localDiscovery , d .queryMetricsProvider , d .queryLogsProvider , d .sandboxLogsReader , d .featureFlags , d .config ), nil
220+ }
221+
189222 authOrgID := ""
190223 if cluster .AuthOrgID != nil {
191224 authOrgID = * cluster .AuthOrgID
192225 }
193226
194- c , err = newRemoteCluster (
195- context .WithoutCancel (ctx ),
227+ return newRemoteCluster (
196228 d .tel ,
197229 cluster .Endpoint ,
198230 cluster .EndpointTls ,
@@ -201,14 +233,89 @@ func (d clustersSyncStore) PoolInsert(ctx context.Context, cluster queries.Clust
201233 cluster .SandboxProxyDomain ,
202234 authOrgID ,
203235 )
236+ }
237+
238+ func (p * Pool ) initialInventory (lifecycleCtx context.Context , ctx context.Context ) ([]queries.Cluster , error ) {
239+ clusters , err := p .store .SourceList (ctx )
204240 if err != nil {
205- logger .L ().Error (ctx , "Initializing remote cluster failed" , zap .Error (err ), logger .WithClusterID (cluster .ID ))
241+ return nil , err
242+ }
243+
244+ targets := make ([]queries.Cluster , 0 , len (clusters ))
245+ for _ , cluster := range clusters {
246+ if cluster .ID != consts .LocalClusterID {
247+ targets = append (targets , cluster )
248+
249+ continue
250+ }
251+
252+ local , localErr := p .store .newCluster (cluster )
253+ if localErr != nil {
254+ return nil , localErr
255+ }
256+ if err := p .activateCluster (lifecycleCtx , local , syncImmediately ); err != nil {
257+ return nil , err
258+ }
259+ }
260+
261+ return targets , nil
262+ }
263+
264+ func (p * Pool ) startInitialRemoteCluster (lifecycleCtx context.Context , ctx context.Context , cluster queries.Cluster ) error {
265+ c , err := p .store .newCluster (cluster )
266+ if err != nil {
267+ return err
268+ }
269+
270+ if err := c .SyncInstances (ctx ); err != nil {
271+ _ = c .Close (ctx )
272+
273+ return err
274+ }
275+
276+ // The snapshot is already applied, so the periodic loop starts at its first
277+ // scheduled tick rather than repeating discovery now.
278+ if err := p .activateCluster (lifecycleCtx , c , syncOnNextTick ); err != nil {
279+ _ = c .Close (ctx )
280+
281+ return err
282+ }
283+
284+ return nil
285+ }
286+
287+ // activateCluster publishes a constructed cluster: it starts the cluster's
288+ // periodic discovery and adds it to the pool. Both happen under lifecycleMu,
289+ // the mutex Close holds while cancelling, so a cluster can never start or
290+ // become visible after shutdown began. Callers construct the cluster and run
291+ // any network calls outside this method, both to keep those off the mutex and
292+ // because the caller owns the cluster until activation succeeds.
293+ func (p * Pool ) activateCluster (lifecycleCtx context.Context , cluster * Cluster , mode initialSyncMode ) error {
294+ p .lifecycleMu .Lock ()
295+ defer p .lifecycleMu .Unlock ()
296+
297+ if err := lifecycleCtx .Err (); err != nil {
298+ return err
299+ }
300+
301+ cluster .Start (lifecycleCtx , mode )
302+ p .clusters .Insert (cluster .ID .String (), cluster )
303+
304+ return nil
305+ }
306+
307+ func (p * Pool ) finishStartup (lifecycleCtx context.Context ) {
308+ p .lifecycleMu .Lock ()
309+ defer p .lifecycleMu .Unlock ()
206310
311+ if lifecycleCtx .Err () != nil {
207312 return
208313 }
209314
210- d .clusters .Insert (clusterID , c )
211- logger .L ().Info (ctx , "Remote cluster initialized successfully" , logger .WithClusterID (cluster .ID ))
315+ p .startupReadyOnce .Do (func () {
316+ close (p .startupReady )
317+ go p .synchronization .Start (lifecycleCtx , clustersSyncInterval , clusterSyncTimeout , false )
318+ })
212319}
213320
214321func (d clustersSyncStore ) PoolUpdate (_ context.Context , _ * Cluster ) {
0 commit comments