Merge pull request #246 from nyannyacha/fix-pool-sep

fix: make major worker tasks managed separately from the user worker

Author: laktek

crates/base/src/deno_runtime.rs

@@ -522,6 +522,7 @@ mod test {
         MainWorkerRuntimeOpts, UserWorkerMsgs, UserWorkerRuntimeOpts, WorkerContextInitOpts,
         WorkerRuntimeOpts,
     };
+    use serial_test::serial;
     use std::collections::HashMap;
     use std::fs;
     use std::fs::File;
@@ -532,6 +533,7 @@ mod test {
     use tokio::sync::{mpsc, watch};
 
     #[tokio::test]
+    #[serial]
     async fn test_module_code_no_eszip() {
         let (worker_pool_tx, _) = mpsc::unbounded_channel::<UserWorkerMsgs>();
         let mut rt = DenoRuntime::new(WorkerContextInitOpts {
@@ -559,6 +561,7 @@ mod test {
     }
 
     #[tokio::test]
+    #[serial]
     #[allow(clippy::arc_with_non_send_sync)]
     async fn test_eszip_with_source_file() {
         let (worker_pool_tx, _) = mpsc::unbounded_channel::<UserWorkerMsgs>();
@@ -616,6 +619,7 @@ mod test {
     }
 
     #[tokio::test]
+    #[serial]
     #[allow(clippy::arc_with_non_send_sync)]
     async fn test_create_eszip_from_graph() {
         let (worker_pool_tx, _) = mpsc::unbounded_channel::<UserWorkerMsgs>();
@@ -707,6 +711,7 @@ mod test {
 
     // Main Runtime should have access to `EdgeRuntime`
     #[tokio::test]
+    #[serial]
     async fn test_main_runtime_creation() {
         let mut runtime = create_runtime(None, None, None).await;
 
@@ -726,6 +731,7 @@ mod test {
 
     // User Runtime Should not have access to EdgeRuntime
     #[tokio::test]
+    #[serial]
     async fn test_user_runtime_creation() {
         let mut runtime = create_runtime(
             None,
@@ -749,6 +755,7 @@ mod test {
     }
 
     #[tokio::test]
+    #[serial]
     async fn test_main_rt_fs() {
         let mut main_rt = create_runtime(None, Some(std::env::vars().collect()), None).await;
 
@@ -798,6 +805,7 @@ mod test {
     // }
 
     #[tokio::test]
+    #[serial]
     async fn test_os_ops() {
         let mut user_rt = create_runtime(
             None,
@@ -920,6 +928,7 @@ mod test {
     }
 
     #[tokio::test]
+    #[serial]
     async fn test_os_env_vars() {
         std::env::set_var("Supa_Test", "Supa_Value");
         let mut main_rt = create_runtime(None, Some(std::env::vars().collect()), None).await;
@@ -1016,6 +1025,7 @@ mod test {
     }
 
     #[tokio::test]
+    #[serial]
     async fn test_read_file_user_rt() {
         let mut user_rt = create_basic_user_runtime("./test_cases/readFile", 20, 1000).await;
         let (_tx, unix_stream_rx) =
@@ -1033,6 +1043,7 @@ mod test {
     }
 
     #[tokio::test]
+    #[serial]
     async fn test_array_buffer_allocation_below_limit() {
         let mut user_rt = create_basic_user_runtime("./test_cases/array_buffers", 20, 1000).await;
         let (_tx, unix_stream_rx) =
@@ -1042,6 +1053,7 @@ mod test {
     }
 
     #[tokio::test]
+    #[serial]
     async fn test_array_buffer_allocation_above_limit() {
         let mut user_rt = create_basic_user_runtime("./test_cases/array_buffers", 15, 1000).await;
         let (_tx, unix_stream_rx) =

crates/base/src/rt_worker/rt.rs

@@ -1,5 +1,10 @@
+use std::num::NonZeroUsize;
+
 use once_cell::sync::Lazy;
 
+pub const DEFAULT_PRIMARY_WORKER_POOL_SIZE: usize = 2;
+pub const DEFAULT_USER_WORKER_POOL_SIZE: usize = 1;
+
 pub static SUPERVISOR_RT: Lazy<tokio::runtime::Runtime> = Lazy::new(|| {
     tokio::runtime::Builder::new_multi_thread()
         .enable_all()
@@ -8,14 +13,46 @@ pub static SUPERVISOR_RT: Lazy<tokio::runtime::Runtime> = Lazy::new(|| {
         .unwrap()
 });
 
-pub static WORKER_RT: Lazy<tokio_util::task::LocalPoolHandle> = Lazy::new(|| {
+// NOTE: This pool is for the main and event workers. The reason why they should
+// separate from the user worker pool is they can starve them if user workers
+// are saturated.
+pub static PRIMARY_WORKER_RT: Lazy<tokio_util::task::LocalPoolHandle> = Lazy::new(|| {
+    let maybe_pool_size = std::env::var("EDGE_RUNTIME_PRIMARY_WORKER_POOL_SIZE")
+        .ok()
+        .and_then(|it| it.parse::<usize>().ok())
+        .map(|it| {
+            if it < DEFAULT_PRIMARY_WORKER_POOL_SIZE {
+                DEFAULT_PRIMARY_WORKER_POOL_SIZE
+            } else {
+                it
+            }
+        });
+
+    tokio_util::task::LocalPoolHandle::new(
+        maybe_pool_size.unwrap_or(DEFAULT_PRIMARY_WORKER_POOL_SIZE),
+    )
+});
+
+pub static USER_WORKER_RT: Lazy<tokio_util::task::LocalPoolHandle> = Lazy::new(|| {
     let maybe_pool_size = std::env::var("EDGE_RUNTIME_WORKER_POOL_SIZE")
         .ok()
-        .and_then(|it| it.parse::<usize>().ok());
+        .and_then(|it| it.parse::<usize>().ok())
+        .map(|it| {
+            if it < DEFAULT_USER_WORKER_POOL_SIZE {
+                DEFAULT_USER_WORKER_POOL_SIZE
+            } else {
+                it
+            }
+        });
 
     tokio_util::task::LocalPoolHandle::new(if cfg!(debug_assertions) {
-        maybe_pool_size.unwrap_or(1)
+        maybe_pool_size.unwrap_or(DEFAULT_USER_WORKER_POOL_SIZE)
     } else {
-        maybe_pool_size.unwrap_or(std::thread::available_parallelism().unwrap().get())
+        maybe_pool_size.unwrap_or(
+            std::thread::available_parallelism()
+                .ok()
+                .map(NonZeroUsize::get)
+                .unwrap_or(DEFAULT_USER_WORKER_POOL_SIZE),
+        )
     })
 });

crates/base/src/rt_worker/worker.rs

@@ -98,8 +98,13 @@ impl Worker {
         let is_user_worker = opts.conf.is_user_worker();
 
         let cancel = self.cancel.clone();
+        let rt = if is_user_worker {
+            &rt::USER_WORKER_RT
+        } else {
+            &rt::PRIMARY_WORKER_RT
+        };
 
-        let _worker_handle = rt::WORKER_RT.spawn_pinned(move || {
+        let _worker_handle = rt.spawn_pinned(move || {
             tokio::task::spawn_local(async move {
                 let (maybe_cpu_usage_metrics_tx, maybe_cpu_usage_metrics_rx) = is_user_worker
                     .then(unbounded_channel::<CPUUsageMetrics>)