From 03a2d8427a83fdc893a965cc3b5c5a59b5c0509d Mon Sep 17 00:00:00 2001
From: Graham Christensen <graham@grahamc.com>
Date: Thu, 8 Feb 2018 10:06:08 -0500
Subject: [PATCH] Support specifying initial heap size for evaluations

---
 ofborg/src/config.rs              |  2 ++
 ofborg/src/nix.rs                 | 60 ++++++++++++++++++++-----------
 ofborg/src/tasks/build.rs         |  2 +-
 ofborg/src/tasks/massrebuilder.rs |  2 +-
 4 files changed, 44 insertions(+), 22 deletions(-)

diff --git a/ofborg/src/config.rs b/ofborg/src/config.rs
index 5381b81..0928f7e 100644
--- a/ofborg/src/config.rs
+++ b/ofborg/src/config.rs
@@ -42,6 +42,7 @@ pub struct NixConfig {
     pub system: String,
     pub remote: String,
     pub build_timeout_seconds: u16,
+    pub initial_heap_size: Option<String>
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -105,6 +106,7 @@ impl Config {
             self.nix.system.clone(),
             self.nix.remote.clone(),
             self.nix.build_timeout_seconds,
+            self.nix.initial_heap_size.clone(),
         );
     }
 }
diff --git a/ofborg/src/nix.rs b/ofborg/src/nix.rs
index 8b590b7..4ad4958 100644
--- a/ofborg/src/nix.rs
+++ b/ofborg/src/nix.rs
@@ -13,43 +13,36 @@ pub struct Nix {
     remote: String,
     build_timeout: u16,
     limit_supported_systems: bool,
+    initial_heap_size: Option<String>,
 }
 
 impl Nix {
-    pub fn new(system: String, remote: String, build_timeout: u16) -> Nix {
+    pub fn new(system: String, remote: String, build_timeout: u16, initial_heap_size: Option<String>) -> Nix {
         return Nix {
             system: system,
             remote: remote,
             build_timeout: build_timeout,
+            initial_heap_size: initial_heap_size,
             limit_supported_systems: true,
         };
     }
 
     pub fn with_system(&self, system: String) -> Nix {
-        return Nix {
-            system: system,
-            remote: self.remote.clone(),
-            build_timeout: self.build_timeout,
-            limit_supported_systems: self.limit_supported_systems,
-        };
+        let mut n = self.clone();
+        n.system = system;
+        return n;
     }
 
     pub fn with_limited_supported_systems(&self) -> Nix {
-        return Nix {
-            system: self.system.clone(),
-            remote: self.remote.clone(),
-            build_timeout: self.build_timeout,
-            limit_supported_systems: true,
-        };
+        let mut n = self.clone();
+        n.limit_supported_systems = true;
+        return n;
     }
 
     pub fn without_limited_supported_systems(&self) -> Nix {
-        return Nix {
-            system: self.system.clone(),
-            remote: self.remote.clone(),
-            build_timeout: self.build_timeout,
-            limit_supported_systems: false,
-        };
+        let mut n = self.clone();
+        n.limit_supported_systems = false;
+        return n;
     }
 
     pub fn safely_build_attrs(
@@ -132,6 +125,10 @@ impl Nix {
         command.env("NIX_PATH", nixpath);
         command.env("NIX_REMOTE", &self.remote);
 
+        if let Some(ref initial_heap_size) = self.initial_heap_size {
+            command.env("GC_INITIAL_HEAP_SIZE", &initial_heap_size);
+        }
+
         let path = env::var("PATH").unwrap();
         command.env("PATH", path);
 
@@ -165,7 +162,7 @@ impl Nix {
 #[cfg(test)]
 mod tests {
     fn nix() -> Nix {
-        Nix::new("x86_64-linux".to_owned(), "daemon".to_owned(), 1800)
+        Nix::new("x86_64-linux".to_owned(), "daemon".to_owned(), 1800, None)
     }
 
     fn build_path() -> PathBuf {
@@ -295,6 +292,29 @@ mod tests {
         );
     }
 
+    #[test]
+    fn safe_command_custom_gc() {
+        let nix = Nix::new("x86_64-linux".to_owned(), "daemon".to_owned(), 1800, Some("4g".to_owned()));
+
+        let ret: Result<File, File> =
+            nix.run(
+                nix.safe_command("./environment.sh", build_path().as_path(), vec![]),
+                true,
+            );
+
+        assert_run(
+            ret,
+            Expect::Pass,
+            vec![
+                "HOME=/homeless-shelter",
+                "NIX_PATH=nixpkgs=",
+                "NIX_REMOTE=",
+                "PATH=",
+                "GC_INITIAL_HEAP_SIZE=4g",
+            ],
+        );
+    }
+
     #[test]
     fn safe_command_options() {
         let nix = nix();
diff --git a/ofborg/src/tasks/build.rs b/ofborg/src/tasks/build.rs
index dd28b87..09bf604 100644
--- a/ofborg/src/tasks/build.rs
+++ b/ofborg/src/tasks/build.rs
@@ -353,7 +353,7 @@ mod tests {
     use ofborg::test_scratch::TestScratch;
 
     fn nix() -> nix::Nix {
-        nix::Nix::new("x86_64-linux".to_owned(), "daemon".to_owned(), 1800)
+        nix::Nix::new("x86_64-linux".to_owned(), "daemon".to_owned(), 1800, None)
     }
 
     fn tpath(component: &str) -> PathBuf {
diff --git a/ofborg/src/tasks/massrebuilder.rs b/ofborg/src/tasks/massrebuilder.rs
index 0032e7f..3f7f18b 100644
--- a/ofborg/src/tasks/massrebuilder.rs
+++ b/ofborg/src/tasks/massrebuilder.rs
@@ -717,7 +717,7 @@ mod tests {
 
     #[test]
     fn stdenv_checking() {
-        let nix = Nix::new(String::from("x86_64-linux"), String::from("daemon"), 1200);
+        let nix = Nix::new(String::from("x86_64-linux"), String::from("daemon"), 1200, None);
         let mut stdenv =
             Stdenvs::new(
                 nix.clone(),