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
19pub struct Unconfigured;
24
25pub struct Built;
27
28pub struct Started;
30
31pub struct Stopped;
33
34pub 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 #[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 #[must_use]
64 pub fn with_graceful_timeout(mut self, timeout: Duration) -> Self {
65 self.graceful_timeout = timeout;
66 self
67 }
68
69 #[must_use]
71 pub fn with_component(mut self, component: Arc<dyn Component>) -> Self {
72 self.components.push(component);
73 self
74 }
75
76 #[must_use]
78 pub fn with_container(mut self, container: Arc<Container>) -> Self {
79 self.container = container;
80 self
81 }
82
83 #[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 #[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 #[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 #[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 #[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 #[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 #[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 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
178pub 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 pub fn builder(config: C) -> AppBuilder<C> {
201 AppBuilder::new(config)
202 }
203}
204
205impl<C> App<Built, C> {
206 #[must_use]
208 pub fn config(&self) -> Arc<C> {
209 Arc::clone(&self.config)
210 }
211
212 #[must_use]
214 pub fn container(&self) -> Arc<Container> {
215 Arc::clone(&self.container)
216 }
217
218 #[must_use]
220 pub fn lifecycle_event_bus(&self) -> EventBus<LifecycleEvent> {
221 self.lifecycle_events.clone()
222 }
223
224 #[must_use]
226 pub fn shutdown_token(&self) -> CancellationToken {
227 self.shutdown_token.clone()
228 }
229
230 #[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 #[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 #[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 #[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 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 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 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 #[must_use]
408 pub fn config(&self) -> Arc<C> {
409 Arc::clone(&self.config)
410 }
411
412 #[must_use]
414 pub fn container(&self) -> Arc<Container> {
415 Arc::clone(&self.container)
416 }
417
418 #[must_use]
420 pub fn lifecycle_event_bus(&self) -> EventBus<LifecycleEvent> {
421 self.lifecycle_events.clone()
422 }
423
424 #[must_use]
426 pub fn shutdown_token(&self) -> CancellationToken {
427 self.shutdown_token.clone()
428 }
429
430 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 #[must_use]
500 pub fn config(&self) -> Arc<C> {
501 Arc::clone(&self.config)
502 }
503
504 #[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}