[antlir2][image_test] log booted test stdout/err as tpx artifact

Summary:
Instead of simply copying this when the test finishes (successfully or failed),
also upload it as a TPX artifact so that it's always visible in the UI on any
kind of test failure.

Test Plan:
```name="Make a test that times out while printing stuff"
❯ hg diff
 diff --git a/fbcode/antlir/antlir2/testing/tests/test.py b/fbcode/antlir/antlir2/testing/tests/test.py
 --- a/fbcode/antlir/antlir2/testing/tests/test.py
+++ b/fbcode/antlir/antlir2/testing/tests/test.py
@@ -8,6 +8,8 @@
 import os
 import pwd
 import stat
+import sys
+import time
 import unittest

@@ -33,3 +35,9 @@
             os.makedev(1, 3),
             "/dev/null device number is wrong",
         )
+
+    def test_timeout(self) -> None:
+        while True:
+            print("running...")
+            print("running...", file=sys.stderr)
+            time.sleep(1)

```

```
❯ buck2 test fbcode//antlir/antlir2/testing/tests:test-py-boot-centos9 -- timeout --timeout=5
✉ Timeout: fbcode//antlir/antlir2/testing/tests:test-py-boot-centos9 - test_timeout (antlir.antlir2.testing.tests.test.Test) (5.0s)
Your test timed out because its execution time exceeded 5 seconds. Profile your test to understand what makes it slow
stdout: from_execution_error::timeout::Buck2TpxTimeout
stderr:
Buck UI: https://www.internalfb.com/buck2/79de5265-756f-4e6e-b737-ec829da3cb40
Test UI: https://www.internalfb.com/intern/testinfra/testrun/5910974838205890
Network: Up: 0B  Down: 0B
Analyzing targets. Remaining     0/153
Executing actions. Remaining     0/364
Command: test.
Time elapsed: 7.8s
Test execution completed but the tests failed
Tests finished: Pass 0. Fail 1. Fatal 0. Skip 0. Infra Failure 0. Build failure 0
1 TESTS FAILED
  ✗ fbcode//antlir/antlir2/testing/tests:test-py-boot-centos9 - test_timeout (antlir.antlir2.testing.tests.test.Test)
```
But it's visible in the UI:
https://pxl.cl/8wvxq

Reviewed By: justintrudell

Differential Revision: D88067169

fbshipit-source-id: cade40784d38c285e542a50c9d79766ca5ba9ad8

Author: vmagro

antlir/antlir2/testing/image_test/BUCK

