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