use std::{future::Future, mem::take, sync::Arc, time::Instant}; use tokio::{process::Command as TokioCommand, select, task::JoinHandle}; use tracing::{instrument, trace, trace_span, Instrument}; use watchexec_signals::Signal; use crate::{ command::Command, errors::{sync_io_error, SyncIoError}, flag::Flag, job::priority::Timer, }; use super::{ job::Job, messages::{Control, ControlMessage}, priority, state::CommandState, }; /// Spawn a job task and return a [`Job`] handle and a [`JoinHandle`]. /// /// The job task immediately starts in the background: it does not need polling. #[must_use] #[instrument(level = "trace")] pub fn start_job(command: Arc) -> (Job, JoinHandle<()>) { let (sender, mut receiver) = priority::new(); let gone = Flag::default(); let done = gone.clone(); ( Job { command: command.clone(), control_queue: sender, gone, }, tokio::spawn(async move { let mut error_handler = ErrorHandler::None; let mut spawn_hook = SpawnHook::None; let mut command_state = CommandState::Pending; let mut previous_run = None; let mut stop_timer = None; let mut on_end: Vec = Vec::new(); let mut on_end_restart: Option = None; enum Loop { Normally, Skip, Break, } 'main: loop { select! { result = command_state.wait(), if command_state.is_running() => { trace!(?result, ?command_state, "got wait result"); match async { #[cfg(test)] eprintln!("[{:?}] waited: {result:?}", Instant::now()); match result { Err(err) => { let fut = error_handler.call(sync_io_error(err)); fut.await; return Loop::Skip; } Ok(true) => { trace!(existing=?stop_timer, "erasing stop timer"); stop_timer = None; trace!(count=%on_end.len(), "raising all pending end flags"); for done in take(&mut on_end).into_iter() { done.raise(); } if let Some(flag) = on_end_restart.take() { trace!("continuing a graceful restart"); let mut spawnable = command.to_spawnable(); previous_run = Some(command_state.reset()); spawn_hook .call( &mut spawnable, &JobTaskContext { command: command.clone(), current: &command_state, previous: previous_run.as_ref(), }, ) .await; if let Err(err) = command_state.spawn(command.clone(), spawnable).await { let fut = error_handler.call(sync_io_error(err)); fut.await; return Loop::Skip; } trace!("raising graceful restart's flag"); flag.raise(); } } Ok(false) => { trace!("child wasn't running, ignoring wait result"); } } Loop::Normally }.instrument(trace_span!("handle wait result")).await { Loop::Normally => {} Loop::Skip => { trace!("skipping to next event"); continue 'main; } Loop::Break => { trace!("breaking out of main loop"); break 'main; } } } Some(ControlMessage { control, done }) = receiver.recv(&mut stop_timer) => { match async { trace!(?control, ?command_state, "got control message"); #[cfg(test)] eprintln!("[{:?}] control: {control:?}", Instant::now()); macro_rules! try_with_handler { ($erroring:expr) => { match $erroring { Err(err) => { let fut = error_handler.call(sync_io_error(err)); fut.await; trace!("raising done flag for this control after error"); done.raise(); return Loop::Normally; } Ok(value) => value, } }; } match control { Control::Start => { if command_state.is_running() { trace!("child is running, skip"); } else { let mut spawnable = command.to_spawnable(); previous_run = Some(command_state.reset()); spawn_hook .call( &mut spawnable, &JobTaskContext { command: command.clone(), current: &command_state, previous: previous_run.as_ref(), }, ) .await; try_with_handler!(command_state.spawn(command.clone(), spawnable).await); } } Control::Stop => { if let CommandState::Running { child, started, .. } = &mut command_state { trace!("stopping child"); try_with_handler!(child.kill().await); trace!("waiting on child"); let status = try_with_handler!(child.wait().await); trace!(?status, "got child end status"); command_state = CommandState::Finished { status: status.into(), started: *started, finished: Instant::now(), }; trace!(count=%on_end.len(), "raising all pending end flags"); for done in take(&mut on_end).into_iter() { done.raise(); } } else { trace!("child isn't running, skip"); } } Control::GracefulStop { signal, grace } => { if let CommandState::Running { child, .. } = &mut command_state { try_with_handler!(signal_child(signal, child).await); trace!(?grace, "setting up graceful stop timer"); stop_timer.replace(Timer::stop(grace, done)); return Loop::Skip; } else { trace!("child isn't running, skip"); } } Control::TryRestart => { if let CommandState::Running { child, started, .. } = &mut command_state { trace!("stopping child"); try_with_handler!(child.kill().await); trace!("waiting on child"); let status = try_with_handler!(child.wait().await); trace!(?status, "got child end status"); command_state = CommandState::Finished { status: status.into(), started: *started, finished: Instant::now(), }; previous_run = Some(command_state.reset()); trace!(count=%on_end.len(), "raising all pending end flags"); for done in take(&mut on_end).into_iter() { done.raise(); } let mut spawnable = command.to_spawnable(); spawn_hook .call( &mut spawnable, &JobTaskContext { command: command.clone(), current: &command_state, previous: previous_run.as_ref(), }, ) .await; try_with_handler!(command_state.spawn(command.clone(), spawnable).await); } else { trace!("child isn't running, skip"); } } Control::TryGracefulRestart { signal, grace } => { if let CommandState::Running { child, .. } = &mut command_state { try_with_handler!(signal_child(signal, child).await); trace!(?grace, "setting up graceful stop timer"); stop_timer.replace(Timer::restart(grace, done.clone())); trace!("setting up graceful restart flag"); on_end_restart = Some(done); return Loop::Skip; } else { trace!("child isn't running, skip"); } } Control::ContinueTryGracefulRestart => { trace!("continuing a graceful try-restart"); if let CommandState::Running { child, started, .. } = &mut command_state { trace!("stopping child forcefully"); try_with_handler!(child.kill().await); trace!("waiting on child"); let status = try_with_handler!(child.wait().await); trace!(?status, "got child end status"); command_state = CommandState::Finished { status: status.into(), started: *started, finished: Instant::now(), }; trace!(count=%on_end.len(), "raising all pending end flags"); for done in take(&mut on_end).into_iter() { done.raise(); } } let mut spawnable = command.to_spawnable(); previous_run = Some(command_state.reset()); spawn_hook .call( &mut spawnable, &JobTaskContext { command: command.clone(), current: &command_state, previous: previous_run.as_ref(), }, ) .await; try_with_handler!(command_state.spawn(command.clone(), spawnable).await); } Control::Signal(signal) => { if let CommandState::Running { child, .. } = &mut command_state { try_with_handler!(signal_child(signal, child).await); } else { trace!("child isn't running, skip"); } } Control::Delete => { trace!("raising done flag immediately"); done.raise(); return Loop::Break; } Control::NextEnding => { if !matches!(command_state, CommandState::Finished { .. }) { trace!("queue end flag"); on_end.push(done); return Loop::Skip; } else { trace!("child is finished, raise done flag immediately"); done.raise(); return Loop::Normally; } } Control::SyncFunc(f) => { f(&JobTaskContext { command: command.clone(), current: &command_state, previous: previous_run.as_ref(), }); } Control::AsyncFunc(f) => { Box::into_pin(f(&JobTaskContext { command: command.clone(), current: &command_state, previous: previous_run.as_ref(), })) .await; } Control::SetSyncErrorHandler(f) => { trace!("setting sync error handler"); error_handler = ErrorHandler::Sync(f); } Control::SetAsyncErrorHandler(f) => { trace!("setting async error handler"); error_handler = ErrorHandler::Async(f); } Control::UnsetErrorHandler => { trace!("unsetting error handler"); error_handler = ErrorHandler::None; } Control::SetSyncSpawnHook(f) => { trace!("setting sync spawn hook"); spawn_hook = SpawnHook::Sync(f); } Control::SetAsyncSpawnHook(f) => { trace!("setting async spawn hook"); spawn_hook = SpawnHook::Async(f); } Control::UnsetSpawnHook => { trace!("unsetting spawn hook"); spawn_hook = SpawnHook::None; } } trace!("raising control done flag"); done.raise(); Loop::Normally }.instrument(trace_span!("handle control message")).await { Loop::Normally => {} Loop::Skip => { trace!("skipping to next event (without raising done flag)"); continue 'main; } Loop::Break => { trace!("breaking out of main loop"); break 'main; } } } } } trace!("raising job done flag"); done.raise(); }), ) } macro_rules! sync_async_callbox { ($name:ident, $synct:ty, $asynct:ty, ($($argname:ident : $argtype:ty),*)) => { pub enum $name { None, Sync($synct), Async($asynct), } impl $name { #[instrument(level = "trace", skip(self, $($argname),*))] pub async fn call(&self, $($argname: $argtype),*) { match self { $name::None => (), $name::Sync(f) => { ::tracing::trace!("calling sync {:?}", stringify!($name)); f($($argname),*) } $name::Async(f) => { ::tracing::trace!("calling async {:?}", stringify!($name)); Box::into_pin(f($($argname),*)).await } } } } }; } /// Job task internals exposed via hooks. #[derive(Debug)] pub struct JobTaskContext<'task> { /// The job's [`Command`]. pub command: Arc, /// The current state of the job. pub current: &'task CommandState, /// The state of the previous iteration of the job, if any. /// /// This is generally [`CommandState::Finished`], but may be other states in rare cases. pub previous: Option<&'task CommandState>, } pub type SyncFunc = Box) + Send + Sync + 'static>; pub type AsyncFunc = Box< dyn (FnOnce(&JobTaskContext<'_>) -> Box + Send + Sync>) + Send + Sync + 'static, >; pub type SyncSpawnHook = Arc) + Send + Sync + 'static>; pub type AsyncSpawnHook = Arc< dyn (Fn(&mut TokioCommand, &JobTaskContext<'_>) -> Box + Send + Sync>) + Send + Sync + 'static, >; sync_async_callbox!(SpawnHook, SyncSpawnHook, AsyncSpawnHook, (command: &mut TokioCommand, context: &JobTaskContext<'_>)); pub type SyncErrorHandler = Arc; pub type AsyncErrorHandler = Arc< dyn (Fn(SyncIoError) -> Box + Send + Sync>) + Send + Sync + 'static, >; sync_async_callbox!(ErrorHandler, SyncErrorHandler, AsyncErrorHandler, (error: SyncIoError)); #[cfg_attr(not(windows), allow(clippy::needless_pass_by_ref_mut))] // needed for start_kill() #[instrument(level = "trace")] async fn signal_child( signal: Signal, #[cfg(test)] child: &mut super::TestChild, #[cfg(not(test))] child: &mut command_group::tokio::ErasedChild, ) -> std::io::Result<()> { #[cfg(unix)] { let sig = signal .to_nix() .or_else(|| Signal::Terminate.to_nix()) .expect("UNWRAP: guaranteed for Signal::Terminate default"); trace!(signal=?sig, "sending signal"); child.signal(sig)?; } #[cfg(windows)] if signal == Signal::ForceStop { trace!("starting kill, without waiting"); child.start_kill()?; } else { trace!(?signal, "ignoring unsupported signal"); } Ok(()) }