use std::{
future::Future,
sync::Arc,
sync::atomic::{AtomicBool, Ordering},
};
use saddle_core::{ComponentLifecycle, ErrorKind, Result, SaddleError};
use crate::RequestLifecycle;
static RUNTIME_STARTED: AtomicBool = AtomicBool::new(false);
pub struct Application {
components: Vec<Arc<dyn ComponentLifecycle>>,
requests: RequestLifecycle,
}
impl Application {
pub fn new() -> Self {
Self {
components: Vec::new(),
requests: RequestLifecycle::new(),
}
}
pub fn request_lifecycle(&self) -> RequestLifecycle {
self.requests.clone()
}
pub fn register<C>(&mut self, component: C) -> Result<()>
where
C: ComponentLifecycle + 'static,
{
self.register_shared(Arc::new(component))
}
pub fn register_shared(&mut self, component: Arc<dyn ComponentLifecycle>) -> Result<()> {
if self
.components
.iter()
.any(|registered| registered.name() == component.name())
{
return Err(SaddleError::new(
ErrorKind::Conflict,
"runtime.duplicate_component",
format!("component '{}' is already registered", component.name()),
));
}
self.components.push(component);
Ok(())
}
pub fn run(self) -> Result<()> {
Self::run_with(|| async move { Ok(self) })
}
pub fn run_with<F, Fut>(bootstrap: F) -> Result<()>
where
F: FnOnce() -> Fut + Send + 'static,
Fut: Future<Output = Result<Self>> + Send + 'static,
{
if RUNTIME_STARTED
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.is_err()
{
return Err(SaddleError::new(
ErrorKind::Conflict,
"runtime.already_started",
"the Saddle runtime has already started in this process",
));
}
let runtime = build_runtime()?;
runtime.block_on(async {
let signal = ShutdownSignal::register()?;
bootstrap_and_run(bootstrap, signal.wait()).await
})
}
async fn run_until_shutdown<F>(self, shutdown: F) -> Result<()>
where
F: Future<Output = Result<()>>,
{
tokio::pin!(shutdown);
let signal_before_start = tokio::select! {
biased;
signal_result = &mut shutdown => Some(signal_result),
_ = std::future::ready(()) => None,
};
if let Some(signal_result) = signal_before_start {
self.requests.begin_draining();
self.requests.wait_until_drained().await;
self.requests.mark_stopped();
return signal_result;
}
let mut started = 0;
for component in &self.components {
let start = component.start();
tokio::pin!(start);
let mut shutdown_during_start = None;
let start_result = tokio::select! {
biased;
signal_result = &mut shutdown => {
shutdown_during_start = Some(signal_result);
start.await
}
start_result = &mut start => start_result,
};
if let Err(error) = start_result {
self.requests.begin_draining();
self.requests.wait_until_drained().await;
let _ = self.shutdown_components(started).await;
self.requests.mark_stopped();
return Err(error);
}
started += 1;
if let Some(signal_result) = shutdown_during_start {
self.requests.begin_draining();
self.requests.wait_until_drained().await;
let shutdown_result = self.shutdown_components(started).await;
self.requests.mark_stopped();
return signal_result.and(shutdown_result);
}
}
self.requests.mark_ready();
let signal_result = shutdown.await;
self.requests.begin_draining();
self.requests.wait_until_drained().await;
let shutdown_result = self.shutdown_components(started).await;
self.requests.mark_stopped();
signal_result.and(shutdown_result)
}
async fn shutdown_components(&self, started: usize) -> Result<()> {
let mut first_error = None;
for component in self.components[..started].iter().rev() {
if let Err(error) = component.shutdown().await
&& first_error.is_none()
{
first_error = Some(error);
}
}
first_error.map_or(Ok(()), Err)
}
}
async fn bootstrap_and_run<F, Fut, S>(bootstrap: F, shutdown: S) -> Result<()>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<Application>>,
S: Future<Output = Result<()>>,
{
let application = bootstrap().await?;
application.run_until_shutdown(shutdown).await
}
impl Default for Application {
fn default() -> Self {
Self::new()
}
}
fn build_runtime() -> Result<tokio::runtime::Runtime> {
tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
.map_err(|_| {
SaddleError::new(
ErrorKind::Infrastructure,
"runtime.initialization_failed",
"failed to initialize the Saddle async runtime",
)
})
}
#[cfg(unix)]
struct ShutdownSignal {
interrupt: tokio::signal::unix::Signal,
terminate: tokio::signal::unix::Signal,
}
#[cfg(unix)]
impl ShutdownSignal {
fn register() -> Result<Self> {
Ok(Self {
interrupt: tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt())
.map_err(|_| signal_error())?,
terminate: tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
.map_err(|_| signal_error())?,
})
}
async fn wait(mut self) -> Result<()> {
tokio::select! {
_ = self.interrupt.recv() => Ok(()),
_ = self.terminate.recv() => Ok(()),
}
}
}
#[cfg(windows)]
struct ShutdownSignal {
ctrl_c: tokio::signal::windows::CtrlC,
ctrl_break: tokio::signal::windows::CtrlBreak,
}
#[cfg(windows)]
impl ShutdownSignal {
fn register() -> Result<Self> {
Ok(Self {
ctrl_c: tokio::signal::windows::ctrl_c().map_err(|_| signal_error())?,
ctrl_break: tokio::signal::windows::ctrl_break().map_err(|_| signal_error())?,
})
}
async fn wait(mut self) -> Result<()> {
tokio::select! {
_ = self.ctrl_c.recv() => Ok(()),
_ = self.ctrl_break.recv() => Ok(()),
}
}
}
fn signal_error() -> SaddleError {
SaddleError::new(
ErrorKind::Infrastructure,
"runtime.signal_registration_failed",
"failed to register the application shutdown signal",
)
}
#[cfg(test)]
mod tests {
use std::sync::Mutex;
use saddle_core::LifecycleFuture;
use super::*;
struct RecordingComponent {
name: &'static str,
events: Arc<Mutex<Vec<String>>>,
start_error: bool,
shutdown_error: bool,
}
struct BlockingStartComponent {
events: Arc<Mutex<Vec<String>>>,
started: Mutex<Option<tokio::sync::oneshot::Sender<()>>>,
release: Mutex<Option<tokio::sync::oneshot::Receiver<()>>>,
}
impl ComponentLifecycle for RecordingComponent {
fn name(&self) -> &'static str {
self.name
}
fn start(&self) -> LifecycleFuture<'_> {
Box::pin(async move {
self.events
.lock()
.unwrap()
.push(format!("start:{}", self.name));
if self.start_error {
Err(test_error("start failed"))
} else {
Ok(())
}
})
}
fn shutdown(&self) -> LifecycleFuture<'_> {
Box::pin(async move {
self.events
.lock()
.unwrap()
.push(format!("shutdown:{}", self.name));
if self.shutdown_error {
Err(test_error("shutdown failed"))
} else {
Ok(())
}
})
}
}
impl ComponentLifecycle for BlockingStartComponent {
fn name(&self) -> &'static str {
"blocking"
}
fn start(&self) -> LifecycleFuture<'_> {
Box::pin(async move {
self.events
.lock()
.unwrap()
.push("start:blocking".to_owned());
let started = self.started.lock().unwrap().take().unwrap();
let release = self.release.lock().unwrap().take().unwrap();
started.send(()).unwrap();
release.await.unwrap();
Ok(())
})
}
fn shutdown(&self) -> LifecycleFuture<'_> {
Box::pin(async move {
self.events
.lock()
.unwrap()
.push("shutdown:blocking".to_owned());
Ok(())
})
}
}
fn component(name: &'static str, events: &Arc<Mutex<Vec<String>>>) -> RecordingComponent {
RecordingComponent {
name,
events: Arc::clone(events),
start_error: false,
shutdown_error: false,
}
}
fn test_error(message: &'static str) -> SaddleError {
SaddleError::new(ErrorKind::Infrastructure, "test.failure", message)
}
fn test_runtime() -> tokio::runtime::Runtime {
tokio::runtime::Builder::new_current_thread()
.build()
.expect("test runtime must build")
}
async fn shutdown_when_ready(requests: RequestLifecycle) -> Result<()> {
while requests.phase() != crate::ApplicationPhase::Ready {
tokio::task::yield_now().await;
}
Ok(())
}
#[test]
fn components_start_in_order_and_shutdown_in_reverse() {
let events = Arc::new(Mutex::new(Vec::new()));
let mut application = Application::new();
application.register(component("db", &events)).unwrap();
application.register(component("service", &events)).unwrap();
let shutdown = shutdown_when_ready(application.request_lifecycle());
test_runtime()
.block_on(application.run_until_shutdown(shutdown))
.unwrap();
assert_eq!(
*events.lock().unwrap(),
[
"start:db",
"start:service",
"shutdown:service",
"shutdown:db"
]
);
}
#[test]
fn async_bootstrap_runs_before_early_shutdown_prevents_component_start() {
let events = Arc::new(Mutex::new(Vec::new()));
let bootstrap_events = Arc::clone(&events);
test_runtime()
.block_on(bootstrap_and_run(
move || async move {
bootstrap_events
.lock()
.unwrap()
.push("bootstrap".to_owned());
let mut application = Application::new();
application.register(component("component", &bootstrap_events))?;
Ok(application)
},
std::future::ready(Ok(())),
))
.unwrap();
assert_eq!(*events.lock().unwrap(), ["bootstrap"]);
}
#[test]
fn failed_async_bootstrap_does_not_start_components() {
let error = test_runtime()
.block_on(bootstrap_and_run(
|| async { Err(test_error("bootstrap failed")) },
std::future::pending(),
))
.unwrap_err();
assert_eq!(error.message(), "bootstrap failed");
}
#[test]
fn startup_failure_rolls_back_only_started_components() {
let events = Arc::new(Mutex::new(Vec::new()));
let mut application = Application::new();
application.register(component("first", &events)).unwrap();
let mut failing = component("failing", &events);
failing.start_error = true;
application.register(failing).unwrap();
application.register(component("never", &events)).unwrap();
let shutdown = shutdown_when_ready(application.request_lifecycle());
let error = test_runtime()
.block_on(application.run_until_shutdown(shutdown))
.unwrap_err();
assert_eq!(error.message(), "start failed");
assert_eq!(
*events.lock().unwrap(),
["start:first", "start:failing", "shutdown:first"]
);
}
#[test]
fn shutdown_continues_after_a_component_error() {
let events = Arc::new(Mutex::new(Vec::new()));
let mut application = Application::new();
application.register(component("first", &events)).unwrap();
let mut failing = component("second", &events);
failing.shutdown_error = true;
application.register(failing).unwrap();
let shutdown = shutdown_when_ready(application.request_lifecycle());
let error = test_runtime()
.block_on(application.run_until_shutdown(shutdown))
.unwrap_err();
assert_eq!(error.message(), "shutdown failed");
assert_eq!(
*events.lock().unwrap(),
[
"start:first",
"start:second",
"shutdown:second",
"shutdown:first"
]
);
}
#[test]
fn duplicate_component_names_are_rejected() {
let events = Arc::new(Mutex::new(Vec::new()));
let mut application = Application::new();
application.register(component("db", &events)).unwrap();
let error = application.register(component("db", &events)).unwrap_err();
assert_eq!(error.code(), "runtime.duplicate_component");
}
#[test]
fn application_shutdown_waits_for_an_admitted_request() {
test_runtime().block_on(async {
let application = Application::new();
let requests = application.request_lifecycle();
let (release, released) = tokio::sync::oneshot::channel();
let shutdown = async move {
shutdown_when_ready(requests.clone()).await?;
let request = requests
.try_accept()
.expect("application is ready before waiting for shutdown");
tokio::spawn(async move {
released.await.unwrap();
drop(request);
});
Ok(())
};
let running = tokio::spawn(application.run_until_shutdown(shutdown));
tokio::task::yield_now().await;
assert!(!running.is_finished());
release.send(()).unwrap();
running.await.unwrap().unwrap();
});
}
#[test]
fn signal_failure_before_start_prevents_component_startup() {
let events = Arc::new(Mutex::new(Vec::new()));
let mut application = Application::new();
application.register(component("service", &events)).unwrap();
let error = test_runtime()
.block_on(application.run_until_shutdown(async { Err(signal_error()) }))
.unwrap_err();
assert_eq!(error.code(), "runtime.signal_registration_failed");
assert!(events.lock().unwrap().is_empty());
}
#[test]
fn shutdown_during_startup_stops_starting_and_rolls_back() {
test_runtime().block_on(async {
let events = Arc::new(Mutex::new(Vec::new()));
let (started_tx, started_rx) = tokio::sync::oneshot::channel();
let (release_tx, release_rx) = tokio::sync::oneshot::channel();
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
let mut application = Application::new();
application
.register(BlockingStartComponent {
events: Arc::clone(&events),
started: Mutex::new(Some(started_tx)),
release: Mutex::new(Some(release_rx)),
})
.unwrap();
application.register(component("never", &events)).unwrap();
let running = tokio::spawn(application.run_until_shutdown(async move {
shutdown_rx.await.unwrap();
Ok(())
}));
started_rx.await.unwrap();
shutdown_tx.send(()).unwrap();
tokio::task::yield_now().await;
release_tx.send(()).unwrap();
running.await.unwrap().unwrap();
assert_eq!(
*events.lock().unwrap(),
["start:blocking", "shutdown:blocking"]
);
});
}
#[test]
fn managed_runtime_provides_an_async_io_driver() {
build_runtime()
.unwrap()
.block_on(async {
tokio::net::TcpListener::bind((std::net::Ipv4Addr::LOCALHOST, 0)).await
})
.expect("service listeners require the managed async I/O driver");
}
}