From d2a80dda88317d090f4c8bd65e39cf706a7daa5e Mon Sep 17 00:00:00 2001
From: Ilan Joselevich <personal@ilanjoselevich.com>
Date: Mon, 12 Aug 2024 00:02:16 +0300
Subject: feat(tvix/cli): Add derivation file dumping functionality

Provides a derivation file dumping functionality for tvix-cli that can
be used when passing the --drv-dumpdir CLI arg to tvix-cli.

This will dump all the known derivation files into the specified
directory, making it easier to debug derivation divergences between Tvix
generated drvs and the drvs generated by Nix.

Supersedes: https://cl.tvl.fyi/c/depot/+/11265

Change-Id: I0e10b26eba22032b84ac543af0d4150ad87aed3e
Reviewed-on: https://cl.tvl.fyi/c/depot/+/12192
Autosubmit: Ilan Joselevich <personal@ilanjoselevich.com>
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>
---
 tvix/cli/src/args.rs           | 11 +++++++++++
 tvix/cli/src/lib.rs            | 17 ++++++++++++++++-
 tvix/glue/src/known_paths.rs   | 21 +++++++++++++++++++++
 tvix/glue/src/tvix_store_io.rs |  2 +-
 4 files changed, 49 insertions(+), 2 deletions(-)

(limited to 'tvix')

diff --git a/tvix/cli/src/args.rs b/tvix/cli/src/args.rs
index 5b8f6b48b134..36f9a6a262dd 100644
--- a/tvix/cli/src/args.rs
+++ b/tvix/cli/src/args.rs
@@ -72,4 +72,15 @@ pub struct Args {
 
     #[arg(long, env, default_value = "dummy://")]
     pub build_service_addr: String,
+
+    /// An optional path in which Derivations encountered during evaluation
+    /// are dumped into, after evaluation. If it doesn't exist, the directory is created.
+    ///
+    /// Files dumped there are named like they would show up in `/nix/store`,
+    /// if produced by Nix. Existing files are not overwritten.
+    ///
+    /// This is only for debugging and diffing purposes for post-eval inspection;
+    /// Tvix does not read from these.
+    #[clap(long)]
+    pub drv_dumpdir: Option<PathBuf>,
 }
diff --git a/tvix/cli/src/lib.rs b/tvix/cli/src/lib.rs
index 060665480f1e..beb4c505207c 100644
--- a/tvix/cli/src/lib.rs
+++ b/tvix/cli/src/lib.rs
@@ -112,7 +112,7 @@ pub fn evaluate(
             eval_builder = eval_builder.add_builtins(impure_builtins());
             eval_builder = add_derivation_builtins(eval_builder, Rc::clone(&tvix_store_io));
             eval_builder = add_fetcher_builtins(eval_builder, Rc::clone(&tvix_store_io));
-            eval_builder = add_import_builtins(eval_builder, tvix_store_io);
+            eval_builder = add_import_builtins(eval_builder, Rc::clone(&tvix_store_io));
         }
     };
     eval_builder = configure_nix_path(eval_builder, &args.nix_search_path);
@@ -173,6 +173,21 @@ pub fn evaluate(
         }
     }
 
+    if let Some(dumpdir) = &args.drv_dumpdir {
+        // Dump all known derivations files to `dumpdir`.
+        std::fs::create_dir_all(dumpdir).expect("failed to create drv dumpdir");
+        tvix_store_io
+            .known_paths
+            .borrow()
+            .get_derivations()
+            // Skip already dumped derivations.
+            .filter(|(drv_path, _)| !dumpdir.join(drv_path.to_string()).exists())
+            .for_each(|(drv_path, drv)| {
+                std::fs::write(dumpdir.join(drv_path.to_string()), drv.to_aterm_bytes())
+                    .expect("failed to write drv to dumpdir");
+            })
+    }
+
     Ok(EvalResult {
         globals,
         value: result.value,
diff --git a/tvix/glue/src/known_paths.rs b/tvix/glue/src/known_paths.rs
index edc57c38f247..b303b8a13520 100644
--- a/tvix/glue/src/known_paths.rs
+++ b/tvix/glue/src/known_paths.rs
@@ -129,6 +129,11 @@ impl KnownPaths {
             .get(output_path)
             .map(|(name, fetch)| (name.to_owned(), fetch.to_owned()))
     }
+
+    /// Returns an iterator over all known derivations and their store path.
+    pub fn get_derivations(&self) -> impl Iterator<Item = (&StorePath, &Derivation)> {
+        self.derivations.iter().map(|(k, v)| (k, &v.1))
+    }
 }
 
 #[cfg(test)]
@@ -269,5 +274,21 @@ mod tests {
         );
     }
 
+    #[test]
+    fn get_derivations_working() {
+        let mut known_paths = KnownPaths::default();
+
+        // Add BAR_DRV
+        known_paths.add_derivation(BAR_DRV_PATH.clone(), BAR_DRV.clone());
+
+        // We should be able to find BAR_DRV_PATH and BAR_DRV as a pair in get_derivations.
+        assert_eq!(
+            Some((&BAR_DRV_PATH.clone(), &BAR_DRV.clone())),
+            known_paths
+                .get_derivations()
+                .find(|(s, d)| (*s, *d) == (&BAR_DRV_PATH, &BAR_DRV))
+        );
+    }
+
     // TODO: add test panicking about missing digest
 }
diff --git a/tvix/glue/src/tvix_store_io.rs b/tvix/glue/src/tvix_store_io.rs
index c25194d50fda..54e3550fc940 100644
--- a/tvix/glue/src/tvix_store_io.rs
+++ b/tvix/glue/src/tvix_store_io.rs
@@ -67,7 +67,7 @@ pub struct TvixStoreIO {
     >,
 
     // Paths known how to produce, by building or fetching.
-    pub(crate) known_paths: RefCell<KnownPaths>,
+    pub known_paths: RefCell<KnownPaths>,
 }
 
 impl TvixStoreIO {
-- 
cgit 1.4.1