Add a config option to set the default build memory limit

Nemo157 · syphar · commit 3277ddd3d87a · 2023-09-01T05:39:04.000+02:00
diff --git a/src/config.rs b/src/config.rs
@@ -101,6 +101,7 @@ pub struct Config {
     pub(crate) docker_image: Option<String>,
     pub(crate) toolchain: String,
     pub(crate) build_cpu_limit: Option<u32>,
+    pub(crate) build_default_memory_limit: Option<usize>,
     pub(crate) include_default_targets: bool,
     pub(crate) disable_memory_limit: bool,
 }
@@ -203,6 +204,7 @@ impl Config {
                 .or(maybe_env("DOCSRS_DOCKER_IMAGE")?),
             toolchain: env("DOCSRS_TOOLCHAIN", "nightly".to_string())?,
             build_cpu_limit: maybe_env("DOCSRS_BUILD_CPU_LIMIT")?,
+            build_default_memory_limit: maybe_env("DOCSRS_BUILD_DEFAULT_MEMORY_LIMIT")?,
             include_default_targets: env("DOCSRS_INCLUDE_DEFAULT_TARGETS", true)?,
             disable_memory_limit: env("DOCSRS_DISABLE_MEMORY_LIMIT", false)?,
         })
diff --git a/src/docbuilder/limits.rs b/src/docbuilder/limits.rs
@@ -1,9 +1,10 @@
-use crate::db::Overrides;
-use crate::error::Result;
+use crate::{db::Overrides, error::Result, Config};
 use postgres::Client;
 use serde::Serialize;
 use std::time::Duration;
 