@@ -9,14 +9,19 @@ rust_library(
     srcs = ["src/lib.rs"],
     visibility = ["PUBLIC"],
     deps = [
+        "anyhow",
         "clap",
+        "tempfile",
         "thiserror",
     ],
 )
 
 rust_binary(
     name = "image-test",
-    srcs = glob(["src/**/*.rs"]),
+    srcs = glob(
+        ["src/**/*.rs"],
+        exclude = ["src/lib.rs"],
+    ),
     crate_root = "src/main.rs",
     resources = {
         "antlir2_image_test_shell.conf": "antlir2_image_test_shell.conf",

antlir/antlir2/testing/image_test/src/lib.rs

@@ -8,11 +8,18 @@
 use std::collections::HashSet;
 use std::ffi::OsStr;
 use std::ffi::OsString;
+use std::fs::File;
+use std::fs::OpenOptions;
+use std::os::fd::AsRawFd;
+use std::os::fd::FromRawFd;
 use std::path::Path;
 use std::path::PathBuf;
 use std::str::FromStr;
 
+use anyhow::Context;
+use anyhow::Result;
 use clap::Parser;
+use tempfile::NamedTempFile;
 use thiserror::Error;
 
 #[derive(Parser, Clone, Debug)]
@@ -225,6 +232,112 @@ impl KvPair {
     }
 }
 
+/// A file that tpx will upload as an artifact on failing test instances. If not
+/// run under tpx, this will be some other fd (a regular file in /tmp, stderr, etc)
+pub struct TpxArtifact {
+    file: LogFile,
+    path: PathBuf,
+}
+
+enum LogFile {
+    Stderr,
+    File(File),
+    Tmp(NamedTempFile),
+}
+
+impl TpxArtifact {
+    /// Create a file to record additional text logs into. When invoked under
+    /// tpx, this will be uploaded as an artifact. The artifact metadata is set
+    /// up before running the test so that it still gets uploaded even in case
+    /// of a timeout.
+    /// If not running under tpx, this will be sent to stderr
+    fn new_tpx_or_none(name: &str) -> Result<Option<Self>> {
+        // if tpx has provided this artifacts dir, put the logs there so they get
+        // uploaded along with the test results
+        if let Some(artifacts_dir) = std::env::var_os("TEST_RESULT_ARTIFACTS_DIR") {
+            std::fs::create_dir_all(&artifacts_dir).with_context(|| {
+                format!("while creating artifacts dir {}", artifacts_dir.display())
+            })?;
+            let dst = Path::new(&artifacts_dir).join(name);
+            if let Some(annotations_dir) = std::env::var_os("TEST_RESULT_ARTIFACT_ANNOTATIONS_DIR")
+            {
+                std::fs::create_dir_all(&annotations_dir)?;
+                std::fs::write(
+                    Path::new(&annotations_dir).join(format!("{name}.annotation")),
+                    r#"{"type": {"generic_text_log": {}}, "description": "test logs"}"#,
+                )?;
+            }
+            let file = OpenOptions::new()
+                .create(true)
+                .truncate(true)
+                .read(true)
+                .write(true)
+                .open(&dst)
+                .with_context(|| format!("while creating {}", dst.display()))?;
+            Ok(Some(Self {
+                file: LogFile::File(file),
+                path: dst,
+            }))
+        } else {
+            Ok(None)
+        }
+    }
+
+    /// Create a file to record additional text logs into. When invoked under
+    /// tpx, this will be uploaded as an artifact. The artifact metadata is set
+    /// up before running the test so that it still gets uploaded even in case
+    /// of a timeout.
+    /// If not running under tpx, this will be sent to a temporary file.
+    pub fn new_log_file(name: &str) -> Result<Self> {
+        match Self::new_tpx_or_none(name)? {
+            Some(s) => Ok(s),
+            None => {
+                let tmpfile = tempfile::NamedTempFile::new()?;
+                Ok(Self {
+                    path: tmpfile.path().to_owned(),
+                    file: LogFile::Tmp(tmpfile),
+                })
+            }
+        }
+    }
+
+    /// Same as [TpxArtifact::new_log_file], but if not running under tpx, this
+    /// will be sent to stderr
+    pub fn new_log_file_or_stderr(name: &str) -> Result<Self> {
+        match Self::new_tpx_or_none(name)? {
+            Some(s) => Ok(s),
+            None => Ok(Self {
+                file: LogFile::Stderr,
+                path: "/dev/stderr".into(),
+            }),
+        }
+    }
+
+    pub fn as_file(&self) -> std::io::Result<File> {
+        match &self.file {
+            LogFile::Stderr => Ok(unsafe { File::from_raw_fd(std::io::stderr().as_raw_fd()) }),
+            LogFile::File(f) => f.try_clone(),
+            LogFile::Tmp(f) => f.as_file().try_clone(),
+        }
+    }
+
+    pub fn into_file(self) -> File {
+        match self.file {
+            LogFile::Stderr => unsafe { File::from_raw_fd(std::io::stderr().as_raw_fd()) },
+            LogFile::File(f) => f,
+            LogFile::Tmp(f) => f.into_file(),
+        }
+    }
+
+    pub fn is_stderr(&self) -> bool {
+        matches!(self.file, LogFile::Stderr)
+    }
+
+    pub fn path(&self) -> &Path {
+        &self.path
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::env;

antlir/antlir2/testing/image_test/src/spawn_common.rs

@@ -7,11 +7,9 @@
 
 use std::collections::BTreeMap;
 use std::collections::HashMap;
-use std::fs::File;
 use std::fs::Permissions;
+use std::io::Seek;
 use std::io::Write;
-use std::os::fd::AsRawFd;
-use std::os::fd::FromRawFd;
 use std::os::unix::fs::PermissionsExt;
 use std::os::unix::process::CommandExt;
 use std::path::Path;
@@ -27,17 +25,14 @@ use anyhow::Result;
 use anyhow::ensure;
 use bon::builder;
 use image_test_lib::Test;
+use image_test_lib::TpxArtifact;
 use tempfile::NamedTempFile;
 use tracing::debug;
 use tracing::trace;
 
 use crate::exec;
 use crate::runtime;
 
-fn make_log_files(_base: &str) -> Result<(NamedTempFile, NamedTempFile)> {
-    Ok((NamedTempFile::new()?, NamedTempFile::new()?))
-}
-
 #[builder]
 pub(crate) fn run(
     spec: runtime::Spec,
@@ -143,8 +138,9 @@ pub(crate) fn run(
 
     match spec.boot {
         Some(boot) => {
-            let container_stdout = container_stdout_file()?;
-            let (mut test_stdout, mut test_stderr) = make_log_files("test")?;
+            let container_stdout = TpxArtifact::new_log_file_or_stderr("container-stdout.txt")?;
+            let test_stdout = TpxArtifact::new_log_file("test-stdout.txt")?;
+            let test_stderr = TpxArtifact::new_log_file("test-stderr.txt")?;
 
             let mut test_unit_dropin = NamedTempFile::new()?;
             writeln!(test_unit_dropin, "[Unit]")?;
@@ -255,18 +251,22 @@ pub(crate) fn run(
                 .context("while spawning systemd-nspawn")?;
             let res = child.wait().context("while waiting for systemd-nspawn")?;
 
+            let mut test_stdout = test_stdout.into_file();
+            let mut test_stderr = test_stderr.into_file();
+            test_stdout.rewind()?;
+            test_stderr.rewind()?;
             std::io::copy(&mut test_stdout, &mut std::io::stdout())?;
             std::io::copy(&mut test_stderr, &mut std::io::stderr())?;
 
             if !res.success() {
                 // if the container stdout is not already being dumped to
                 // stdout/err, then print out the path where it can be found
-                if let ContainerConsoleOutput::File { path, .. } = &container_stdout {
+                if !container_stdout.is_stderr() {
                     eprintln!(
                         "full container console output can be found at: '{}'",
-                        path.display()
+                        container_stdout.path().display()
                     );
-                    eprintln!("{}", path.display());
+                    eprintln!("{}", container_stdout.path().display());
                 }
                 std::process::exit(res.code().unwrap_or(255))
             } else {
@@ -289,42 +289,3 @@ pub(crate) fn run(
         }
     }
 }
-
-enum ContainerConsoleOutput {
-    File { path: PathBuf, file: File },
-    Stderr,
-}
-
-impl ContainerConsoleOutput {
-    fn as_file(&self) -> Result<File> {
-        match self {
-            Self::File { file, .. } => file.try_clone().context("while cloning file fd"),
-            Self::Stderr => Ok(unsafe { File::from_raw_fd(std::io::stderr().as_raw_fd()) }),
-        }
-    }
-}
-
-/// Create a file to record container stdout into. When invoked under tpx, this
-/// will be uploaded as an artifact. The artifact metadata is set up before
-/// running the test so that it still gets uploaded even in case of a timeout
-fn container_stdout_file() -> Result<ContainerConsoleOutput> {
-    // if tpx has provided this artifacts dir, put the logs there so they get
-    // uploaded along with the test results
-    if let Some(artifacts_dir) = std::env::var_os("TEST_RESULT_ARTIFACTS_DIR") {
-        std::fs::create_dir_all(&artifacts_dir)?;
-        let dst = Path::new(&artifacts_dir).join("container-stdout.txt");
-        if let Some(annotations_dir) = std::env::var_os("TEST_RESULT_ARTIFACT_ANNOTATIONS_DIR") {
-            std::fs::create_dir_all(&annotations_dir)?;
-            std::fs::write(
-                Path::new(&annotations_dir).join("container-stdout.txt.annotation"),
-                r#"{"type": {"generic_text_log": {}}, "description": "systemd logs"}"#,
-            )?;
-        }
-        File::create(&dst)
-            .with_context(|| format!("while creating {}", dst.display()))
-            .map(|file| ContainerConsoleOutput::File { path: dst, file })
-    } else {
-        // otherwise, have it go right to stderr
-        Ok(ContainerConsoleOutput::Stderr)
-    }
-}