fix: sync drop timing of unix stream sender to timing of an isolate drops

Author: nyannyacha

crates/base/src/deno_runtime.rs

@@ -1,4 +1,5 @@
 use crate::rt_worker::supervisor::{CPUUsage, CPUUsageMetrics};
+use crate::rt_worker::worker::UnixStreamEntry;
 use crate::utils::units::mib_to_bytes;
 
 use anyhow::{anyhow, bail, Context, Error};
@@ -25,7 +26,6 @@ use std::fmt;
 use std::os::fd::RawFd;
 use std::sync::Arc;
 use std::task::Poll;
-use tokio::net::UnixStream;
 use tokio::sync::{mpsc, watch};
 
 use crate::snapshot;
@@ -372,17 +372,14 @@ impl DenoRuntime {
 
     pub async fn run(
         &mut self,
-        unix_stream_rx: mpsc::UnboundedReceiver<(UnixStream, Option<watch::Receiver<ConnSync>>)>,
+        unix_stream_rx: mpsc::UnboundedReceiver<UnixStreamEntry>,
         maybe_cpu_usage_metrics_tx: Option<mpsc::UnboundedSender<CPUUsageMetrics>>,
         name: Option<String>,
     ) -> (Result<(), Error>, i64) {
         {
             let op_state_rc = self.js_runtime.op_state();
             let mut op_state = op_state_rc.borrow_mut();
-            op_state
-                .put::<mpsc::UnboundedReceiver<(UnixStream, Option<watch::Receiver<ConnSync>>)>>(
-                    unix_stream_rx,
-                );
+            op_state.put::<mpsc::UnboundedReceiver<UnixStreamEntry>>(unix_stream_rx);
 
             if self.conf.is_main_worker() {
                 op_state.put::<mpsc::UnboundedSender<UserWorkerMsgs>>(
@@ -447,7 +444,7 @@ impl DenoRuntime {
                 cx,
                 PollEventLoopOptions {
                     wait_for_inspector: false,
-                    pump_v8_message_loop: true,
+                    pump_v8_message_loop: !is_termination_requested.is_raised(),
                 },
             );
 
@@ -533,8 +530,8 @@ fn set_v8_flags() {
 #[cfg(test)]
 mod test {
     use crate::deno_runtime::DenoRuntime;
+    use crate::rt_worker::worker::UnixStreamEntry;
     use deno_core::{FastString, ModuleCode, PollEventLoopOptions};
-    use sb_core::conn_sync::ConnSync;
     use sb_graph::emitter::EmitterFactory;
     use sb_graph::{generate_binary_eszip, EszipPayloadKind};
     use sb_workers::context::{
@@ -548,8 +545,7 @@ mod test {
     use std::io::Write;
     use std::path::PathBuf;
     use std::sync::Arc;
-    use tokio::net::UnixStream;
-    use tokio::sync::{mpsc, watch};
+    use tokio::sync::mpsc;
 
     #[tokio::test]
     #[serial]
@@ -1059,8 +1055,7 @@ mod 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) =
-            mpsc::unbounded_channel::<(UnixStream, Option<watch::Receiver<ConnSync>>)>();
+        let (_tx, unix_stream_rx) = mpsc::unbounded_channel::<UnixStreamEntry>();
 
         let (result, _) = user_rt.run(unix_stream_rx, None, None).await;
         match result {
@@ -1077,8 +1072,7 @@ mod 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) =
-            mpsc::unbounded_channel::<(UnixStream, Option<watch::Receiver<ConnSync>>)>();
+        let (_tx, unix_stream_rx) = mpsc::unbounded_channel::<UnixStreamEntry>();
         let (result, _) = user_rt.run(unix_stream_rx, None, None).await;
         assert!(result.is_ok(), "expected no errors");
     }
@@ -1087,8 +1081,7 @@ mod 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) =
-            mpsc::unbounded_channel::<(UnixStream, Option<watch::Receiver<ConnSync>>)>();
+        let (_tx, unix_stream_rx) = mpsc::unbounded_channel::<UnixStreamEntry>();
         let (result, _) = user_rt.run(unix_stream_rx, None, None).await;
         match result {
             Err(err) => {

crates/base/src/rt_worker/implementation/default_handler.rs

@@ -1,15 +1,12 @@
 use crate::deno_runtime::DenoRuntime;
 use crate::rt_worker::supervisor::CPUUsageMetrics;
-use crate::rt_worker::worker::{HandleCreationType, Worker, WorkerHandler};
+use crate::rt_worker::worker::{HandleCreationType, UnixStreamEntry, Worker, WorkerHandler};
 use anyhow::Error;
 use event_worker::events::{BootFailureEvent, PseudoEvent, UncaughtExceptionEvent, WorkerEvents};
 use log::error;
-use sb_core::conn_sync::ConnSync;
 use std::any::Any;
-use tokio::net::UnixStream;
 use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
 use tokio::sync::oneshot::Receiver;
-use tokio::sync::watch;
 
 impl WorkerHandler for Worker {
     fn handle_error(&self, error: Error) -> Result<WorkerEvents, Error> {
@@ -22,7 +19,7 @@ impl WorkerHandler for Worker {
     fn handle_creation(
         &self,
         mut created_rt: DenoRuntime,
-        unix_stream_rx: UnboundedReceiver<(UnixStream, Option<watch::Receiver<ConnSync>>)>,
+        unix_stream_rx: UnboundedReceiver<UnixStreamEntry>,
         termination_event_rx: Receiver<WorkerEvents>,
         maybe_cpu_usage_metrics_tx: Option<UnboundedSender<CPUUsageMetrics>>,
         name: Option<String>,

crates/base/src/rt_worker/worker.rs

@@ -41,13 +41,14 @@ pub struct Worker {
 }
 
 pub type HandleCreationType = Pin<Box<dyn Future<Output = Result<WorkerEvents, Error>>>>;
+pub type UnixStreamEntry = (UnixStream, Option<watch::Receiver<ConnSync>>);
 
 pub trait WorkerHandler: Send {
     fn handle_error(&self, error: Error) -> Result<WorkerEvents, Error>;
     fn handle_creation(
         &self,
         created_rt: DenoRuntime,
-        unix_stream_rx: UnboundedReceiver<(UnixStream, Option<watch::Receiver<ConnSync>>)>,
+        unix_stream_rx: UnboundedReceiver<UnixStreamEntry>,
         termination_event_rx: Receiver<WorkerEvents>,
         maybe_cpu_metrics_tx: Option<UnboundedSender<CPUUsageMetrics>>,
         name: Option<String>,
@@ -57,10 +58,10 @@ pub trait WorkerHandler: Send {
 
 impl Worker {
     pub fn new(init_opts: &WorkerContextInitOpts) -> Result<Self, Error> {
-        let (worker_key, pool_msg_tx, events_msg_tx, cancel, thread_name) =
+        let (worker_key, pool_msg_tx, events_msg_tx, cancel, worker_name) =
             parse_worker_conf(&init_opts.conf);
-        let event_metadata = get_event_metadata(&init_opts.conf);
 
+        let event_metadata = get_event_metadata(&init_opts.conf);
         let worker_boot_start_time = Instant::now();
 
         Ok(Self {
@@ -71,7 +72,7 @@ impl Worker {
             cancel,
             event_metadata,
             worker_key,
-            worker_name: thread_name,
+            worker_name,
         })
     }
 
@@ -82,7 +83,10 @@ impl Worker {
     pub fn start(
         &self,
         mut opts: WorkerContextInitOpts,
-        unix_channel_rx: UnboundedReceiver<(UnixStream, Option<watch::Receiver<ConnSync>>)>,
+        unix_stream_pair: (
+            UnboundedSender<UnixStreamEntry>,
+            UnboundedReceiver<UnixStreamEntry>,
+        ),
         booter_signal: Sender<Result<(), Error>>,
         termination_token: Option<TerminationToken>,
     ) {
@@ -91,6 +95,7 @@ impl Worker {
         let event_metadata = self.event_metadata.clone();
         let supervisor_policy = self.supervisor_policy.unwrap_or_default();
 
+        let (unix_stream_tx, unix_stream_rx) = unix_stream_pair;
         let events_msg_tx = self.events_msg_tx.clone();
         let pool_msg_tx = self.pool_msg_tx.clone();
 
@@ -191,7 +196,7 @@ impl Worker {
 
                         let data = method_cloner.handle_creation(
                             new_runtime,
-                            unix_channel_rx,
+                            unix_stream_rx,
                             termination_event_rx,
                             maybe_cpu_usage_metrics_tx,
                             Some(worker_name),
@@ -216,6 +221,8 @@ impl Worker {
                     }
                 };
 
+                drop(unix_stream_tx);
+
                 match result {
                     Ok(event) => {
                         match event {

crates/base/src/rt_worker/worker_ctx.rs

@@ -29,6 +29,7 @@ use uuid::Uuid;
 
 use super::rt;
 use super::supervisor::{self, CPUTimerParam, CPUUsageMetrics};
+use super::worker::UnixStreamEntry;
 use super::worker_pool::{SupervisorPolicy, WorkerPoolPolicy};
 
 #[derive(Clone)]
@@ -75,7 +76,7 @@ impl TerminationToken {
 }
 
 async fn handle_request(
-    unix_stream_tx: mpsc::UnboundedSender<(UnixStream, Option<watch::Receiver<ConnSync>>)>,
+    unix_stream_tx: mpsc::UnboundedSender<UnixStreamEntry>,
     msg: WorkerRequestMsg,
 ) -> Result<(), Error> {
     // create a unix socket pair
@@ -310,8 +311,7 @@ pub async fn create_worker<Opt: Into<CreateWorkerArgs>>(
     init_opts: Opt,
 ) -> Result<mpsc::UnboundedSender<WorkerRequestMsg>, Error> {
     let (worker_boot_result_tx, worker_boot_result_rx) = oneshot::channel::<Result<(), Error>>();
-    let (unix_stream_tx, unix_stream_rx) =
-        mpsc::unbounded_channel::<(UnixStream, Option<watch::Receiver<ConnSync>>)>();
+    let (unix_stream_tx, unix_stream_rx) = mpsc::unbounded_channel::<UnixStreamEntry>();
 
     let CreateWorkerArgs(init_opts, maybe_supervisor_policy, maybe_termination_token) =
         init_opts.into();
@@ -331,27 +331,31 @@ pub async fn create_worker<Opt: Into<CreateWorkerArgs>>(
     if let Some(worker_struct_ref) = downcast_reference {
         worker_struct_ref.start(
             init_opts,
-            unix_stream_rx,
+            (unix_stream_tx.clone(), unix_stream_rx),
             worker_boot_result_tx,
             maybe_termination_token.clone(),
         );
 
         // create an async task waiting for requests for worker
         let (worker_req_tx, mut worker_req_rx) = mpsc::unbounded_channel::<WorkerRequestMsg>();
 
-        let worker_req_handle: tokio::task::JoinHandle<Result<(), Error>> =
-            tokio::task::spawn(async move {
+        let worker_req_handle: tokio::task::JoinHandle<Result<(), Error>> = tokio::task::spawn({
+            let stream_tx = unix_stream_tx;
+            async move {
                 while let Some(msg) = worker_req_rx.recv().await {
-                    let unix_stream_tx_clone = unix_stream_tx.clone();
-                    tokio::task::spawn(async move {
-                        if let Err(err) = handle_request(unix_stream_tx_clone, msg).await {
-                            error!("worker failed to handle request: {:?}", err);
+                    tokio::task::spawn({
+                        let stream_tx_inner = stream_tx.clone();
+                        async move {
+                            if let Err(err) = handle_request(stream_tx_inner, msg).await {
+                                error!("worker failed to handle request: {:?}", err);
+                            }
                         }
                     });
                 }
 
                 Ok(())
-            });
+            }
+        });
 
         // wait for worker to be successfully booted
         let worker_boot_result = worker_boot_result_rx.await?;