feat(zkvms_guest_io): Track and output memory usage

We're using benchexec's runexec to track how much memory our program
uses. Because of nix compatibility troubles, for now we've disabled
runexec's container feature. This could result in less accurate data.

2 files changed, 28 insertions, 10 deletions

diff --git a/zkvms_guest_io/default.nix b/zkvms_guest_io/default.nix
index 20c75a7..bce5a9c 100644
--- a/zkvms_guest_io/default.nix
+++ b/zkvms_guest_io/default.nix
@@ -1,11 +1,11 @@
-{ craneLib-default, guest, zkvms, hostPackages, lib, }:
+{ craneLib-default, guest, zkvms, hostPackages, lib, benchexec, }:
 let
   commonArgs = {
     name = "${guest}";
 
-    buildInputs =
-      lib.foldr (zkvm: accum: accum ++ [ hostPackages."${zkvm}/${guest}" ]) [ ]
-      zkvms;
+    buildInputs = [ benchexec ] ++
+      (lib.foldr (zkvm: accum: accum ++ [ hostPackages."${zkvm}/${guest}" ]) [ ]
+      zkvms);
 
     src = lib.fileset.toSource {
       root = ./.;
@@ -15,6 +15,10 @@ let
     PROGRAMS = lib.foldr (zkvm: accum:
       hostPackages."${zkvm}/${guest}" + "/bin/${zkvm}_${guest}," + accum) ""
       zkvms;
+
+    postPatch = ''
+      sed -i 's|"runexec"|"${benchexec}/bin/runexec"|' ./src/main.rs
+    '';
   };
 
   cargoArtifacts = craneLib-default.buildDepsOnly commonArgs;
diff --git a/zkvms_guest_io/src/main.rs b/zkvms_guest_io/src/main.rs
index 04d33d1..ab20a32 100644
--- a/zkvms_guest_io/src/main.rs
+++ b/zkvms_guest_io/src/main.rs
@@ -30,14 +30,21 @@ struct Cli {
 }
 
 fn run_command(zkvm_guest_command: &str, operation: &str) -> Result<std::process::Output, Error> {
-    Command::new(zkvm_guest_command)
-        .arg(operation)
+    Command::new("runexec")
+        .args(["--no-container", "--"])
+        .args([zkvm_guest_command, operation])
         .arg("--benchmark")
         .args([ "--metrics-output", "/tmp/current_metrics" ])
         .stdout(Stdio::piped())
         .output()
 }
 
+fn get_runexec_value(output: &String, name: &str, end: char) -> String {
+    let start_bytes = output.find(name).unwrap();
+    let right_half = &output[start_bytes + name.len() + 1..];
+    right_half[..right_half.find(end).unwrap()].to_string()
+}
+
 fn main() {
     let cli = Cli::parse();
 
@@ -65,7 +72,7 @@ fn main() {
 
             let output = run_command(zkvm_guest_command, operation);
 
-            // Couldn't run the command
+            // Couldn't run runexec
             if let Err(msg) = output {
                 println!("Failed to run command!");
                 println!("{msg}");
@@ -75,10 +82,10 @@ fn main() {
                 continue;
             }
 
-            // The command ran and therefore produced some output
+            // runexec ran and therefore produced some output
             let output = output.unwrap();
 
-            // The command ran but exited with non-zero status code
+            // runexec exited with non-zero status code
             if !output.status.success() {
                 println!("Command failed!");
             }
@@ -93,15 +100,22 @@ fn main() {
                 );
             }
 
-            // The command ran but exited with non-zero status code
+            // runexec ran but exited with non-zero status code
             if !output.status.success() {
                 break 'guest_iter;
             }
 
+            // The guest program ran but exited with non-zero status code
+            if get_runexec_value(&stdout, "returnvalue", '\n') != "0" {
+                run[operation] = Null;
+                continue;
+            }
+
             let raw_data = &read_to_string("/tmp/current_metrics")
                 .ok()
                 .unwrap();
             run[operation] = json::parse(raw_data).unwrap();
+            run[operation]["memory"] = get_runexec_value(&stdout, "memory", 'B').parse::<u64>().unwrap().into();
         }
 
         runs["benchmarking"].push(run);