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
25const 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 let ready = |_| async { Ok(()) };
144 match run_with_input_limit(
145 command,
146 None,
147 launch.output_limit,
148 cancellation.clone(),
149 ready,
150 )
151 .await
152 {
153 Ok(output) => Ok(Finished {
154 status: output.status,
155 stdout: output.stdout,
156 stderr: output.stderr,
157 truncated: output.truncated,
158 }),
159 Err(Error::Timeout) => Err(LaunchError::TimedOut),
160 Err(Error::Cancelled) => Err(LaunchError::Cancelled),
161 Err(error) => Err(LaunchError::Failed(
162 error
163 .argument_refusal()
164 .unwrap_or_else(|| error.to_string()),
165 )),
166 }
167 })
168 }
169}
170
171#[derive(Debug, Error)]
172pub enum Error {
173 #[error("process could not be spawned: {0}")]
174 Spawn(#[source] io::Error),
175 #[error("child registration failed: {0}")]
176 Registration(#[source] io::Error),
177 #[error("process supervision failed: {0}")]
178 Io(#[from] io::Error),
179 #[error("process cleanup could not be confirmed: {0}")]
180 Cleanup(#[source] io::Error),
181 #[error("process supervision task failed: {0}")]
182 Task(#[from] tokio::task::JoinError),
183 #[error("process timed out")]
184 Timeout,
185 #[error("process was cancelled")]
186 Cancelled,
187 #[error("process output pipes did not close after cleanup")]
188 OutputDidNotClose,
189}
190
191impl Error {
192 pub(crate) fn cleanup_confirmed(&self) -> bool {
195 match self {
196 Self::Cleanup(_) | Self::Task(_) => false,
197 Self::Spawn(_)
198 | Self::Registration(_)
199 | Self::Io(_)
200 | Self::Timeout
201 | Self::Cancelled
202 | Self::OutputDidNotClose => true,
203 }
204 }
205
206 pub(crate) fn argument_refusal(&self) -> Option<String> {
207 match self {
208 Self::Spawn(error) if error.kind() == io::ErrorKind::InvalidInput => Some(format!(
209 "Process arguments could not be passed safely: {error}"
210 )),
211 _ => None,
212 }
213 }
214}
215
216pub async fn run<F, R>(
220 command: Command,
221 cancellation: CancellationToken,
222 register: F,
223) -> Result<Output, Error>
224where
225 F: FnOnce(ChildIdentity) -> R + Send + 'static,
226 R: Future<Output = io::Result<()>> + Send + 'static,
227{
228 run_with_input(command, None, cancellation, register).await
229}
230
231pub(crate) async fn run_with_input<F, R>(
232 command: Command,
233 input: Option<Vec<u8>>,
234 cancellation: CancellationToken,
235 register: F,
236) -> Result<Output, Error>
237where
238 F: FnOnce(ChildIdentity) -> R + Send + 'static,
239 R: Future<Output = io::Result<()>> + Send + 'static,
240{
241 run_with_input_limit(command, input, OUTPUT_LIMIT, cancellation, register).await
242}
243
244pub(crate) async fn run_with_input_limit<F, R>(
245 command: Command,
246 input: Option<Vec<u8>>,
247 output_limit: usize,
248 cancellation: CancellationToken,
249 register: F,
250) -> Result<Output, Error>
251where
252 F: FnOnce(ChildIdentity) -> R + Send + 'static,
253 R: Future<Output = io::Result<()>> + Send + 'static,
254{
255 let cancellation = cancellation.child_token();
256 let _cancel_on_drop = cancellation.clone().drop_guard();
257 crate::task::joined(
258 tokio::spawn(run_inner(
259 command,
260 input,
261 output_limit,
262 Capture::default(),
263 cancellation,
264 register,
265 ))
266 .await,
267 )?
268}
269
270#[derive(Default)]
273struct Capture {
274 logs: Option<PathBuf>,
275 lines: Option<(Lines, Lines)>,
276}
277
278pub(crate) async fn run_watched(
281 command: Command,
282 cancellation: CancellationToken,
283 run: &Run,
284) -> Result<Output, Error> {
285 let cancellation = cancellation.child_token();
286 let _cancel_on_drop = cancellation.clone().drop_guard();
287 let capture = Capture {
288 logs: None,
289 lines: Some((run.lines(Stream::Stdout), run.lines(Stream::Stderr))),
290 };
291 crate::task::joined(
292 tokio::spawn(run_inner(
293 command,
294 None,
295 OUTPUT_LIMIT,
296 capture,
297 cancellation,
298 |_| async { Ok(()) },
299 ))
300 .await,
301 )?
302}
303
304async fn run_inner<F, R>(
305 command: Command,
306 input: Option<Vec<u8>>,
307 output_limit: usize,
308 Capture { logs, lines }: Capture,
309 cancellation: CancellationToken,
310 register: F,
311) -> Result<Output, Error>
312where
313 F: FnOnce(ChildIdentity) -> R,
314 R: Future<Output = io::Result<()>>,
315{
316 let started = Instant::now();
317 let deadline = started + command.timeout;
318 if cancellation.is_cancelled() {
319 return Err(Error::Cancelled);
320 }
321 if command.timeout.is_zero() {
322 return Err(Error::Timeout);
323 }
324
325 let mut child = spawn_owned(
326 &command,
327 input.is_some(),
328 cancellation.clone(),
329 deadline,
330 register,
331 )
332 .await?;
333
334 let stdout = child.stdout().take().expect("stdout is piped");
335 let stderr = child.stderr().take().expect("stderr is piped");
336 let stdin = child.stdin().take();
337 let reaped = CancellationToken::new();
338 let feed = async {
339 if let (Some(mut stdin), Some(input)) = (stdin, input) {
340 tokio::select! {
341 _ = reaped.cancelled() => {},
342 _ = cancellation.cancelled() => {},
343 _ = sleep_until(deadline) => {},
344 result = stdin.write_all(&input) => {
345 if let Err(error) = result && error.kind() != io::ErrorKind::BrokenPipe {
347 return Err(Error::Io(error));
348 }
349 }
350 }
351 }
352 Ok(())
353 };
354 let wait = async {
355 let result = tokio::select! {
356 biased;
357 _ = cancellation.cancelled() => Err(Error::Cancelled),
358 _ = sleep_until(deadline) => Err(Error::Timeout),
359 status = child.wait() => status.map_err(Error::Io),
360 };
361 let cleanup = if result.is_err() {
362 terminate(&mut child).await
363 } else {
364 child.cleanup().await.map(|_| ())
366 };
367 reaped.cancel();
368 cleanup.map_err(Error::Cleanup)?;
369 result
370 };
371 let (stdout_lines, stderr_lines) = lines.unzip();
372 let capture = async {
373 let read = async {
374 tokio::try_join!(
375 capture(
376 stdout,
377 output_limit,
378 logs.as_ref().map(|dir| dir.join("runtime.stdout")),
379 stdout_lines,
380 ),
381 capture(
382 stderr,
383 output_limit,
384 logs.as_ref().map(|dir| dir.join("runtime.stderr")),
385 stderr_lines,
386 )
387 )
388 };
389 tokio::pin!(read);
390 tokio::select! {
391 result = &mut read => result.map_err(Error::Io),
392 _ = reaped.cancelled() => {
393 timeout(CLEANUP_GRACE, read).await
394 .map_err(|_| Error::OutputDidNotClose)?
395 .map_err(Error::Io)
396 }
397 }
398 };
399 let (status, output, input) = tokio::join!(wait, capture, feed);
400 let status = status?;
401 input?;
402 let ((stdout, stdout_truncated), (stderr, stderr_truncated)) = output?;
403 Ok(Output {
404 status,
405 stdout,
406 stderr,
407 truncated: stdout_truncated || stderr_truncated,
408 duration: started.elapsed(),
409 })
410}
411
412pub(crate) async fn spawn_owned<F, R>(
415 command: &Command,
416 input: bool,
417 cancellation: CancellationToken,
418 deadline: Instant,
419 register: F,
420) -> Result<Child, Error>
421where
422 F: FnOnce(ChildIdentity) -> R,
423 R: Future<Output = io::Result<()>>,
424{
425 let program = platform::program::resolve_with(&command.program, &command.cwd, |name| {
426 environment_value(&command.env, name).map(ToOwned::to_owned)
427 })
428 .map_err(Error::Spawn)?;
429 let mut child = tokio::process::Command::new(&program);
430 child
431 .args(&command.args)
432 .current_dir(&command.cwd)
433 .env_clear()
434 .envs(&command.env)
435 .stdin(if input { Stdio::piped() } else { Stdio::null() })
436 .stdout(Stdio::piped())
437 .stderr(Stdio::piped())
438 .kill_on_drop(true);
439 let (gate, mut spawning) = platform::spawn_gated(child).map_err(Error::Spawn)?;
440 let admission = tokio::select! {
441 biased;
442 _ = cancellation.cancelled() => Err(Error::Cancelled),
443 _ = sleep_until(deadline) => Err(Error::Timeout),
444 result = async {
445 let pid = gate.pid().await?;
446 let child = child_identity(pid).map_err(Error::Registration)?;
447 register(child).await.map_err(Error::Registration)?;
448 Ok::<_, Error>(())
449 } => result,
450 };
451 let admission = admission.and_then(|()| {
453 if cancellation.is_cancelled() {
454 Err(Error::Cancelled)
455 } else if Instant::now() >= deadline {
456 Err(Error::Timeout)
457 } else {
458 gate.admit().map_err(Error::Io)
459 }
460 });
461 drop(gate);
462 let spawned = crate::task::joined((&mut spawning).await)?;
463 let mut child = match spawned {
464 Ok(child) => child,
465 Err(error) => {
466 return Err(match admission {
467 Err(error @ (Error::Registration(_) | Error::Cancelled | Error::Timeout)) => error,
468 _ => Error::Spawn(error),
469 });
470 }
471 };
472 if let Err(error) = admission {
473 terminate(&mut child).await.map_err(Error::Cleanup)?;
474 return Err(error);
475 }
476
477 Ok(child)
478}
479
480fn environment_value<'a>(
481 environment: &'a BTreeMap<OsString, OsString>,
482 name: &str,
483) -> Option<&'a std::ffi::OsStr> {
484 crate::platform::environment::lookup(environment, name)
485}
486
487pub(crate) fn launch_detached(
489 program: &std::ffi::OsStr,
490 args: &[String],
491 cwd: &std::path::Path,
492) -> Result<(), Error> {
493 let program = platform::program::resolve(program, cwd).map_err(Error::Spawn)?;
494 let mut command = tokio::process::Command::new(program);
495 command
496 .args(args)
497 .current_dir(cwd)
498 .stdin(Stdio::null())
499 .stdout(Stdio::null())
500 .stderr(Stdio::null())
501 .kill_on_drop(false);
502 let mut child = platform::spawn_detached(command).map_err(Error::Spawn)?;
503 tokio::spawn(async move {
504 let _ = child.wait().await;
505 });
506 Ok(())
507}
508
509pub(crate) fn launch_watched(
514 program: &std::ffi::OsStr,
515 args: &[String],
516 cwd: &std::path::Path,
517 run: Run,
518) -> Result<(), (Error, Run)> {
519 if !run.watched() {
520 return match launch_detached(program, args, cwd) {
521 Ok(()) => {
522 run.end(End::Launched);
523 Ok(())
524 }
525 Err(error) => Err((error, run)),
526 };
527 }
528 let program = match platform::program::resolve(program, cwd) {
529 Ok(program) => program,
530 Err(error) => return Err((Error::Spawn(error), run)),
531 };
532 let mut command = tokio::process::Command::new(program);
533 command
534 .args(args)
535 .current_dir(cwd)
536 .stdin(Stdio::null())
537 .stdout(Stdio::piped())
538 .stderr(Stdio::piped())
539 .kill_on_drop(false);
540 let mut child = match platform::spawn_watched(command) {
541 Ok(child) => child,
542 Err(error) => return Err((Error::Spawn(error), run)),
543 };
544 let stdout = child.stdout.take().expect("stdout is piped");
545 let stderr = child.stderr.take().expect("stderr is piped");
546 tokio::spawn(async move {
547 let (status, _, _) = tokio::join!(
549 child.wait(),
550 capture(stdout, 0, None, Some(run.lines(Stream::Stdout))),
551 capture(stderr, 0, None, Some(run.lines(Stream::Stderr))),
552 );
553 run.end(match status {
554 Ok(status) => End::exited(&status),
555 Err(error) => End::Failed(error.to_string()),
556 });
557 });
558 Ok(())
559}
560
561pub(crate) fn artifactize_executable() -> Result<std::path::PathBuf, String> {
563 let current = std::env::current_exe().map_err(|e| e.to_string())?;
564 for directory in current
565 .parent()
566 .into_iter()
567 .flat_map(|parent| [Some(parent), parent.parent()])
568 .flatten()
569 {
570 if let Some(candidate) = platform::program::candidates(&directory.join("artifactize"))
571 .into_iter()
572 .find(|path| path.is_file())
573 {
574 return Ok(candidate);
575 }
576 }
577 platform::program::resolve(
578 std::ffi::OsStr::new("artifactize"),
579 &std::env::current_dir().map_err(|e| e.to_string())?,
580 )
581 .map_err(|e| e.to_string())
582}
583
584#[derive(Clone, Copy)]
586pub(crate) enum RunEntry {
587 Supervisor,
588 Executor,
589}
590pub(crate) fn spawn_run(
591 program: &std::path::Path,
592 id: &crate::types::RunId,
593 repo: &std::path::Path,
594 state: &std::path::Path,
595 entry: RunEntry,
596) -> Result<(ChildIdentity, tokio::process::Child), Error> {
597 let mut options = platform::private_options();
598 options.write(true);
599 match entry {
600 RunEntry::Supervisor => {
601 options.create_new(true);
602 }
603 RunEntry::Executor => {
604 options.append(true);
605 }
606 }
607 let directory = state.join("runs").join(id.as_str());
608 let stderr = options
609 .open(directory.join("worker.stderr"))
610 .map_err(Error::Spawn)?;
611 let program = platform::program::resolve(program.as_os_str(), repo).map_err(Error::Spawn)?;
612 let mut command = tokio::process::Command::new(program);
613 command
614 .arg("--repo")
615 .arg(repo)
616 .arg("--state-dir")
617 .arg(state)
618 .arg(match entry {
619 RunEntry::Supervisor => "__verify",
620 RunEntry::Executor => "__execute",
621 })
622 .arg(id.as_str())
623 .current_dir(directory)
626 .stdin(Stdio::null())
627 .stdout(Stdio::null())
628 .stderr(Stdio::from(stderr))
629 .kill_on_drop(false);
630 let child = platform::spawn_run(command).map_err(Error::Spawn)?;
631 let identity = child_identity(child.id().expect("spawned worker has a PID"))
632 .map_err(Error::Registration)?;
633 Ok((identity, child))
634}
635pub(crate) fn launch_run(
636 program: &std::path::Path,
637 id: &crate::types::RunId,
638 repo: &std::path::Path,
639 state: &std::path::Path,
640) -> Result<ChildIdentity, Error> {
641 let (identity, mut child) = spawn_run(program, id, repo, state, RunEntry::Supervisor)?;
642 tokio::spawn(async move {
643 let _ = child.wait().await;
644 });
645 Ok(identity)
646}
647
648pub(crate) fn child_identity(pid: u32) -> io::Result<ChildIdentity> {
649 let start_time = platform::process_start_time(pid)?;
650 Ok(ChildIdentity {
651 pid: ProcessId::new(pid),
652 start_time,
653 })
654}
655
656pub(crate) fn is_alive(owner: Option<ChildIdentity>) -> io::Result<bool> {
659 let Some(owner) = owner else {
660 return Ok(false);
661 };
662 match platform::running_process_start_time(owner.pid.get()) {
663 Ok(start_time) => Ok(start_time == owner.start_time),
664 Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(false),
665 Err(error) => Err(error),
666 }
667}
668
669async fn terminate(child: &mut Child) -> io::Result<()> {
670 let _ = child.interrupt();
671 tokio::select! {
672 _ = sleep(CLEANUP_GRACE) => {},
673 _ = child.wait() => {},
674 }
675 child.cleanup().await?;
676 Ok(())
677}
678
679pub(crate) async fn run_logged<F, R>(
682 command: Command,
683 directory: PathBuf,
684 cancellation: CancellationToken,
685 register: F,
686) -> Result<Output, Error>
687where
688 F: FnOnce(ChildIdentity) -> R + Send + 'static,
689 R: Future<Output = io::Result<()>> + Send + 'static,
690{
691 let cancellation = cancellation.child_token();
692 let _cancel_on_drop = cancellation.clone().drop_guard();
693 crate::task::joined(
694 tokio::spawn(run_inner(
695 command,
696 None,
697 OUTPUT_LIMIT,
698 Capture {
699 logs: Some(directory),
700 lines: None,
701 },
702 cancellation,
703 register,
704 ))
705 .await,
706 )?
707}
708
709async fn capture(
710 mut pipe: impl AsyncRead + Unpin,
711 limit: usize,
712 path: Option<PathBuf>,
713 mut lines: Option<Lines>,
714) -> io::Result<(Vec<u8>, bool)> {
715 let mut log = path
716 .as_ref()
717 .map(|path| {
718 let mut options = platform::private_options();
719 options.write(true).create_new(true);
720 options.open(path).map(tokio::fs::File::from_std)
721 })
722 .transpose()?;
723 let mut output = Vec::new();
724 let mut buffer = [0; CAPTURE_BUFFER_BYTES];
725 let mut truncated = false;
726 loop {
727 let count = pipe.read(&mut buffer).await?;
728 if count == 0 {
729 if let Some(lines) = &mut lines {
730 lines.finish();
731 }
732 if let (Some(log), Some(path)) = (&mut log, &path) {
733 log.flush().await?;
734 let mut options = platform::private_options();
735 options.write(true).create_new(true);
736 let extension = path.extension().expect("stream extension");
737 options.open(
738 path.with_extension(format!("{}.complete", extension.to_string_lossy())),
739 )?;
740 }
741 return Ok((output, truncated));
742 }
743 if let Some(log) = &mut log {
744 log.write_all(&buffer[..count]).await?;
745 log.flush().await?;
747 }
748 if let Some(lines) = &mut lines {
749 lines.push(&buffer[..count]);
750 }
751 let retained = count.min(limit - output.len());
752 output.extend_from_slice(&buffer[..retained]);
753 truncated |= retained < count;
754 }
755}