use std::path::PathBuf;
use std::time::Duration;
use crate::banner::Banner;
use crate::banner::emit;
use crate::container::{Container, seed_builtins};
use crate::error::{Error, Result, combine_results, with_cleanup};
use crate::graph::render_dot;
use crate::invoke::InvokeOut;
use crate::invoke::ScopedInvoker;
use crate::lifecycle::Lifecycle;
use crate::module::Module;
use crate::option::ModOption;
use crate::provide_group::init_group;
use crate::provide_group::require_group;
use crate::scope::ScopeId;
use crate::shutdown::Shutdowner;
use crate::timeout::DEFAULT_TIMEOUT;
use crate::trace;
use crate::trace::dot_graph_written;
use crate::trace::emit_unfinished;
use crate::trace::hooks_abandoned;
use crate::trace::invoke_cancelled;
use crate::trace::invoke_failed;
use crate::trace::invoke_panicked;
use crate::trace::invoking;
use crate::trace::rollback_failed;
use crate::trace::rolled_back;
use crate::trace::rolling_back;
use crate::trace::running_app_dropped;
use crate::trace::start_failed;
use crate::trace::stop_failed;
use crate::trace::stopped;
use crate::wiring::{impl_group_wiring_methods, impl_wiring_methods};
use std::borrow::Cow;
use std::fmt;
use std::fs::write;
use std::future::Future;
use std::mem::take;
use std::time::Instant;
use tokio::time::timeout;
#[derive(Debug)]
pub struct Modrun;
pub struct ModrunBuilder {
options: Vec<Box<dyn ModOption>>,
banner: Banner,
build_timeout: Option<Duration>,
start_timeout: Option<Duration>,
stop_timeout: Option<Duration>,
dot_graph_path: Option<PathBuf>,
}
impl Default for ModrunBuilder {
fn default() -> Self {
Self {
options: Vec::new(),
banner: Banner::default(),
build_timeout: Some(DEFAULT_TIMEOUT),
start_timeout: Some(DEFAULT_TIMEOUT),
stop_timeout: Some(DEFAULT_TIMEOUT),
dot_graph_path: None,
}
}
}
impl fmt::Debug for ModrunBuilder {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ModrunBuilder")
.field("options", &self.options.len())
.field("banner", &self.banner)
.field("build_timeout", &self.build_timeout)
.field("start_timeout", &self.start_timeout)
.field("stop_timeout", &self.stop_timeout)
.field("dot_graph_path", &self.dot_graph_path)
.finish()
}
}
impl Modrun {
#[must_use]
pub fn builder() -> ModrunBuilder {
ModrunBuilder::default()
}
}
impl ModrunBuilder {
fn push_option(&mut self, option: Box<dyn ModOption>) {
self.options.push(option);
}
#[must_use]
pub fn module(mut self, module: Module) -> Self {
self.push_option(module.into_option());
self
}
#[must_use]
pub fn build_timeout(mut self, duration: Duration) -> Self {
self.build_timeout = Some(duration);
self
}
#[must_use]
pub fn no_build_timeout(mut self) -> Self {
self.build_timeout = None;
self
}
#[must_use]
pub fn start_timeout(mut self, duration: Duration) -> Self {
self.start_timeout = Some(duration);
self
}
#[must_use]
pub fn no_start_timeout(mut self) -> Self {
self.start_timeout = None;
self
}
#[must_use]
pub fn stop_timeout(mut self, duration: Duration) -> Self {
self.stop_timeout = Some(duration);
self
}
#[must_use]
pub fn no_stop_timeout(mut self) -> Self {
self.stop_timeout = None;
self
}
#[must_use]
pub fn banner(mut self, text: impl Into<Cow<'static, str>>) -> Self {
self.banner = Banner::Custom(text.into());
self
}
#[must_use]
pub fn no_banner(mut self) -> Self {
self.banner = Banner::Off;
self
}
pub fn render_dot(self) -> Result<String> {
let state = self.prepare_build_state()?;
Ok(render_dot(&state.container, &state.invokers))
}
#[must_use]
pub fn dot_graph(mut self, path: impl Into<PathBuf>) -> Self {
self.dot_graph_path = Some(path.into());
self
}
fn print_banner(&self) {
emit(&self.banner);
}
pub async fn run(self) -> Result<()> {
self.print_banner();
let shutdown = Shutdowner::new();
let lifecycle = Lifecycle::new(shutdown.clone());
let mut signals = SignalWatch::install()?;
let mut state = self.into_build_state(lifecycle.clone(), shutdown.clone())?;
let stop_timeout = state.stop_timeout;
let app = tokio::select! {
biased;
result = run_invokers(&mut state) => match result {
Ok(()) => built_app_from_state(state),
Err(e) => {
return finish_run(
Err(e),
unwind_registered(&lifecycle, stop_timeout).await,
);
}
},
_ = shutdown.wait() => {
crate::trace::shutdown_requested();
return finish_interrupt(
&shutdown,
unwind_registered(&lifecycle, stop_timeout).await,
);
}
signal = signals.recv() => {
crate::trace::received_signal(signal);
shutdown.shutdown();
crate::trace::shutdown_requested();
return finish_run(
Ok(()),
unwind_registered(&lifecycle, stop_timeout).await,
);
}
};
let started = tokio::select! {
biased;
result = app.start_hooks() => result,
_ = shutdown.wait() => {
crate::trace::shutdown_requested();
crate::trace::rolling_back_after_shutdown();
return finish_interrupt(&shutdown, graceful_unwind(&app).await);
}
signal = signals.recv() => {
crate::trace::received_signal(signal);
shutdown.shutdown();
crate::trace::shutdown_requested();
crate::trace::rolling_back_after_shutdown();
return finish_run(Ok(()), graceful_unwind(&app).await);
}
};
match started {
Ok(()) => trace::started(),
Err(err) => {
rolling_back(&err);
let cleanup = graceful_unwind(&app).await;
start_failed(&err);
return finish_run(Err(err), cleanup);
}
}
tokio::select! {
_ = shutdown.wait() => {
crate::trace::shutdown_requested();
}
signal = signals.recv() => {
crate::trace::received_signal(signal);
shutdown.shutdown();
crate::trace::shutdown_requested();
}
}
app.stop().await
}
pub async fn start(self) -> Result<RunningApp> {
self.print_banner();
let shutdown = Shutdowner::new();
let lifecycle = Lifecycle::new(shutdown.clone());
let mut state = self.into_build_state(lifecycle.clone(), shutdown)?;
let stop_timeout = state.stop_timeout;
if let Err(e) = run_invokers(&mut state).await {
return Err(with_cleanup(
e,
unwind_registered(&lifecycle, stop_timeout).await,
));
}
let app = built_app_from_state(state);
app.start().await?;
Ok(RunningApp { inner: Some(app) })
}
fn into_build_state(
mut self,
lifecycle: Lifecycle,
shutdown: Shutdowner,
) -> Result<BuildState> {
let dot_graph_path = self.dot_graph_path.take();
let state = self.prepare_build_state_with(lifecycle, shutdown)?;
if let Some(path) = dot_graph_path {
let dot = render_dot(&state.container, &state.invokers);
write(&path, dot).map_err(|source| Error::io("write dot graph", source))?;
dot_graph_written(path.as_os_str().to_string_lossy().as_ref());
}
Ok(state)
}
fn prepare_build_state(self) -> Result<BuildState> {
let shutdown = Shutdowner::new();
let lifecycle = Lifecycle::new(shutdown.clone());
self.prepare_build_state_with(lifecycle, shutdown)
}
fn prepare_build_state_with(
self,
lifecycle: Lifecycle,
shutdown: Shutdowner,
) -> Result<BuildState> {
let mut state = BuildState {
container: Container::new(),
invokers: Vec::new(),
lifecycle,
build_timeout: self.build_timeout,
start_timeout: self.start_timeout,
stop_timeout: self.stop_timeout,
current_scope: ScopeId::ROOT,
};
seed_builtins(&mut state.container, state.lifecycle.clone(), shutdown)?;
for opt in self.options {
opt.apply(&mut state)?;
}
let invoker_deps: Vec<_> = state
.invokers
.iter()
.map(|scoped| (scoped.scope, scoped.invoker.dep_types()))
.collect();
state.container.validate(&invoker_deps)?;
Ok(state)
}
}
async fn run_invokers(state: &mut BuildState) -> Result<()> {
let invokers = take(&mut state.invokers);
let budget = state.build_timeout;
let invoke = async {
for scoped in invokers {
let ScopedInvoker { scope, invoker } = scoped;
let function = invoker.name();
let deps = invoker.dep_list();
let module = state.container.scopes().name(scope);
invoking(function, deps.as_slice(), module);
let mut invoke_guard = InflightInvoke::new(function, module);
let previous = state.container.enter_scope(scope);
let result = async {
if !deps.as_slice().is_empty() {
state.container.ensure_built(deps.as_slice()).await?;
}
match invoker.call(&state.container) {
InvokeOut::Done(r) => r,
InvokeOut::Fut(fut) => fut.await,
}
}
.await;
state.container.leave_scope(previous);
invoke_guard.finish();
if let Err(ref err) = result {
invoke_failed(function, deps.as_slice(), module, err);
}
result?;
}
Ok(())
};
with_timeout(
budget,
invoke,
Error::BuildTimeout(budget.unwrap_or(Duration::ZERO)),
)
.await
}
struct InflightInvoke {
function: &'static str,
module: &'static str,
finished: bool,
}
impl InflightInvoke {
fn new(function: &'static str, module: &'static str) -> Self {
Self {
function,
module,
finished: false,
}
}
fn finish(&mut self) {
self.finished = true;
}
}
impl Drop for InflightInvoke {
fn drop(&mut self) {
emit_unfinished(
self.finished,
|| invoke_panicked(self.function, self.module),
|| invoke_cancelled(self.function, self.module),
);
}
}
fn built_app_from_state(state: BuildState) -> BuiltApp {
BuiltApp {
lifecycle: state.lifecycle,
start_timeout: state.start_timeout,
stop_timeout: state.stop_timeout,
}
}
fn finish_run(phase: Result<()>, cleanup: Result<()>) -> Result<()> {
match phase {
Ok(()) => cleanup,
Err(e) => Err(with_cleanup(e, cleanup)),
}
}
fn finish_interrupt(shutdown: &Shutdowner, cleanup: Result<()>) -> Result<()> {
if shutdown.is_failure() {
match cleanup {
Ok(()) => Err(Error::TaskFailedDuringStart),
Err(e) if matches!(e, Error::UnwindTimeout(_) | Error::StopTimeout(_)) => {
Err(with_cleanup(Error::TaskFailedDuringStart, Err(e)))
}
Err(e) => Err(e),
}
} else {
cleanup
}
}
fn report_unwind_result(lifecycle: &Lifecycle, result: Result<()>) -> Result<()> {
match result {
Ok(()) => {
rolled_back();
Ok(())
}
Err(err) => {
rollback_failed(&err);
let leftover = lifecycle.pending_stops();
if leftover > 0 {
hooks_abandoned(leftover);
}
Err(err)
}
}
}
async fn unwind_lifecycle(lifecycle: &Lifecycle, stop_timeout: Option<Duration>) -> Result<()> {
lifecycle.prepare_for_unwind();
report_unwind_result(
lifecycle,
with_timeout(
stop_timeout,
lifecycle.unwind_started(),
Error::UnwindTimeout(stop_timeout.unwrap_or(Duration::ZERO)),
)
.await,
)
}
async fn unwind_registered(lifecycle: &Lifecycle, stop_timeout: Option<Duration>) -> Result<()> {
unwind_lifecycle(lifecycle, stop_timeout).await
}
impl_wiring_methods!(ModrunBuilder);
impl_group_wiring_methods!(ModrunBuilder);
impl ModrunBuilder {
#[must_use]
pub fn init_group<T: Clone + Send + Sync + 'static>(mut self) -> Self {
self.push_option(init_group::<T>());
self
}
#[must_use]
pub fn require_group<T: Clone + Send + Sync + 'static>(mut self) -> Self {
self.push_option(require_group::<T>());
self
}
}
pub(crate) struct BuildState {
pub(crate) container: Container,
pub(crate) invokers: Vec<ScopedInvoker>,
pub(crate) lifecycle: Lifecycle,
pub(crate) build_timeout: Option<Duration>,
pub(crate) start_timeout: Option<Duration>,
pub(crate) stop_timeout: Option<Duration>,
pub(crate) current_scope: ScopeId,
}
struct BuiltApp {
lifecycle: Lifecycle,
start_timeout: Option<Duration>,
stop_timeout: Option<Duration>,
}
async fn graceful_unwind(app: &BuiltApp) -> Result<()> {
unwind_lifecycle(&app.lifecycle, app.stop_timeout).await
}
async fn with_timeout(
budget: Option<Duration>,
fut: impl Future<Output = Result<()>>,
timed_out: Error,
) -> Result<()> {
match budget {
None => fut.await,
Some(d) => {
let started = Instant::now();
match timeout(d, fut).await {
Ok(Err(err)) => Err(err),
Ok(Ok(())) if started.elapsed() >= d => Err(timed_out),
Ok(Ok(())) => Ok(()),
Err(_) => Err(timed_out),
}
}
}
}
impl BuiltApp {
async fn start_hooks(&self) -> Result<()> {
with_timeout(
self.start_timeout,
self.lifecycle.start(),
Error::StartTimeout(self.start_timeout.unwrap_or(Duration::ZERO)),
)
.await
}
async fn start(&self) -> Result<()> {
match self.start_hooks().await {
Ok(()) => {
trace::started();
Ok(())
}
Err(err) => {
rolling_back(&err);
let unwound = graceful_unwind(self).await;
start_failed(&err);
combine_results(Err(err), unwound)
}
}
}
async fn stop(&self) -> Result<()> {
let result = with_timeout(
self.stop_timeout,
self.lifecycle.stop(),
Error::StopTimeout(self.stop_timeout.unwrap_or(Duration::ZERO)),
)
.await;
match &result {
Err(err) => {
stop_failed(err);
let leftover = self.lifecycle.pending_stops();
if leftover > 0 {
hooks_abandoned(leftover);
}
}
Ok(()) => stopped(),
}
result
}
}
struct SignalWatch {
#[cfg(all(feature = "signal", unix))]
sigint: tokio::signal::unix::Signal,
#[cfg(all(feature = "signal", unix))]
sigterm: tokio::signal::unix::Signal,
#[cfg(all(feature = "signal", windows))]
ctrl_c: tokio::signal::windows::CtrlC,
#[cfg(all(feature = "signal", windows))]
ctrl_break: tokio::signal::windows::CtrlBreak,
#[cfg(all(feature = "signal", windows))]
ctrl_close: tokio::signal::windows::CtrlClose,
#[cfg(all(feature = "signal", windows))]
ctrl_shutdown: tokio::signal::windows::CtrlShutdown,
}
impl SignalWatch {
fn install() -> Result<Self> {
#[cfg(all(feature = "signal", unix))]
{
use tokio::signal::unix::{SignalKind, signal};
let sigint = signal(SignalKind::interrupt()).map_err(Error::SigintListen)?;
let sigterm = signal(SignalKind::terminate()).map_err(Error::SigtermListen)?;
Ok(Self { sigint, sigterm })
}
#[cfg(all(feature = "signal", windows))]
{
use tokio::signal::windows::{ctrl_break, ctrl_c, ctrl_close, ctrl_shutdown};
Ok(Self {
ctrl_c: ctrl_c().map_err(Error::SignalListen)?,
ctrl_break: ctrl_break().map_err(Error::SignalListen)?,
ctrl_close: ctrl_close().map_err(Error::SignalListen)?,
ctrl_shutdown: ctrl_shutdown().map_err(Error::SignalListen)?,
})
}
#[cfg(not(any(all(feature = "signal", unix), all(feature = "signal", windows))))]
{
Ok(Self {})
}
}
async fn recv(&mut self) -> &'static str {
#[cfg(all(feature = "signal", unix))]
{
tokio::select! {
Some(()) = self.sigint.recv() => "SIGINT",
Some(()) = self.sigterm.recv() => "SIGTERM",
else => std::future::pending().await,
}
}
#[cfg(all(feature = "signal", windows))]
{
tokio::select! {
Some(()) = self.ctrl_c.recv() => "CTRL_C",
Some(()) = self.ctrl_break.recv() => "CTRL_BREAK",
Some(()) = self.ctrl_close.recv() => "CTRL_CLOSE",
Some(()) = self.ctrl_shutdown.recv() => "CTRL_SHUTDOWN",
else => std::future::pending().await,
}
}
#[cfg(not(any(all(feature = "signal", unix), all(feature = "signal", windows))))]
{
let _ = self;
std::future::pending::<&'static str>().await
}
}
}
#[must_use = "dropping a RunningApp skips all OnStop hooks; call stop() instead"]
pub struct RunningApp {
inner: Option<BuiltApp>,
}
impl fmt::Debug for RunningApp {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("RunningApp").finish_non_exhaustive()
}
}
impl RunningApp {
pub async fn stop(mut self) -> Result<()> {
let inner = self
.inner
.take()
.expect("RunningApp always contains a built application");
inner.stop().await
}
}
impl Drop for RunningApp {
fn drop(&mut self) {
if self.inner.is_some() {
running_app_dropped();
#[cfg(debug_assertions)]
eprintln!("modrun: dropping RunningApp without stop(); OnStop hooks will not run");
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn finish_interrupt_retains_failure_on_unwind_timeout() {
let shutdown = Shutdowner::new();
shutdown.fail();
let err = finish_interrupt(
&shutdown,
Err(Error::UnwindTimeout(Duration::from_millis(30))),
)
.unwrap_err();
match err {
Error::CleanupAfterFailure { cleanup, earlier } => {
assert!(
matches!(*cleanup, Error::UnwindTimeout(_)),
"cleanup was {cleanup}"
);
assert!(
matches!(*earlier, Error::TaskFailedDuringStart),
"earlier was {earlier}"
);
}
other => panic!("expected CleanupAfterFailure, got {other}"),
}
}
#[test]
fn finish_interrupt_keeps_join_error() {
let shutdown = Shutdowner::new();
shutdown.fail();
let err = finish_interrupt(&shutdown, Err(Error::hook("died"))).unwrap_err();
let msg = err.to_string();
assert!(msg.contains("died"), "{msg}");
assert!(!msg.contains("cleanup failed"), "{msg}");
}
#[test]
fn finish_interrupt_graceful_ok_is_ok() {
let shutdown = Shutdowner::new();
shutdown.shutdown();
finish_interrupt(&shutdown, Ok(())).unwrap();
}
}