nh-nixos: batch nix path-info calls in nh os info

crates/nh-nixos/src/generations.rs

@@ -1,4 +1,9 @@
-use std::{collections::HashMap, fs, path::Path, process};
+use std::{
+  collections::HashMap,
+  fs,
+  path::{Path, PathBuf},
+  process,
+};
 
 use chrono::{DateTime, Local, TimeZone, Utc};
 use clap::ValueEnum;
@@ -95,6 +100,92 @@ pub fn from_dir(generation_dir: &Path) -> Option<u64> {
     })
 }
 
+fn closure_size_from_json(
+  json: &serde_json::Value,
+  store_path_str: &str,
+) -> Option<u64> {
+  json.as_array().map_or_else(
+    || {
+      json.as_object().and_then(|obj| {
+        obj
+          .get(&*store_path_str)
+          .and_then(|value| value.get("closureSize"))
+          .and_then(serde_json::Value::as_u64)
+      })
+    },
+    |arr| {
+      arr.iter().find_map(|entry| {
+        let path = entry.get("path")?.as_str()?;
+        let size = entry.get("closureSize")?.as_u64()?;
+        (path == store_path_str).then_some(size)
+      })
+    },
+  )
+}
+
+#[expect(clippy::cast_precision_loss)]
+fn bytes_to_gb_string(bytes: u64) -> String {
+  format!("{:.1} GB", bytes as f64 / 1_073_741_824.0)
+}
+
+/// Get closure sizes for all given generation directories in a single
+/// `nix path-info` invocation.
+///
+/// # Returns
+///
+/// A map from generation directory path to formatted closure size
+/// string.
+#[must_use]
+pub fn get_closure_sizes_batch(
+  generation_dirs: &[&Path],
+) -> HashMap<PathBuf, String> {
+  if generation_dirs.is_empty() {
+    return HashMap::new();
+  }
+
+  let store_paths: Vec<PathBuf> = generation_dirs
+    .iter()
+    .map(|p| p.read_link().unwrap_or_else(|_| p.to_path_buf()))
+    .collect();
+
+  let output = match process::Command::new("nix")
+    .args(["path-info", "-Sh", "--json"])
+    .args(generation_dirs)
+    .output()
+  {
+    Ok(out) => out,
+    Err(e) => {
+      debug!("get_closure_sizes_batch: failed to run nix path-info: {e:?}");
+      return HashMap::new();
+    },
+  };
+
+  let output_str = String::from_utf8_lossy(&output.stdout);
+
+  let json: serde_json::Value =
+    match serde_json::from_str::<serde_json::Value>(&output_str) {
+      Ok(j) => j,
+      Err(e) => {
+        debug!(
+          "get_closure_sizes_batch: failed to parse JSON: {e} output: \
+           {output_str}"
+        );
+        return HashMap::new();
+      },
+    };
+
+  generation_dirs
+    .iter()
+    .zip(store_paths.iter())
+    .map(|(gen_dir, store_path)| {
+      let store_path_str = store_path.to_string_lossy();
+      let size_str = closure_size_from_json(&json, &store_path_str)
+        .map_or_else(|| "Unknown".to_string(), bytes_to_gb_string);
+      (gen_dir.to_path_buf(), size_str)
+    })
+    .collect()
+}
+
 #[must_use]
 pub fn get_closure_size(generation_dir: &Path) -> String {
   let store_path = generation_dir
@@ -108,7 +199,6 @@ pub fn get_closure_size(generation_dir: &Path) -> String {
     .output()
   {
     Ok(out) => out,
-
     Err(e) => {
       debug!("get_closure_size: failed to run nix path-info: {e:?}");
       return "Unknown".to_string();
@@ -120,7 +210,6 @@ pub fn get_closure_size(generation_dir: &Path) -> String {
   let json: serde_json::Value =
     match serde_json::from_str::<serde_json::Value>(&output_str) {
       Ok(j) => j,
-
       Err(e) => {
         debug!(
           "get_closure_size: failed to parse JSON: {e} output: {output_str}"
@@ -129,57 +218,26 @@ pub fn get_closure_size(generation_dir: &Path) -> String {
       },
     };
 
-  let closure_size = json.as_array().map_or_else(
+  closure_size_from_json(&json, &store_path_str).map_or_else(
     || {
-      json.as_object().and_then(|obj| {
-        obj
-          .iter()
-          .find(|(path, _)| path.as_str() == store_path_str)
-          .and_then(|(_, value)| value.get("closureSize"))
-          .and_then(serde_json::Value::as_u64)
-      })
-    },
-    |arr| {
-      arr.iter().find_map(|entry| {
-        let path = entry.get("path")?.as_str()?;
-        let size = entry.get("closureSize")?.as_u64()?;
-        (path == store_path_str).then_some(size)
-      })
-    },
-  );
-
-  closure_size.map_or_else(
-    || {
-      let paths: Vec<String> = json.as_array().map_or_else(
-        || {
-          json
-            .as_object()
-            .map_or_else(Vec::new, |obj| obj.keys().cloned().collect())
-        },
-        |arr| {
-          arr
-            .iter()
-            .filter_map(|e| e.get("path")?.as_str().map(ToString::to_string))
-            .collect()
-        },
-      );
-
       debug!(
         "get_closure_size: store_path not found or closureSize missing. \
-         store_path: {store_path_str}, json paths: {:?}, output: {}",
-        paths, output_str
+         store_path: {store_path_str}, output: {output_str}"
       );
       "Unknown".to_string()
     },
-    #[expect(clippy::cast_precision_loss)]
-    |bytes| format!("{:.1} GB", bytes as f64 / 1_073_741_824.0),
+    bytes_to_gb_string,
   )
 }
 
 #[must_use]
-pub fn describe(generation_dir: &Path) -> Option<GenerationInfo> {
+pub fn describe(
+  generation_dir: &Path,
+  closure_size: Option<String>,
+) -> Option<GenerationInfo> {
   let generation_number = from_dir(generation_dir)?;
-  let closure_size = get_closure_size(generation_dir);
+  let closure_size =
+    closure_size.unwrap_or_else(|| get_closure_size(generation_dir));
   // Get metadata once and reuse for both date and existence checks
   let metadata = fs::metadata(generation_dir).ok()?;
   let build_date = metadata

crates/nh-nixos/src/nixos.rs

@@ -1227,7 +1227,7 @@ fn list_generations() -> Result<Vec<generations::GenerationInfo>> {
     if let Some(name) = path.file_name().and_then(|s| s.to_str())
       && name.starts_with("system-")
       && name.ends_with("-link")
-      && let Some(gen_info) = generations::describe(&path)
+      && let Some(gen_info) = generations::describe(&path, None)
     {
       generations.push(gen_info);
     }
@@ -1376,9 +1376,16 @@ impl OsGenerationsArgs {
       })
       .collect();
 
+    let gen_dir_refs: Vec<&std::path::Path> =
+      generations.iter().map(PathBuf::as_path).collect();
+    let closure_sizes = generations::get_closure_sizes_batch(&gen_dir_refs);
+
     let descriptions: Vec<generations::GenerationInfo> = generations
       .iter()
-      .filter_map(|gen_dir| generations::describe(gen_dir))
+      .filter_map(|gen_dir| {
+        let size = closure_sizes.get(gen_dir).cloned();
+        generations::describe(gen_dir, size)
+      })
       .collect();
 
     generations::print_info(descriptions, self.fields.as_deref())?;