1use std::{
2 future::Future,
3 sync::Arc,
4 sync::atomic::{AtomicBool, Ordering},
5};
6
7use saddle_core::{ComponentLifecycle, ErrorKind, Result, SaddleError};
8
9use crate::RequestLifecycle;
10
11static RUNTIME_STARTED: AtomicBool = AtomicBool::new(false);
12const WORKER_THREADS: usize = 2;
13const MAX_IO_EVENTS_PER_TICK: usize = 5;
14
15pub struct Application {
20 components: Vec<Arc<dyn ComponentLifecycle>>,
21 requests: RequestLifecycle,
22 deployment_resource_budget: Option<saddle_admission::DeploymentResourceBudget>,
23 ingress_bridge_issued: AtomicBool,
24 #[cfg(all(target_arch = "x86_64", target_os = "linux"))]
25 pending_driver_finalizer: crate::post_driver::PendingDriverFinalizerSlot,
26}
27
28impl Application {
29 pub fn new() -> Self {
31 Self {
32 components: Vec::new(),
33 requests: RequestLifecycle::new(),
34 deployment_resource_budget: None,
35 ingress_bridge_issued: AtomicBool::new(false),
36 #[cfg(all(target_arch = "x86_64", target_os = "linux"))]
37 pending_driver_finalizer: crate::post_driver::PendingDriverFinalizerSlot::new(),
38 }
39 }
40
41 #[cfg(all(target_arch = "x86_64", target_os = "linux"))]
42 pub(crate) fn install_prevalidated_components(
43 &mut self,
44 components: Vec<Arc<dyn ComponentLifecycle>>,
45 ) {
46 debug_assert!(self.components.is_empty());
47 self.components = components;
48 }
49
50 #[cfg(all(target_arch = "x86_64", target_os = "linux"))]
51 #[doc(hidden)]
52 pub fn pending_driver_finalizer(&self) -> crate::post_driver::PendingDriverFinalizerSlot {
53 self.pending_driver_finalizer.clone()
54 }
55
56 #[cfg(all(test, target_arch = "x86_64", target_os = "linux"))]
57 pub(crate) fn post_driver_is_unarmed_for_test(&self) -> bool {
58 self.pending_driver_finalizer.is_unarmed_for_test()
59 }
60
61 #[cfg(all(target_arch = "x86_64", target_os = "linux"))]
62 #[doc(hidden)]
63 #[allow(clippy::result_large_err)]
64 pub fn commit_post_driver_install(
65 &self,
66 binding: saddle_admission::VerifiedPostDriverInstallBinding,
67 ) {
68 self.pending_driver_finalizer
69 .commit_verified_install(binding)
70 }
71
72 pub(crate) fn reserved_post_driver_submit(
73 &self,
74 ) -> crate::post_driver::MustSubmitDriverFinalizer {
75 self.pending_driver_finalizer.reserved_submit_handle()
76 }
77
78 pub fn request_lifecycle(&self) -> RequestLifecycle {
80 self.requests.clone()
81 }
82
83 pub fn health(&self) -> crate::ApplicationHealth {
86 self.requests.health()
87 }
88
89 #[doc(hidden)]
92 pub fn managed_ingress_bridge(
93 &self,
94 capacity: usize,
95 ) -> Option<crate::alpha1_ingress::ManagedIngressBridge> {
96 if capacity == 0 {
97 return None;
98 }
99 if self
100 .ingress_bridge_issued
101 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
102 .is_err()
103 {
104 return None;
105 }
106 crate::alpha1_ingress::ManagedIngressBridge::new(self.requests.clone(), capacity)
107 }
108
109 pub fn register<C>(&mut self, component: C) -> Result<()>
113 where
114 C: ComponentLifecycle + 'static,
115 {
116 self.register_shared(Arc::new(component))
117 }
118
119 pub fn register_shared(&mut self, component: Arc<dyn ComponentLifecycle>) -> Result<()> {
121 if self
122 .components
123 .iter()
124 .any(|registered| registered.name() == component.name())
125 {
126 return Err(SaddleError::new(
127 ErrorKind::Conflict,
128 "runtime.duplicate_component",
129 format!("component '{}' is already registered", component.name()),
130 ));
131 }
132 self.components.push(component);
133 Ok(())
134 }
135
136 pub fn run(self) -> Result<()> {
142 Self::run_with(|| async move { Ok(self) })
143 }
144
145 pub fn run_with<F, Fut>(bootstrap: F) -> Result<()>
152 where
153 F: FnOnce() -> Fut + Send + 'static,
154 Fut: Future<Output = Result<Self>> + Send + 'static,
155 {
156 if RUNTIME_STARTED
157 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
158 .is_err()
159 {
160 return Err(SaddleError::new(
161 ErrorKind::Conflict,
162 "runtime.already_started",
163 "the Saddle runtime has already started in this process",
164 ));
165 }
166
167 let runtime = build_runtime()?;
168
169 #[cfg(all(target_arch = "x86_64", target_os = "linux"))]
170 {
171 Self::run_with_owned_runtime(runtime, bootstrap)
172 }
173
174 #[cfg(not(all(target_arch = "x86_64", target_os = "linux")))]
175 runtime.block_on(async {
176 let signal = ShutdownSignal::register()?;
177 bootstrap_and_run(bootstrap, signal.wait()).await
178 })
179 }
180
181 #[doc(hidden)]
184 pub fn run_with_deployment_resource_budget<F, Fut>(
185 budget: saddle_admission::DeploymentResourceBudget,
186 bootstrap: F,
187 ) -> Result<()>
188 where
189 F: FnOnce() -> Fut + Send + 'static,
190 Fut: Future<Output = Result<Self>> + Send + 'static,
191 {
192 Self::run_with(move || async move {
193 let mut application = bootstrap().await?;
194 if application.deployment_resource_budget.is_some() {
195 return Err(SaddleError::new(
196 ErrorKind::Conflict,
197 "runtime.deployment_resource_budget_already_installed",
198 "the deployment resource budget was already installed",
199 ));
200 }
201 application.deployment_resource_budget = Some(budget);
202 Ok(application)
203 })
204 }
205
206 #[cfg(all(target_arch = "x86_64", target_os = "linux"))]
207 pub(crate) fn claim_process_runtime() -> Result<()> {
208 if RUNTIME_STARTED
209 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
210 .is_err()
211 {
212 return Err(SaddleError::new(
213 ErrorKind::Conflict,
214 "runtime.already_started",
215 "the Saddle runtime has already started in this process",
216 ));
217 }
218 Ok(())
219 }
220
221 #[cfg(all(target_arch = "x86_64", target_os = "linux"))]
222 pub(crate) fn run_with_owned_runtime<F, Fut>(
223 runtime: tokio::runtime::Runtime,
224 bootstrap: F,
225 ) -> Result<()>
226 where
227 F: FnOnce() -> Fut,
228 Fut: Future<Output = Result<Self>>,
229 {
230 let outcome = runtime.block_on(async {
231 let signal = ShutdownSignal::register()?;
232 let application = bootstrap().await?;
233 let finalizer = application.pending_driver_finalizer();
234 let result = application.run_until_shutdown(signal.wait()).await;
235 Ok::<_, SaddleError>((finalizer, result))
236 });
237 match outcome {
238 Ok((finalizer, result)) => finalizer.finish(runtime, result),
239 Err(error) => {
240 drop(runtime);
241 Err(error)
242 }
243 }
244 }
245
246 pub(crate) async fn run_until_shutdown<F>(self, shutdown: F) -> Result<()>
247 where
248 F: Future<Output = Result<()>>,
249 {
250 tokio::pin!(shutdown);
251 let signal_before_start = tokio::select! {
252 biased;
253 signal_result = &mut shutdown => Some(signal_result),
254 _ = std::future::ready(()) => None,
255 };
256 if let Some(signal_result) = signal_before_start {
257 self.requests.begin_draining();
258 self.requests.wait_until_drained().await;
259 self.requests.mark_stopped();
260 return signal_result;
261 }
262
263 let mut started = 0;
264
265 for component in &self.components {
266 let start = component.start();
267 tokio::pin!(start);
268 let mut shutdown_during_start = None;
269 let start_result = tokio::select! {
270 biased;
271 signal_result = &mut shutdown => {
272 shutdown_during_start = Some(signal_result);
273 start.await
277 }
278 start_result = &mut start => start_result,
279 };
280
281 if let Err(error) = start_result {
282 self.requests.begin_draining();
283 self.requests.wait_until_drained().await;
284 let _ = self.shutdown_components(started).await;
285 self.requests.mark_stopped();
286 return Err(error);
287 }
288 started += 1;
289
290 if let Some(signal_result) = shutdown_during_start {
291 self.requests.begin_draining();
292 self.requests.wait_until_drained().await;
293 let shutdown_result = self.shutdown_components(started).await;
294 self.requests.mark_stopped();
295 return signal_result.and(shutdown_result);
296 }
297 }
298
299 self.requests.mark_ready();
300 let signal_result = shutdown.await;
301 self.requests.begin_draining();
302 self.requests.wait_until_drained().await;
303 let shutdown_result = self.shutdown_components(started).await;
304 self.requests.mark_stopped();
305
306 signal_result.and(shutdown_result)
307 }
308
309 async fn shutdown_components(&self, started: usize) -> Result<()> {
310 let mut first_error = None;
311 for component in self.components[..started].iter().rev() {
312 if let Err(error) = component.shutdown().await {
313 if first_error.is_none() {
314 first_error = Some(error);
315 }
316 }
317 }
318 first_error.map_or(Ok(()), Err)
319 }
320}
321
322#[cfg(any(test, not(all(target_arch = "x86_64", target_os = "linux"))))]
323async fn bootstrap_and_run<F, Fut, S>(bootstrap: F, shutdown: S) -> Result<()>
324where
325 F: FnOnce() -> Fut,
326 Fut: Future<Output = Result<Application>>,
327 S: Future<Output = Result<()>>,
328{
329 let application = bootstrap().await?;
330 application.run_until_shutdown(shutdown).await
331}
332
333impl Default for Application {
334 fn default() -> Self {
335 Self::new()
336 }
337}
338
339fn build_runtime() -> Result<tokio::runtime::Runtime> {
340 tokio::runtime::Builder::new_multi_thread()
341 .worker_threads(WORKER_THREADS)
342 .max_io_events_per_tick(MAX_IO_EVENTS_PER_TICK)
343 .enable_all()
344 .build()
345 .map_err(|_| {
346 SaddleError::new(
347 ErrorKind::Infrastructure,
348 "runtime.initialization_failed",
349 "failed to initialize the Saddle async runtime",
350 )
351 })
352}
353
354#[cfg(all(target_arch = "x86_64", target_os = "linux"))]
355pub(crate) fn claim_owned_runtime() -> Result<tokio::runtime::Runtime> {
356 Application::claim_process_runtime()?;
357 build_runtime()
358}
359
360#[cfg(unix)]
361pub(crate) struct ShutdownSignal {
362 interrupt: tokio::signal::unix::Signal,
363 terminate: tokio::signal::unix::Signal,
364}
365
366#[cfg(unix)]
367impl ShutdownSignal {
368 pub(crate) fn register() -> Result<Self> {
370 Ok(Self {
371 interrupt: tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt())
372 .map_err(|_| signal_error())?,
373 terminate: tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
374 .map_err(|_| signal_error())?,
375 })
376 }
377
378 pub(crate) async fn wait(mut self) -> Result<()> {
379 tokio::select! {
380 _ = self.interrupt.recv() => Ok(()),
381 _ = self.terminate.recv() => Ok(()),
382 }
383 }
384}
385
386#[cfg(windows)]
387struct ShutdownSignal {
388 ctrl_c: tokio::signal::windows::CtrlC,
389 ctrl_break: tokio::signal::windows::CtrlBreak,
390}
391
392#[cfg(windows)]
393impl ShutdownSignal {
394 fn register() -> Result<Self> {
396 Ok(Self {
397 ctrl_c: tokio::signal::windows::ctrl_c().map_err(|_| signal_error())?,
398 ctrl_break: tokio::signal::windows::ctrl_break().map_err(|_| signal_error())?,
399 })
400 }
401
402 async fn wait(mut self) -> Result<()> {
403 tokio::select! {
404 _ = self.ctrl_c.recv() => Ok(()),
405 _ = self.ctrl_break.recv() => Ok(()),
406 }
407 }
408}
409
410fn signal_error() -> SaddleError {
411 SaddleError::new(
412 ErrorKind::Infrastructure,
413 "runtime.signal_registration_failed",
414 "failed to register the application shutdown signal",
415 )
416}
417
418#[cfg(test)]
419mod tests {
420 use std::sync::Mutex;
421
422 use saddle_core::LifecycleFuture;
423
424 use super::*;
425 use crate::ApplicationPhase;
426
427 struct RecordingComponent {
428 name: &'static str,
429 events: Arc<Mutex<Vec<String>>>,
430 start_error: bool,
431 shutdown_error: bool,
432 }
433
434 struct BlockingStartComponent {
435 events: Arc<Mutex<Vec<String>>>,
436 started: Mutex<Option<tokio::sync::oneshot::Sender<()>>>,
437 release: Mutex<Option<tokio::sync::oneshot::Receiver<()>>>,
438 }
439
440 struct HealthAwareListener {
441 health: crate::ApplicationHealth,
442 events: Arc<Mutex<Vec<String>>>,
443 }
444
445 impl ComponentLifecycle for RecordingComponent {
446 fn name(&self) -> &'static str {
447 self.name
448 }
449
450 fn start(&self) -> LifecycleFuture<'_> {
451 Box::pin(async move {
452 self.events
453 .lock()
454 .unwrap()
455 .push(format!("start:{}", self.name));
456 if self.start_error {
457 Err(test_error("start failed"))
458 } else {
459 Ok(())
460 }
461 })
462 }
463
464 fn shutdown(&self) -> LifecycleFuture<'_> {
465 Box::pin(async move {
466 self.events
467 .lock()
468 .unwrap()
469 .push(format!("shutdown:{}", self.name));
470 if self.shutdown_error {
471 Err(test_error("shutdown failed"))
472 } else {
473 Ok(())
474 }
475 })
476 }
477 }
478
479 impl ComponentLifecycle for BlockingStartComponent {
480 fn name(&self) -> &'static str {
481 "blocking"
482 }
483
484 fn start(&self) -> LifecycleFuture<'_> {
485 Box::pin(async move {
486 self.events
487 .lock()
488 .unwrap()
489 .push("start:blocking".to_owned());
490 let started = self.started.lock().unwrap().take().unwrap();
491 let release = self.release.lock().unwrap().take().unwrap();
492 started.send(()).unwrap();
493 release.await.unwrap();
494 Ok(())
495 })
496 }
497
498 fn shutdown(&self) -> LifecycleFuture<'_> {
499 Box::pin(async move {
500 self.events
501 .lock()
502 .unwrap()
503 .push("shutdown:blocking".to_owned());
504 Ok(())
505 })
506 }
507 }
508
509 impl ComponentLifecycle for HealthAwareListener {
510 fn name(&self) -> &'static str {
511 "health-aware-listener"
512 }
513
514 fn start(&self) -> LifecycleFuture<'_> {
515 Box::pin(async move {
516 let snapshot = self.health.snapshot();
517 assert!(snapshot.is_live());
518 assert!(!snapshot.is_ready());
519 assert_eq!(snapshot.phase(), ApplicationPhase::Starting);
520 self.events
521 .lock()
522 .unwrap()
523 .push("listener:accepting".into());
524 Ok(())
525 })
526 }
527
528 fn shutdown(&self) -> LifecycleFuture<'_> {
529 Box::pin(async move {
530 let snapshot = self.health.snapshot();
531 assert!(snapshot.is_live());
532 assert!(!snapshot.is_ready());
533 assert_eq!(snapshot.phase(), ApplicationPhase::Draining);
534 self.events.lock().unwrap().push("listener:stopped".into());
535 Ok(())
536 })
537 }
538 }
539
540 fn component(name: &'static str, events: &Arc<Mutex<Vec<String>>>) -> RecordingComponent {
541 RecordingComponent {
542 name,
543 events: Arc::clone(events),
544 start_error: false,
545 shutdown_error: false,
546 }
547 }
548
549 fn test_error(message: &'static str) -> SaddleError {
550 SaddleError::new(ErrorKind::Infrastructure, "test.failure", message)
551 }
552
553 fn test_runtime() -> tokio::runtime::Runtime {
554 tokio::runtime::Builder::new_current_thread()
555 .build()
556 .expect("test runtime must build")
557 }
558
559 async fn shutdown_when_ready(requests: RequestLifecycle) -> Result<()> {
560 while requests.phase() != crate::ApplicationPhase::Ready {
561 tokio::task::yield_now().await;
562 }
563 Ok(())
564 }
565
566 #[test]
567 fn components_start_in_order_and_shutdown_in_reverse() {
568 let events = Arc::new(Mutex::new(Vec::new()));
569 let mut application = Application::new();
570 application.register(component("db", &events)).unwrap();
571 application.register(component("service", &events)).unwrap();
572 let shutdown = shutdown_when_ready(application.request_lifecycle());
573
574 test_runtime()
575 .block_on(application.run_until_shutdown(shutdown))
576 .unwrap();
577
578 assert_eq!(
579 *events.lock().unwrap(),
580 [
581 "start:db",
582 "start:service",
583 "shutdown:service",
584 "shutdown:db"
585 ]
586 );
587 }
588
589 #[test]
590 fn health_uses_the_unique_lifecycle_and_clears_ready_before_listener_shutdown() {
591 test_runtime().block_on(async {
592 let events = Arc::new(Mutex::new(Vec::new()));
593 let mut application = Application::new();
594 let health = application.health();
595 let initial = health.snapshot();
596 assert!(initial.is_live());
597 assert!(!initial.is_ready());
598 assert_eq!(initial.phase(), ApplicationPhase::Starting);
599
600 application
601 .register(HealthAwareListener {
602 health: health.clone(),
603 events: Arc::clone(&events),
604 })
605 .unwrap();
606 let shutdown_health = health.clone();
607 application
608 .run_until_shutdown(async move {
609 loop {
610 let snapshot = shutdown_health.snapshot();
611 if snapshot.is_ready() {
612 assert!(snapshot.is_live());
613 assert_eq!(snapshot.phase(), ApplicationPhase::Ready);
614 return Ok(());
615 }
616 tokio::task::yield_now().await;
617 }
618 })
619 .await
620 .unwrap();
621
622 let stopped = health.snapshot();
623 assert!(!stopped.is_live());
624 assert!(!stopped.is_ready());
625 assert_eq!(stopped.phase(), ApplicationPhase::Stopped);
626 assert_eq!(
627 *events.lock().unwrap(),
628 ["listener:accepting", "listener:stopped"]
629 );
630 });
631 }
632
633 #[test]
634 fn async_bootstrap_runs_before_early_shutdown_prevents_component_start() {
635 let events = Arc::new(Mutex::new(Vec::new()));
636 let bootstrap_events = Arc::clone(&events);
637
638 test_runtime()
639 .block_on(bootstrap_and_run(
640 move || async move {
641 bootstrap_events
642 .lock()
643 .unwrap()
644 .push("bootstrap".to_owned());
645 let mut application = Application::new();
646 application.register(component("component", &bootstrap_events))?;
647 Ok(application)
648 },
649 std::future::ready(Ok(())),
650 ))
651 .unwrap();
652
653 assert_eq!(*events.lock().unwrap(), ["bootstrap"]);
654 }
655
656 #[test]
657 fn failed_async_bootstrap_does_not_start_components() {
658 let error = test_runtime()
659 .block_on(bootstrap_and_run(
660 || async { Err(test_error("bootstrap failed")) },
661 std::future::pending(),
662 ))
663 .unwrap_err();
664 assert_eq!(error.message(), "bootstrap failed");
665 }
666
667 #[test]
668 fn startup_failure_rolls_back_only_started_components() {
669 let events = Arc::new(Mutex::new(Vec::new()));
670 let mut application = Application::new();
671 application.register(component("first", &events)).unwrap();
672 let mut failing = component("failing", &events);
673 failing.start_error = true;
674 application.register(failing).unwrap();
675 application.register(component("never", &events)).unwrap();
676 let shutdown = shutdown_when_ready(application.request_lifecycle());
677
678 let error = test_runtime()
679 .block_on(application.run_until_shutdown(shutdown))
680 .unwrap_err();
681
682 assert_eq!(error.message(), "start failed");
683 assert_eq!(
684 *events.lock().unwrap(),
685 ["start:first", "start:failing", "shutdown:first"]
686 );
687 }
688
689 #[test]
690 fn shutdown_continues_after_a_component_error() {
691 let events = Arc::new(Mutex::new(Vec::new()));
692 let mut application = Application::new();
693 application.register(component("first", &events)).unwrap();
694 let mut failing = component("second", &events);
695 failing.shutdown_error = true;
696 application.register(failing).unwrap();
697 let shutdown = shutdown_when_ready(application.request_lifecycle());
698
699 let error = test_runtime()
700 .block_on(application.run_until_shutdown(shutdown))
701 .unwrap_err();
702
703 assert_eq!(error.message(), "shutdown failed");
704 assert_eq!(
705 *events.lock().unwrap(),
706 [
707 "start:first",
708 "start:second",
709 "shutdown:second",
710 "shutdown:first"
711 ]
712 );
713 }
714
715 #[test]
716 fn duplicate_component_names_are_rejected() {
717 let events = Arc::new(Mutex::new(Vec::new()));
718 let mut application = Application::new();
719 application.register(component("db", &events)).unwrap();
720
721 let error = application.register(component("db", &events)).unwrap_err();
722 assert_eq!(error.code(), "runtime.duplicate_component");
723 }
724
725 #[test]
726 fn application_shutdown_waits_for_an_admitted_request() {
727 test_runtime().block_on(async {
728 let application = Application::new();
729 let requests = application.request_lifecycle();
730 let (release, released) = tokio::sync::oneshot::channel();
731
732 let shutdown = async move {
733 shutdown_when_ready(requests.clone()).await?;
734 let request = requests
735 .try_accept()
736 .expect("application is ready before waiting for shutdown");
737 tokio::spawn(async move {
738 released.await.unwrap();
739 drop(request);
740 });
741 Ok(())
742 };
743 let running = tokio::spawn(application.run_until_shutdown(shutdown));
744
745 tokio::task::yield_now().await;
746 assert!(!running.is_finished());
747 release.send(()).unwrap();
748 running.await.unwrap().unwrap();
749 });
750 }
751
752 #[test]
753 fn signal_failure_before_start_prevents_component_startup() {
754 let events = Arc::new(Mutex::new(Vec::new()));
755 let mut application = Application::new();
756 application.register(component("service", &events)).unwrap();
757
758 let error = test_runtime()
759 .block_on(application.run_until_shutdown(async { Err(signal_error()) }))
760 .unwrap_err();
761
762 assert_eq!(error.code(), "runtime.signal_registration_failed");
763 assert!(events.lock().unwrap().is_empty());
764 }
765
766 #[test]
767 fn shutdown_during_startup_stops_starting_and_rolls_back() {
768 test_runtime().block_on(async {
769 let events = Arc::new(Mutex::new(Vec::new()));
770 let (started_tx, started_rx) = tokio::sync::oneshot::channel();
771 let (release_tx, release_rx) = tokio::sync::oneshot::channel();
772 let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
773 let mut application = Application::new();
774 application
775 .register(BlockingStartComponent {
776 events: Arc::clone(&events),
777 started: Mutex::new(Some(started_tx)),
778 release: Mutex::new(Some(release_rx)),
779 })
780 .unwrap();
781 application.register(component("never", &events)).unwrap();
782
783 let running = tokio::spawn(application.run_until_shutdown(async move {
784 shutdown_rx.await.unwrap();
785 Ok(())
786 }));
787 started_rx.await.unwrap();
788 shutdown_tx.send(()).unwrap();
789 tokio::task::yield_now().await;
790 release_tx.send(()).unwrap();
791
792 running.await.unwrap().unwrap();
793 assert_eq!(
794 *events.lock().unwrap(),
795 ["start:blocking", "shutdown:blocking"]
796 );
797 });
798 }
799
800 #[test]
801 fn managed_runtime_provides_an_async_io_driver() {
802 build_runtime()
803 .unwrap()
804 .block_on(async {
805 tokio::net::TcpListener::bind((std::net::Ipv4Addr::LOCALHOST, 0)).await
806 })
807 .expect("service listeners require the managed async I/O driver");
808 }
809}