Merge pull request #3135 from fibonacci1729/limit-mem-usage

Add support for limiting instance memory usage in spin up

Author: fibonacci1729

crates/runtime-config/src/lib.rs

@@ -45,6 +45,8 @@ pub struct ResolvedRuntimeConfig<T> {
     ///
     /// `None` is used for an "unset" log directory.
     pub log_dir: Option<PathBuf>,
+    /// The maximum memory allocation limit.
+    pub max_instance_memory: Option<usize>,
     /// The input TOML, for informational summaries.
     pub toml: toml::Table,
 }
@@ -137,12 +139,15 @@ where
 
         let toml = toml_resolver.toml();
         let log_dir = toml_resolver.log_dir()?;
+        let max_instance_memory = toml_resolver.max_instance_memory()?;
+
         let source = TomlRuntimeConfigSource::new(
             toml_resolver,
             &key_value_resolver,
             tls_resolver.as_ref(),
             &sqlite_resolver,
         );
+
         // Note: all valid fields in the runtime config must have been referenced at
         // this point or the finalizer will fail due to `validate_all_keys_used`
         // not passing.
@@ -154,6 +159,7 @@ where
             sqlite_resolver,
             state_dir,
             log_dir,
+            max_instance_memory,
             toml,
         })
     }
@@ -167,6 +173,11 @@ where
     pub fn log_dir(&self) -> Option<PathBuf> {
         self.log_dir.clone()
     }
+
+    /// The maximum memory allocation limit.
+    pub fn max_instance_memory(&self) -> Option<usize> {
+        self.max_instance_memory
+    }
 }
 
 #[derive(Clone, Debug)]
@@ -261,6 +272,16 @@ impl<'a> TomlResolver<'a> {
         }
     }
 
+    /// Get the configured maximum memory allocation limit.
+    pub fn max_instance_memory(&self) -> anyhow::Result<Option<usize>> {
+        self.table
+            .get("max_instance_memory")
+            .and_then(|v| v.as_integer())
+            .map(|toml_value| toml_value.try_into())
+            .transpose()
+            .map_err(Into::into)
+    }
+
     /// Validate that all keys in the TOML file have been used.
     pub fn validate_all_keys_used(&self) -> spin_factors::Result<()> {
         self.table.validate_all_keys_used()

crates/runtime-factors/src/build.rs

@@ -6,8 +6,9 @@ use anyhow::Context as _;
 use spin_factors_executor::FactorsExecutor;
 use spin_runtime_config::ResolvedRuntimeConfig;
 use spin_trigger::cli::{
-    FactorsConfig, InitialKvSetterHook, KeyValueDefaultStoreSummaryHook, RuntimeFactorsBuilder,
-    SqlStatementExecutorHook, SqliteDefaultStoreSummaryHook, StdioLoggingExecutorHooks,
+    FactorsConfig, InitialKvSetterHook, KeyValueDefaultStoreSummaryHook, MaxInstanceMemoryHook,
+    RuntimeFactorsBuilder, SqlStatementExecutorHook, SqliteDefaultStoreSummaryHook,
+    StdioLoggingExecutorHooks,
 };
 
 /// A [`RuntimeFactorsBuilder`] for [`TriggerFactors`].
@@ -56,6 +57,16 @@ impl RuntimeFactorsBuilder for FactorsBuilder {
         executor.add_hooks(InitialKvSetterHook::new(args.key_values.clone()));
         executor.add_hooks(SqliteDefaultStoreSummaryHook);
         executor.add_hooks(KeyValueDefaultStoreSummaryHook);
+
+        let max_instance_memory = args
+            .max_instance_memory
+            .or(runtime_config.max_instance_memory());
+
+        // Only add the hook if a max instance memory size is specified via flag or runtime config.
+        if let Some(max_instance_memory) = max_instance_memory {
+            executor.add_hooks(MaxInstanceMemoryHook::new(max_instance_memory));
+        }
+
         Ok(())
     }
 }

crates/runtime-factors/src/lib.rs

@@ -100,6 +100,10 @@ pub struct TriggerAppArgs {
     /// To run from a file, prefix the filename with @ e.g. spin up --sqlite @migration.sql
     #[clap(long = "sqlite")]
     pub sqlite_statements: Vec<String>,
+
+    /// Sets the maxmimum memory allocation limit for an instance in bytes.
+    #[clap(long, env = "SPIN_MAX_INSTANCE_MEMORY")]
+    pub max_instance_memory: Option<usize>,
 }
 
 impl From<ResolvedRuntimeConfig<TriggerFactorsRuntimeConfig>> for TriggerFactorsRuntimeConfig {

crates/trigger/src/cli.rs

@@ -1,5 +1,6 @@
 mod initial_kv_setter;
 mod launch_metadata;
+mod max_instance_memory;
 mod sqlite_statements;
 mod stdio;
 mod summary;
@@ -19,6 +20,7 @@ use spin_factors_executor::{ComponentLoader, FactorsExecutor};
 use crate::{loader::ComponentLoader as ComponentLoaderImpl, Trigger, TriggerApp};
 pub use initial_kv_setter::InitialKvSetterHook;
 pub use launch_metadata::LaunchMetadata;
+pub use max_instance_memory::MaxInstanceMemoryHook;
 pub use sqlite_statements::SqlStatementExecutorHook;
 use stdio::FollowComponents;
 pub use stdio::StdioLoggingExecutorHooks;

crates/trigger/src/cli/max_instance_memory.rs

@@ -0,0 +1,26 @@
+use spin_core::async_trait;
+use spin_factors::RuntimeFactors;
+use spin_factors_executor::{ExecutorHooks, FactorsInstanceBuilder};
+
+/// An [`ExecutorHooks`] that sets the maximum memory allocation limit.
+pub struct MaxInstanceMemoryHook {
+    max_instance_memory: usize,
+}
+
+impl MaxInstanceMemoryHook {
+    pub fn new(max_instance_memory: usize) -> Self {
+        Self {
+            max_instance_memory,
+        }
+    }
+}
+
+#[async_trait]
+impl<F: RuntimeFactors, U> ExecutorHooks<F, U> for MaxInstanceMemoryHook {
+    fn prepare_instance(&self, builder: &mut FactorsInstanceBuilder<F, U>) -> anyhow::Result<()> {
+        builder
+            .store_builder()
+            .max_memory_size(self.max_instance_memory);
+        Ok(())
+    }
+}