[antlir2][vmtest] upload env values

Summary:
Test envs are only evaluated when running `buck test`. To aid
reproducing test environment with interactive shell, dump the envs for
debugging purpose and upload to tpx.

Test Plan:
Run test and see env uploaded.
```
$ buck test antlir/antlir2/antlir2_vm/tests:rust-test
File changed: fbcode//antlir/antlir2/antlir2_vm/src/main.rs
Buck UI: https://www.internalfb.com/buck2/2a847842-ee86-41a9-a571-eaeee9effea7
Test UI: https://www.internalfb.com/intern/testinfra/testrun/10414574143836274
Network: Up: 7.2KiB  Down: 0B  (reSessionID-3b689d95-be91-4bbc-bf8d-565e3f5469f9)
Jobs completed: 12. Time elapsed: 55.1s.
Cache hits: 0%. Commands: 1 (cached: 0, remote: 0, local: 1)
Tests finished: Pass 3. Fail 0. Fatal 0. Skip 0. Build failure 0
```
The `Logs` tab on the left now has env.txt:
https://www.internalfb.com/intern/testinfra/diagnostics/10414574143836274.281475080786032.1696628348/

Reviewed By: epilatow

Differential Revision: D50039958

fbshipit-source-id: 494441db526ac4444584cb72de076b15ab92e9aa

antlir/antlir2/antlir2_vm/BUCK

@@ -21,6 +21,7 @@ rust_binary(
         "derive_builder",
         "once_cell",
         "serde",
+        "serde_json",
         "tempfile",
         "thiserror",
         "tracing",

antlir/antlir2/antlir2_vm/src/main.rs

@@ -27,6 +27,7 @@ use anyhow::Context;
 use clap::Args;
 use clap::Parser;
 use clap::Subcommand;
+use image_test_lib::KvPair;
 use image_test_lib::Test;
 use json_arg::JsonFile;
 use tempfile::tempdir;
@@ -41,7 +42,7 @@ use crate::runtime::set_runtime;
 use crate::types::MachineOpts;
 use crate::types::RuntimeOpts;
 use crate::types::VMArgs;
-use crate::utils::console_output_path_for_tpx;
+use crate::utils::create_tpx_logs;
 use crate::utils::env_names_to_kvpairs;
 use crate::utils::log_command;
 use crate::vm::VM;
@@ -147,6 +148,28 @@ struct ValidatedVMArgs {
     is_list: bool,
 }
 
+/// Record and upload envs for debugging purpose
+#[cfg(not(test))]
+fn record_envs(envs: &[KvPair]) -> Result<()> {
+    let env_file = create_tpx_logs("env", "env vars")?;
+    if let Some(file) = env_file {
+        std::fs::write(
+            file,
+            envs.iter()
+                .map(|s| s.to_os_string().to_string_lossy().to_string())
+                .collect::<Vec<_>>()
+                .join("\n"),
+        )?;
+    }
+    Ok(())
+}
+
+/// We only want envs for actual VM test, not unit tests here.
+#[cfg(test)]
+fn record_envs(_envs: &[KvPair]) -> Result<()> {
+    Ok(())
+}
+
 /// Further validate `VMArgs` parsed by clap and generate a new `VMArgs` with
 /// content specific to test execution.
 fn get_test_vm_args(orig_args: &VMArgs, cli_envs: Vec<String>) -> Result<ValidatedVMArgs> {
@@ -168,6 +191,7 @@ fn get_test_vm_args(orig_args: &VMArgs, cli_envs: Vec<String>) -> Result<Validat
         }
     }
     let envs = env_names_to_kvpairs(env_names);
+    record_envs(&envs)?;
 
     #[derive(Debug, Parser)]
     struct TestArgsParser {
@@ -189,7 +213,7 @@ fn get_test_vm_args(orig_args: &VMArgs, cli_envs: Vec<String>) -> Result<Validat
     vm_args.output_dirs = test_args.test.output_dirs().into_iter().collect();
     vm_args.mode.command = Some(test_args.test.into_inner_cmd());
     vm_args.command_envs = envs;
-    vm_args.console_output_file = console_output_path_for_tpx()?;
+    vm_args.console_output_file = create_tpx_logs("console", "console logs")?;
     Ok(ValidatedVMArgs {
         inner: vm_args,
         is_list,
@@ -282,7 +306,6 @@ fn main() -> Result<()> {
 
 #[cfg(test)]
 mod test {
-    use image_test_lib::KvPair;
 
     use super::*;
     use crate::types::VMModeArgs;

antlir/antlir2/antlir2_vm/src/utils.rs

@@ -13,6 +13,7 @@ use std::path::PathBuf;
 use std::process::Command;
 
 use image_test_lib::KvPair;
+use serde_json::json;
 use tracing::debug;
 use tracing::error;
 
@@ -64,21 +65,30 @@ pub(crate) fn run_command_capture_output(command: &mut Command) -> Result<(), st
     Ok(())
 }
 
-/// Return a path to record VM console output. When invoked under tpx, this
-/// will be uploaded as an artifact.
-pub(crate) fn console_output_path_for_tpx() -> Result<Option<PathBuf>, std::io::Error> {
+/// Return a path to record debugging data. When invoked under tpx, this will be
+/// uploaded as an artifact.
+pub(crate) fn create_tpx_logs(
+    name: &str,
+    description: &str,
+) -> Result<Option<PathBuf>, std::io::Error> {
     // 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") {
         fs::create_dir_all(&artifacts_dir)?;
-        let dst = Path::new(&artifacts_dir).join("console.txt");
+        let dst = Path::new(&artifacts_dir).join(format!("{}.txt", name));
         // The artifact metadata is set up before running the test so that it
         // still gets uploaded even in case of a timeout
         if let Some(annotations_dir) = std::env::var_os("TEST_RESULT_ARTIFACT_ANNOTATIONS_DIR") {
             fs::create_dir_all(&annotations_dir)?;
             fs::write(
-                Path::new(&annotations_dir).join("console.txt.annotation"),
-                r#"{"type": {"generic_text_log": {}}, "description": "console logs"}"#,
+                Path::new(&annotations_dir).join(format!("{}.txt.annotation", name)),
+                json!({
+                    "type": {
+                        "generic_text_log": {},
+                    },
+                    "description": description,
+                })
+                .to_string(),
             )?;
         }
         Ok(Some(dst))

antlir/antlir2/docs/docs/internals/vm-tests.md

@@ -89,6 +89,15 @@ command executed. This could be helpful when you want to run the test inside the
 VM shell. You can find more details in the
 [example section](#putting-it-together-an-investigation-example)
 
+Note: If your test target uses `env`, they won't be present for the interactive
+debugging sub targets. This is due to limitation in how envs are populated
+during tests, which are not fully available for `buck run`. One workaround is to
+run `buck test <test> -- --env RUST_LOG=debug` first, and look for the ssh
+command spawning the test in the failure output. It should contain a full list
+of envs that you can copy into your interactive shell. For Meta users, there are
+[additional integration](fb/vm-tests.md#more-internal-debugging-tips) to provide
+you the envs.
+
 ### Logging
 
 By default, the logging level is `info`. It only prints basic information like