Impl memory limit for flake evaluation

Cargo.lock

@@ -768,10 +768,12 @@ name = "nix-interop"
 version = "0.0.0"
 dependencies = [
  "anyhow",
+ "rustix 0.38.3",
  "serde",
  "serde_json",
  "serde_repr",
  "syntax",
+ "thiserror",
  "tokio",
 ]
 

crates/nil/src/config.rs

@@ -74,6 +74,8 @@ pub struct Config {
     pub formatting_command: Option<Vec<String>>,
     #[parse("/nix/binary", default = "nix".into())]
     pub nix_binary: PathBuf,
+    #[parse("/nix/maxMemoryMB", default = Some(2048))]
+    pub nix_max_memory_mb: Option<u64>,
     #[parse("/nix/flake/autoArchive")]
     pub nix_flake_auto_archive: Option<bool>,
     #[parse("/nix/flake/autoEvalInputs")]
@@ -98,4 +100,8 @@ impl Config {
         ensure!(v != Some(Vec::new()), "command must not be empty");
         Ok(v)
     }
+
+    pub fn nix_max_memory(&self) -> Option<u64> {
+        self.nix_max_memory_mb?.checked_mul(1 << 20)
+    }
 }

crates/nil/src/server.rs

@@ -672,6 +672,7 @@ impl Server {
                 &flake_url,
                 Some(watcher_tx),
                 include_legacy,
+                config.nix_max_memory(),
             ));
             let ret = loop {
                 match tokio::time::timeout(PROGRESS_REPORT_PERIOD, eval_fut.as_mut()).await {
@@ -684,13 +685,13 @@ impl Server {
                 Ok(output) => output,
                 Err(err) => {
                     // Don't spam on configuration errors (eg. bad Nix path).
-                    if error_cnt == 0 {
+                    error_cnt += 1;
+                    if error_cnt <= 3 {
                         client.show_message_ext(
                             MessageType::ERROR,
                             format!("Flake input {input_name:?} cannot be evaluated: {err:#}"),
                         );
                     }
-                    error_cnt += 1;
                     continue;
                 }
             };

crates/nix-interop/Cargo.toml

@@ -11,7 +11,11 @@ serde = { version = "1.0.152", features = ["derive"] }
 serde_json = "1.0.91"
 serde_repr = "0.1.10"
 syntax = { path = "../syntax" }
+thiserror = "1.0.43"
 tokio = { version = "1.27.0", features = ["io-util", "macros", "process", "sync"] }
 
+[target.'cfg(unix)'.dependencies]
+rustix = { version = "0.38.3", default-features = false, features = ["process"] }
+
 [dev-dependencies]
 tokio = { version = "1.27.0", features = ["macros", "parking_lot", "rt", "sync"] } # parking_lot is required for `OnceCell::const_new`.

crates/nix-interop/src/flake_output.rs

@@ -8,15 +8,17 @@ use tokio::io::{AsyncBufReadExt, BufReader};
 use tokio::process::Command;
 use tokio::sync::watch;
 
-use crate::FlakeUrl;
+use crate::{FlakeUrl, NixOutOfMemory};
 
 pub async fn eval_flake_output(
     nix_command: &Path,
     flake_url: &FlakeUrl,
     watcher_tx: Option<watch::Sender<String>>,
     legacy: bool,
+    memory_limit: Option<u64>,
 ) -> Result<FlakeOutput> {
-    let mut child = Command::new(nix_command)
+    let mut command = Command::new(nix_command);
+    command
         .kill_on_drop(true)
         .args([
             "flake",
@@ -30,12 +32,34 @@ pub async fn eval_flake_output(
         .arg(flake_url)
         .stdin(Stdio::null())
         .stdout(Stdio::piped())
-        .stderr(Stdio::piped())
-        .spawn()
-        .context("Failed to spawn `nix`")?;
+        .stderr(Stdio::piped());
+
+    #[cfg(unix)]
+    unsafe {
+        if let Some(limit) = memory_limit {
+            use rustix::process::{setrlimit, Resource, Rlimit};
+            command.pre_exec(move || {
+                // NB. RSS limit has no effect on modern Linux. We set DATA limit instead.
+                setrlimit(
+                    Resource::Data,
+                    Rlimit {
+                        current: Some(limit),
+                        maximum: Some(limit),
+                    },
+                )
+                .map_err(|err| std::io::Error::from_raw_os_error(err.raw_os_error()))
+            });
+        }
+    }
+
+    #[cfg(not(unix))]
+    let _unused = memory_limit;
+
+    let mut child = command.spawn().context("Failed to spawn `nix`")?;
 
     let stderr = child.stderr.take().expect("Piped");
     let mut error_msg = String::new();
+    let mut oom = false;
     let consume_stderr_fut = async {
         let mut stderr = BufReader::new(stderr);
         let mut line = String::new();
@@ -43,7 +67,8 @@ pub async fn eval_flake_output(
             line.clear();
             matches!(stderr.read_line(&mut line).await, Ok(n) if n != 0)
         } {
-            if let Some(inner) = line.trim().strip_prefix("evaluating '") {
+            let line = line.trim();
+            if let Some(inner) = line.strip_prefix("evaluating '") {
                 if let Some(inner) = inner.strip_suffix("'...") {
                     if let Some(tx) = &watcher_tx {
                         tx.send_modify(|buf| {
@@ -53,13 +78,20 @@ pub async fn eval_flake_output(
                     }
                 }
             } else {
-                error_msg.push_str(&line);
+                if line == "error: out of memory" {
+                    oom = true;
+                }
+                error_msg.push_str(line);
             }
         }
     };
     let wait_fut = child.wait_with_output();
     let output = tokio::join!(consume_stderr_fut, wait_fut).1?;
 
+    if oom {
+        return Err(NixOutOfMemory.into());
+    }
+
     ensure!(
         output.status.success(),
         "`nix flake show {}` failed with {}. Stderr:\n{}",
@@ -123,7 +155,7 @@ mod tests {
     async fn eval_outputs() {
         let flake_url = FlakeUrl::new_path("./tests/test_flake");
         let (tx, rx) = watch::channel(String::new());
-        let output = eval_flake_output("nix".as_ref(), &flake_url, Some(tx), false)
+        let output = eval_flake_output("nix".as_ref(), &flake_url, Some(tx), false, None)
             .await
             .unwrap();
         // Even if the system is omitted, the attrpath is still printed in progress.
@@ -138,4 +170,16 @@ mod tests {
         assert_eq!(leaf.name.as_ref().unwrap(), "hello-1.2.3");
         assert_eq!(leaf.description.as_deref(), Some("A test derivation"));
     }
+
+    #[tokio::test]
+    #[ignore = "requires calling 'nix'"]
+    async fn memory_limit() {
+        let flake_url = FlakeUrl::new_path("./tests/oom_flake");
+        // 64MiB. This should be large enough to start the Nix evaluator itself without crash.
+        let limit = 64 << 20;
+        let err = eval_flake_output("nix".as_ref(), &flake_url, None, false, Some(limit))
+            .await
+            .unwrap_err();
+        assert!(err.is::<NixOutOfMemory>(), "expect OOM but got: {err}");
+    }
 }

crates/nix-interop/src/lib.rs

@@ -13,6 +13,10 @@ pub const DEFAULT_IMPORT_FILE: &str = "default.nix";
 pub const FLAKE_FILE: &str = "flake.nix";
 pub const FLAKE_LOCK_FILE: &str = "flake.lock";
 
+#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
+#[error("Nix exceeds memory limit")]
+pub struct NixOutOfMemory;
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct FlakeUrl(String);
 

crates/nix-interop/tests/oom_flake/flake.nix

@@ -0,0 +1,7 @@
+{
+  outputs = { self }:
+  # Allocates a list with ~1G elements.
+  # NB. This evaluates to false, so it will never be cached.
+  assert builtins.length (builtins.head (builtins.genList (x: x) 1000000000)) == 42;
+  { };
+}

dev/vim-coc.nix

@@ -100,6 +100,7 @@ let
     nil.formatting.command = [ "${pkgs.nixpkgs-fmt}/bin/nixpkgs-fmt" ];
     nil.diagnostics.excludedFiles = [ "generated.nix" ];
     nil.nix.flake.autoEvalInputs = true;
+    nil.nix.maxMemoryMB = 2048;
   };
 
 

docs/configuration.md

@@ -45,6 +45,14 @@ Default configuration:
       // Type: string
       // Example: "/run/current-system/sw/bin/nix"
       "binary": "nix",
+      // The virtual memory limit in MiB for `nix` evaluation.
+      // Currently it only applies to flake evaluation when `autoEvalInputs` is
+      // enabled, on *NIX platforms. Other `nix` invocations may be also
+      // applied in the future. `null` means no limit.
+      //
+      // Type: number | null
+      // Example: 1024
+      "maxMemoryMB": 2048,
       "flake": {
         // Auto-archiving behavior which may use network.
         //