feat(memory): add reserved_memory_bytes opt (#16433)

src/cmd_all/src/standalone.rs

@@ -359,6 +359,7 @@ mod test {
                             connector_rpc_sink_payload_format: None,
                             config_path: "src/config/test.toml",
                             total_memory_bytes: 34359738368,
+                            reserved_memory_bytes: None,
                             parallelism: 10,
                             role: Both,
                             metrics_level: None,

src/compute/src/lib.rs

@@ -90,6 +90,13 @@ pub struct ComputeNodeOpts {
     #[clap(long, env = "RW_TOTAL_MEMORY_BYTES", default_value_t = default_total_memory_bytes())]
     pub total_memory_bytes: usize,
 
+    /// Reserved memory for the compute node in bytes.
+    /// If not set, a portion (default to 30%) for the total_memory_bytes will be used as the reserved memory.
+    ///
+    /// The total memory compute and storage can use is `total_memory_bytes` - `reserved_memory_bytes`.
+    #[clap(long, env = "RW_RESERVED_MEMORY_BYTES")]
+    pub reserved_memory_bytes: Option<usize>,
+
     /// The parallelism that the compute node will register to the scheduler of the meta service.
     #[clap(long, env = "RW_PARALLELISM", default_value_t = default_parallelism())]
     #[override_opts(if_absent, path = streaming.actor_runtime_worker_threads_num)]

src/compute/src/memory/config.rs

@@ -19,6 +19,8 @@ use risingwave_common::config::{
 };
 use risingwave_common::util::pretty_bytes::convert;
 
+use crate::ComputeNodeOpts;
+
 /// The minimal memory requirement of computing tasks in megabytes.
 pub const MIN_COMPUTE_MEMORY_MB: usize = 512;
 /// The memory reserved for system usage (stack and code segment of processes, allocation
@@ -41,20 +43,24 @@ const STORAGE_SHARED_BUFFER_MEMORY_PROPORTION: f64 = 0.3;
 /// Each compute node reserves some memory for stack and code segment of processes, allocation
 /// overhead, network buffer, etc. based on `SYSTEM_RESERVED_MEMORY_PROPORTION`. The reserve memory
 /// size must be larger than `MIN_SYSTEM_RESERVED_MEMORY_MB`
-pub fn reserve_memory_bytes(total_memory_bytes: usize) -> (usize, usize) {
-    if total_memory_bytes < MIN_COMPUTE_MEMORY_MB << 20 {
+pub fn reserve_memory_bytes(opts: &ComputeNodeOpts) -> (usize, usize) {
+    if opts.total_memory_bytes < MIN_COMPUTE_MEMORY_MB << 20 {
         panic!(
             "The total memory size ({}) is too small. It must be at least {} MB.",
-            convert(total_memory_bytes as _),
+            convert(opts.total_memory_bytes as _),
             MIN_COMPUTE_MEMORY_MB
         );
     }
 
-    let reserved = std::cmp::max(
-        (total_memory_bytes as f64 * SYSTEM_RESERVED_MEMORY_PROPORTION).ceil() as usize,
-        MIN_SYSTEM_RESERVED_MEMORY_MB << 20,
-    );
-    (reserved, total_memory_bytes - reserved)
+    // If `reserved_memory_bytes` is not set, use `SYSTEM_RESERVED_MEMORY_PROPORTION` * `total_memory_bytes`.
+    let reserved = opts.reserved_memory_bytes.unwrap_or_else(|| {
+        (opts.total_memory_bytes as f64 * SYSTEM_RESERVED_MEMORY_PROPORTION).ceil() as usize
+    });
+
+    // Should have at least `MIN_SYSTEM_RESERVED_MEMORY_MB` for reserved memory.
+    let reserved = std::cmp::max(reserved, MIN_SYSTEM_RESERVED_MEMORY_MB << 20);
+
+    (reserved, opts.total_memory_bytes - reserved)
 }
 
 /// Decide the memory limit for each storage cache. If not specified in `StorageConfig`, memory
@@ -230,21 +236,41 @@ pub fn storage_memory_config(
 
 #[cfg(test)]
 mod tests {
+    use clap::Parser;
     use risingwave_common::config::StorageConfig;
 
     use super::{reserve_memory_bytes, storage_memory_config};
+    use crate::ComputeNodeOpts;
 
     #[test]
     fn test_reserve_memory_bytes() {
         // at least 512 MB
-        let (reserved, non_reserved) = reserve_memory_bytes(1536 << 20);
-        assert_eq!(reserved, 512 << 20);
-        assert_eq!(non_reserved, 1024 << 20);
+        {
+            let mut opts = ComputeNodeOpts::parse_from(vec![] as Vec<String>);
+            opts.total_memory_bytes = 1536 << 20;
+            let (reserved, non_reserved) = reserve_memory_bytes(&opts);
+            assert_eq!(reserved, 512 << 20);
+            assert_eq!(non_reserved, 1024 << 20);
+        }
 
         // reserve based on proportion
-        let (reserved, non_reserved) = reserve_memory_bytes(10 << 30);
-        assert_eq!(reserved, 3 << 30);
-        assert_eq!(non_reserved, 7 << 30);
+        {
+            let mut opts = ComputeNodeOpts::parse_from(vec![] as Vec<String>);
+            opts.total_memory_bytes = 10 << 30;
+            let (reserved, non_reserved) = reserve_memory_bytes(&opts);
+            assert_eq!(reserved, 3 << 30);
+            assert_eq!(non_reserved, 7 << 30);
+        }
+
+        // reserve based on opts
+        {
+            let mut opts = ComputeNodeOpts::parse_from(vec![] as Vec<String>);
+            opts.total_memory_bytes = 10 << 30;
+            opts.reserved_memory_bytes = Some(2 << 30);
+            let (reserved, non_reserved) = reserve_memory_bytes(&opts);
+            assert_eq!(reserved, 2 << 30);
+            assert_eq!(non_reserved, 8 << 30);
+        }
     }
 
     #[test]

src/compute/src/server.rs

@@ -103,7 +103,7 @@ pub async fn compute_node_serve(
 
     // Register to the cluster. We're not ready to serve until activate is called.
     let (meta_client, system_params) = MetaClient::register_new(
-        opts.meta_address,
+        opts.meta_address.clone(),
         WorkerType::ComputeNode,
         &advertise_addr,
         Property {
@@ -122,8 +122,7 @@ pub async fn compute_node_serve(
     let embedded_compactor_enabled =
         embedded_compactor_enabled(state_store_url, config.storage.disable_remote_compactor);
 
-    let (reserved_memory_bytes, non_reserved_memory_bytes) =
-        reserve_memory_bytes(opts.total_memory_bytes);
+    let (reserved_memory_bytes, non_reserved_memory_bytes) = reserve_memory_bytes(&opts);
     let storage_memory_config = storage_memory_config(
         non_reserved_memory_bytes,
         embedded_compactor_enabled,