use std::{
collections::BTreeMap,
ffi::OsString,
future::Future,
io,
path::PathBuf,
process::{ExitStatus, Stdio},
time::Duration,
};
use thiserror::Error;
use tokio::{
io::{AsyncRead, AsyncReadExt, AsyncWriteExt},
time::{Instant, sleep, sleep_until, timeout},
};
use tokio_util::sync::CancellationToken;
use crate::{
platform::{self, Child},
tools::human::log::{End, Lines, Run, Stream},
};
const OUTPUT_LIMIT: usize = 128 * 1024;
const CLEANUP_GRACE: Duration = Duration::from_secs(1);
const CAPTURE_BUFFER_BYTES: usize = 8192;
#[derive(Debug, Clone)]
pub struct Command {
pub program: OsString,
pub args: Vec<OsString>,
pub cwd: PathBuf,
pub env: BTreeMap<OsString, OsString>,
pub timeout: Duration,
}
#[derive(
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize, serde::Deserialize,
)]
#[serde(transparent)]
pub struct ProcessId(std::num::NonZeroU32);
impl ProcessId {
pub fn new(pid: u32) -> Self {
Self(std::num::NonZeroU32::new(pid).expect("a real process id is never 0"))
}
pub fn get(self) -> u32 {
self.0.get()
}
}
impl From<std::num::NonZeroU32> for ProcessId {
fn from(pid: std::num::NonZeroU32) -> Self {
Self(pid)
}
}
impl std::fmt::Display for ProcessId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
self.0.fmt(f)
}
}
impl rusqlite::types::ToSql for ProcessId {
fn to_sql(&self) -> rusqlite::Result<rusqlite::types::ToSqlOutput<'_>> {
Ok(rusqlite::types::ToSqlOutput::from(i64::from(self.0.get())))
}
}
impl rusqlite::types::FromSql for ProcessId {
fn column_result(value: rusqlite::types::ValueRef<'_>) -> rusqlite::types::FromSqlResult<Self> {
let pid = u32::try_from(value.as_i64()?)
.map_err(|error| rusqlite::types::FromSqlError::Other(error.into()))?;
std::num::NonZeroU32::new(pid).map(Self).ok_or_else(|| {
rusqlite::types::FromSqlError::Other(
"a saved owner_pid of 0 is not a process id".into(),
)
})
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ChildIdentity {
pub pid: ProcessId,
pub start_time: u64,
}
#[derive(Debug)]
pub struct Output {
pub status: ExitStatus,
pub stdout: Vec<u8>,
pub stderr: Vec<u8>,
pub truncated: bool,
pub duration: Duration,
}
pub(crate) fn tool_output<'a>(
program: &str,
args: impl IntoIterator<Item = &'a std::ffi::OsStr>,
) -> Option<Vec<u8>> {
platform::program::tool_output(program, args)
}
pub(crate) struct Launcher {
pub env: BTreeMap<OsString, OsString>,
}
impl artifactize_tools::launch::Launcher for Launcher {
fn run<'a>(
&'a self,
launch: artifactize_tools::launch::Launch,
cancellation: &'a CancellationToken,
) -> artifactize_tools::launch::Running<'a> {
use artifactize_tools::launch::{Finished, LaunchError};
Box::pin(async move {
let command = Command {
program: launch.program.into(),
args: launch.args,
cwd: launch.cwd,
env: self.env.clone(),
timeout: launch.timeout,
};
let ready = |_| async { Ok(()) };
match run_with_input_limit(
command,
None,
launch.output_limit,
cancellation.clone(),
ready,
)
.await
{
Ok(output) => Ok(Finished {
status: output.status,
stdout: output.stdout,
stderr: output.stderr,
truncated: output.truncated,
}),
Err(Error::Timeout) => Err(LaunchError::TimedOut),
Err(Error::Cancelled) => Err(LaunchError::Cancelled),
Err(error) => Err(LaunchError::Failed(
error
.argument_refusal()
.unwrap_or_else(|| error.to_string()),
)),
}
})
}
}
#[derive(Debug, Error)]
pub enum Error {
#[error("process could not be spawned: {0}")]
Spawn(#[source] io::Error),
#[error("child registration failed: {0}")]
Registration(#[source] io::Error),
#[error("process supervision failed: {0}")]
Io(#[from] io::Error),
#[error("process cleanup could not be confirmed: {0}")]
Cleanup(#[source] io::Error),
#[error("process supervision task failed: {0}")]
Task(#[from] tokio::task::JoinError),
#[error("process timed out")]
Timeout,
#[error("process was cancelled")]
Cancelled,
#[error("process output pipes did not close after cleanup")]
OutputDidNotClose,
}
impl Error {
pub(crate) fn cleanup_confirmed(&self) -> bool {
match self {
Self::Cleanup(_) | Self::Task(_) => false,
Self::Spawn(_)
| Self::Registration(_)
| Self::Io(_)
| Self::Timeout
| Self::Cancelled
| Self::OutputDidNotClose => true,
}
}
pub(crate) fn argument_refusal(&self) -> Option<String> {
match self {
Self::Spawn(error) if error.kind() == io::ErrorKind::InvalidInput => Some(format!(
"Process arguments could not be passed safely: {error}"
)),
_ => None,
}
}
}
pub async fn run<F, R>(
command: Command,
cancellation: CancellationToken,
register: F,
) -> Result<Output, Error>
where
F: FnOnce(ChildIdentity) -> R + Send + 'static,
R: Future<Output = io::Result<()>> + Send + 'static,
{
run_with_input(command, None, cancellation, register).await
}
pub(crate) async fn run_with_input<F, R>(
command: Command,
input: Option<Vec<u8>>,
cancellation: CancellationToken,
register: F,
) -> Result<Output, Error>
where
F: FnOnce(ChildIdentity) -> R + Send + 'static,
R: Future<Output = io::Result<()>> + Send + 'static,
{
run_with_input_limit(command, input, OUTPUT_LIMIT, cancellation, register).await
}
pub(crate) async fn run_with_input_limit<F, R>(
command: Command,
input: Option<Vec<u8>>,
output_limit: usize,
cancellation: CancellationToken,
register: F,
) -> Result<Output, Error>
where
F: FnOnce(ChildIdentity) -> R + Send + 'static,
R: Future<Output = io::Result<()>> + Send + 'static,
{
let cancellation = cancellation.child_token();
let _cancel_on_drop = cancellation.clone().drop_guard();
crate::task::joined(
tokio::spawn(run_inner(
command,
input,
output_limit,
Capture::default(),
cancellation,
register,
))
.await,
)?
}
#[derive(Default)]
struct Capture {
logs: Option<PathBuf>,
lines: Option<(Lines, Lines)>,
}
pub(crate) async fn run_watched(
command: Command,
cancellation: CancellationToken,
run: &Run,
) -> Result<Output, Error> {
let cancellation = cancellation.child_token();
let _cancel_on_drop = cancellation.clone().drop_guard();
let capture = Capture {
logs: None,
lines: Some((run.lines(Stream::Stdout), run.lines(Stream::Stderr))),
};
crate::task::joined(
tokio::spawn(run_inner(
command,
None,
OUTPUT_LIMIT,
capture,
cancellation,
|_| async { Ok(()) },
))
.await,
)?
}
async fn run_inner<F, R>(
command: Command,
input: Option<Vec<u8>>,
output_limit: usize,
Capture { logs, lines }: Capture,
cancellation: CancellationToken,
register: F,
) -> Result<Output, Error>
where
F: FnOnce(ChildIdentity) -> R,
R: Future<Output = io::Result<()>>,
{
let started = Instant::now();
let deadline = started + command.timeout;
if cancellation.is_cancelled() {
return Err(Error::Cancelled);
}
if command.timeout.is_zero() {
return Err(Error::Timeout);
}
let mut child = spawn_owned(
&command,
input.is_some(),
cancellation.clone(),
deadline,
register,
)
.await?;
let stdout = child.stdout().take().expect("stdout is piped");
let stderr = child.stderr().take().expect("stderr is piped");
let stdin = child.stdin().take();
let reaped = CancellationToken::new();
let feed = async {
if let (Some(mut stdin), Some(input)) = (stdin, input) {
tokio::select! {
_ = reaped.cancelled() => {},
_ = cancellation.cancelled() => {},
_ = sleep_until(deadline) => {},
result = stdin.write_all(&input) => {
if let Err(error) = result && error.kind() != io::ErrorKind::BrokenPipe {
return Err(Error::Io(error));
}
}
}
}
Ok(())
};
let wait = async {
let result = tokio::select! {
biased;
_ = cancellation.cancelled() => Err(Error::Cancelled),
_ = sleep_until(deadline) => Err(Error::Timeout),
status = child.wait() => status.map_err(Error::Io),
};
let cleanup = if result.is_err() {
terminate(&mut child).await
} else {
child.cleanup().await.map(|_| ())
};
reaped.cancel();
cleanup.map_err(Error::Cleanup)?;
result
};
let (stdout_lines, stderr_lines) = lines.unzip();
let capture = async {
let read = async {
tokio::try_join!(
capture(
stdout,
output_limit,
logs.as_ref().map(|dir| dir.join("runtime.stdout")),
stdout_lines,
),
capture(
stderr,
output_limit,
logs.as_ref().map(|dir| dir.join("runtime.stderr")),
stderr_lines,
)
)
};
tokio::pin!(read);
tokio::select! {
result = &mut read => result.map_err(Error::Io),
_ = reaped.cancelled() => {
timeout(CLEANUP_GRACE, read).await
.map_err(|_| Error::OutputDidNotClose)?
.map_err(Error::Io)
}
}
};
let (status, output, input) = tokio::join!(wait, capture, feed);
let status = status?;
input?;
let ((stdout, stdout_truncated), (stderr, stderr_truncated)) = output?;
Ok(Output {
status,
stdout,
stderr,
truncated: stdout_truncated || stderr_truncated,
duration: started.elapsed(),
})
}
pub(crate) async fn spawn_owned<F, R>(
command: &Command,
input: bool,
cancellation: CancellationToken,
deadline: Instant,
register: F,
) -> Result<Child, Error>
where
F: FnOnce(ChildIdentity) -> R,
R: Future<Output = io::Result<()>>,
{
let program = platform::program::resolve_with(&command.program, &command.cwd, |name| {
environment_value(&command.env, name).map(ToOwned::to_owned)
})
.map_err(Error::Spawn)?;
let mut child = tokio::process::Command::new(&program);
child
.args(&command.args)
.current_dir(&command.cwd)
.env_clear()
.envs(&command.env)
.stdin(if input { Stdio::piped() } else { Stdio::null() })
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true);
let (gate, mut spawning) = platform::spawn_gated(child).map_err(Error::Spawn)?;
let admission = tokio::select! {
biased;
_ = cancellation.cancelled() => Err(Error::Cancelled),
_ = sleep_until(deadline) => Err(Error::Timeout),
result = async {
let pid = gate.pid().await?;
let child = child_identity(pid).map_err(Error::Registration)?;
register(child).await.map_err(Error::Registration)?;
Ok::<_, Error>(())
} => result,
};
let admission = admission.and_then(|()| {
if cancellation.is_cancelled() {
Err(Error::Cancelled)
} else if Instant::now() >= deadline {
Err(Error::Timeout)
} else {
gate.admit().map_err(Error::Io)
}
});
drop(gate);
let spawned = crate::task::joined((&mut spawning).await)?;
let mut child = match spawned {
Ok(child) => child,
Err(error) => {
return Err(match admission {
Err(error @ (Error::Registration(_) | Error::Cancelled | Error::Timeout)) => error,
_ => Error::Spawn(error),
});
}
};
if let Err(error) = admission {
terminate(&mut child).await.map_err(Error::Cleanup)?;
return Err(error);
}
Ok(child)
}
fn environment_value<'a>(
environment: &'a BTreeMap<OsString, OsString>,
name: &str,
) -> Option<&'a std::ffi::OsStr> {
crate::platform::environment::lookup(environment, name)
}
pub(crate) fn launch_detached(
program: &std::ffi::OsStr,
args: &[String],
cwd: &std::path::Path,
) -> Result<(), Error> {
let program = platform::program::resolve(program, cwd).map_err(Error::Spawn)?;
let mut command = tokio::process::Command::new(program);
command
.args(args)
.current_dir(cwd)
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.kill_on_drop(false);
let mut child = platform::spawn_detached(command).map_err(Error::Spawn)?;
tokio::spawn(async move {
let _ = child.wait().await;
});
Ok(())
}
pub(crate) fn launch_watched(
program: &std::ffi::OsStr,
args: &[String],
cwd: &std::path::Path,
run: Run,
) -> Result<(), (Error, Run)> {
if !run.watched() {
return match launch_detached(program, args, cwd) {
Ok(()) => {
run.end(End::Launched);
Ok(())
}
Err(error) => Err((error, run)),
};
}
let program = match platform::program::resolve(program, cwd) {
Ok(program) => program,
Err(error) => return Err((Error::Spawn(error), run)),
};
let mut command = tokio::process::Command::new(program);
command
.args(args)
.current_dir(cwd)
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(false);
let mut child = match platform::spawn_watched(command) {
Ok(child) => child,
Err(error) => return Err((Error::Spawn(error), run)),
};
let stdout = child.stdout.take().expect("stdout is piped");
let stderr = child.stderr.take().expect("stderr is piped");
tokio::spawn(async move {
let (status, _, _) = tokio::join!(
child.wait(),
capture(stdout, 0, None, Some(run.lines(Stream::Stdout))),
capture(stderr, 0, None, Some(run.lines(Stream::Stderr))),
);
run.end(match status {
Ok(status) => End::exited(&status),
Err(error) => End::Failed(error.to_string()),
});
});
Ok(())
}
pub(crate) fn artifactize_executable() -> Result<std::path::PathBuf, String> {
let current = std::env::current_exe().map_err(|e| e.to_string())?;
for directory in current
.parent()
.into_iter()
.flat_map(|parent| [Some(parent), parent.parent()])
.flatten()
{
if let Some(candidate) = platform::program::candidates(&directory.join("artifactize"))
.into_iter()
.find(|path| path.is_file())
{
return Ok(candidate);
}
}
platform::program::resolve(
std::ffi::OsStr::new("artifactize"),
&std::env::current_dir().map_err(|e| e.to_string())?,
)
.map_err(|e| e.to_string())
}
#[derive(Clone, Copy)]
pub(crate) enum RunEntry {
Supervisor,
Executor,
}
pub(crate) fn spawn_run(
program: &std::path::Path,
id: &crate::types::RunId,
repo: &std::path::Path,
state: &std::path::Path,
entry: RunEntry,
) -> Result<(ChildIdentity, tokio::process::Child), Error> {
let mut options = platform::private_options();
options.write(true);
match entry {
RunEntry::Supervisor => {
options.create_new(true);
}
RunEntry::Executor => {
options.append(true);
}
}
let directory = state.join("runs").join(id.as_str());
let stderr = options
.open(directory.join("worker.stderr"))
.map_err(Error::Spawn)?;
let program = platform::program::resolve(program.as_os_str(), repo).map_err(Error::Spawn)?;
let mut command = tokio::process::Command::new(program);
command
.arg("--repo")
.arg(repo)
.arg("--state-dir")
.arg(state)
.arg(match entry {
RunEntry::Supervisor => "__verify",
RunEntry::Executor => "__execute",
})
.arg(id.as_str())
.current_dir(directory)
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::from(stderr))
.kill_on_drop(false);
let child = platform::spawn_run(command).map_err(Error::Spawn)?;
let identity = child_identity(child.id().expect("spawned worker has a PID"))
.map_err(Error::Registration)?;
Ok((identity, child))
}
pub(crate) fn launch_run(
program: &std::path::Path,
id: &crate::types::RunId,
repo: &std::path::Path,
state: &std::path::Path,
) -> Result<ChildIdentity, Error> {
let (identity, mut child) = spawn_run(program, id, repo, state, RunEntry::Supervisor)?;
tokio::spawn(async move {
let _ = child.wait().await;
});
Ok(identity)
}
pub(crate) fn child_identity(pid: u32) -> io::Result<ChildIdentity> {
let start_time = platform::process_start_time(pid)?;
Ok(ChildIdentity {
pid: ProcessId::new(pid),
start_time,
})
}
pub(crate) fn is_alive(owner: Option<ChildIdentity>) -> io::Result<bool> {
let Some(owner) = owner else {
return Ok(false);
};
match platform::running_process_start_time(owner.pid.get()) {
Ok(start_time) => Ok(start_time == owner.start_time),
Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(false),
Err(error) => Err(error),
}
}
async fn terminate(child: &mut Child) -> io::Result<()> {
let _ = child.interrupt();
tokio::select! {
_ = sleep(CLEANUP_GRACE) => {},
_ = child.wait() => {},
}
child.cleanup().await?;
Ok(())
}
pub(crate) async fn run_logged<F, R>(
command: Command,
directory: PathBuf,
cancellation: CancellationToken,
register: F,
) -> Result<Output, Error>
where
F: FnOnce(ChildIdentity) -> R + Send + 'static,
R: Future<Output = io::Result<()>> + Send + 'static,
{
let cancellation = cancellation.child_token();
let _cancel_on_drop = cancellation.clone().drop_guard();
crate::task::joined(
tokio::spawn(run_inner(
command,
None,
OUTPUT_LIMIT,
Capture {
logs: Some(directory),
lines: None,
},
cancellation,
register,
))
.await,
)?
}
async fn capture(
mut pipe: impl AsyncRead + Unpin,
limit: usize,
path: Option<PathBuf>,
mut lines: Option<Lines>,
) -> io::Result<(Vec<u8>, bool)> {
let mut log = path
.as_ref()
.map(|path| {
let mut options = platform::private_options();
options.write(true).create_new(true);
options.open(path).map(tokio::fs::File::from_std)
})
.transpose()?;
let mut output = Vec::new();
let mut buffer = [0; CAPTURE_BUFFER_BYTES];
let mut truncated = false;
loop {
let count = pipe.read(&mut buffer).await?;
if count == 0 {
if let Some(lines) = &mut lines {
lines.finish();
}
if let (Some(log), Some(path)) = (&mut log, &path) {
log.flush().await?;
let mut options = platform::private_options();
options.write(true).create_new(true);
let extension = path.extension().expect("stream extension");
options.open(
path.with_extension(format!("{}.complete", extension.to_string_lossy())),
)?;
}
return Ok((output, truncated));
}
if let Some(log) = &mut log {
log.write_all(&buffer[..count]).await?;
log.flush().await?;
}
if let Some(lines) = &mut lines {
lines.push(&buffer[..count]);
}
let retained = count.min(limit - output.len());
output.extend_from_slice(&buffer[..retained]);
truncated |= retained < count;
}
}