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