From de90c121a5b3e8d816a19f14aa7cf4c6ccccf1d0 Mon Sep 17 00:00:00 2001 From: Junjie Wu Date: Tue, 10 Oct 2023 12:41:55 -0700 Subject: [PATCH] [antlir2][vm] read env values instead of passing them in Summary: Currently, vmtest pass in `--setenv` with values from `ExternalRunnerTestInfo.env`. This mostly works except when the env values contain macros like `$(location)`. Based on testing, those macros will not be evaluated at the time our provider generates the full command for spawning the test. Thus we get them passed in verbatim into container and VM, which obviously won't work. Instead, defer reading of the env values until our VM process spawns, at which point the macros must have been evaluated. This does mean that [shell] sub target will not be able to reproduce the same environment as the real test. `buck run` simply won't set any of those test envs. However, we don't really have a choice given `buck run` will not get the evaluated value anyway. More context: https://fb.workplace.com/groups/askbuck/posts/25235126456109298 https://fb.workplace.com/groups/antlirusers/posts/1780377695727348 https://www.internalfb.com/diff/D49933136 Test Plan: Enhance unit test and vmtests to cover this case. CI signal should cover all these tests. ```name=Verify the failing test $ buck2 test 'fbcode//mode/dbgo-cov' '-c' 'fbcode.cxx_coverage_only=strobelight' fbcode//strobelight/server/vmtest:pylatency-bpf-verification-vm-test-5.19.0-0_fbk4_10711_g9cbb08e4dfe0 -- --exact 'strobelight/server/vmtest:pylatency-bpf-verification-vm-test-5.19.0-0_fbk4_10711_g9cbb08e4dfe0 - PyLatencyVerificationTest.Load' --collect-coverage '--code-coverage-session=test_session' --force-tpx File changed: fbsource//default.profraw File changed: fbcode//strobelight/server/vmtest/strobelight_test_hhvm_usdt.py File changed: fbcode//strobelight/server/test/USDTDummyLoopRRM.cpp 12 additional file change events Buck UI: https://www.internalfb.com/buck2/28bf261e-14d0-4f8b-8c07-3e6fb2186a53 Test UI: https://www.internalfb.com/intern/testinfra/testrun/4503599830155826 Network: Up: 174B Down: 11MiB (reSessionID-c122a95b-c397-4508-aa43-ce8ccb5cd48c) Jobs completed: 149. Time elapsed: 52.2s. Cache hits: 100%. Commands: 1 (cached: 1, remote: 0, local: 0) Tests finished: Pass 1. Fail 0. Fatal 0. Skip 0. Build failure 0 ``` Reviewed By: epilatow Differential Revision: D49981391 fbshipit-source-id: db7304bad30c6180fd1a0ffc5c411ef9bfa1a971 --- antlir/antlir2/antlir2_vm/bzl/test.bzl | 5 +- antlir/antlir2/antlir2_vm/src/isolation.rs | 34 ------ antlir/antlir2/antlir2_vm/src/main.rs | 83 ++++++-------- antlir/antlir2/antlir2_vm/src/utils.rs | 123 +++++++++++++++++++++ antlir/antlir2/antlir2_vm/src/vm.rs | 12 +- antlir/antlir2/antlir2_vm/tests/BUCK | 19 ++-- antlir/antlir2/antlir2_vm/tests/test.cpp | 7 ++ antlir/antlir2/antlir2_vm/tests/test.py | 5 + antlir/antlir2/antlir2_vm/tests/test.rs | 9 ++ antlir/antlir2/antlir2_vm/tests/test.sh | 8 ++ 10 files changed, 197 insertions(+), 108 deletions(-) diff --git a/antlir/antlir2/antlir2_vm/bzl/test.bzl b/antlir/antlir2/antlir2_vm/bzl/test.bzl index 920bdb967cd..d221a552c5f 100644 --- a/antlir/antlir2/antlir2_vm/bzl/test.bzl +++ b/antlir/antlir2/antlir2_vm/bzl/test.bzl @@ -22,10 +22,7 @@ def _impl(ctx: AnalysisContext) -> list[Provider]: cmd_args(ctx.attrs.vm_host[VMHostInfo].image[LayerInfo].subvol_symlink, format = "--image={}"), cmd_args(ctx.attrs.vm_host[VMHostInfo].machine_spec, format = "--machine-spec={}"), cmd_args(ctx.attrs.vm_host[VMHostInfo].runtime_spec, format = "--runtime-spec={}"), - cmd_args( - ["{}={}".format(k, v) for k, v in ctx.attrs.test[ExternalRunnerTestInfo].env.items()], - format = "--setenv={}", - ), + cmd_args([k for k in ctx.attrs.test[ExternalRunnerTestInfo].env], format = "--passenv={}"), ) test_args = cmd_args( ctx.attrs.test[ExternalRunnerTestInfo].test_type, diff --git a/antlir/antlir2/antlir2_vm/src/isolation.rs b/antlir/antlir2/antlir2_vm/src/isolation.rs index 3017dcb6fd2..97d51532aaf 100644 --- a/antlir/antlir2/antlir2_vm/src/isolation.rs +++ b/antlir/antlir2/antlir2_vm/src/isolation.rs @@ -8,7 +8,6 @@ use std::collections::BTreeMap; use std::collections::HashSet; use std::env; -use std::ffi::OsString; use std::path::PathBuf; use std::process::Command; @@ -87,21 +86,6 @@ impl Platform { } } -/// If these env exist, always pass them through. -const PASSTHROUGH_ENVS: &[&str] = &["RUST_LOG", "ANTLIR_BUCK"]; - -/// Generate default passthrough env vars -pub(crate) fn default_passthrough_envs() -> Vec { - PASSTHROUGH_ENVS - .iter() - .filter(|x| env::var(*x).is_ok()) - .map(|x| KvPair { - key: x.to_string(), - value: OsString::from(env::var(x).expect("must exist")), - }) - .collect() -} - /// Return IsolatedContext ready for executing a command inside isolation /// # Arguments /// * `image` - container image that would be used to run the VM @@ -136,21 +120,3 @@ pub(crate) fn is_isolated() -> Result { debug!("systemd-detect-virt returned: {}", virt); Ok(virt == "systemd-nspawn") } - -#[cfg(test)] -mod test { - use super::*; - - #[test] - fn test_default_envs() { - assert_eq!(default_passthrough_envs(), Vec::new()); - env::set_var("RUST_LOG", "hello"); - assert_eq!( - default_passthrough_envs(), - vec![KvPair { - key: "RUST_LOG".into(), - value: "hello".into(), - }], - ); - } -} diff --git a/antlir/antlir2/antlir2_vm/src/main.rs b/antlir/antlir2/antlir2_vm/src/main.rs index e6f4d6d5471..51a38818747 100644 --- a/antlir/antlir2/antlir2_vm/src/main.rs +++ b/antlir/antlir2/antlir2_vm/src/main.rs @@ -27,7 +27,6 @@ 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; @@ -35,7 +34,6 @@ use tracing::debug; use tracing_subscriber::filter::LevelFilter; use tracing_subscriber::prelude::*; -use crate::isolation::default_passthrough_envs; use crate::isolation::is_isolated; use crate::isolation::isolated; use crate::isolation::Platform; @@ -44,6 +42,7 @@ use crate::types::MachineOpts; use crate::types::RuntimeOpts; use crate::types::VMArgs; use crate::utils::console_output_path_for_tpx; +use crate::utils::env_names_to_kvpairs; use crate::utils::log_command; use crate::vm::VM; @@ -84,10 +83,9 @@ struct IsolateCmdArgs { /// Path to container image. #[arg(long)] image: PathBuf, - /// Set these env vars in the container. If VM executes a command, these - /// env vars will also be prepended to the command. + /// Pass through these environment variables into the container and VM, if they exist. #[arg(long)] - setenv: Vec, + passenv: Vec, /// Args for run command #[clap(flatten)] run_cmd_args: RunCmdArgs, @@ -111,9 +109,8 @@ fn run(args: &RunCmdArgs) -> Result<()> { /// Enter isolated container and then respawn itself inside it with `run` /// command and its parameters. fn respawn(args: &IsolateCmdArgs) -> Result<()> { - let mut envs = default_passthrough_envs(); - envs.extend(args.setenv.clone()); let mut vm_args = args.run_cmd_args.vm_args.clone(); + let envs = env_names_to_kvpairs(args.passenv.clone()); vm_args.command_envs = envs.clone(); // Let's always capture console output unless it's console mode @@ -142,21 +139,6 @@ fn respawn(args: &IsolateCmdArgs) -> Result<()> { Ok(()) } -/// Merge all sources of our envs into final list of env vars we should use -/// everywhere for tests. Dedup is handled by functions that use the result. -fn get_test_envs(from_cli: &[KvPair]) -> Vec { - // This handles common envs like RUST_LOG - let mut envs = default_passthrough_envs(); - envs.extend_from_slice(from_cli); - // forward test runner env vars to the inner test - for (key, val) in std::env::vars() { - if key.starts_with("TEST_PILOT") { - envs.push((key, OsString::from(val)).into()); - } - } - envs -} - /// Validated `VMArgs` and other necessary metadata for tests. struct ValidatedVMArgs { /// VMArgs that will be passed into the VM with modified fields @@ -167,7 +149,7 @@ struct ValidatedVMArgs { /// 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: &[KvPair]) -> Result { +fn get_test_vm_args(orig_args: &VMArgs, cli_envs: Vec) -> Result { if orig_args.timeout_secs.is_none() { return Err(anyhow!("Test command must specify --timeout-secs.")); } @@ -177,7 +159,15 @@ fn get_test_vm_args(orig_args: &VMArgs, cli_envs: &[KvPair]) -> Result R /// inputs instead of allowing caller to pass them in. Some inputs are parsed /// from the test command. fn test(args: &IsolateCmdArgs) -> Result<()> { - let validated_args = get_test_vm_args(&args.run_cmd_args.vm_args, &args.setenv)?; + let validated_args = get_test_vm_args(&args.run_cmd_args.vm_args, args.passenv.clone())?; let mut command = if validated_args.is_list { list_test_command(args, &validated_args) } else { @@ -292,24 +282,11 @@ fn main() -> Result<()> { #[cfg(test)] mod test { + use image_test_lib::KvPair; + use super::*; use crate::types::VMModeArgs; - #[test] - fn test_get_test_envs() { - env::set_var("RUST_LOG", "hello"); - env::set_var("TEST_PILOT_A", "A"); - let from_cli = vec![KvPair::from(("foo", "bar"))]; - assert_eq!( - get_test_envs(&from_cli), - vec![ - KvPair::from(("RUST_LOG", "hello")), - KvPair::from(("foo", "bar")), - KvPair::from(("TEST_PILOT_A", "A")), - ], - ) - } - #[test] fn test_get_test_vm_args() { let valid = VMArgs { @@ -320,25 +297,35 @@ mod test { }, ..Default::default() }; - let empty_env = Vec::::new(); let mut expected = valid.clone(); expected.mode.command = Some(vec![OsString::from("whatever")]); - let parsed = get_test_vm_args(&valid, &empty_env).expect("Parsing should succeed"); - assert_eq!(parsed.inner, expected); + let parsed = get_test_vm_args(&valid, vec![]).expect("Parsing should succeed"); + assert_eq!(parsed.inner.mode, expected.mode); assert!(!parsed.is_list); let mut timeout = valid.clone(); timeout.timeout_secs = None; - assert!(get_test_vm_args(&timeout, &empty_env).is_err()); + assert!(get_test_vm_args(&timeout, vec![]).is_err()); let mut output_dirs = valid.clone(); output_dirs.output_dirs = vec![PathBuf::from("/some")]; - assert!(get_test_vm_args(&output_dirs, &empty_env).is_err()); + assert!(get_test_vm_args(&output_dirs, vec![]).is_err()); - let mut command = valid; + let mut command = valid.clone(); command.mode.command = None; - assert!(get_test_vm_args(&command, &empty_env).is_err()); + assert!(get_test_vm_args(&command, vec![]).is_err()); command.mode.command = Some(vec![OsString::from("invalid")]); - assert!(get_test_vm_args(&command, &empty_env).is_err()); + assert!(get_test_vm_args(&command, vec![]).is_err()); + + let env_var_test = valid; + std::env::set_var("TEST_PILOT_A", "A"); + let parsed = get_test_vm_args(&env_var_test, vec![]).expect("Parsing should succeed"); + assert!( + parsed + .inner + .command_envs + .iter() + .any(|x| *x == KvPair::from(("TEST_PILOT_A", "A"))) + ); } } diff --git a/antlir/antlir2/antlir2_vm/src/utils.rs b/antlir/antlir2/antlir2_vm/src/utils.rs index 2647b12075b..44d8f940b1f 100644 --- a/antlir/antlir2/antlir2_vm/src/utils.rs +++ b/antlir/antlir2/antlir2_vm/src/utils.rs @@ -5,12 +5,14 @@ * LICENSE file in the root directory of this source tree. */ +use std::collections::HashSet; use std::fs; use std::io::ErrorKind; use std::path::Path; use std::path::PathBuf; use std::process::Command; +use image_test_lib::KvPair; use tracing::debug; use tracing::error; @@ -85,6 +87,24 @@ pub(crate) fn console_output_path_for_tpx() -> Result, std::io:: } } +/// Convert a list of env names into KvPair with its values +pub(crate) fn env_names_to_kvpairs(env_names: Vec) -> Vec { + let mut names: HashSet<_> = env_names.into_iter().collect(); + // If these env exist, always pass them through. + ["RUST_LOG", "RUST_BACKTRACE", "ANTLIR_BUCK"] + .iter() + .for_each(|name| { + names.insert(name.to_string()); + }); + names + .iter() + .filter_map(|name| match std::env::var(name) { + Ok(value) => Some(KvPair::from((name, value))), + _ => None, + }) + .collect() +} + #[cfg(test)] /// Helper function for converting qemu args to a single string for asserting in tests. /// This is usually only needed for string only functions like `contains`. @@ -97,6 +117,10 @@ pub(crate) fn qemu_args_to_string(args: &[std::ffi::OsString]) -> String { #[cfg(test)] mod test { + use std::collections::HashMap; + use std::env; + use std::ffi::OsString; + use super::*; #[test] @@ -110,4 +134,103 @@ mod test { format!("Program: `hello`. Args: `{:?}`", vec!["world"]), ); } + + struct EnvTest { + envs: Vec<(&'static str, &'static str)>, + passenv: Vec<&'static str>, + result: HashMap, + } + + const ALLOWED_ENV_NAMES: &[&str] = &[ + "RUST_LOG", + "RUST_BACKTRACE", + "ANTLIR_BUCK", + "TEST_PILOT_A", + "OTHER", + ]; + + impl EnvTest { + fn new( + envs: Vec<(&'static str, &'static str)>, + passenv: Vec<&'static str>, + result: Vec<(&'static str, &'static str)>, + ) -> Self { + // We have to clear all envs across tests, so we must know the ones to clear. + envs.iter().for_each(|(name, _)| { + assert!( + ALLOWED_ENV_NAMES.contains(name), + "{name} not allowed for testing" + ) + }); + Self { + envs, + passenv, + result: result + .into_iter() + .map(|(k, v)| (k.into(), v.into())) + .collect(), + } + } + + fn test(&self) { + ALLOWED_ENV_NAMES.iter().for_each(|name| { + env::remove_var(name); + }); + self.envs.iter().for_each(|(name, val)| { + env::set_var(name, val); + }); + let result: HashMap<_, _> = + env_names_to_kvpairs(self.passenv.iter().map(|s| s.to_string()).collect()) + .drain(..) + .map(|pair| (pair.key, pair.value)) + .collect(); + assert_eq!(result, self.result); + } + } + + #[test] + fn test_env_names_to_kvpairs() { + [ + EnvTest::new(vec![], vec![], vec![]), + // Reads nothing + EnvTest::new( + vec![], + vec!["RUST_LOG", "ANTLIR_BUCK", "TEST_PILOT_A", "OTHER"], + vec![], + ), + // Always pass through + EnvTest::new( + vec![ + ("RUST_LOG", "info"), + ("ANTLIR_BUCK", "1"), + ("TEST_PILOT_A", "A"), + ("OTHER", "other"), + ], + vec![], + vec![("RUST_LOG", "info"), ("ANTLIR_BUCK", "1")], + ), + // Selection + EnvTest::new( + vec![("TEST_PILOT_A", "A"), ("OTHER", "other")], + vec!["TEST_PILOT_A"], + vec![("TEST_PILOT_A", "A")], + ), + // Mixed + EnvTest::new( + vec![ + ("RUST_LOG", "info"), + ("TEST_PILOT_A", "A"), + ("OTHER", "other"), + ], + vec!["TEST_PILOT_A"], + vec![("TEST_PILOT_A", "A"), ("RUST_LOG", "info")], + ), + ] + .iter() + .enumerate() + .for_each(|(i, test)| { + println!("Running test #{i}"); + test.test(); + }); + } } diff --git a/antlir/antlir2/antlir2_vm/src/vm.rs b/antlir/antlir2/antlir2_vm/src/vm.rs index 6ab0ba625c0..d494c577591 100644 --- a/antlir/antlir2/antlir2_vm/src/vm.rs +++ b/antlir/antlir2/antlir2_vm/src/vm.rs @@ -5,7 +5,6 @@ * LICENSE file in the root directory of this source tree. */ -use std::collections::HashMap; use std::collections::HashSet; use std::ffi::OsString; use std::fs; @@ -26,7 +25,6 @@ use std::thread::JoinHandle; use std::time::Duration; use std::time::Instant; -use image_test_lib::KvPair; use thiserror::Error; use tracing::debug; use tracing::error; @@ -524,16 +522,8 @@ impl VM { // Just wait for the human that's trying to debug with console self.wait_for_timeout::<()>(f.into_inner(), start_ts, None)?; } else if !self.args.mode.container { - // Open ssh shell or run command over ssh - let dedup_envs: HashMap<_, _> = self - .args - .command_envs - .iter() - .cloned() - .map(|p| (p.key, p.value)) - .collect(); let mut ssh_cmd = GuestSSHCommand::new()?.ssh_cmd(); - dedup_envs.iter().map(KvPair::from).for_each(|kv| { + self.args.command_envs.iter().for_each(|kv| { ssh_cmd.arg(kv.to_os_string()); }); if let Some(command) = &self.args.mode.command { diff --git a/antlir/antlir2/antlir2_vm/tests/BUCK b/antlir/antlir2/antlir2_vm/tests/BUCK index 9aa763b8296..e4fd549a821 100644 --- a/antlir/antlir2/antlir2_vm/tests/BUCK +++ b/antlir/antlir2/antlir2_vm/tests/BUCK @@ -9,7 +9,9 @@ vm.cpp_test( srcs = ["test.cpp"], env = { "ANTLIR2_TEST": "1", + "ENV_ARTIFACT": "$(exe :sidecar)", }, + run_as_bundle = True, vm_host = get_vm(), ) @@ -18,7 +20,9 @@ vm.python_test( srcs = ["test.py"], env = { "ANTLIR2_TEST": "1", + "ENV_ARTIFACT": "$(exe :sidecar)", }, + run_as_bundle = True, vm_host = get_vm(), ) @@ -29,7 +33,9 @@ vm.rust_test( crate_root = "test.rs", env = { "ANTLIR2_TEST": "1", + "ENV_ARTIFACT": "$(exe :sidecar)", }, + run_as_bundle = True, vm_host = get_vm(), ) @@ -37,6 +43,7 @@ vm.sh_test( name = "sh-test", env = { "ANTLIR2_TEST": "1", + "ENV_ARTIFACT": "$(exe :sidecar)", }, test = "test.sh", vm_host = get_vm(), @@ -51,23 +58,13 @@ vm.sh_test( crate_root = "test.rs", env = { "ANTLIR2_TEST": "1", + "ENV_ARTIFACT": "$(exe :sidecar)", }, vm_host = target, ) for name, target in PRECONFIGURED_VM.items() ] -# Verify bundle -vm.cpp_test( - name = "cpp-test-bundle", - srcs = ["test.cpp"], - env = { - "ANTLIR2_TEST": "1", - }, - run_as_bundle = True, - vm_host = get_vm(), -) - # Test for sidecar rust_binary( name = "sidecar", diff --git a/antlir/antlir2/antlir2_vm/tests/test.cpp b/antlir/antlir2/antlir2_vm/tests/test.cpp index 02386e102c6..46db47bea86 100644 --- a/antlir/antlir2/antlir2_vm/tests/test.cpp +++ b/antlir/antlir2/antlir2_vm/tests/test.cpp @@ -7,6 +7,7 @@ #include #include +#include TEST(CppTest, TestIsRoot) { EXPECT_EQ(getuid(), 0); @@ -15,3 +16,9 @@ TEST(CppTest, TestIsRoot) { TEST(CppTest, TestEnvPropagated) { ASSERT_STREQ(std::getenv("ANTLIR2_TEST"), "1"); } + +TEST(CppTest, TestEnvArtifactExists) { + auto artifact = std::getenv("ENV_ARTIFACT"); + ASSERT_NE(nullptr, artifact); + ASSERT_TRUE(std::filesystem::exists(artifact)); +} diff --git a/antlir/antlir2/antlir2_vm/tests/test.py b/antlir/antlir2/antlir2_vm/tests/test.py index 9bf17b48e62..e3b3f2785ed 100644 --- a/antlir/antlir2/antlir2_vm/tests/test.py +++ b/antlir/antlir2/antlir2_vm/tests/test.py @@ -14,3 +14,8 @@ def test_is_root(self) -> None: def test_env_propagated(self) -> None: self.assertEqual("1", os.getenv("ANTLIR2_TEST")) + + def test_env_artifact_exists(self) -> None: + artifact = os.getenv("ENV_ARTIFACT") + self.assertNotEqual(None, artifact) + self.assertTrue(os.path.exists(artifact)) diff --git a/antlir/antlir2/antlir2_vm/tests/test.rs b/antlir/antlir2/antlir2_vm/tests/test.rs index fee064a1622..6677aed215d 100644 --- a/antlir/antlir2/antlir2_vm/tests/test.rs +++ b/antlir/antlir2/antlir2_vm/tests/test.rs @@ -5,6 +5,8 @@ * LICENSE file in the root directory of this source tree. */ +use std::path::Path; + #[link(name = "c")] extern "C" { fn geteuid() -> u32; @@ -19,3 +21,10 @@ fn is_root() { fn env_propagated() { assert_eq!("1", std::env::var("ANTLIR2_TEST").expect("env var missing")); } + +#[test] +fn env_artifact_exists() { + let artifact = std::env::var("ENV_ARTIFACT"); + assert!(artifact.is_ok()); + assert!(Path::new(&artifact.expect("Already checked")).exists()); +} diff --git a/antlir/antlir2/antlir2_vm/tests/test.sh b/antlir/antlir2/antlir2_vm/tests/test.sh index 2fdcff11490..c690b7fb13b 100755 --- a/antlir/antlir2/antlir2_vm/tests/test.sh +++ b/antlir/antlir2/antlir2_vm/tests/test.sh @@ -15,3 +15,11 @@ if [ "${ANTLIR2_TEST}" != "1" ]; then echo "Env var ANTLIR2_TEST is wrong or missing!" exit 2 fi + +if [ -z "${ENV_ARTIFACT}" ]; then + echo "Env var ENV_ARTIFACT is missing!" + exit 3 +elif ! [ -f "${ENV_ARTIFACT}" ]; then + echo "Env var ENV_ARTIFACT does not point to a valid target!" + exit 3 +fi