Skip to main content

rskit_bootstrap/
app.rs

1use std::marker::PhantomData;
2use std::sync::Arc;
3use std::time::Duration;
4
5use rskit_component::{Component, Registry};
6use rskit_config::AppConfig;
7use rskit_di::Container;
8use rskit_errors::{AppError, AppResult, ErrorCode};
9use rskit_hook::{EventBus, EventBusConfig};
10use rskit_provider::Provider;
11use tokio_util::sync::CancellationToken;
12
13use crate::hooks::LifecycleEvent;
14use crate::lifecycle::{
15    AsyncHook, LifecycleHooks, LifecyclePhase, make_hook, record_stop_error, run_hooks,
16    wait_for_signal,
17};
18
19/// Typestate marker retained only as the unbuildable initial marker.
20///
21/// Applications are constructed through [`AppBuilder`], which validates config
22/// and returns [`App<Built, C>`].
23pub struct Unconfigured;
24
25/// Typestate marker for a built application whose components are not running.
26pub struct Built;
27
28/// Typestate marker for an application with started components.
29pub struct Started;
30
31/// Typestate marker for an application that has completed shutdown.
32pub struct Stopped;
33
34/// Builder for [`App`].
35///
36/// The builder is the composition root for configuration validation,
37/// component registration, dependency injection, provider registration, and
38/// lifecycle hook ordering.
39pub struct AppBuilder<C: AppConfig> {
40    config: C,
41    graceful_timeout: Duration,
42    components: Vec<Arc<dyn Component>>,
43    container: Arc<Container>,
44    lifecycle_events: EventBus<LifecycleEvent>,
45    hooks: LifecycleHooks,
46}
47
48impl<C: AppConfig> AppBuilder<C> {
49    /// Create a new builder with the given application configuration.
50    #[must_use]
51    pub fn new(config: C) -> Self {
52        Self {
53            config,
54            graceful_timeout: Duration::from_secs(30),
55            components: Vec::new(),
56            container: Arc::new(Container::new()),
57            lifecycle_events: EventBus::new(EventBusConfig::default()),
58            hooks: LifecycleHooks::default(),
59        }
60    }
61
62    /// Set the graceful shutdown timeout.
63    #[must_use]
64    pub fn with_graceful_timeout(mut self, timeout: Duration) -> Self {
65        self.graceful_timeout = timeout;
66        self
67    }
68
69    /// Register a component to be started and stopped by the app lifecycle.
70    #[must_use]
71    pub fn with_component(mut self, component: Arc<dyn Component>) -> Self {
72        self.components.push(component);
73        self
74    }
75
76    /// Use an existing typed DI container as the application container.
77    #[must_use]
78    pub fn with_container(mut self, container: Arc<Container>) -> Self {
79        self.container = container;
80        self
81    }
82
83    /// Use an existing bounded lifecycle event bus.
84    #[must_use]
85    pub fn with_lifecycle_event_bus(mut self, lifecycle_events: EventBus<LifecycleEvent>) -> Self {
86        self.lifecycle_events = lifecycle_events;
87        self
88    }
89
90    /// Register a typed dependency in the application DI container.
91    #[must_use]
92    pub fn with_dependency<T>(self, dependency: Arc<T>) -> Self
93    where
94        T: Send + Sync + 'static,
95    {
96        self.container.register(dependency);
97        self
98    }
99
100    /// Register a provider implementation in the application DI container.
101    #[must_use]
102    pub fn with_provider<P>(self, provider: Arc<P>) -> Self
103    where
104        P: Provider + Send + Sync + 'static,
105    {
106        self.container.register(provider);
107        self
108    }
109
110    /// Register a hook called before components start.
111    #[must_use]
112    pub fn before_start<F, Fut>(mut self, hook: F) -> Self
113    where
114        F: Fn(CancellationToken) -> Fut + Send + Sync + 'static,
115        Fut: std::future::Future<Output = AppResult<()>> + Send + 'static,
116    {
117        self.hooks.before_start.push(make_hook(hook));
118        self
119    }
120
121    /// Register a hook called after components start and readiness checks pass.
122    #[must_use]
123    pub fn after_start<F, Fut>(mut self, hook: F) -> Self
124    where
125        F: Fn(CancellationToken) -> Fut + Send + Sync + 'static,
126        Fut: std::future::Future<Output = AppResult<()>> + Send + 'static,
127    {
128        self.hooks.after_start.push(make_hook(hook));
129        self
130    }
131
132    /// Register a hook called before components stop.
133    #[must_use]
134    pub fn before_stop<F, Fut>(mut self, hook: F) -> Self
135    where
136        F: Fn(CancellationToken) -> Fut + Send + Sync + 'static,
137        Fut: std::future::Future<Output = AppResult<()>> + Send + 'static,
138    {
139        self.hooks.before_stop.push(make_hook(hook));
140        self
141    }
142
143    /// Register a hook called after components stop.
144    #[must_use]
145    pub fn after_stop<F, Fut>(mut self, hook: F) -> Self
146    where
147        F: Fn(CancellationToken) -> Fut + Send + Sync + 'static,
148        Fut: std::future::Future<Output = AppResult<()>> + Send + 'static,
149    {
150        self.hooks.after_stop.push(make_hook(hook));
151        self
152    }
153
154    /// Validate config and build the application in the [`Built`] state.
155    pub fn build(self) -> AppResult<App<Built, C>> {
156        self.config
157            .validate()
158            .map_err(|error| AppError::invalid_input("config", error.to_string()))?;
159
160        let mut registry = Registry::new();
161        for component in self.components {
162            registry.register(component);
163        }
164
165        Ok(App {
166            _state: PhantomData,
167            config: Arc::new(self.config),
168            container: self.container,
169            registry,
170            lifecycle_events: self.lifecycle_events,
171            hooks: self.hooks,
172            graceful_timeout: self.graceful_timeout,
173            shutdown_token: CancellationToken::new(),
174        })
175    }
176}
177
178/// Application orchestrator with typestate lifecycle.
179///
180/// Valid transitions are compile-time scoped:
181///
182/// ```text
183/// AppBuilder::build() -> App<Built, C>
184/// App<Built, C>::start() -> App<Started, C>
185/// App<Started, C>::stop() -> App<Stopped, C>
186/// ```
187pub struct App<S, C> {
188    _state: PhantomData<S>,
189    config: Arc<C>,
190    container: Arc<Container>,
191    registry: Registry,
192    lifecycle_events: EventBus<LifecycleEvent>,
193    hooks: LifecycleHooks,
194    graceful_timeout: Duration,
195    shutdown_token: CancellationToken,
196}
197
198impl<C: AppConfig> App<Built, C> {
199    /// Create an [`AppBuilder`].
200    pub fn builder(config: C) -> AppBuilder<C> {
201        AppBuilder::new(config)
202    }
203}
204
205impl<C> App<Built, C> {
206    /// Return a clone of the shared application configuration.
207    #[must_use]
208    pub fn config(&self) -> Arc<C> {
209        Arc::clone(&self.config)
210    }
211
212    /// Return the typed application DI container.
213    #[must_use]
214    pub fn container(&self) -> Arc<Container> {
215        Arc::clone(&self.container)
216    }
217
218    /// Return the bounded lifecycle event bus used by this app.
219    #[must_use]
220    pub fn lifecycle_event_bus(&self) -> EventBus<LifecycleEvent> {
221        self.lifecycle_events.clone()
222    }
223
224    /// Clone the shutdown token to share with long-running tasks.
225    #[must_use]
226    pub fn shutdown_token(&self) -> CancellationToken {
227        self.shutdown_token.clone()
228    }
229
230    /// Register a hook called before components start.
231    #[must_use]
232    pub fn before_start<F, Fut>(mut self, hook: F) -> Self
233    where
234        F: Fn(CancellationToken) -> Fut + Send + Sync + 'static,
235        Fut: std::future::Future<Output = AppResult<()>> + Send + 'static,
236    {
237        self.hooks.before_start.push(make_hook(hook));
238        self
239    }
240
241    /// Register a hook called after components start and readiness checks pass.
242    #[must_use]
243    pub fn after_start<F, Fut>(mut self, hook: F) -> Self
244    where
245        F: Fn(CancellationToken) -> Fut + Send + Sync + 'static,
246        Fut: std::future::Future<Output = AppResult<()>> + Send + 'static,
247    {
248        self.hooks.after_start.push(make_hook(hook));
249        self
250    }
251
252    /// Register a hook called before components stop.
253    #[must_use]
254    pub fn before_stop<F, Fut>(mut self, hook: F) -> Self
255    where
256        F: Fn(CancellationToken) -> Fut + Send + Sync + 'static,
257        Fut: std::future::Future<Output = AppResult<()>> + Send + 'static,
258    {
259        self.hooks.before_stop.push(make_hook(hook));
260        self
261    }
262
263    /// Register a hook called after components stop.
264    #[must_use]
265    pub fn after_stop<F, Fut>(mut self, hook: F) -> Self
266    where
267        F: Fn(CancellationToken) -> Fut + Send + Sync + 'static,
268        Fut: std::future::Future<Output = AppResult<()>> + Send + 'static,
269    {
270        self.hooks.after_stop.push(make_hook(hook));
271        self
272    }
273}
274
275impl<C: AppConfig> App<Built, C> {
276    /// Start components and transition to [`Started`].
277    pub async fn start(self) -> AppResult<App<Started, C>> {
278        crate::summary::print_startup(self.config.service_config(), self.registry.len());
279        self.run_hooks(LifecyclePhase::BeforeStart, &self.hooks.before_start)
280            .await?;
281
282        if let Err(error) = self.registry.start_all().await {
283            return Err(self
284                .rollback_startup_failure(error.context("component startup failed"), false)
285                .await);
286        }
287
288        if let Err(error) = self.ready_check() {
289            return Err(self
290                .rollback_startup_failure(error.context("component readiness failed"), true)
291                .await);
292        }
293
294        if let Err(error) = self
295            .run_hooks(LifecyclePhase::AfterStart, &self.hooks.after_start)
296            .await
297        {
298            return Err(self
299                .rollback_startup_failure(error.context("after_start hooks failed"), true)
300                .await);
301        }
302
303        Ok(App {
304            _state: PhantomData,
305            config: self.config,
306            container: self.container,
307            registry: self.registry,
308            lifecycle_events: self.lifecycle_events,
309            hooks: self.hooks,
310            graceful_timeout: self.graceful_timeout,
311            shutdown_token: self.shutdown_token,
312        })
313    }
314
315    /// Run as a long-running service until a signal or shutdown token fires.
316    pub async fn run(self) -> AppResult<()> {
317        let started = self.start().await?;
318        wait_for_signal(started.shutdown_token.clone()).await;
319        started.stop().await.map(|_| ())
320    }
321
322    /// Run a finite task after startup, then stop the app.
323    pub async fn run_task<F, Fut>(self, task: F) -> AppResult<()>
324    where
325        F: FnOnce(Arc<C>, CancellationToken) -> Fut,
326        Fut: std::future::Future<Output = AppResult<()>>,
327    {
328        let started = self.start().await?;
329        let task_result = tokio::select! {
330            result = task(Arc::clone(&started.config), started.shutdown_token.clone()) => result,
331            _ = wait_for_signal_owned(started.shutdown_token.clone()) => {
332                Ok(())
333            }
334        };
335        let stop_result = started.stop().await.map(|_| ());
336        match (task_result, stop_result) {
337            (Ok(()), Ok(())) => Ok(()),
338            (Ok(()), Err(stop_error)) => Err(stop_error),
339            (Err(task_error), Ok(())) => Err(task_error),
340            (Err(task_error), Err(stop_error)) => Err(task_error
341                .context("application shutdown also failed after task error")
342                .with_detail("shutdown_error", stop_error.to_string())),
343        }
344    }
345
346    fn ready_check(&self) -> AppResult<()> {
347        let unhealthy = self
348            .registry
349            .health_all()
350            .into_iter()
351            .filter(|health| !health.is_healthy())
352            .map(|health| match health.message {
353                Some(message) => format!("{}={}({message})", health.name, health.status),
354                None => format!("{}={}", health.name, health.status),
355            })
356            .collect::<Vec<_>>();
357
358        if unhealthy.is_empty() {
359            Ok(())
360        } else {
361            Err(AppError::service_unavailable(format!(
362                "unhealthy components: {}",
363                unhealthy.join(", ")
364            )))
365        }
366    }
367
368    async fn rollback_startup_failure(
369        &self,
370        mut error: AppError,
371        stop_components: bool,
372    ) -> AppError {
373        if let Err(stop_hook_error) = self
374            .run_hooks(LifecyclePhase::BeforeStop, &self.hooks.before_stop)
375            .await
376        {
377            tracing::warn!(error = %stop_hook_error, "before_stop hook failed during startup rollback");
378            error = error.context(format!(
379                "startup rollback before_stop hooks failed: {stop_hook_error}"
380            ));
381        }
382
383        if stop_components && let Err(stop_error) = self.registry.stop_all().await {
384            tracing::warn!(error = %stop_error, "component rollback error");
385            error = error.context(format!(
386                "startup rollback component stop failed: {stop_error}"
387            ));
388        }
389        error
390    }
391}
392
393impl<C> App<Built, C> {
394    async fn run_hooks(&self, phase: LifecyclePhase, hooks: &[AsyncHook]) -> AppResult<()> {
395        run_hooks(
396            phase,
397            hooks,
398            self.shutdown_token.clone(),
399            &self.lifecycle_events,
400        )
401        .await
402    }
403}
404
405impl<C> App<Started, C> {
406    /// Return a clone of the shared application configuration.
407    #[must_use]
408    pub fn config(&self) -> Arc<C> {
409        Arc::clone(&self.config)
410    }
411
412    /// Return the typed application DI container.
413    #[must_use]
414    pub fn container(&self) -> Arc<Container> {
415        Arc::clone(&self.container)
416    }
417
418    /// Return the bounded lifecycle event bus used by this app.
419    #[must_use]
420    pub fn lifecycle_event_bus(&self) -> EventBus<LifecycleEvent> {
421        self.lifecycle_events.clone()
422    }
423
424    /// Clone the shutdown token to share with long-running tasks.
425    #[must_use]
426    pub fn shutdown_token(&self) -> CancellationToken {
427        self.shutdown_token.clone()
428    }
429
430    /// Stop components and transition to [`Stopped`].
431    pub async fn stop(self) -> AppResult<App<Stopped, C>> {
432        let mut stop_error: Option<AppError> = None;
433        let stop_result = tokio::time::timeout(
434            self.graceful_timeout,
435            run_hooks(
436                LifecyclePhase::BeforeStop,
437                &self.hooks.before_stop,
438                self.shutdown_token.clone(),
439                &self.lifecycle_events,
440            ),
441        )
442        .await;
443
444        match stop_result {
445            Ok(Ok(())) => {}
446            Ok(Err(error)) => {
447                tracing::warn!(error = %error, "before_stop hook error");
448                stop_error = Some(error);
449            }
450            Err(_) => {
451                tracing::warn!("graceful shutdown timeout elapsed during before_stop hooks");
452                stop_error = Some(AppError::new(
453                    ErrorCode::Timeout,
454                    "graceful shutdown timeout elapsed during before_stop hooks",
455                ));
456            }
457        }
458
459        if let Err(error) = self.registry.stop_all().await {
460            tracing::warn!(error = %error, "component shutdown error");
461            record_stop_error(&mut stop_error, error, "component shutdown failed");
462        }
463
464        if let Err(error) = run_hooks(
465            LifecyclePhase::AfterStop,
466            &self.hooks.after_stop,
467            self.shutdown_token.clone(),
468            &self.lifecycle_events,
469        )
470        .await
471        {
472            record_stop_error(&mut stop_error, error, "after_stop hooks failed");
473        }
474
475        if let Err(error) = self.container.close().await {
476            record_stop_error(&mut stop_error, error, "container close failed");
477        }
478        tracing::info!("shutdown complete");
479
480        if let Some(error) = stop_error {
481            return Err(error);
482        }
483
484        Ok(App {
485            _state: PhantomData,
486            config: self.config,
487            container: self.container,
488            registry: self.registry,
489            lifecycle_events: self.lifecycle_events,
490            hooks: self.hooks,
491            graceful_timeout: self.graceful_timeout,
492            shutdown_token: self.shutdown_token,
493        })
494    }
495}
496
497impl<C> App<Stopped, C> {
498    /// Return a clone of the shared application configuration.
499    #[must_use]
500    pub fn config(&self) -> Arc<C> {
501        Arc::clone(&self.config)
502    }
503
504    /// Return the bounded lifecycle event bus used by this app.
505    #[must_use]
506    pub fn lifecycle_event_bus(&self) -> EventBus<LifecycleEvent> {
507        self.lifecycle_events.clone()
508    }
509}
510
511async fn wait_for_signal_owned(token: CancellationToken) {
512    wait_for_signal(token).await;
513}
514
515#[cfg(test)]
516mod tests {
517    use std::sync::Arc;
518    use std::sync::atomic::{AtomicUsize, Ordering};
519
520    use parking_lot::Mutex;
521    use rskit_config::ServiceConfig;
522    use tokio_util::sync::CancellationToken;
523
524    use super::AppBuilder;
525
526    #[derive(Debug, Default, serde::Deserialize)]
527    struct TestCfg {
528        #[serde(default)]
529        service: rskit_config::ServiceConfig,
530    }
531
532    impl rskit_validation::Validate for TestCfg {
533        fn validate(&self) -> Result<(), validator::ValidationErrors> {
534            rskit_validation::Validate::validate(&self.service)
535        }
536    }
537
538    impl rskit_config::AppConfig for TestCfg {
539        fn apply_defaults(&mut self) {}
540
541        fn service_config(&self) -> &ServiceConfig {
542            &self.service
543        }
544    }
545
546    #[tokio::test]
547    async fn app_builder_builds_successfully() {
548        let result = AppBuilder::new(TestCfg::default()).build();
549        assert!(result.is_ok());
550    }
551
552    #[tokio::test]
553    async fn lifecycle_hooks_run_in_order() {
554        let order = Arc::new(Mutex::new(Vec::new()));
555        let before_start = Arc::clone(&order);
556        let after_start = Arc::clone(&order);
557        let before_stop = Arc::clone(&order);
558        let after_stop = Arc::clone(&order);
559
560        let app = AppBuilder::new(TestCfg::default())
561            .before_start(move |_token| {
562                let order = Arc::clone(&before_start);
563                async move {
564                    order.lock().push("before_start");
565                    Ok(())
566                }
567            })
568            .after_start(move |_token| {
569                let order = Arc::clone(&after_start);
570                async move {
571                    order.lock().push("after_start");
572                    Ok(())
573                }
574            })
575            .before_stop(move |_token| {
576                let order = Arc::clone(&before_stop);
577                async move {
578                    order.lock().push("before_stop");
579                    Ok(())
580                }
581            })
582            .after_stop(move |_token| {
583                let order = Arc::clone(&after_stop);
584                async move {
585                    order.lock().push("after_stop");
586                    Ok(())
587                }
588            })
589            .build()
590            .expect("build should succeed");
591
592        app.run_task(|_cfg: Arc<TestCfg>, _cancel: CancellationToken| async move { Ok(()) })
593            .await
594            .expect("run_task should complete");
595
596        assert_eq!(
597            *order.lock(),
598            vec!["before_start", "after_start", "before_stop", "after_stop"]
599        );
600    }
601
602    #[tokio::test]
603    async fn lifecycle_events_are_published() {
604        let app = AppBuilder::new(TestCfg::default())
605            .build()
606            .expect("build should succeed");
607        let mut subscriber = app.lifecycle_event_bus().subscribe();
608
609        app.run_task(|_cfg: Arc<TestCfg>, _cancel: CancellationToken| async move { Ok(()) })
610            .await
611            .expect("run_task should complete");
612
613        let before_start = subscriber.recv().await.expect("before_start event");
614        let after_start = subscriber.recv().await.expect("after_start event");
615        let before_stop = subscriber.recv().await.expect("before_stop event");
616        let after_stop = subscriber.recv().await.expect("after_stop event");
617        assert_eq!(before_start.kind(), crate::LifecycleEventType::BeforeStart);
618        assert_eq!(after_start.kind(), crate::LifecycleEventType::AfterStart);
619        assert_eq!(before_stop.kind(), crate::LifecycleEventType::BeforeStop);
620        assert_eq!(after_stop.kind(), crate::LifecycleEventType::AfterStop);
621    }
622
623    #[tokio::test]
624    async fn app_run_task_executes_and_exits() {
625        let app = AppBuilder::new(TestCfg::default())
626            .build()
627            .expect("build should succeed");
628        let runs = Arc::new(AtomicUsize::new(0));
629        let run_counter = Arc::clone(&runs);
630
631        app.run_task(move |_cfg: Arc<TestCfg>, _cancel: CancellationToken| {
632            let runs = Arc::clone(&run_counter);
633            async move {
634                runs.fetch_add(1, Ordering::SeqCst);
635                Ok::<(), rskit_errors::AppError>(())
636            }
637        })
638        .await
639        .expect("task should run");
640
641        assert_eq!(runs.load(Ordering::SeqCst), 1);
642    }
643
644    #[tokio::test]
645    async fn hook_failure_stops_startup() {
646        let app = AppBuilder::new(TestCfg::default())
647            .before_start(|_| async {
648                Err::<(), _>(rskit_errors::AppError::new(
649                    rskit_errors::ErrorCode::Internal,
650                    "hard fail",
651                ))
652            })
653            .build()
654            .expect("build should succeed");
655
656        let error = app
657            .run_task(|_cfg: Arc<TestCfg>, _cancel: CancellationToken| async move { Ok(()) })
658            .await
659            .expect_err("hook failure should fail startup");
660        assert!(error.to_string().contains("before_start hook failed"));
661    }
662
663    #[tokio::test]
664    async fn stop_hooks_run_after_shutdown_token_is_cancelled() {
665        let ran = Arc::new(AtomicUsize::new(0));
666        let before_stop_ran = Arc::clone(&ran);
667        let after_stop_ran = Arc::clone(&ran);
668
669        let app = AppBuilder::new(TestCfg::default())
670            .before_stop(move |_token| {
671                let ran = Arc::clone(&before_stop_ran);
672                async move {
673                    ran.fetch_add(1, Ordering::SeqCst);
674                    Ok(())
675                }
676            })
677            .after_stop(move |_token| {
678                let ran = Arc::clone(&after_stop_ran);
679                async move {
680                    ran.fetch_add(1, Ordering::SeqCst);
681                    Ok(())
682                }
683            })
684            .build()
685            .expect("build should succeed")
686            .start()
687            .await
688            .expect("start should succeed");
689
690        app.shutdown_token().cancel();
691        app.stop().await.expect("stop hooks should still run");
692
693        assert_eq!(ran.load(Ordering::SeqCst), 2);
694    }
695}