1use std::future::Future;
11use std::marker::PhantomData;
12use std::pin::Pin;
13
14use tokio_util::sync::CancellationToken;
15
16use crate::{
17 RuntimeComponentBundle, RuntimeComponentKind, RuntimeComponentRegistry, log_shutdown_report,
18 spawn_termination_signal_handler,
19};
20
21type RuntimeHookFuture = Pin<Box<dyn Future<Output = ()> + Send>>;
22type RuntimeHook = Box<dyn FnOnce() -> RuntimeHookFuture + Send>;
23type RuntimeComponentRegistration = Box<dyn FnOnce(&mut RuntimeComponentRegistry) + Send>;
24type ShutdownResourceFuture = Pin<Box<dyn Future<Output = Result<(), String>> + Send>>;
25type ShutdownResourceFn<T> = Box<dyn FnOnce(T) -> ShutdownResourceFuture + Send>;
26
27fn empty_runtime_hook() -> RuntimeHook {
28 Box::new(|| Box::pin(async {}))
29}
30
31pub struct ServiceLifecycle<S> {
37 server: S,
38 shutdown_token: CancellationToken,
39}
40
41impl<S> ServiceLifecycle<S> {
42 pub fn new(server: S, shutdown_token: CancellationToken) -> Self {
44 Self {
45 server,
46 shutdown_token,
47 }
48 }
49}
50
51pub struct RuntimeServiceComponent<S> {
58 component_name: &'static str,
59 kind: RuntimeComponentKind,
60 service: S,
61 shutdown_token: CancellationToken,
62 stop_on_signal: RuntimeHook,
63}
64
65impl<S> RuntimeServiceComponent<S> {
66 pub fn new<F, Fut>(
68 component_name: &'static str,
69 kind: RuntimeComponentKind,
70 service: S,
71 shutdown_token: CancellationToken,
72 stop_on_signal: F,
73 ) -> Self
74 where
75 F: FnOnce() -> Fut + Send + 'static,
76 Fut: Future<Output = ()> + Send + 'static,
77 {
78 Self {
79 component_name,
80 kind,
81 service,
82 shutdown_token,
83 stop_on_signal: Box::new(move || Box::pin(stop_on_signal())),
84 }
85 }
86
87 fn into_parts(self) -> RuntimeServiceParts<S> {
88 RuntimeServiceParts {
89 service: self.service,
90 shutdown_token: self.shutdown_token,
91 stop_on_signal: self.stop_on_signal,
92 }
93 }
94}
95
96struct RuntimeServiceParts<S> {
97 service: S,
98 shutdown_token: CancellationToken,
99 stop_on_signal: RuntimeHook,
100}
101
102pub trait AsterRuntimeComponent<S> {
108 type Output;
110
111 fn apply(self, builder: AsterRuntimeBuilder<S>) -> Self::Output;
113}
114
115pub struct RuntimeComponentBundleRegistration<B> {
117 bundle: B,
118}
119
120pub struct RuntimeComponentWithShutdown<C, F> {
122 build: F,
123 _component: PhantomData<fn() -> C>,
124}
125
126pub struct TryRuntimeComponentWithShutdown<C, F, E> {
128 build: F,
129 _component: PhantomData<fn() -> C>,
130 _error: PhantomData<fn() -> E>,
131}
132
133pub const fn runtime_component<B>(bundle: B) -> RuntimeComponentBundleRegistration<B> {
135 RuntimeComponentBundleRegistration { bundle }
136}
137
138pub fn runtime_component_with_shutdown<C, F>(build: F) -> RuntimeComponentWithShutdown<C, F>
144where
145 F: FnOnce(CancellationToken) -> C,
146{
147 RuntimeComponentWithShutdown {
148 build,
149 _component: PhantomData,
150 }
151}
152
153pub fn try_runtime_component_with_shutdown<C, F, E>(
155 build: F,
156) -> TryRuntimeComponentWithShutdown<C, F, E>
157where
158 F: FnOnce(CancellationToken) -> Result<C, E>,
159{
160 TryRuntimeComponentWithShutdown {
161 build,
162 _component: PhantomData,
163 _error: PhantomData,
164 }
165}
166
167pub struct ShutdownResourceComponent<T> {
174 component_name: &'static str,
175 kind: RuntimeComponentKind,
176 phase_name: &'static str,
177 dependencies: &'static [&'static str],
178 resource: T,
179 shutdown: ShutdownResourceFn<T>,
180}
181
182impl<T> ShutdownResourceComponent<T> {
183 pub fn new<F, Fut>(
185 component_name: &'static str,
186 kind: RuntimeComponentKind,
187 phase_name: &'static str,
188 resource: T,
189 shutdown: F,
190 ) -> Self
191 where
192 F: FnOnce(T) -> Fut + Send + 'static,
193 Fut: Future<Output = Result<(), String>> + Send + 'static,
194 {
195 Self {
196 component_name,
197 kind,
198 phase_name,
199 dependencies: &[],
200 resource,
201 shutdown: Box::new(move |resource| Box::pin(shutdown(resource))),
202 }
203 }
204
205 #[must_use]
207 pub const fn depends_on_all(mut self, dependencies: &'static [&'static str]) -> Self {
208 self.dependencies = dependencies;
209 self
210 }
211}
212
213impl<T> RuntimeComponentBundle for ShutdownResourceComponent<T>
214where
215 T: Send + 'static,
216{
217 fn register(self, registry: &mut RuntimeComponentRegistry) {
218 let Self {
219 component_name,
220 kind,
221 phase_name,
222 dependencies,
223 resource,
224 shutdown,
225 } = self;
226 registry
227 .component(component_name)
228 .kind(kind)
229 .depends_on_all(dependencies)
230 .shutdown_once(phase_name, None, resource, shutdown);
231 }
232}
233
234pub fn shutdown_resource_component<T, F, Fut>(
236 component_name: &'static str,
237 kind: RuntimeComponentKind,
238 phase_name: &'static str,
239 resource: T,
240 shutdown: F,
241) -> RuntimeComponentBundleRegistration<ShutdownResourceComponent<T>>
242where
243 T: Send + 'static,
244 F: FnOnce(T) -> Fut + Send + 'static,
245 Fut: Future<Output = Result<(), String>> + Send + 'static,
246{
247 runtime_component(ShutdownResourceComponent::new(
248 component_name,
249 kind,
250 phase_name,
251 resource,
252 shutdown,
253 ))
254}
255
256pub fn shutdown_resource_component_after<T, F, Fut>(
258 component_name: &'static str,
259 kind: RuntimeComponentKind,
260 phase_name: &'static str,
261 dependencies: &'static [&'static str],
262 resource: T,
263 shutdown: F,
264) -> RuntimeComponentBundleRegistration<ShutdownResourceComponent<T>>
265where
266 T: Send + 'static,
267 F: FnOnce(T) -> Fut + Send + 'static,
268 Fut: Future<Output = Result<(), String>> + Send + 'static,
269{
270 runtime_component(
271 ShutdownResourceComponent::new(component_name, kind, phase_name, resource, shutdown)
272 .depends_on_all(dependencies),
273 )
274}
275
276impl<S, Service> AsterRuntimeComponent<S> for RuntimeServiceComponent<Service> {
277 type Output = AsterRuntimeBuilder<Service>;
278
279 fn apply(self, mut builder: AsterRuntimeBuilder<S>) -> Self::Output {
280 let component_name = self.component_name;
281 let kind = self.kind;
282 let assembly_error = if builder.service.is_some() {
283 Some(AsterRuntimeError::DuplicateService)
284 } else {
285 builder.assembly_error
286 };
287 builder.components.push(Box::new(move |registry| {
288 registry.component(component_name).kind(kind);
289 }));
290
291 AsterRuntimeBuilder {
292 service: Some(self.into_parts()),
293 shutdown_token: builder.shutdown_token,
294 before_shutdown: builder.before_shutdown,
295 components: builder.components,
296 assembly_error,
297 }
298 }
299}
300
301impl<S, B> AsterRuntimeComponent<S> for RuntimeComponentBundleRegistration<B>
302where
303 B: RuntimeComponentBundle + Send + 'static,
304{
305 type Output = AsterRuntimeBuilder<S>;
306
307 fn apply(self, mut builder: AsterRuntimeBuilder<S>) -> Self::Output {
308 builder
309 .components
310 .push(Box::new(move |registry| self.bundle.register(registry)));
311 builder
312 }
313}
314
315impl<S, C, F> AsterRuntimeComponent<S> for RuntimeComponentWithShutdown<C, F>
316where
317 C: AsterRuntimeComponent<S>,
318 F: FnOnce(CancellationToken) -> C,
319{
320 type Output = C::Output;
321
322 fn apply(self, builder: AsterRuntimeBuilder<S>) -> Self::Output {
323 let component = (self.build)(builder.shutdown_token.clone());
324 component.apply(builder)
325 }
326}
327
328impl<S, C, F, E> AsterRuntimeComponent<S> for TryRuntimeComponentWithShutdown<C, F, E>
329where
330 C: AsterRuntimeComponent<S>,
331 F: FnOnce(CancellationToken) -> Result<C, E>,
332{
333 type Output = Result<C::Output, E>;
334
335 fn apply(self, builder: AsterRuntimeBuilder<S>) -> Self::Output {
336 let component = (self.build)(builder.shutdown_token.clone())?;
337 Ok(component.apply(builder))
338 }
339}
340
341impl<B> RuntimeComponentBundle for RuntimeComponentBundleRegistration<B>
342where
343 B: RuntimeComponentBundle,
344{
345 fn register(self, registry: &mut RuntimeComponentRegistry) {
346 self.bundle.register(registry);
347 }
348}
349
350#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
352pub enum AsterRuntimeError {
353 #[error("runtime requires exactly one service component")]
355 MissingService,
356 #[error("runtime only supports one service component")]
358 DuplicateService,
359 #[error("runtime component graph is invalid: {source}")]
361 ComponentGraph {
362 source: crate::RuntimeComponentGraphError,
364 },
365 #[error("runtime startup aborted by a required component phase")]
367 Startup {
368 report: crate::StartupReport,
370 },
371}
372
373pub struct AsterRuntime<S> {
388 service: S,
389 shutdown_token: CancellationToken,
390 stop_on_signal: RuntimeHook,
391 before_shutdown: RuntimeHook,
392 registry: RuntimeComponentRegistry,
393}
394
395impl AsterRuntime<()> {
396 #[must_use]
398 pub fn builder() -> AsterRuntimeBuilder<()> {
399 AsterRuntimeBuilder::new()
400 }
401}
402
403impl<S> AsterRuntime<S>
404where
405 S: Future,
406{
407 pub async fn run(mut self) -> Result<S::Output, AsterRuntimeError> {
418 let startup_report = self.registry.startup().await;
419 if startup_report.aborted() {
420 let shutdown_report = self.registry.shutdown().await;
421 log_shutdown_report(&shutdown_report);
422 return Err(AsterRuntimeError::Startup {
423 report: startup_report,
424 });
425 }
426
427 let _signal_task =
428 spawn_termination_signal_handler(self.shutdown_token, self.stop_on_signal);
429
430 let output = self.service.await;
431 tracing::info!("service stopped");
432
433 (self.before_shutdown)().await;
434
435 let report = self.registry.shutdown().await;
436 log_shutdown_report(&report);
437
438 Ok(output)
439 }
440}
441
442pub struct AsterRuntimeBuilder<S = ()> {
444 service: Option<RuntimeServiceParts<S>>,
445 shutdown_token: CancellationToken,
446 before_shutdown: RuntimeHook,
447 components: Vec<RuntimeComponentRegistration>,
448 assembly_error: Option<AsterRuntimeError>,
449}
450
451impl AsterRuntimeBuilder<()> {
452 fn new() -> Self {
453 Self {
454 service: None,
455 shutdown_token: CancellationToken::new(),
456 before_shutdown: empty_runtime_hook(),
457 components: Vec::new(),
458 assembly_error: None,
459 }
460 }
461}
462
463impl<S> AsterRuntimeBuilder<S> {
464 pub fn shutdown_token(&self) -> &CancellationToken {
470 &self.shutdown_token
471 }
472
473 pub fn component<C>(self, component: C) -> C::Output
475 where
476 C: AsterRuntimeComponent<S>,
477 {
478 component.apply(self)
479 }
480
481 #[must_use]
483 pub fn before_shutdown<F, Fut>(mut self, before_shutdown: F) -> Self
484 where
485 F: FnOnce() -> Fut + Send + 'static,
486 Fut: Future<Output = ()> + Send + 'static,
487 {
488 self.before_shutdown = Box::new(move || Box::pin(before_shutdown()));
489 self
490 }
491
492 pub fn build(self) -> Result<AsterRuntime<S>, AsterRuntimeError> {
499 if let Some(error) = self.assembly_error {
500 return Err(error);
501 }
502
503 let Some(service) = self.service else {
504 return Err(AsterRuntimeError::MissingService);
505 };
506
507 let mut registry = RuntimeComponentRegistry::new();
508 for component in self.components {
509 component(&mut registry);
510 }
511 registry
512 .validate()
513 .map_err(|source| AsterRuntimeError::ComponentGraph { source })?;
514
515 Ok(AsterRuntime {
516 service: service.service,
517 shutdown_token: service.shutdown_token,
518 stop_on_signal: service.stop_on_signal,
519 before_shutdown: self.before_shutdown,
520 registry,
521 })
522 }
523}
524
525impl<S> AsterRuntimeBuilder<S>
526where
527 S: Future,
528{
529 pub async fn run(self) -> Result<S::Output, AsterRuntimeError> {
535 self.build()?.run().await
536 }
537}
538
539impl<S> ServiceLifecycle<S>
540where
541 S: Future,
542{
543 pub async fn run<Stop, StopFut, AfterStop, AfterStopFut>(
545 self,
546 stop_on_signal: Stop,
547 after_stop: AfterStop,
548 ) -> S::Output
549 where
550 Stop: FnOnce() -> StopFut + Send + 'static,
551 StopFut: Future<Output = ()> + Send + 'static,
552 AfterStop: FnOnce() -> AfterStopFut,
553 AfterStopFut: Future<Output = ()>,
554 {
555 let _signal_task = spawn_termination_signal_handler(self.shutdown_token, stop_on_signal);
556
557 let server_result = self.server.await;
558 tracing::info!("server stopped");
559 after_stop().await;
560 server_result
561 }
562}
563
564#[cfg(test)]
565mod tests {
566 use std::sync::{
567 Arc,
568 atomic::{AtomicBool, Ordering},
569 };
570
571 use tokio_util::sync::CancellationToken;
572
573 use super::{AsterRuntime, AsterRuntimeError, RuntimeServiceComponent, ServiceLifecycle};
574 use crate::{RuntimeComponentBundle, RuntimeComponentKind, RuntimeComponentRegistry};
575
576 #[tokio::test]
577 async fn lifecycle_runs_after_stop_and_returns_server_result() {
578 let after_stop_ran = Arc::new(AtomicBool::new(false));
579 let observed_after_stop = Arc::clone(&after_stop_ran);
580
581 let result = ServiceLifecycle::new(
582 async { Ok::<_, &'static str>(42) },
583 CancellationToken::default(),
584 )
585 .run(
586 || async {},
587 move || {
588 let observed_after_stop = Arc::clone(&observed_after_stop);
589 async move {
590 observed_after_stop.store(true, Ordering::SeqCst);
591 }
592 },
593 )
594 .await;
595
596 assert_eq!(result, Ok(42));
597 assert!(after_stop_ran.load(Ordering::SeqCst));
598 }
599
600 struct TestShutdownComponent {
601 events: Arc<std::sync::Mutex<Vec<&'static str>>>,
602 }
603
604 impl RuntimeComponentBundle for TestShutdownComponent {
605 fn register(self, registry: &mut RuntimeComponentRegistry) {
606 registry.component_shutdown(
607 "test",
608 RuntimeComponentKind::Product,
609 "test_shutdown",
610 None,
611 move || {
612 let events = Arc::clone(&self.events);
613 async move {
614 events.lock().unwrap().push("component");
615 Ok(())
616 }
617 },
618 );
619 }
620 }
621
622 #[tokio::test]
623 async fn aster_runtime_runs_before_shutdown_and_registered_components() {
624 let events = Arc::new(std::sync::Mutex::new(Vec::new()));
625 let before_events = Arc::clone(&events);
626 let component_events = Arc::clone(&events);
627
628 let result = AsterRuntime::builder()
629 .component(RuntimeServiceComponent::new(
630 "http",
631 RuntimeComponentKind::Core,
632 async { Ok::<_, &'static str>(7) },
633 CancellationToken::default(),
634 || async {},
635 ))
636 .before_shutdown(move || {
637 let before_events = Arc::clone(&before_events);
638 async move {
639 before_events.lock().unwrap().push("before");
640 }
641 })
642 .component(crate::runtime_component(TestShutdownComponent {
643 events: component_events,
644 }))
645 .run()
646 .await
647 .expect("runtime should run");
648
649 assert_eq!(result, Ok(7));
650 assert_eq!(events.lock().unwrap().as_slice(), ["before", "component"]);
651 }
652
653 #[tokio::test]
654 async fn aster_runtime_runs_component_shutdown_when_startup_aborts() {
655 let events = Arc::new(std::sync::Mutex::new(Vec::new()));
656 let started_events = Arc::clone(&events);
657 let shutdown_events = Arc::clone(&events);
658
659 let result = AsterRuntime::builder()
660 .component(RuntimeServiceComponent::new(
661 "http",
662 RuntimeComponentKind::Core,
663 std::future::pending::<()>(),
664 CancellationToken::default(),
665 || async {},
666 ))
667 .component(crate::runtime_component(
668 move |registry: &mut RuntimeComponentRegistry| {
669 registry.component_startup(
670 "database",
671 RuntimeComponentKind::Core,
672 "database_connect",
673 crate::StartupPhaseFailurePolicy::Required,
674 move || {
675 let started_events = Arc::clone(&started_events);
676 async move {
677 started_events.lock().unwrap().push("database_started");
678 Ok(())
679 }
680 },
681 );
682 registry.component_shutdown(
683 "database",
684 RuntimeComponentKind::Core,
685 "database_close",
686 None,
687 move || {
688 let shutdown_events = Arc::clone(&shutdown_events);
689 async move {
690 shutdown_events.lock().unwrap().push("database_shutdown");
691 Ok(())
692 }
693 },
694 );
695 registry.component_startup(
696 "cache",
697 RuntimeComponentKind::Core,
698 "cache_connect",
699 crate::StartupPhaseFailurePolicy::Required,
700 || async { Err("cache unavailable".to_string()) },
701 );
702 },
703 ))
704 .run()
705 .await;
706
707 assert!(matches!(result, Err(AsterRuntimeError::Startup { .. })));
708 assert_eq!(
711 events.lock().unwrap().as_slice(),
712 ["database_started", "database_shutdown"]
713 );
714 }
715
716 #[tokio::test]
717 async fn aster_runtime_builder_shares_shutdown_token_with_components() {
718 let observed = Arc::new(AtomicBool::new(false));
719 let observed_component = Arc::clone(&observed);
720
721 let result = AsterRuntime::builder()
722 .component(crate::runtime_component_with_shutdown(|shutdown| {
723 let component_token = shutdown.clone();
724 RuntimeServiceComponent::new(
725 "http",
726 RuntimeComponentKind::Core,
727 async move {
728 component_token.cancel();
729 Ok::<_, &'static str>(())
730 },
731 shutdown,
732 || async {},
733 )
734 }))
735 .component(crate::runtime_component_with_shutdown(|shutdown| {
736 crate::runtime_component(move |registry: &mut RuntimeComponentRegistry| {
737 registry.component_shutdown(
738 "observer",
739 RuntimeComponentKind::Product,
740 "observe_shared_shutdown",
741 None,
742 move || {
743 let observed_component = Arc::clone(&observed_component);
744 let shutdown = shutdown.clone();
745 async move {
746 observed_component.store(shutdown.is_cancelled(), Ordering::SeqCst);
747 Ok(())
748 }
749 },
750 );
751 })
752 }))
753 .run()
754 .await
755 .expect("runtime should run");
756
757 assert_eq!(result, Ok(()));
758 assert!(observed.load(Ordering::SeqCst));
759 }
760
761 #[test]
762 fn aster_runtime_requires_service_component() {
763 let result = AsterRuntime::builder().build();
764 assert!(matches!(result, Err(AsterRuntimeError::MissingService)));
765 }
766
767 #[test]
768 fn aster_runtime_rejects_duplicate_service_components() {
769 let result = AsterRuntime::builder()
770 .component(RuntimeServiceComponent::new(
771 "http",
772 RuntimeComponentKind::Core,
773 async {},
774 CancellationToken::default(),
775 || async {},
776 ))
777 .component(RuntimeServiceComponent::new(
778 "worker",
779 RuntimeComponentKind::Core,
780 async {},
781 CancellationToken::default(),
782 || async {},
783 ))
784 .build();
785
786 assert!(matches!(result, Err(AsterRuntimeError::DuplicateService)));
787 }
788
789 #[test]
790 fn aster_runtime_rejects_invalid_component_graph() {
791 let result = AsterRuntime::builder()
792 .component(RuntimeServiceComponent::new(
793 "http",
794 RuntimeComponentKind::Core,
795 async {},
796 CancellationToken::default(),
797 || async {},
798 ))
799 .component(crate::runtime_component(
800 |registry: &mut RuntimeComponentRegistry| {
801 registry.component("database").depends_on("missing");
802 },
803 ))
804 .build();
805
806 assert!(matches!(
807 result,
808 Err(AsterRuntimeError::ComponentGraph {
809 source: crate::RuntimeComponentGraphError::MissingDependency {
810 component: "database",
811 dependency: "missing"
812 }
813 })
814 ));
815 }
816
817 #[test]
818 fn shutdown_resource_component_registers_dependencies_and_shutdown() {
819 let registry = RuntimeComponentRegistry::configured(|registry| {
820 crate::shutdown_resource_component_after(
821 "mail_outbox",
822 RuntimeComponentKind::Mail,
823 "mail_outbox_drain",
824 &["background_tasks"],
825 42_u8,
826 |_| async { Ok(()) },
827 )
828 .register(registry);
829 });
830
831 let descriptor = registry
832 .descriptor("mail_outbox")
833 .expect("shutdown resource component should be registered");
834 assert_eq!(descriptor.kind, RuntimeComponentKind::Mail);
835 assert_eq!(descriptor.dependencies, vec!["background_tasks"]);
836 assert_eq!(
837 descriptor
838 .shutdown
839 .first()
840 .expect("shutdown phase should be registered")
841 .phase_name,
842 "mail_outbox_drain"
843 );
844 }
845}