1use std::{
4 collections::BTreeMap,
5 ffi::OsString,
6 future::Future,
7 io,
8 path::PathBuf,
9 process::{ExitStatus, Stdio},
10 time::Duration,
11};
12
13use thiserror::Error;
14use tokio::{
15 io::AsyncWriteExt,
16 time::{Instant, sleep, sleep_until},
17};
18use tokio_util::sync::CancellationToken;
19
20use crate::{
21 platform::{self, Child},
22 tools::human::log::{End, Lines, Run, Stream},
23};
24
25pub(crate) const OUTPUT_LIMIT: usize = 128 * 1024;
27const CLEANUP_GRACE: Duration = Duration::from_secs(1);
29
30mod capture;
31pub use capture::{CaptureError, StreamFailure};
32#[cfg(test)]
33pub(crate) use capture::{OUTPUT_CLOSE_GRACE, slow_disk};
34
35#[derive(Debug, Clone)]
37pub struct Command {
38 pub program: OsString,
39 pub args: Vec<OsString>,
40 pub cwd: PathBuf,
41 pub env: BTreeMap<OsString, OsString>,
43 pub timeout: Duration,
44}
45
46#[derive(
50 Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize, serde::Deserialize,
51)]
52#[serde(transparent)]
53pub struct ProcessId(std::num::NonZeroU32);
54
55impl ProcessId {
56 pub fn new(pid: u32) -> Self {
58 Self(std::num::NonZeroU32::new(pid).expect("a real process id is never 0"))
59 }
60
61 pub fn get(self) -> u32 {
62 self.0.get()
63 }
64}
65
66impl From<std::num::NonZeroU32> for ProcessId {
67 fn from(pid: std::num::NonZeroU32) -> Self {
68 Self(pid)
69 }
70}
71
72impl std::fmt::Display for ProcessId {
73 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74 self.0.fmt(f)
75 }
76}
77
78impl rusqlite::types::ToSql for ProcessId {
79 fn to_sql(&self) -> rusqlite::Result<rusqlite::types::ToSqlOutput<'_>> {
80 Ok(rusqlite::types::ToSqlOutput::from(i64::from(self.0.get())))
81 }
82}
83
84impl rusqlite::types::FromSql for ProcessId {
85 fn column_result(value: rusqlite::types::ValueRef<'_>) -> rusqlite::types::FromSqlResult<Self> {
86 let pid = u32::try_from(value.as_i64()?)
87 .map_err(|error| rusqlite::types::FromSqlError::Other(error.into()))?;
88 std::num::NonZeroU32::new(pid).map(Self).ok_or_else(|| {
89 rusqlite::types::FromSqlError::Other(
90 "a saved owner_pid of 0 is not a process id".into(),
91 )
92 })
93 }
94}
95
96#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
98pub struct ChildIdentity {
99 pub pid: ProcessId,
100 pub start_time: u64,
103}
104
105#[derive(Debug)]
106pub struct Output {
107 pub status: ExitStatus,
108 pub stdout: Vec<u8>,
109 pub stderr: Vec<u8>,
110 pub truncated: bool,
111 pub duration: Duration,
112}
113
114pub(crate) fn tool_output<'a>(
118 program: &str,
119 args: impl IntoIterator<Item = &'a std::ffi::OsStr>,
120) -> Option<Vec<u8>> {
121 platform::program::tool_output(program, args)
122}
123
124pub(crate) struct Launcher {
128 pub env: BTreeMap<OsString, OsString>,
129}
130
131impl artifactize_tools::launch::Launcher for Launcher {
132 fn run<'a>(
133 &'a self,
134 launch: artifactize_tools::launch::Launch,
135 cancellation: &'a CancellationToken,
136 ) -> artifactize_tools::launch::Running<'a> {
137 use artifactize_tools::launch::{Finished, LaunchError};
138 Box::pin(async move {
139 let command = Command {
140 program: launch.program.into(),
141 args: launch.args,
142 cwd: launch.cwd,
143 env: self.env.clone(),
144 timeout: launch.timeout,
145 };
146 match run_unregistered(command, None, launch.output_limit, cancellation.clone()).await {
147 Ok(output) => Ok(Finished {
148 status: output.status,
149 stdout: output.stdout,
150 stderr: output.stderr,
151 truncated: output.truncated,
152 }),
153 Err(Error::Timeout) => Err(LaunchError::TimedOut),
154 Err(Error::Cancelled) => Err(LaunchError::Cancelled),
155 Err(error) => Err(LaunchError::Failed(
156 error
157 .argument_refusal()
158 .unwrap_or_else(|| error.to_string()),
159 )),
160 }
161 })
162 }
163}
164
165#[derive(Debug, Error)]
166pub enum Error {
167 #[error("process could not be spawned: {0}")]
168 Spawn(#[source] io::Error),
169 #[error("child registration failed: {0}")]
170 Registration(#[source] io::Error),
171 #[error("process supervision failed: {0}")]
172 Io(#[from] io::Error),
173 #[error("process cleanup could not be confirmed: {0}")]
174 Cleanup(#[source] io::Error),
175 #[error("process supervision task failed: {0}")]
176 Task(#[from] tokio::task::JoinError),
177 #[error("process timed out")]
178 Timeout,
179 #[error("process was cancelled")]
180 Cancelled,
181 #[error(transparent)]
184 Capture(Box<CaptureError>),
185}
186
187impl Error {
188 pub(crate) fn cleanup_confirmed(&self) -> bool {
191 match self {
192 Self::Cleanup(_) | Self::Task(_) => false,
193 Self::Spawn(_)
194 | Self::Registration(_)
195 | Self::Io(_)
196 | Self::Timeout
197 | Self::Cancelled
198 | Self::Capture(_) => true,
199 }
200 }
201
202 pub(crate) fn argument_refusal(&self) -> Option<String> {
203 match self {
204 Self::Spawn(error) if error.kind() == io::ErrorKind::InvalidInput => Some(format!(
205 "Process arguments could not be passed safely: {error}"
206 )),
207 _ => None,
208 }
209 }
210}
211
212pub async fn run<F, R>(
216 command: Command,
217 cancellation: CancellationToken,
218 register: F,
219) -> Result<Output, Error>
220where
221 F: FnOnce(ChildIdentity) -> R + Send + 'static,
222 R: Future<Output = io::Result<()>> + Send + 'static,
223{
224 run_with_input(command, None, cancellation, register).await
225}
226
227pub(crate) async fn run_with_input<F, R>(
228 command: Command,
229 input: Option<Vec<u8>>,
230 cancellation: CancellationToken,
231 register: F,
232) -> Result<Output, Error>
233where
234 F: FnOnce(ChildIdentity) -> R + Send + 'static,
235 R: Future<Output = io::Result<()>> + Send + 'static,
236{
237 run_with_input_limit(command, input, OUTPUT_LIMIT, cancellation, register).await
238}
239
240pub(crate) async fn run_with_input_limit<F, R>(
241 command: Command,
242 input: Option<Vec<u8>>,
243 output_limit: usize,
244 cancellation: CancellationToken,
245 register: F,
246) -> Result<Output, Error>
247where
248 F: FnOnce(ChildIdentity) -> R + Send + 'static,
249 R: Future<Output = io::Result<()>> + Send + 'static,
250{
251 let cancellation = cancellation.child_token();
252 let _cancel_on_drop = cancellation.clone().drop_guard();
253 crate::task::joined(
254 tokio::spawn(run_inner(
255 command,
256 input,
257 output_limit,
258 Capture::default(),
259 cancellation,
260 Some(register),
261 ))
262 .await,
263 )?
264}
265
266type Unregistered = fn(ChildIdentity) -> std::future::Ready<io::Result<()>>;
269
270pub(crate) async fn run_unregistered(
273 command: Command,
274 input: Option<Vec<u8>>,
275 output_limit: usize,
276 cancellation: CancellationToken,
277) -> Result<Output, Error> {
278 let cancellation = cancellation.child_token();
279 let _cancel_on_drop = cancellation.clone().drop_guard();
280 crate::task::joined(
281 tokio::spawn(run_inner(
282 command,
283 input,
284 output_limit,
285 Capture::default(),
286 cancellation,
287 None::<Unregistered>,
288 ))
289 .await,
290 )?
291}
292
293#[derive(Default)]
296struct Capture {
297 logs: Option<PathBuf>,
298 lines: Option<(Lines, Lines)>,
299}
300
301pub(crate) async fn run_watched(
304 command: Command,
305 cancellation: CancellationToken,
306 run: &Run,
307) -> Result<Output, Error> {
308 let cancellation = cancellation.child_token();
309 let _cancel_on_drop = cancellation.clone().drop_guard();
310 let capture = Capture {
311 logs: None,
312 lines: Some((run.lines(Stream::Stdout), run.lines(Stream::Stderr))),
313 };
314 crate::task::joined(
315 tokio::spawn(run_inner(
316 command,
317 None,
318 OUTPUT_LIMIT,
319 capture,
320 cancellation,
321 None::<Unregistered>,
322 ))
323 .await,
324 )?
325}
326
327async fn run_inner<F, R>(
328 command: Command,
329 input: Option<Vec<u8>>,
330 output_limit: usize,
331 Capture { logs, lines }: Capture,
332 cancellation: CancellationToken,
333 register: Option<F>,
334) -> Result<Output, Error>
335where
336 F: FnOnce(ChildIdentity) -> R,
337 R: Future<Output = io::Result<()>>,
338{
339 let started = Instant::now();
340 let deadline = started + command.timeout;
341 if cancellation.is_cancelled() {
342 return Err(Error::Cancelled);
343 }
344 if command.timeout.is_zero() {
345 return Err(Error::Timeout);
346 }
347
348 let mut child = spawn_owned(
349 &command,
350 input.is_some(),
351 cancellation.clone(),
352 deadline,
353 register,
354 )
355 .await?;
356
357 let stdout = child.stdout().take().expect("stdout is piped");
358 let stderr = child.stderr().take().expect("stderr is piped");
359 let stdin = child.stdin().take();
360 let gone = CancellationToken::new();
363 let feed = async {
364 if let (Some(mut stdin), Some(input)) = (stdin, input) {
365 tokio::select! {
366 _ = gone.cancelled() => {},
367 _ = cancellation.cancelled() => {},
368 _ = sleep_until(deadline) => {},
369 result = stdin.write_all(&input) => {
370 if let Err(error) = result && error.kind() != io::ErrorKind::BrokenPipe {
372 return Err(Error::Io(error));
373 }
374 }
375 }
376 }
377 Ok(())
378 };
379 let wait = async {
380 let result = tokio::select! {
381 biased;
382 _ = cancellation.cancelled() => Err(Error::Cancelled),
383 _ = sleep_until(deadline) => Err(Error::Timeout),
384 status = child.wait() => status.map_err(Error::Io),
385 };
386 let cleanup = if result.is_err() {
387 terminate(&mut child).await
388 } else {
389 child.cleanup().await.map(|_| ())
391 };
392 gone.cancel();
393 cleanup.map_err(Error::Cleanup)?;
394 result
395 };
396 let (stdout_lines, stderr_lines) = lines.unzip();
397 let evidence = |name| logs.as_ref().map(|directory| directory.join(name));
398 let streams = async {
399 tokio::join!(
400 capture::capture(
401 stdout,
402 output_limit,
403 evidence("runtime.stdout"),
404 stdout_lines,
405 &gone
406 ),
407 capture::capture(
408 stderr,
409 output_limit,
410 evidence("runtime.stderr"),
411 stderr_lines,
412 &gone
413 )
414 )
415 };
416 let (status, (stdout, stderr), input) = tokio::join!(wait, streams, feed);
417 let status = status?;
418 input?;
419 capture::finish(status, stdout, stderr, started.elapsed())
420}
421
422pub(crate) async fn spawn_owned<F, R>(
428 command: &Command,
429 input: bool,
430 cancellation: CancellationToken,
431 deadline: Instant,
432 register: Option<F>,
433) -> Result<Child, Error>
434where
435 F: FnOnce(ChildIdentity) -> R,
436 R: Future<Output = io::Result<()>>,
437{
438 let program = platform::program::resolve_with(&command.program, &command.cwd, |name| {
439 environment_value(&command.env, name).map(ToOwned::to_owned)
440 })
441 .map_err(Error::Spawn)?;
442 let mut child = tokio::process::Command::new(&program);
443 child
444 .args(&command.args)
445 .current_dir(&command.cwd)
446 .env_clear()
447 .envs(&command.env)
448 .stdin(if input { Stdio::piped() } else { Stdio::null() })
449 .stdout(Stdio::piped())
450 .stderr(Stdio::piped())
451 .kill_on_drop(true);
452 let Some(register) = register else {
453 return platform::spawn_now(child).await.map_err(Error::Spawn);
454 };
455 let (gate, mut spawning) = platform::spawn_gated(child).map_err(Error::Spawn)?;
456 let admission = tokio::select! {
457 biased;
458 _ = cancellation.cancelled() => Err(Error::Cancelled),
459 _ = sleep_until(deadline) => Err(Error::Timeout),
460 result = async {
461 let pid = gate.pid().await?;
462 let child = child_identity(pid).map_err(Error::Registration)?;
463 register(child).await.map_err(Error::Registration)?;
464 Ok::<_, Error>(())
465 } => result,
466 };
467 let admission = admission.and_then(|()| {
469 if cancellation.is_cancelled() {
470 Err(Error::Cancelled)
471 } else if Instant::now() >= deadline {
472 Err(Error::Timeout)
473 } else {
474 gate.admit().map_err(Error::Io)
475 }
476 });
477 drop(gate);
478 let spawned = crate::task::joined((&mut spawning).await)?;
479 let mut child = match spawned {
480 Ok(child) => child,
481 Err(error) => {
482 return Err(match admission {
483 Err(error @ (Error::Registration(_) | Error::Cancelled | Error::Timeout)) => error,
484 _ => Error::Spawn(error),
485 });
486 }
487 };
488 if let Err(error) = admission {
489 terminate(&mut child).await.map_err(Error::Cleanup)?;
490 return Err(error);
491 }
492
493 Ok(child)
494}
495
496fn environment_value<'a>(
497 environment: &'a BTreeMap<OsString, OsString>,
498 name: &str,
499) -> Option<&'a std::ffi::OsStr> {
500 crate::platform::environment::lookup(environment, name)
501}
502
503pub(crate) fn launch_detached(
505 program: &std::ffi::OsStr,
506 args: &[String],
507 cwd: &std::path::Path,
508) -> Result<(), Error> {
509 let program = platform::program::resolve(program, cwd).map_err(Error::Spawn)?;
510 let mut command = tokio::process::Command::new(program);
511 command
512 .args(args)
513 .current_dir(cwd)
514 .env_clear()
515 .envs(platform::environment::snapshot())
516 .stdin(Stdio::null())
517 .stdout(Stdio::null())
518 .stderr(Stdio::null())
519 .kill_on_drop(false);
520 let mut child = platform::spawn_detached(command).map_err(Error::Spawn)?;
521 tokio::spawn(async move {
522 let _ = child.wait().await;
523 });
524 Ok(())
525}
526
527pub(crate) fn launch_watched(
532 program: &std::ffi::OsStr,
533 args: &[String],
534 cwd: &std::path::Path,
535 run: Run,
536) -> Result<(), (Error, Run)> {
537 if !run.watched() {
538 return match launch_detached(program, args, cwd) {
539 Ok(()) => {
540 run.end(End::Launched);
541 Ok(())
542 }
543 Err(error) => Err((error, run)),
544 };
545 }
546 let program = match platform::program::resolve(program, cwd) {
547 Ok(program) => program,
548 Err(error) => return Err((Error::Spawn(error), run)),
549 };
550 let mut command = tokio::process::Command::new(program);
551 command
552 .args(args)
553 .current_dir(cwd)
554 .env_clear()
555 .envs(platform::environment::snapshot())
556 .stdin(Stdio::null())
557 .stdout(Stdio::piped())
558 .stderr(Stdio::piped())
559 .kill_on_drop(false);
560 let mut child = match platform::spawn_watched(command) {
561 Ok(child) => child,
562 Err(error) => return Err((Error::Spawn(error), run)),
563 };
564 let stdout = child.stdout.take().expect("stdout is piped");
565 let stderr = child.stderr.take().expect("stderr is piped");
566 tokio::spawn(async move {
567 let held = CancellationToken::new();
570 let (status, _, _) = tokio::join!(
571 child.wait(),
572 capture::capture(stdout, 0, None, Some(run.lines(Stream::Stdout)), &held),
573 capture::capture(stderr, 0, None, Some(run.lines(Stream::Stderr)), &held),
574 );
575 run.end(match status {
576 Ok(status) => End::exited(&status),
577 Err(error) => End::Failed(error.to_string()),
578 });
579 });
580 Ok(())
581}
582
583pub(crate) fn artifactize_executable() -> Result<std::path::PathBuf, String> {
585 let current = std::env::current_exe().map_err(|e| e.to_string())?;
586 for directory in current
587 .parent()
588 .into_iter()
589 .flat_map(|parent| [Some(parent), parent.parent()])
590 .flatten()
591 {
592 if let Some(candidate) = platform::program::candidates(&directory.join("artifactize"))
593 .into_iter()
594 .find(|path| path.is_file())
595 {
596 return Ok(candidate);
597 }
598 }
599 platform::program::resolve(
600 std::ffi::OsStr::new("artifactize"),
601 &std::env::current_dir().map_err(|e| e.to_string())?,
602 )
603 .map_err(|e| e.to_string())
604}
605
606#[derive(Clone, Copy)]
608pub(crate) enum RunEntry {
609 Supervisor,
610 Executor,
611}
612pub(crate) fn spawn_run(
613 program: &std::path::Path,
614 id: &crate::types::RunId,
615 repo: &std::path::Path,
616 state: &std::path::Path,
617 entry: RunEntry,
618) -> Result<(ChildIdentity, tokio::process::Child), Error> {
619 let mut options = platform::private_options();
620 options.write(true);
621 match entry {
622 RunEntry::Supervisor => {
623 options.create_new(true);
624 }
625 RunEntry::Executor => {
626 options.append(true);
627 }
628 }
629 let directory = state.join("runs").join(id.as_str());
630 let stderr = options
631 .open(directory.join("worker.stderr"))
632 .map_err(Error::Spawn)?;
633 let program = platform::program::resolve(program.as_os_str(), repo).map_err(Error::Spawn)?;
634 let mut command = tokio::process::Command::new(program);
635 command
636 .arg("--repo")
637 .arg(repo)
638 .arg("--state-dir")
639 .arg(state)
640 .arg(match entry {
641 RunEntry::Supervisor => "__verify",
642 RunEntry::Executor => "__execute",
643 })
644 .arg(id.as_str())
645 .current_dir(directory)
648 .env_clear()
649 .envs(platform::environment::snapshot())
650 .stdin(Stdio::null())
651 .stdout(Stdio::null())
652 .stderr(Stdio::from(stderr))
653 .kill_on_drop(false);
654 let child = platform::spawn_run(command).map_err(Error::Spawn)?;
655 let identity = child_identity(child.id().expect("spawned worker has a PID"))
656 .map_err(Error::Registration)?;
657 Ok((identity, child))
658}
659pub(crate) fn launch_run(
660 program: &std::path::Path,
661 id: &crate::types::RunId,
662 repo: &std::path::Path,
663 state: &std::path::Path,
664) -> Result<ChildIdentity, Error> {
665 let (identity, mut child) = spawn_run(program, id, repo, state, RunEntry::Supervisor)?;
666 tokio::spawn(async move {
667 let _ = child.wait().await;
668 });
669 Ok(identity)
670}
671
672pub(crate) fn child_identity(pid: u32) -> io::Result<ChildIdentity> {
673 let start_time = platform::process_start_time(pid)?;
674 Ok(ChildIdentity {
675 pid: ProcessId::new(pid),
676 start_time,
677 })
678}
679
680pub(crate) fn is_alive(owner: Option<ChildIdentity>) -> io::Result<bool> {
683 let Some(owner) = owner else {
684 return Ok(false);
685 };
686 match platform::running_process_start_time(owner.pid.get()) {
687 Ok(start_time) => Ok(start_time == owner.start_time),
688 Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(false),
689 Err(error) => Err(error),
690 }
691}
692
693async fn terminate(child: &mut Child) -> io::Result<()> {
694 let _ = child.interrupt();
695 tokio::select! {
696 _ = sleep(CLEANUP_GRACE) => {},
697 _ = child.wait() => {},
698 }
699 child.cleanup().await?;
700 Ok(())
701}
702
703pub(crate) async fn run_logged<F, R>(
706 command: Command,
707 directory: PathBuf,
708 cancellation: CancellationToken,
709 register: F,
710) -> Result<Output, Error>
711where
712 F: FnOnce(ChildIdentity) -> R + Send + 'static,
713 R: Future<Output = io::Result<()>> + Send + 'static,
714{
715 let cancellation = cancellation.child_token();
716 let _cancel_on_drop = cancellation.clone().drop_guard();
717 crate::task::joined(
718 tokio::spawn(run_inner(
719 command,
720 None,
721 OUTPUT_LIMIT,
722 Capture {
723 logs: Some(directory),
724 lines: None,
725 },
726 cancellation,
727 Some(register),
728 ))
729 .await,
730 )?
731}