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 #[doc(hidden)]
84 pub fn managed_ingress_bridge(
85 &self,
86 capacity: usize,
87 ) -> Option<crate::alpha1_ingress::ManagedIngressBridge> {
88 if capacity == 0 {
89 return None;
90 }
91 if self
92 .ingress_bridge_issued
93 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
94 .is_err()
95 {
96 return None;
97 }
98 crate::alpha1_ingress::ManagedIngressBridge::new(self.requests.clone(), capacity)
99 }
100
101 pub fn register<C>(&mut self, component: C) -> Result<()>
105 where
106 C: ComponentLifecycle + 'static,
107 {
108 self.register_shared(Arc::new(component))
109 }
110
111 pub fn register_shared(&mut self, component: Arc<dyn ComponentLifecycle>) -> Result<()> {
113 if self
114 .components
115 .iter()
116 .any(|registered| registered.name() == component.name())
117 {
118 return Err(SaddleError::new(
119 ErrorKind::Conflict,
120 "runtime.duplicate_component",
121 format!("component '{}' is already registered", component.name()),
122 ));
123 }
124 self.components.push(component);
125 Ok(())
126 }
127
128 pub fn run(self) -> Result<()> {
134 Self::run_with(|| async move { Ok(self) })
135 }
136
137 pub fn run_with<F, Fut>(bootstrap: F) -> Result<()>
144 where
145 F: FnOnce() -> Fut + Send + 'static,
146 Fut: Future<Output = Result<Self>> + Send + 'static,
147 {
148 if RUNTIME_STARTED
149 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
150 .is_err()
151 {
152 return Err(SaddleError::new(
153 ErrorKind::Conflict,
154 "runtime.already_started",
155 "the Saddle runtime has already started in this process",
156 ));
157 }
158
159 let runtime = build_runtime()?;
160
161 #[cfg(all(target_arch = "x86_64", target_os = "linux"))]
162 {
163 Self::run_with_owned_runtime(runtime, bootstrap)
164 }
165
166 #[cfg(not(all(target_arch = "x86_64", target_os = "linux")))]
167 runtime.block_on(async {
168 let signal = ShutdownSignal::register()?;
169 bootstrap_and_run(bootstrap, signal.wait()).await
170 })
171 }
172
173 #[cfg(all(target_arch = "x86_64", target_os = "linux"))]
174 pub(crate) fn claim_process_runtime() -> Result<()> {
175 if RUNTIME_STARTED
176 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
177 .is_err()
178 {
179 return Err(SaddleError::new(
180 ErrorKind::Conflict,
181 "runtime.already_started",
182 "the Saddle runtime has already started in this process",
183 ));
184 }
185 Ok(())
186 }
187
188 #[cfg(all(target_arch = "x86_64", target_os = "linux"))]
189 pub(crate) fn run_with_owned_runtime<F, Fut>(
190 runtime: tokio::runtime::Runtime,
191 bootstrap: F,
192 ) -> Result<()>
193 where
194 F: FnOnce() -> Fut,
195 Fut: Future<Output = Result<Self>>,
196 {
197 let outcome = runtime.block_on(async {
198 let signal = ShutdownSignal::register()?;
199 let application = bootstrap().await?;
200 let finalizer = application.pending_driver_finalizer();
201 let result = application.run_until_shutdown(signal.wait()).await;
202 Ok::<_, SaddleError>((finalizer, result))
203 });
204 match outcome {
205 Ok((finalizer, result)) => finalizer.finish(runtime, result),
206 Err(error) => {
207 drop(runtime);
208 Err(error)
209 }
210 }
211 }
212
213 pub(crate) async fn run_until_shutdown<F>(self, shutdown: F) -> Result<()>
214 where
215 F: Future<Output = Result<()>>,
216 {
217 tokio::pin!(shutdown);
218 let signal_before_start = tokio::select! {
219 biased;
220 signal_result = &mut shutdown => Some(signal_result),
221 _ = std::future::ready(()) => None,
222 };
223 if let Some(signal_result) = signal_before_start {
224 self.requests.begin_draining();
225 self.requests.wait_until_drained().await;
226 self.requests.mark_stopped();
227 return signal_result;
228 }
229
230 let mut started = 0;
231
232 for component in &self.components {
233 let start = component.start();
234 tokio::pin!(start);
235 let mut shutdown_during_start = None;
236 let start_result = tokio::select! {
237 biased;
238 signal_result = &mut shutdown => {
239 shutdown_during_start = Some(signal_result);
240 start.await
244 }
245 start_result = &mut start => start_result,
246 };
247
248 if let Err(error) = start_result {
249 self.requests.begin_draining();
250 self.requests.wait_until_drained().await;
251 let _ = self.shutdown_components(started).await;
252 self.requests.mark_stopped();
253 return Err(error);
254 }
255 started += 1;
256
257 if let Some(signal_result) = shutdown_during_start {
258 self.requests.begin_draining();
259 self.requests.wait_until_drained().await;
260 let shutdown_result = self.shutdown_components(started).await;
261 self.requests.mark_stopped();
262 return signal_result.and(shutdown_result);
263 }
264 }
265
266 self.requests.mark_ready();
267 let signal_result = shutdown.await;
268 self.requests.begin_draining();
269 self.requests.wait_until_drained().await;
270 let shutdown_result = self.shutdown_components(started).await;
271 self.requests.mark_stopped();
272
273 signal_result.and(shutdown_result)
274 }
275
276 async fn shutdown_components(&self, started: usize) -> Result<()> {
277 let mut first_error = None;
278 for component in self.components[..started].iter().rev() {
279 if let Err(error) = component.shutdown().await {
280 if first_error.is_none() {
281 first_error = Some(error);
282 }
283 }
284 }
285 first_error.map_or(Ok(()), Err)
286 }
287}
288
289#[cfg(any(test, not(all(target_arch = "x86_64", target_os = "linux"))))]
290async fn bootstrap_and_run<F, Fut, S>(bootstrap: F, shutdown: S) -> Result<()>
291where
292 F: FnOnce() -> Fut,
293 Fut: Future<Output = Result<Application>>,
294 S: Future<Output = Result<()>>,
295{
296 let application = bootstrap().await?;
297 application.run_until_shutdown(shutdown).await
298}
299
300impl Default for Application {
301 fn default() -> Self {
302 Self::new()
303 }
304}
305
306fn build_runtime() -> Result<tokio::runtime::Runtime> {
307 tokio::runtime::Builder::new_multi_thread()
308 .worker_threads(WORKER_THREADS)
309 .max_io_events_per_tick(MAX_IO_EVENTS_PER_TICK)
310 .enable_all()
311 .build()
312 .map_err(|_| {
313 SaddleError::new(
314 ErrorKind::Infrastructure,
315 "runtime.initialization_failed",
316 "failed to initialize the Saddle async runtime",
317 )
318 })
319}
320
321#[cfg(unix)]
322pub(crate) struct ShutdownSignal {
323 interrupt: tokio::signal::unix::Signal,
324 terminate: tokio::signal::unix::Signal,
325}
326
327#[cfg(unix)]
328impl ShutdownSignal {
329 pub(crate) fn register() -> Result<Self> {
331 Ok(Self {
332 interrupt: tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt())
333 .map_err(|_| signal_error())?,
334 terminate: tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
335 .map_err(|_| signal_error())?,
336 })
337 }
338
339 pub(crate) async fn wait(mut self) -> Result<()> {
340 tokio::select! {
341 _ = self.interrupt.recv() => Ok(()),
342 _ = self.terminate.recv() => Ok(()),
343 }
344 }
345}
346
347#[cfg(windows)]
348struct ShutdownSignal {
349 ctrl_c: tokio::signal::windows::CtrlC,
350 ctrl_break: tokio::signal::windows::CtrlBreak,
351}
352
353#[cfg(windows)]
354impl ShutdownSignal {
355 fn register() -> Result<Self> {
357 Ok(Self {
358 ctrl_c: tokio::signal::windows::ctrl_c().map_err(|_| signal_error())?,
359 ctrl_break: tokio::signal::windows::ctrl_break().map_err(|_| signal_error())?,
360 })
361 }
362
363 async fn wait(mut self) -> Result<()> {
364 tokio::select! {
365 _ = self.ctrl_c.recv() => Ok(()),
366 _ = self.ctrl_break.recv() => Ok(()),
367 }
368 }
369}
370
371fn signal_error() -> SaddleError {
372 SaddleError::new(
373 ErrorKind::Infrastructure,
374 "runtime.signal_registration_failed",
375 "failed to register the application shutdown signal",
376 )
377}
378
379#[cfg(test)]
380mod tests {
381 use std::sync::Mutex;
382
383 use saddle_core::LifecycleFuture;
384
385 use super::*;
386
387 struct RecordingComponent {
388 name: &'static str,
389 events: Arc<Mutex<Vec<String>>>,
390 start_error: bool,
391 shutdown_error: bool,
392 }
393
394 struct BlockingStartComponent {
395 events: Arc<Mutex<Vec<String>>>,
396 started: Mutex<Option<tokio::sync::oneshot::Sender<()>>>,
397 release: Mutex<Option<tokio::sync::oneshot::Receiver<()>>>,
398 }
399
400 impl ComponentLifecycle for RecordingComponent {
401 fn name(&self) -> &'static str {
402 self.name
403 }
404
405 fn start(&self) -> LifecycleFuture<'_> {
406 Box::pin(async move {
407 self.events
408 .lock()
409 .unwrap()
410 .push(format!("start:{}", self.name));
411 if self.start_error {
412 Err(test_error("start failed"))
413 } else {
414 Ok(())
415 }
416 })
417 }
418
419 fn shutdown(&self) -> LifecycleFuture<'_> {
420 Box::pin(async move {
421 self.events
422 .lock()
423 .unwrap()
424 .push(format!("shutdown:{}", self.name));
425 if self.shutdown_error {
426 Err(test_error("shutdown failed"))
427 } else {
428 Ok(())
429 }
430 })
431 }
432 }
433
434 impl ComponentLifecycle for BlockingStartComponent {
435 fn name(&self) -> &'static str {
436 "blocking"
437 }
438
439 fn start(&self) -> LifecycleFuture<'_> {
440 Box::pin(async move {
441 self.events
442 .lock()
443 .unwrap()
444 .push("start:blocking".to_owned());
445 let started = self.started.lock().unwrap().take().unwrap();
446 let release = self.release.lock().unwrap().take().unwrap();
447 started.send(()).unwrap();
448 release.await.unwrap();
449 Ok(())
450 })
451 }
452
453 fn shutdown(&self) -> LifecycleFuture<'_> {
454 Box::pin(async move {
455 self.events
456 .lock()
457 .unwrap()
458 .push("shutdown:blocking".to_owned());
459 Ok(())
460 })
461 }
462 }
463
464 fn component(name: &'static str, events: &Arc<Mutex<Vec<String>>>) -> RecordingComponent {
465 RecordingComponent {
466 name,
467 events: Arc::clone(events),
468 start_error: false,
469 shutdown_error: false,
470 }
471 }
472
473 fn test_error(message: &'static str) -> SaddleError {
474 SaddleError::new(ErrorKind::Infrastructure, "test.failure", message)
475 }
476
477 fn test_runtime() -> tokio::runtime::Runtime {
478 tokio::runtime::Builder::new_current_thread()
479 .build()
480 .expect("test runtime must build")
481 }
482
483 async fn shutdown_when_ready(requests: RequestLifecycle) -> Result<()> {
484 while requests.phase() != crate::ApplicationPhase::Ready {
485 tokio::task::yield_now().await;
486 }
487 Ok(())
488 }
489
490 #[test]
491 fn components_start_in_order_and_shutdown_in_reverse() {
492 let events = Arc::new(Mutex::new(Vec::new()));
493 let mut application = Application::new();
494 application.register(component("db", &events)).unwrap();
495 application.register(component("service", &events)).unwrap();
496 let shutdown = shutdown_when_ready(application.request_lifecycle());
497
498 test_runtime()
499 .block_on(application.run_until_shutdown(shutdown))
500 .unwrap();
501
502 assert_eq!(
503 *events.lock().unwrap(),
504 [
505 "start:db",
506 "start:service",
507 "shutdown:service",
508 "shutdown:db"
509 ]
510 );
511 }
512
513 #[test]
514 fn async_bootstrap_runs_before_early_shutdown_prevents_component_start() {
515 let events = Arc::new(Mutex::new(Vec::new()));
516 let bootstrap_events = Arc::clone(&events);
517
518 test_runtime()
519 .block_on(bootstrap_and_run(
520 move || async move {
521 bootstrap_events
522 .lock()
523 .unwrap()
524 .push("bootstrap".to_owned());
525 let mut application = Application::new();
526 application.register(component("component", &bootstrap_events))?;
527 Ok(application)
528 },
529 std::future::ready(Ok(())),
530 ))
531 .unwrap();
532
533 assert_eq!(*events.lock().unwrap(), ["bootstrap"]);
534 }
535
536 #[test]
537 fn failed_async_bootstrap_does_not_start_components() {
538 let error = test_runtime()
539 .block_on(bootstrap_and_run(
540 || async { Err(test_error("bootstrap failed")) },
541 std::future::pending(),
542 ))
543 .unwrap_err();
544 assert_eq!(error.message(), "bootstrap failed");
545 }
546
547 #[test]
548 fn startup_failure_rolls_back_only_started_components() {
549 let events = Arc::new(Mutex::new(Vec::new()));
550 let mut application = Application::new();
551 application.register(component("first", &events)).unwrap();
552 let mut failing = component("failing", &events);
553 failing.start_error = true;
554 application.register(failing).unwrap();
555 application.register(component("never", &events)).unwrap();
556 let shutdown = shutdown_when_ready(application.request_lifecycle());
557
558 let error = test_runtime()
559 .block_on(application.run_until_shutdown(shutdown))
560 .unwrap_err();
561
562 assert_eq!(error.message(), "start failed");
563 assert_eq!(
564 *events.lock().unwrap(),
565 ["start:first", "start:failing", "shutdown:first"]
566 );
567 }
568
569 #[test]
570 fn shutdown_continues_after_a_component_error() {
571 let events = Arc::new(Mutex::new(Vec::new()));
572 let mut application = Application::new();
573 application.register(component("first", &events)).unwrap();
574 let mut failing = component("second", &events);
575 failing.shutdown_error = true;
576 application.register(failing).unwrap();
577 let shutdown = shutdown_when_ready(application.request_lifecycle());
578
579 let error = test_runtime()
580 .block_on(application.run_until_shutdown(shutdown))
581 .unwrap_err();
582
583 assert_eq!(error.message(), "shutdown failed");
584 assert_eq!(
585 *events.lock().unwrap(),
586 [
587 "start:first",
588 "start:second",
589 "shutdown:second",
590 "shutdown:first"
591 ]
592 );
593 }
594
595 #[test]
596 fn duplicate_component_names_are_rejected() {
597 let events = Arc::new(Mutex::new(Vec::new()));
598 let mut application = Application::new();
599 application.register(component("db", &events)).unwrap();
600
601 let error = application.register(component("db", &events)).unwrap_err();
602 assert_eq!(error.code(), "runtime.duplicate_component");
603 }
604
605 #[test]
606 fn application_shutdown_waits_for_an_admitted_request() {
607 test_runtime().block_on(async {
608 let application = Application::new();
609 let requests = application.request_lifecycle();
610 let (release, released) = tokio::sync::oneshot::channel();
611
612 let shutdown = async move {
613 shutdown_when_ready(requests.clone()).await?;
614 let request = requests
615 .try_accept()
616 .expect("application is ready before waiting for shutdown");
617 tokio::spawn(async move {
618 released.await.unwrap();
619 drop(request);
620 });
621 Ok(())
622 };
623 let running = tokio::spawn(application.run_until_shutdown(shutdown));
624
625 tokio::task::yield_now().await;
626 assert!(!running.is_finished());
627 release.send(()).unwrap();
628 running.await.unwrap().unwrap();
629 });
630 }
631
632 #[test]
633 fn signal_failure_before_start_prevents_component_startup() {
634 let events = Arc::new(Mutex::new(Vec::new()));
635 let mut application = Application::new();
636 application.register(component("service", &events)).unwrap();
637
638 let error = test_runtime()
639 .block_on(application.run_until_shutdown(async { Err(signal_error()) }))
640 .unwrap_err();
641
642 assert_eq!(error.code(), "runtime.signal_registration_failed");
643 assert!(events.lock().unwrap().is_empty());
644 }
645
646 #[test]
647 fn shutdown_during_startup_stops_starting_and_rolls_back() {
648 test_runtime().block_on(async {
649 let events = Arc::new(Mutex::new(Vec::new()));
650 let (started_tx, started_rx) = tokio::sync::oneshot::channel();
651 let (release_tx, release_rx) = tokio::sync::oneshot::channel();
652 let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
653 let mut application = Application::new();
654 application
655 .register(BlockingStartComponent {
656 events: Arc::clone(&events),
657 started: Mutex::new(Some(started_tx)),
658 release: Mutex::new(Some(release_rx)),
659 })
660 .unwrap();
661 application.register(component("never", &events)).unwrap();
662
663 let running = tokio::spawn(application.run_until_shutdown(async move {
664 shutdown_rx.await.unwrap();
665 Ok(())
666 }));
667 started_rx.await.unwrap();
668 shutdown_tx.send(()).unwrap();
669 tokio::task::yield_now().await;
670 release_tx.send(()).unwrap();
671
672 running.await.unwrap().unwrap();
673 assert_eq!(
674 *events.lock().unwrap(),
675 ["start:blocking", "shutdown:blocking"]
676 );
677 });
678 }
679
680 #[test]
681 fn managed_runtime_provides_an_async_io_driver() {
682 build_runtime()
683 .unwrap()
684 .block_on(async {
685 tokio::net::TcpListener::bind((std::net::Ipv4Addr::LOCALHOST, 0)).await
686 })
687 .expect("service listeners require the managed async I/O driver");
688 }
689}