+const GB: usize = 1024 * 1024 * 1024;
+
 #[derive(Debug, Clone, PartialEq, Eq, Serialize)]
 pub(crate) struct Limits {
     memory: usize,
@@ -13,21 +14,20 @@ pub(crate) struct Limits {
     max_log_size: usize,
 }
 
-impl Default for Limits {
-    fn default() -> Self {
+impl Limits {
+    pub(crate) fn new(config: &Config) -> Self {
         Self {
-            memory: 3 * 1024 * 1024 * 1024,        // 3 GB
+            // 3 GB default default
+            memory: config.build_default_memory_limit.unwrap_or(3 * GB),
             timeout: Duration::from_secs(15 * 60), // 15 minutes
             targets: 10,
             networking: false,
             max_log_size: 100 * 1024, // 100 KB
         }
     }
-}
 
-impl Limits {
-    pub(crate) fn for_crate(conn: &mut Client, name: &str) -> Result<Self> {
-        let default = Self::default();
+    pub(crate) fn for_crate(config: &Config, conn: &mut Client, name: &str) -> Result<Self> {
+        let default = Self::new(config);
         let overrides = Overrides::for_crate(conn, name)?.unwrap_or_default();
         Ok(Self {
             memory: overrides.memory.unwrap_or(default.memory),
@@ -74,8 +74,8 @@ mod test {
 
             let krate = "hexponent";
             // limits work if no crate has limits set
-            let hexponent = Limits::for_crate(&mut db.conn(), krate)?;
-            assert_eq!(hexponent, Limits::default());
+            let hexponent = Limits::for_crate(&env.config(), &mut db.conn(), krate)?;
+            assert_eq!(hexponent, Limits::new(&env.config()));
 
             Overrides::save(
                 &mut db.conn(),
@@ -86,12 +86,12 @@ mod test {
                 },
             )?;
             // limits work if crate has limits set
-            let hexponent = Limits::for_crate(&mut db.conn(), krate)?;
+            let hexponent = Limits::for_crate(&env.config(), &mut db.conn(), krate)?;
             assert_eq!(
                 hexponent,
                 Limits {
                     targets: 15,
-                    ..Limits::default()
+                    ..Limits::new(&env.config())
                 }
             );
 
@@ -101,7 +101,7 @@ mod test {
                 memory: 100_000,
                 timeout: Duration::from_secs(300),
                 targets: 1,
-                ..Limits::default()
+                ..Limits::new(&env.config())
             };
             Overrides::save(
                 &mut db.conn(),
@@ -112,7 +112,10 @@ mod test {
                     timeout: Some(limits.timeout),
                 },
             )?;
-            assert_eq!(limits, Limits::for_crate(&mut db.conn(), krate)?);
+            assert_eq!(
+                limits,
+                Limits::for_crate(&env.config(), &mut db.conn(), krate)?
+            );
             Ok(())
         });
     }
@@ -130,10 +133,26 @@ mod test {
                     ..Overrides::default()
                 },
             )?;
-            let limits = Limits::for_crate(&mut db.conn(), krate)?;
+            let limits = Limits::for_crate(&env.config(), &mut db.conn(), krate)?;
             assert_eq!(limits.targets, 1);
 
             Ok(())
         });
     }
+
+    #[test]
+    fn config_default_memory_limit() {
+        wrapper(|env| {
+            env.override_config(|config| {
+                config.build_default_memory_limit = Some(6 * GB);
+            });
+
+            let db = env.db();
+
+            let limits = Limits::for_crate(&env.config(), &mut db.conn(), "krate")?;
+            assert_eq!(limits.memory, 6 * GB);
+
+            Ok(())
+        });
+    }
 }
diff --git a/src/docbuilder/rustwide_builder.rs b/src/docbuilder/rustwide_builder.rs
@@ -228,7 +228,7 @@ impl RustwideBuilder {
         info!("building a dummy crate to get essential files");
 
         let mut conn = self.db.get()?;
-        let limits = Limits::for_crate(&mut conn, DUMMY_CRATE_NAME)?;
+        let limits = Limits::for_crate(&self.config, &mut conn, DUMMY_CRATE_NAME)?;
 
         let mut build_dir = self
             .workspace
@@ -342,7 +342,7 @@ impl RustwideBuilder {
             return Ok(false);
         }
 
-        let limits = Limits::for_crate(&mut conn, name)?;
+        let limits = Limits::for_crate(&self.config, &mut conn, name)?;
         #[cfg(target_os = "linux")]
         if !self.config.disable_memory_limit {
             use anyhow::Context;
diff --git a/src/web/builds.rs b/src/web/builds.rs
@@ -5,6 +5,7 @@ use crate::{
     impl_axum_webpage,
     utils::spawn_blocking,
     web::{error::AxumResult, match_version_axum, MetaData},
+    Config,
 };
 use anyhow::Result;
 use axum::{
@@ -15,6 +16,7 @@ use axum::{
 };
 use chrono::{DateTime, Utc};
 use serde::Serialize;
+use std::sync::Arc;
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize)]
 pub(crate) struct Build {
@@ -40,6 +42,7 @@ impl_axum_webpage! {
 pub(crate) async fn build_list_handler(
     Path((name, req_version)): Path<(String, String)>,
     Extension(pool): Extension<Pool>,
+    Extension(config): Extension<Arc<Config>>,
 ) -> AxumResult<impl IntoResponse> {
     let (version, version_or_latest) = match match_version_axum(&pool, &name, Some(&req_version))
         .await?
@@ -62,7 +65,7 @@ pub(crate) async fn build_list_handler(
         move || {
             let mut conn = pool.get()?;
             Ok((
-                Limits::for_crate(&mut conn, &name)?,
+                Limits::for_crate(&config, &mut conn, &name)?,
                 get_builds(&mut conn, &name, &version)?,
                 MetaData::from_crate(&mut conn, &name, &version, &version_or_latest)?,
             ))
diff --git a/src/web/sitemap.rs b/src/web/sitemap.rs
@@ -7,6 +7,7 @@ use crate::{
         error::{AxumNope, AxumResult},
         AxumErrorPage,
     },
+    Config,
 };
 use axum::{
     extract::{Extension, Path},
@@ -15,6 +16,7 @@ use axum::{
 };
 use chrono::{DateTime, TimeZone, Utc};
 use serde::Serialize;
+use std::sync::Arc;
 
 /// sitemap index
 #[derive(Debug, Clone, PartialEq, Eq, Serialize)]
@@ -70,9 +72,9 @@ pub(crate) async fn sitemap_handler(
                     MAX(releases.release_time) as release_time
              FROM crates
              INNER JOIN releases ON releases.crate_id = crates.id
-             WHERE 
-                rustdoc_status = true AND 
-                crates.name ILIKE $1 
+             WHERE
+                rustdoc_status = true AND
+                crates.name ILIKE $1
              GROUP BY crates.name, releases.target_name
              ",
             &[&format!("{letter}%")],
@@ -112,6 +114,7 @@ impl_axum_webpage!(AboutBuilds = "core/about/builds.html");
 
 pub(crate) async fn about_builds_handler(
     Extension(pool): Extension<Pool>,
+    Extension(config): Extension<Arc<Config>>,
 ) -> AxumResult<impl IntoResponse> {
     let rustc_version = spawn_blocking(move || {
         let mut conn = pool.get()?;
@@ -121,7 +124,7 @@ pub(crate) async fn about_builds_handler(
 
     Ok(AboutBuilds {
         rustc_version,
-        limits: Limits::default(),
+        limits: Limits::new(&config),
         active_tab: "builds",
     })
 }
