fix: make initial/max heap size configurable for main/event worker (#568)

* fix: make initial/max heap size configurable for main/event worker

* stamp: polishing

Author: nyannyacha

cli/src/env.rs

@@ -1,11 +1,15 @@
 use base::deno_runtime;
-use clap::builder::{BoolishValueParser, TypedValueParser};
+use clap::{
+    builder::{BoolishValueParser, TypedValueParser},
+    value_parser,
+};
 use once_cell::sync::OnceCell;
 
 pub(super) fn resolve_deno_runtime_env() {
     let boolish_parser = BoolishValueParser::new();
+    let u64_parser = value_parser!(u64);
     let dumb_command = clap::Command::new(env!("CARGO_BIN_NAME"));
-    let resolve_boolish_env = move |key: &'static str, cell: &'static OnceCell<bool>| {
+    let resolve_boolish_env = |key: &'static str, cell: &'static OnceCell<bool>| {
         cell.get_or_init(|| {
             std::env::var_os(key)
                 .map(|it| {
@@ -17,6 +21,18 @@ pub(super) fn resolve_deno_runtime_env() {
         })
     };
 
+    let resolve_u64_env = |key: &'static str, cell: &'static OnceCell<u64>| {
+        cell.get_or_init(|| {
+            std::env::var_os(key)
+                .map(|it| {
+                    u64_parser
+                        .parse_ref(&dumb_command, None, &it)
+                        .unwrap_or_default()
+                })
+                .unwrap_or_default()
+        })
+    };
+
     deno_runtime::MAYBE_DENO_VERSION.get_or_init(|| deno_manifest::version().to_string());
 
     resolve_boolish_env(
@@ -33,4 +49,24 @@ pub(super) fn resolve_deno_runtime_env() {
         "EDGE_RUNTIME_INCLUDE_MALLOCED_MEMORY_ON_MEMCHECK",
         &deno_runtime::SHOULD_INCLUDE_MALLOCED_MEMORY_ON_MEMCHECK,
     );
+
+    resolve_u64_env(
+        "EDGE_RUNTIME_MAIN_WORKER_INITIAL_HEAP_SIZE_MIB",
+        &deno_runtime::MAIN_WORKER_INITIAL_HEAP_SIZE_MIB,
+    );
+
+    resolve_u64_env(
+        "EDGE_RUNTIME_MAIN_WORKER_MAX_HEAP_SIZE_MIB",
+        &deno_runtime::MAIN_WORKER_MAX_HEAP_SIZE_MIB,
+    );
+
+    resolve_u64_env(
+        "EDGE_RUNTIME_EVENT_WORKER_INITIAL_HEAP_SIZE_MIB",
+        &deno_runtime::EVENT_WORKER_INITIAL_HEAP_SIZE_MIB,
+    );
+
+    resolve_u64_env(
+        "EDGE_RUNTIME_EVENT_WORKER_MAX_HEAP_SIZE_MIB",
+        &deno_runtime::EVENT_WORKER_MAX_HEAP_SIZE_MIB,
+    );
 }

crates/base/src/deno_runtime.rs

@@ -78,7 +78,7 @@ use sb_event_worker::events::{EventMetadata, WorkerEventWithMetadata};
 use sb_event_worker::js_interceptors::sb_events_js_interceptors;
 use sb_event_worker::sb_user_event_worker;
 use sb_node::deno_node;
-use sb_workers::context::{UserWorkerMsgs, WorkerContextInitOpts, WorkerRuntimeOpts};
+use sb_workers::context::{UserWorkerMsgs, WorkerContextInitOpts, WorkerKind, WorkerRuntimeOpts};
 use sb_workers::sb_user_workers;
 
 const DEFAULT_ALLOC_CHECK_INT_MSEC: u64 = 1000;
@@ -107,6 +107,11 @@ pub static SHOULD_USE_VERBOSE_DEPRECATED_API_WARNING: OnceCell<bool> = OnceCell:
 pub static SHOULD_INCLUDE_MALLOCED_MEMORY_ON_MEMCHECK: OnceCell<bool> = OnceCell::new();
 pub static MAYBE_DENO_VERSION: OnceCell<String> = OnceCell::new();
 
+pub static MAIN_WORKER_INITIAL_HEAP_SIZE_MIB: OnceCell<u64> = OnceCell::new();
+pub static MAIN_WORKER_MAX_HEAP_SIZE_MIB: OnceCell<u64> = OnceCell::new();
+pub static EVENT_WORKER_INITIAL_HEAP_SIZE_MIB: OnceCell<u64> = OnceCell::new();
+pub static EVENT_WORKER_MAX_HEAP_SIZE_MIB: OnceCell<u64> = OnceCell::new();
+
 thread_local! {
     // NOTE: Suppose we have met `.await` points while initializing a
     // DenoRuntime. In that case, the current v8 isolate's thread-local state can be
@@ -761,37 +766,62 @@ where
         let beforeunload_cpu_threshold = ArcSwapOption::<u64>::from_pointee(None);
         let beforeunload_mem_threshold = ArcSwapOption::<u64>::from_pointee(None);
 
-        if conf.is_user_worker() {
-            let conf = maybe_user_conf.unwrap();
-            let memory_limit_bytes = mib_to_bytes(conf.memory_limit_mb) as usize;
-
-            beforeunload_mem_threshold.store(
-                flags
-                    .beforeunload_memory_pct
-                    .and_then(|it| percentage_value(memory_limit_bytes as u64, it))
-                    .map(Arc::new),
-            );
+        match conf.to_worker_kind() {
+            WorkerKind::UserWorker => {
+                let conf = maybe_user_conf.unwrap();
+                let memory_limit_bytes = mib_to_bytes(conf.memory_limit_mb) as usize;
 
-            if conf.cpu_time_hard_limit_ms > 0 {
-                beforeunload_cpu_threshold.store(
+                beforeunload_mem_threshold.store(
                     flags
-                        .beforeunload_cpu_pct
-                        .and_then(|it| percentage_value(conf.cpu_time_hard_limit_ms, it))
+                        .beforeunload_memory_pct
+                        .and_then(|it| percentage_value(memory_limit_bytes as u64, it))
                         .map(Arc::new),
                 );
-            }
 
-            let allocator = CustomAllocator::new(memory_limit_bytes);
+                if conf.cpu_time_hard_limit_ms > 0 {
+                    beforeunload_cpu_threshold.store(
+                        flags
+                            .beforeunload_cpu_pct
+                            .and_then(|it| percentage_value(conf.cpu_time_hard_limit_ms, it))
+                            .map(Arc::new),
+                    );
+                }
 
-            allocator.set_waker(mem_check.waker.clone());
+                let allocator = CustomAllocator::new(memory_limit_bytes);
 
-            mem_check.limit = Some(memory_limit_bytes);
-            create_params = Some(
-                v8::CreateParams::default()
-                    .heap_limits(mib_to_bytes(0) as usize, memory_limit_bytes)
-                    .array_buffer_allocator(allocator.into_v8_allocator()),
-            )
-        };
+                allocator.set_waker(mem_check.waker.clone());
+
+                mem_check.limit = Some(memory_limit_bytes);
+                create_params = Some(
+                    v8::CreateParams::default()
+                        .heap_limits(mib_to_bytes(0) as usize, memory_limit_bytes)
+                        .array_buffer_allocator(allocator.into_v8_allocator()),
+                )
+            }
+            kind => {
+                assert_ne!(kind, WorkerKind::UserWorker);
+                let initial_heap_size = match kind {
+                    WorkerKind::MainWorker => &MAIN_WORKER_INITIAL_HEAP_SIZE_MIB,
+                    WorkerKind::EventsWorker => &EVENT_WORKER_INITIAL_HEAP_SIZE_MIB,
+                    _ => unreachable!(),
+                };
+                let max_heap_size = match kind {
+                    WorkerKind::MainWorker => &MAIN_WORKER_MAX_HEAP_SIZE_MIB,
+                    WorkerKind::EventsWorker => &EVENT_WORKER_MAX_HEAP_SIZE_MIB,
+                    _ => unreachable!(),
+                };
+
+                let initial_heap_size = initial_heap_size.get().cloned().unwrap_or_default();
+                let max_heap_size = max_heap_size.get().cloned().unwrap_or_default();
+
+                if max_heap_size > 0 {
+                    create_params = Some(v8::CreateParams::default().heap_limits(
+                        mib_to_bytes(initial_heap_size) as usize,
+                        mib_to_bytes(max_heap_size) as usize,
+                    ));
+                }
+            }
+        }
 
         let mem_check = Arc::new(mem_check);
         let runtime_options = RuntimeOptions {