feat(single-node): add memory allocation (#19895)

Author: fuyufjh

src/cmd_all/src/single_node.rs

@@ -15,7 +15,9 @@
 use clap::Parser;
 use home::home_dir;
 use risingwave_common::config::MetaBackend;
+use risingwave_common::util::resource_util::memory::system_memory_available_bytes;
 use risingwave_compactor::CompactorOpts;
+use risingwave_compute::memory::config::gradient_reserve_memory_bytes;
 use risingwave_compute::ComputeNodeOpts;
 use risingwave_frontend::FrontendOpts;
 use risingwave_meta_node::MetaNodeOpts;
@@ -85,7 +87,7 @@ impl SingleNodeOpts {
 #[derive(Eq, PartialOrd, PartialEq, Debug, Clone, Parser)]
 pub struct NodeSpecificOpts {
     // ------- Compute Node Options -------
-    /// Total available memory for the compute node in bytes. Used by both computing and storage.
+    /// Total available memory for all the nodes
     #[clap(long)]
     pub total_memory_bytes: Option<usize>,
 
@@ -207,10 +209,21 @@ pub fn map_single_node_opts_to_standalone_opts(opts: SingleNodeOpts) -> ParsedSt
     frontend_opts.meta_addr = meta_addr.parse().unwrap();
     compactor_opts.meta_address = meta_addr.parse().unwrap();
 
+    // Allocate memory for each node
+    let total_memory_bytes = if let Some(total_memory_bytes) = opts.node_opts.total_memory_bytes {
+        total_memory_bytes
+    } else {
+        system_memory_available_bytes()
+    };
+    frontend_opts.frontend_total_memory_bytes = memory_for_frontend(total_memory_bytes);
+    compactor_opts.compactor_total_memory_bytes = memory_for_compactor(total_memory_bytes);
+    compute_opts.total_memory_bytes = total_memory_bytes
+        - memory_for_frontend(total_memory_bytes)
+        - memory_for_compactor(total_memory_bytes);
+    compute_opts.memory_manager_target_bytes =
+        Some(gradient_reserve_memory_bytes(total_memory_bytes));
+
     // Apply node-specific options
-    if let Some(total_memory_bytes) = opts.node_opts.total_memory_bytes {
-        compute_opts.total_memory_bytes = total_memory_bytes;
-    }
     if let Some(parallelism) = opts.node_opts.parallelism {
         compute_opts.parallelism = parallelism;
     }
@@ -234,3 +247,19 @@ pub fn map_single_node_opts_to_standalone_opts(opts: SingleNodeOpts) -> ParsedSt
         compactor_opts: Some(compactor_opts),
     }
 }
+
+fn memory_for_frontend(total_memory_bytes: usize) -> usize {
+    if total_memory_bytes <= (16 << 30) {
+        total_memory_bytes / 8
+    } else {
+        (total_memory_bytes - (16 << 30)) / 16 + (16 << 30) / 8
+    }
+}
+
+fn memory_for_compactor(total_memory_bytes: usize) -> usize {
+    if total_memory_bytes <= (16 << 30) {
+        total_memory_bytes / 8
+    } else {
+        (total_memory_bytes - (16 << 30)) / 16 + (16 << 30) / 8
+    }
+}

src/compute/src/lib.rs

@@ -87,7 +87,8 @@ pub struct ComputeNodeOpts {
     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.
+    /// If not set, a portion (default to 30% for the first 16GB and 20% for the rest)
+    /// 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")]
@@ -97,6 +98,13 @@ pub struct ComputeNodeOpts {
     /// If not set, the default value is `total_memory_bytes` - `reserved_memory_bytes`
     ///
     /// It's strongly recommended to set it for standalone deployment.
+    ///
+    /// ## Why need this?
+    ///
+    /// Our [`crate::memory::manager::MemoryManager`] works by reading the memory statistics from
+    /// Jemalloc. This is fine when running the compute node alone; while for standalone mode,
+    /// the memory usage of **all nodes** are counted. Thus, we need to pass a reasonable total
+    /// usage so that the memory is kept around this value.
     #[clap(long, env = "RW_MEMORY_MANAGER_TARGET_BYTES")]
     pub memory_manager_target_bytes: Option<usize>,
 

src/compute/src/memory/config.rs

@@ -80,7 +80,7 @@ pub fn reserve_memory_bytes(opts: &ComputeNodeOpts) -> (usize, usize) {
 /// The reserved memory size is calculated based on the following gradient:
 /// - 30% of the first 16GB
 /// - 20% of the rest
-fn gradient_reserve_memory_bytes(total_memory_bytes: usize) -> usize {
+pub fn gradient_reserve_memory_bytes(total_memory_bytes: usize) -> usize {
     let mut total_memory_bytes = total_memory_bytes;
     let mut reserved = 0;
     for i in 0..RESERVED_MEMORY_LEVELS.len() {