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