test: Add support for A/B-Tests that need to compile old versions of FC

Needed for the vulnerability tests, as they need to run commands inside
of microvms.

Signed-off-by: Patrick Roy <[email protected]>

tests/framework/ab_test.py

@@ -23,6 +23,7 @@
 """
 import contextlib
 import os
+import shutil
 import statistics
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -31,6 +32,10 @@
 import scipy
 
 from framework import utils
+from framework.defs import FC_WORKSPACE_DIR
+from framework.microvm import Microvm
+from framework.with_filelock import with_filelock
+from host_tools.cargo_build import get_firecracker_binaries
 
 # Locally, this will always compare against main, even if we try to merge into, say, a feature branch.
 # We might want to do a more sophisticated way to determine a "parent" branch here.
@@ -100,14 +105,19 @@ def git_ab_test(
         return result_a, result_b, comparison
 
 
+def is_pr() -> bool:
+    """Returns ``True` iff we are executing in the context of a build kite run on a pull request"""
+    return os.environ.get("BUILDKITE_PULL_REQUEST", "false") != "false"
+
+
 def git_ab_test_command_if_pr(command: str):
     """Runs the given bash command as an A/B-Test if we're in a pull request context (asserting that its stdout and
     stderr did not change across the PR). Otherwise runs the command, asserting it returns a zero exit code
     """
-    if os.environ.get("BUILDKITE_PULL_REQUEST", "false") == "false":
-        utils.run_cmd(command)
-    else:
+    if is_pr():
         git_ab_test_command(command)
+    else:
+        utils.run_cmd(command)
 
 
 def git_ab_test_command(command: str):
@@ -123,6 +133,74 @@ def git_ab_test_command(command: str):
     ), f"The output of running command `{command}` changed:\nOld:\nstdout:\n{old_out}\nstderr\n{old_err}\n\nNew:\nstdout:\n{new_out}\nstderr:\n{new_err}"
 
 
+def git_ab_test_with_binaries(
+    test_runner: Callable[[Path, Path], T],
+    comparator: Callable[[T, T], U] = default_comparator,
+    *,
+    a_revision: str = DEFAULT_A_REVISION,
+    b_revision: Optional[str] = None,
+) -> (T, T, U):
+    """Similar to `git_ab_test`, with the only difference being that this function compiles firecracker at the specified
+    revisions and only passes the firecracker binaries to the test_runner. Maintains a cache of previously compiled
+    revisions, to prevent excessive recompilation across different tests of the same revision
+    """
+
+    @with_filelock
+    def grab_binaries(checkout: Path):
+        with chdir(checkout):
+            revision = utils.run_cmd("git rev-parse HEAD").stdout.strip()
+
+        revision_store = FC_WORKSPACE_DIR / "build" / revision
+        if not revision_store.exists():
+            firecracker, jailer = get_firecracker_binaries(workspace_dir=checkout)
+
+            revision_store.mkdir(parents=True, exist_ok=True)
+            shutil.copy(firecracker, revision_store / "firecracker")
+            shutil.copy(jailer, revision_store / "jailer")
+
+        return (
+            revision_store / "firecracker",
+            revision_store / "jailer",
+        )
+
+    return git_ab_test(
+        lambda checkout, _is_a: test_runner(*grab_binaries(checkout)),
+        comparator,
+        a_revision=a_revision,
+        b_revision=b_revision,
+    )
+
+
+def git_ab_test_ssh_command(
+    microvm_factory: Callable[[Path, Path], Microvm], command: str
+):
+    """The same as git_ab_test_command, but via SSH"""
+
+    def test_runner(firecracker, jailer):
+        microvm = microvm_factory(firecracker, jailer)
+        return microvm.ssh.run(command)
+
+    (old_out, old_err), (new_out, new_err), the_same = git_ab_test_with_binaries(
+        test_runner
+    )
+
+    assert (
+        the_same
+    ), f"The output of running command `{command}` changed:\nOld:\nstdout:\n{old_out}\nstderr\n{old_err}\n\nNew:\nstdout:\n{new_out}\nstderr:\n{new_err}"
+
+
+def git_ab_test_ssh_command_if_pr(
+    microvm_factory: Callable[[Path, Path], Microvm], command: str
+):
+    """The same as git_ab_test_command_if_pr, but via SSH"""
+    if is_pr():
+        git_ab_test_ssh_command(microvm_factory, command)
+    else:
+        microvm = microvm_factory(*get_firecracker_binaries())
+        ecode, stdout, stderr = microvm.ssh.run(command)
+        assert ecode == 0, f"stdout:\n{stdout}\nstderr:\n{stderr}\n"
+
+
 def check_regression(a_samples: List[float], b_samples: List[float]):
     """Checks for a regression by performing a permutation test. A permutation test is a non-parametric test that takes
     three parameters: Two populations (sets of samples) and a function computing a "statistic" based on two populations.
@@ -167,6 +245,10 @@ def temporary_checkout(revision: str):
 
         yield Path(tmp_dir)
 
+        # If we compiled firecracker inside the checkout, python's recursive shutil.rmdir will
+        # run incredibly long. Thus, remove manually.
+        utils.run_cmd(f"rm -rf {tmp_dir}")
+
 
 # Once we upgrade to python 3.11, this will be in contextlib:
 # https://docs.python.org/3/library/contextlib.html#contextlib.chdir

tests/host_tools/cargo_build.py

@@ -71,7 +71,7 @@ def get_binary(name, *, workspace_dir=FC_WORKSPACE_DIR):
     """Build a binary"""
     target = DEFAULT_BUILD_TARGET
     target_dir = workspace_dir / "build" / "cargo_target"
-    bin_path = target_dir / target / name
+    bin_path = target_dir / target / "release" / name
     if not bin_path.exists():
         env = {"RUSTFLAGS": get_rustflags()}
         cargo(