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