From 20ddc77133938f7bf863b2c976e43807c02efd6d Mon Sep 17 00:00:00 2001
From: Eric Fu <eric@singularity-data.com>
Date: Wed, 20 Nov 2024 14:40:16 +0800
Subject: [PATCH] allow override memory_manager_target_bytes

---
 src/cmd_all/src/standalone.rs |  1 +
 src/compute/src/lib.rs        |  7 +++++++
 src/compute/src/server.rs     | 15 +++++++--------
 3 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/src/cmd_all/src/standalone.rs b/src/cmd_all/src/standalone.rs
index 6eb62492999a0..34d9749ff0e7f 100644
--- a/src/cmd_all/src/standalone.rs
+++ b/src/cmd_all/src/standalone.rs
@@ -509,6 +509,7 @@ mod test {
                             enable_barrier_read: None,
                             temp_secret_file_dir: "./frontend/secrets/",
                             frontend_total_memory_bytes: 34359738368,
+                            memory_manager_target_bytes: None,
                         },
                     ),
                     compactor_opts: None,
diff --git a/src/compute/src/lib.rs b/src/compute/src/lib.rs
index 1336a84980cea..ef4b5c5e32d3b 100644
--- a/src/compute/src/lib.rs
+++ b/src/compute/src/lib.rs
@@ -92,6 +92,13 @@ pub struct ComputeNodeOpts {
     #[clap(long, env = "RW_RESERVED_MEMORY_BYTES")]
     pub reserved_memory_bytes: Option<usize>,
 
+    /// Target memory usage for Memory Manager.
+    /// If not set, the default value is `total_memory_bytes` - `reserved_memory_bytes`
+    ///
+    /// It's strongly recommended to set it for standalone deployment.
+    #[clap(long, env = "RW_MEMORY_MANAGER_TARGET_BYTES")]
+    pub memory_manager_target_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)]
diff --git a/src/compute/src/server.rs b/src/compute/src/server.rs
index 52b6718115776..8de74cd45645b 100644
--- a/src/compute/src/server.rs
+++ b/src/compute/src/server.rs
@@ -288,15 +288,14 @@ pub async fn compute_node_serve(
         batch_mem_limit(compute_memory_bytes, opts.role.for_serving()),
     ));
 
-    // NOTE: Due to some limits, we use `compute_memory_bytes + storage_memory_bytes` as
-    // `total_compute_memory_bytes` for memory control. This is just a workaround for some
-    // memory control issues and should be modified as soon as we figure out a better solution.
-    //
-    // Related issues:
-    // - https://github.com/risingwavelabs/risingwave/issues/8696
-    // - https://github.com/risingwavelabs/risingwave/issues/8822
+    let target_memory = if let Some(v) = opts.memory_manager_target_bytes {
+        v
+    } else {
+        compute_memory_bytes + storage_memory_bytes
+    };
+
     let memory_mgr = MemoryManager::new(MemoryManagerConfig {
-        target_memory: compute_memory_bytes + storage_memory_bytes,
+        target_memory,
         threshold_aggressive: config
             .streaming
             .developer