Add cli interface to sandbox limit overrides

Author: Nemo157

Cargo.lock

@@ -107,6 +107,7 @@ chrono = { version = "0.4.11", default-features = false, features = ["clock", "s
 
 # Transitive dependencies we don't use directly but need to have specific versions of
 thread_local = "1.1.3"
+humantime = "2.1.0"
 
 [dependencies.postgres]
 version = "0.19"

Cargo.toml

@@ -7,7 +7,7 @@ use std::sync::Arc;
 use anyhow::{anyhow, Context as _, Error, Result};
 use clap::{Parser, Subcommand, ValueEnum};
 use docs_rs::cdn::CdnBackend;
-use docs_rs::db::{self, add_path_into_database, Pool, PoolClient};
+use docs_rs::db::{self, add_path_into_database, Overrides, Pool, PoolClient};
 use docs_rs::repositories::RepositoryStatsUpdater;
 use docs_rs::utils::{
     get_config, queue_builder, remove_crate_priority, set_crate_priority, ConfigName,
@@ -16,6 +16,7 @@ use docs_rs::{
     start_web_server, BuildQueue, Config, Context, Index, Metrics, PackageKind, RustwideBuilder,
     Storage,
 };
+use humantime::Duration;
 use once_cell::sync::OnceCell;
 use tokio::runtime::{Builder, Runtime};
 use tracing_log::LogTracer;
@@ -417,6 +418,12 @@ enum DatabaseSubcommand {
         command: BlacklistSubcommand,
     },
 
+    /// Limit overrides operations
+    Limits {
+        #[command(subcommand)]
+        command: LimitsSubcommand,
+    },
+
     /// Compares the database with the index and resolves inconsistencies
     #[cfg(feature = "consistency_check")]
     Synchronize {
@@ -473,6 +480,8 @@ impl DatabaseSubcommand {
             .context("failed to delete the crate")?,
             Self::Blacklist { command } => command.handle_args(ctx)?,
 
+            Self::Limits { command } => command.handle_args(ctx)?,
+
             #[cfg(feature = "consistency_check")]
             Self::Synchronize { dry_run } => {
                 docs_rs::utils::consistency::run_check(&ctx, dry_run)?;
@@ -482,6 +491,72 @@ impl DatabaseSubcommand {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Subcommand)]
+enum LimitsSubcommand {
+    /// Get sandbox limit overrides for a crate
+    Get { crate_name: String },
+
+    /// List sandbox limit overrides for all crates
+    List,
+
+    /// Set sandbox limits overrides for a crate
+    Set {
+        crate_name: String,
+        #[arg(long)]
+        memory: Option<usize>,
+        #[arg(long)]
+        targets: Option<usize>,
+        #[arg(long)]
+        timeout: Option<Duration>,
+    },
+
+    /// Remove sandbox limits overrides for a crate
+    Remove { crate_name: String },
+}
+
+impl LimitsSubcommand {
+    fn handle_args(self, ctx: BinContext) -> Result<()> {
+        let conn = &mut *ctx.conn()?;
+        match self {
+            Self::Get { crate_name } => {
+                let overrides = Overrides::for_crate(conn, &crate_name)?;
+                println!("sandbox limit overrides for {crate_name} = {overrides:?}");
+            }
+
+            Self::List => {
+                for (crate_name, overrides) in Overrides::all(conn)? {
+                    println!("sandbox limit overrides for {crate_name} = {overrides:?}");
+                }
+            }
+
+            Self::Set {
+                crate_name,
+                memory,
+                targets,
+                timeout,
+            } => {
+                let overrides = Overrides::for_crate(conn, &crate_name)?;
+                println!("previous sandbox limit overrides for {crate_name} = {overrides:?}");
+                let overrides = Overrides {
+                    memory,
+                    targets,
+                    timeout: timeout.map(Into::into),
+                };
+                Overrides::save(conn, &crate_name, overrides)?;
+                let overrides = Overrides::for_crate(conn, &crate_name)?;
+                println!("new sandbox limit overrides for {crate_name} = {overrides:?}");
+            }
+
+            Self::Remove { crate_name } => {
+                let overrides = Overrides::for_crate(conn, &crate_name)?;
+                println!("previous overrides for {crate_name} = {overrides:?}");
+                Overrides::remove(conn, &crate_name)?;
+            }
+        }
+        Ok(())
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Subcommand)]
 enum BlacklistSubcommand {
     /// List all crates on the blacklist

src/bin/cratesfyi.rs

@@ -1,18 +1,22 @@
 //! Database operations
 
-pub use self::add_package::update_crate_data_in_database;
 pub(crate) use self::add_package::{
     add_build_into_database, add_doc_coverage, add_package_into_database,
 };
-pub use self::delete::{delete_crate, delete_version};
-pub use self::file::{add_path_into_database, add_path_into_remote_archive};
-pub use self::migrate::migrate;
-pub use self::pool::{Pool, PoolClient, PoolError};
+pub use self::{
+    add_package::update_crate_data_in_database,
+    delete::{delete_crate, delete_version},
+    file::{add_path_into_database, add_path_into_remote_archive},
+    migrate::migrate,
+    overrides::Overrides,
+    pool::{Pool, PoolClient, PoolError},
+};
 
 mod add_package;
 pub mod blacklist;
 pub mod delete;
 pub(crate) mod file;
 mod migrate;
+mod overrides;
 mod pool;
 pub(crate) mod types;

src/db/mod.rs

@@ -0,0 +1,129 @@
+use crate::error::Result;
+use postgres::Client;
+use std::time::Duration;
+
+#[derive(Default, Debug, Clone, Copy, Eq, PartialEq)]
+pub struct Overrides {
+    pub memory: Option<usize>,
+    pub targets: Option<usize>,
+    pub timeout: Option<Duration>,
+}
+
+impl Overrides {
+    pub fn all(conn: &mut Client) -> Result<Vec<(String, Self)>> {
+        Ok(conn
+            .query("SELECT * FROM sandbox_overrides", &[])?
+            .into_iter()
+            .map(|row| (row.get("crate_name"), Self::from_row(row)))
+            .collect())
+    }
+
+    pub fn for_crate(conn: &mut Client, krate: &str) -> Result<Option<Self>> {
+        Ok(conn
+            .query_opt(
+                "SELECT * FROM sandbox_overrides WHERE crate_name = $1",
+                &[&krate],
+            )?
+            .map(Self::from_row))
+    }
+
+    fn from_row(row: postgres::Row) -> Self {
+        Self {
+            memory: row
+                .get::<_, Option<i64>>("max_memory_bytes")
+                .map(|i| i as usize),
+            targets: row.get::<_, Option<i32>>("max_targets").map(|i| i as usize),
+            timeout: row
+                .get::<_, Option<i32>>("timeout_seconds")
+                .map(|i| Duration::from_secs(i as u64)),
+        }
+    }
+
+    pub fn save(conn: &mut Client, krate: &str, overrides: Self) -> Result<()> {
+        if overrides.timeout.is_some() && overrides.targets.is_none() {
+            tracing::warn!("setting `Overrides::timeout` implies a default `Overrides::targets = 1`, prefer setting this explicitly");
+        }
+        conn.execute(
+            "
+                INSERT INTO sandbox_overrides (
+                    crate_name, max_memory_bytes, max_targets, timeout_seconds
+                )
+                VALUES ($1, $2, $3, $4)
+                ON CONFLICT (crate_name) DO UPDATE
+                    SET
+                        max_memory_bytes = $2,
+                        max_targets = $3,
+                        timeout_seconds = $4
+                ",
+            &[
+                &krate,
+                &overrides.memory.map(|i| i as i64),
+                &overrides.targets.map(|i| i as i32),
+                &overrides.timeout.map(|d| d.as_secs() as i32),
+            ],
+        )?;
+        Ok(())
+    }
+
+    pub fn remove(conn: &mut Client, krate: &str) -> Result<()> {
+        conn.execute(
+            "DELETE FROM sandbox_overrides WHERE crate_name = $1",
+            &[&krate],
+        )?;
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::{db::Overrides, test::*};
+    use std::time::Duration;
+
+    #[test]
+    fn retrieve_overrides() {
+        wrapper(|env| {
+            let db = env.db();
+
+            let krate = "hexponent";
+
+            // no overrides
+            let actual = Overrides::for_crate(&mut db.conn(), krate)?;
+            assert_eq!(actual, None);
+
+            // add partial overrides
+            let expected = Overrides {
+                targets: Some(1),
+                ..Overrides::default()
+            };
+            Overrides::save(&mut db.conn(), krate, expected)?;
+            let actual = Overrides::for_crate(&mut db.conn(), krate)?;
+            assert_eq!(actual, Some(expected));
+
+            // overwrite with full overrides
+            let expected = Overrides {
+                memory: Some(100_000),
+                targets: Some(1),
+                timeout: Some(Duration::from_secs(300)),
+            };
+            Overrides::save(&mut db.conn(), krate, expected)?;
+            let actual = Overrides::for_crate(&mut db.conn(), krate)?;
+            assert_eq!(actual, Some(expected));
+
+            // overwrite with partial overrides
+            let expected = Overrides {
+                memory: Some(1),
+                ..Overrides::default()
+            };
+            Overrides::save(&mut db.conn(), krate, expected)?;
+            let actual = Overrides::for_crate(&mut db.conn(), krate)?;
+            assert_eq!(actual, Some(expected));
+
+            // remove overrides
+            Overrides::remove(&mut db.conn(), krate)?;
+            let actual = Overrides::for_crate(&mut db.conn(), krate)?;
+            assert_eq!(actual, None);
+
+            Ok(())
+        });
+    }
+}

src/db/overrides.rs

@@ -1,3 +1,4 @@
+use crate::db::Overrides;
 use crate::error::Result;
 use postgres::Client;
 use serde::Serialize;
@@ -26,29 +27,18 @@ impl Default for Limits {
 
 impl Limits {
     pub(crate) fn for_crate(conn: &mut Client, name: &str) -> Result<Self> {
-        let mut limits = Self::default();
-
-        let res = conn.query(
-            "SELECT * FROM sandbox_overrides WHERE crate_name = $1;",
-            &[&name],
-        )?;
-        if !res.is_empty() {
-            let row = &res[0];
-            if let Some(memory) = row.get::<_, Option<i64>>("max_memory_bytes") {
-                limits.memory = memory as usize;
-            }
-            let timeout = row.get::<_, Option<i32>>("timeout_seconds");
-            if let Some(timeout) = timeout {
-                limits.timeout = Duration::from_secs(timeout as u64);
-            }
-            if let Some(targets) = row.get::<_, Option<i32>>("max_targets") {
-                limits.targets = targets as usize;
-            } else if timeout.is_some() {
-                limits.targets = 1;
-            }
-        }
-
-        Ok(limits)
+        let default = Self::default();
+        let overrides = Overrides::for_crate(conn, name)?.unwrap_or_default();
+        Ok(Self {
+            memory: overrides.memory.unwrap_or(default.memory),
+            targets: overrides
+                .targets
+                .or(overrides.timeout.map(|_| 1))
+                .unwrap_or(default.targets),
+            timeout: overrides.timeout.unwrap_or(default.timeout),
+            networking: default.networking,
+            max_log_size: default.max_log_size,
+        })
     }
 
     pub(crate) fn memory(&self) -> usize {
@@ -87,9 +77,13 @@ mod test {
             let hexponent = Limits::for_crate(&mut db.conn(), krate)?;
             assert_eq!(hexponent, Limits::default());
 
-            db.conn().query(
-                "INSERT INTO sandbox_overrides (crate_name, max_targets) VALUES ($1, 15)",
-                &[&krate],
+            Overrides::save(
+                &mut db.conn(),
+                krate,
+                Overrides {
+                    targets: Some(15),
+                    ..Overrides::default()
+                },
             )?;
             // limits work if crate has limits set
             let hexponent = Limits::for_crate(&mut db.conn(), krate)?;
@@ -109,10 +103,14 @@ mod test {
                 targets: 1,
                 ..Limits::default()
             };
-            db.conn().query(
-                "INSERT INTO sandbox_overrides (crate_name, max_memory_bytes, timeout_seconds, max_targets)
-                 VALUES ($1, $2, $3, $4)",
-                &[&krate, &(limits.memory as i64), &(limits.timeout.as_secs() as i32), &(limits.targets as i32)]
+            Overrides::save(
+                &mut db.conn(),
+                krate,
+                Overrides {
+                    memory: Some(limits.memory),
+                    targets: Some(limits.targets),
+                    timeout: Some(limits.timeout),
+                },
             )?;
             assert_eq!(limits, Limits::for_crate(&mut db.conn(), krate)?);
             Ok(())
@@ -124,9 +122,13 @@ mod test {
         wrapper(|env| {
             let db = env.db();
             let krate = "hexponent";
-            db.conn().query(
-                "INSERT INTO sandbox_overrides (crate_name, timeout_seconds) VALUES ($1, 20*60);",
-                &[&krate],
+            Overrides::save(
+                &mut db.conn(),
+                krate,
+                Overrides {
+                    timeout: Some(Duration::from_secs(20 * 60)),
+                    ..Overrides::default()
+                },
             )?;
             let limits = Limits::for_crate(&mut db.conn(), krate)?;
             assert_eq!(limits.targets, 1);