fix(storage): do not acquire quota for shared buffer in serving mode (#…

…16753) Signed-off-by: Little-Wallace <[email protected]>
risingwavelabs · May 22, 2024 · 68e9b03 · 68e9b03
1 parent 2f44407
commit 68e9b03
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 7 deletions.
diff --git a/src/compute/src/memory/config.rs b/src/compute/src/memory/config.rs
@@ -69,6 +69,7 @@ pub fn storage_memory_config(
     non_reserved_memory_bytes: usize,
     embedded_compactor_enabled: bool,
     storage_config: &StorageConfig,
+    is_serving: bool,
 ) -> StorageMemoryConfig {
     let (storage_memory_proportion, compactor_memory_proportion) = if embedded_compactor_enabled {
         (STORAGE_MEMORY_PROPORTION, COMPACTOR_MEMORY_PROPORTION)
@@ -110,14 +111,18 @@ pub fn storage_memory_config(
         * STORAGE_SHARED_BUFFER_MEMORY_PROPORTION)
         .ceil() as usize)
         >> 20;
-    let shared_buffer_capacity_mb =
+    let mut shared_buffer_capacity_mb =
         storage_config
             .shared_buffer_capacity_mb
             .unwrap_or(std::cmp::min(
                 default_shared_buffer_capacity_mb,
                 STORAGE_SHARED_BUFFER_MAX_MEMORY_MB,
             ));
-    if shared_buffer_capacity_mb != default_shared_buffer_capacity_mb {
+    if is_serving {
+        default_block_cache_capacity_mb += default_shared_buffer_capacity_mb;
+        // set 1 to pass internal check
+        shared_buffer_capacity_mb = 1;
+    } else if shared_buffer_capacity_mb != default_shared_buffer_capacity_mb {
         default_block_cache_capacity_mb += default_shared_buffer_capacity_mb;
         default_block_cache_capacity_mb =
             default_block_cache_capacity_mb.saturating_sub(shared_buffer_capacity_mb);
@@ -279,17 +284,36 @@ mod tests {
 
         let total_non_reserved_memory_bytes = 8 << 30;
 
-        let memory_config =
-            storage_memory_config(total_non_reserved_memory_bytes, true, &storage_config);
+        let memory_config = storage_memory_config(
+            total_non_reserved_memory_bytes,
+            true,
+            &storage_config,
+            false,
+        );
         assert_eq!(memory_config.block_cache_capacity_mb, 737);
         assert_eq!(memory_config.meta_cache_capacity_mb, 860);
         assert_eq!(memory_config.shared_buffer_capacity_mb, 737);
         assert_eq!(memory_config.compactor_memory_limit_mb, 819);
 
+        let memory_config = storage_memory_config(
+            total_non_reserved_memory_bytes,
+            false,
+            &storage_config,
+            true,
+        );
+        assert_eq!(memory_config.block_cache_capacity_mb, 1966);
+        assert_eq!(memory_config.meta_cache_capacity_mb, 1146);
+        assert_eq!(memory_config.shared_buffer_capacity_mb, 1);
+        assert_eq!(memory_config.compactor_memory_limit_mb, 0);
+
         storage_config.data_file_cache.dir = "data".to_string();
         storage_config.meta_file_cache.dir = "meta".to_string();
-        let memory_config =
-            storage_memory_config(total_non_reserved_memory_bytes, true, &storage_config);
+        let memory_config = storage_memory_config(
+            total_non_reserved_memory_bytes,
+            true,
+            &storage_config,
+            false,
+        );
         assert_eq!(memory_config.block_cache_capacity_mb, 737);
         assert_eq!(memory_config.meta_cache_capacity_mb, 860);
         assert_eq!(memory_config.shared_buffer_capacity_mb, 737);
@@ -299,7 +323,7 @@ mod tests {
         storage_config.cache.meta_cache_capacity_mb = Some(128);
         storage_config.shared_buffer_capacity_mb = Some(1024);
         storage_config.compactor_memory_limit_mb = Some(512);
-        let memory_config = storage_memory_config(0, true, &storage_config);
+        let memory_config = storage_memory_config(0, true, &storage_config, false);
         assert_eq!(memory_config.block_cache_capacity_mb, 512);
         assert_eq!(memory_config.meta_cache_capacity_mb, 128);
         assert_eq!(memory_config.shared_buffer_capacity_mb, 1024);

diff --git a/src/compute/src/server.rs b/src/compute/src/server.rs
@@ -127,6 +127,7 @@ pub async fn compute_node_serve(
         non_reserved_memory_bytes,
         embedded_compactor_enabled,
         &config.storage,
+        !opts.role.for_streaming(),
     );
 
     let storage_memory_bytes = total_storage_memory_limit_bytes(&storage_memory_config);