fix: introduce termination token for all worker control hierarchy

nyannyacha · nyannyacha · commit f9e8a6a34dfa · 2024-01-17T08:52:36.000+09:00
diff --git a/crates/base/src/commands.rs b/crates/base/src/commands.rs
@@ -1,5 +1,5 @@
 use crate::{
-    rt_worker::worker_pool::WorkerPoolPolicy,
+    rt_worker::{worker_ctx::TerminationToken, worker_pool::WorkerPoolPolicy},
     server::{Server, ServerCodes, WorkerEntrypoints},
 };
 use anyhow::Error;
@@ -16,6 +16,7 @@ pub async fn start_server(
     no_module_cache: bool,
     callback_tx: Option<Sender<ServerCodes>>,
     entrypoints: WorkerEntrypoints,
+    termination_token: Option<TerminationToken>,
 ) -> Result<(), Error> {
     let mut server = Server::new(
         ip,
@@ -27,7 +28,9 @@ pub async fn start_server(
         no_module_cache,
         callback_tx,
         entrypoints,
+        termination_token,
     )
     .await?;
+
     server.listen().await
 }
diff --git a/crates/base/src/macros/test_macros.rs b/crates/base/src/macros/test_macros.rs
@@ -1,25 +1,25 @@
 #[macro_export]
 macro_rules! integration_test {
-    ($main_file:expr, $port:expr, $url:expr, $function: expr) => {
-        let (tx, mut rx) = mpsc::channel::<ServerCodes>(1);
+    ($main_file:expr, $port:expr, $url:expr, ($($function:tt)+) $(, $termination_token: expr)?) => {
+        let (tx, mut rx) = tokio::sync::mpsc::channel::<base::server::ServerCodes>(1);
 
         let signal = tokio::spawn(async move {
-            while let Some(ServerCodes::Listening) = rx.recv().await {
-                let req = reqwest::get(format!("http://localhost:{}/{}", $port, $url)).await;
-                return Some(req);
+            while let Some(base::server::ServerCodes::Listening) = rx.recv().await {
+                integration_test!(@req $port, $url, ($($function)+));
             }
             None
         });
 
-        select! {
+        tokio::select! {
             resp = signal => {
                 if let Ok(maybe_response_from_server) = resp {
-                    $function(maybe_response_from_server.unwrap()).await;
+                    let resp = maybe_response_from_server.unwrap();
+                    integration_test!(@resp resp, ($($function)+)).await;
                 } else {
                     panic!("Request thread had a heart attack");
                 }
             }
-            _ = start_server(
+            _ = base::commands::start_server(
                 "0.0.0.0",
                 $port,
                 String::from($main_file),
@@ -31,10 +31,36 @@ macro_rules! integration_test {
                 $crate::server::WorkerEntrypoints {
                     main: None,
                     events: None,
-                }
+                },
+                integration_test!(@term $(, $termination_token)?)
             ) => {
                 panic!("This one should not end first");
             }
         }
     };
+
+    (@term , $termination_token:expr) => {
+        Some($termination_token)
+    };
+
+    (@term) => {
+        None
+    };
+
+    (@req $port:expr, $url:expr, ($req:expr, $_:expr)) => {
+        return $req($port, $url).await;
+    };
+
+    (@req $port:expr, $url:expr, $_:expr) => {
+        let req = reqwest::get(format!("http://localhost:{}/{}", $port, $url)).await;
+        return Some(req);
+    };
+
+    (@resp $var:ident, ($_:expr, $resp:expr)) => {
+        $resp($var)
+    };
+
+    (@resp $var:ident, $resp:expr) => {
+        $resp($var)
+    };
 }
diff --git a/crates/base/src/rt_worker/worker_ctx.rs b/crates/base/src/rt_worker/worker_ctx.rs
@@ -267,9 +267,9 @@ impl From<(WorkerContextInitOpts, SupervisorPolicy)> for CreateWorkerArgs {
     }
 }
 
-impl From<(WorkerContextInitOpts, TerminationToken)> for CreateWorkerArgs {
-    fn from(val: (WorkerContextInitOpts, TerminationToken)) -> Self {
-        CreateWorkerArgs(val.0, None, Some(val.1))
+impl<T: Into<Option<TerminationToken>>> From<(WorkerContextInitOpts, T)> for CreateWorkerArgs {
+    fn from(val: (WorkerContextInitOpts, T)) -> Self {
+        CreateWorkerArgs(val.0, None, val.1.into())
     }
 }
 
@@ -417,6 +417,7 @@ pub async fn create_main_worker(
     no_module_cache: bool,
     user_worker_msgs_tx: mpsc::UnboundedSender<UserWorkerMsgs>,
     maybe_entrypoint: Option<String>,
+    termination_token: Option<TerminationToken>,
 ) -> Result<mpsc::UnboundedSender<WorkerRequestMsg>, Error> {
     let mut service_path = main_worker_path.clone();
     let mut maybe_eszip = None;
@@ -427,20 +428,23 @@ pub async fn create_main_worker(
         }
     }
 
-    let main_worker_req_tx = create_worker(WorkerContextInitOpts {
-        service_path,
-        import_map_path,
-        no_module_cache,
-        events_rx: None,
-        timing: None,
-        maybe_eszip,
-        maybe_entrypoint,
-        maybe_module_code: None,
-        conf: WorkerRuntimeOpts::MainWorker(MainWorkerRuntimeOpts {
-            worker_pool_tx: user_worker_msgs_tx,
-        }),
-        env_vars: std::env::vars().collect(),
-    })
+    let main_worker_req_tx = create_worker((
+        WorkerContextInitOpts {
+            service_path,
+            import_map_path,
+            no_module_cache,
+            events_rx: None,
+            timing: None,
+            maybe_eszip,
+            maybe_entrypoint,
+            maybe_module_code: None,
+            conf: WorkerRuntimeOpts::MainWorker(MainWorkerRuntimeOpts {
+                worker_pool_tx: user_worker_msgs_tx,
+            }),
+            env_vars: std::env::vars().collect(),
+        },
+        termination_token,
+    ))
     .await
     .map_err(|err| anyhow!("main worker boot error: {}", err))?;
 
@@ -452,6 +456,7 @@ pub async fn create_events_worker(
     import_map_path: Option<String>,
     no_module_cache: bool,
     maybe_entrypoint: Option<String>,
+    termination_token: Option<TerminationToken>,
 ) -> Result<mpsc::UnboundedSender<WorkerEventWithMetadata>, Error> {
     let (events_tx, events_rx) = mpsc::unbounded_channel::<WorkerEventWithMetadata>();
 
@@ -466,18 +471,21 @@ pub async fn create_events_worker(
         }
     }
 
-    let _ = create_worker(WorkerContextInitOpts {
-        service_path,
-        no_module_cache,
-        import_map_path,
-        env_vars: std::env::vars().collect(),
-        events_rx: Some(events_rx),
-        timing: None,
-        maybe_eszip,
-        maybe_entrypoint,
-        maybe_module_code: None,
-        conf: WorkerRuntimeOpts::EventsWorker(EventWorkerRuntimeOpts {}),
-    })
+    let _ = create_worker((
+        WorkerContextInitOpts {
+            service_path,
+            no_module_cache,
+            import_map_path,
+            env_vars: std::env::vars().collect(),
+            events_rx: Some(events_rx),
+            timing: None,
+            maybe_eszip,
+            maybe_entrypoint,
+            maybe_module_code: None,
+            conf: WorkerRuntimeOpts::EventsWorker(EventWorkerRuntimeOpts {}),
+        },
+        termination_token,
+    ))
     .await
     .map_err(|err| anyhow!("events worker boot error: {}", err))?;
 
diff --git a/crates/base/src/server.rs b/crates/base/src/server.rs
@@ -1,5 +1,5 @@
 use crate::rt_worker::worker_ctx::{
-    create_events_worker, create_main_worker, create_user_worker_pool,
+    create_events_worker, create_main_worker, create_user_worker_pool, TerminationToken,
 };
 use crate::rt_worker::worker_pool::WorkerPoolPolicy;
 use anyhow::Error;
@@ -159,6 +159,7 @@ pub struct Server {
     port: u16,
     main_worker_req_tx: mpsc::UnboundedSender<WorkerRequestMsg>,
     callback_tx: Option<Sender<ServerCodes>>,
+    termination_token: TerminationToken,
 }
 
 impl Server {
@@ -173,10 +174,12 @@ impl Server {
         no_module_cache: bool,
         callback_tx: Option<Sender<ServerCodes>>,
         entrypoints: WorkerEntrypoints,
+        termination_token: Option<TerminationToken>,
     ) -> Result<Self, Error> {
         let mut worker_events_sender: Option<mpsc::UnboundedSender<WorkerEventWithMetadata>> = None;
         let maybe_events_entrypoint = entrypoints.events;
         let maybe_main_entrypoint = entrypoints.main;
+        let termination_token = termination_token.unwrap_or_default();
 
         // Create Event Worker
         if let Some(events_service_path) = maybe_events_service_path {
@@ -188,6 +191,7 @@ impl Server {
                 import_map_path.clone(),
                 no_module_cache,
                 maybe_events_entrypoint,
+                Some(termination_token.child_token()),
             )
             .await?;
 
@@ -210,6 +214,7 @@ impl Server {
             no_module_cache,
             user_worker_msgs_tx,
             maybe_main_entrypoint,
+            Some(termination_token.child_token()),
         )
         .await?;
 
@@ -219,12 +224,19 @@ impl Server {
             port,
             main_worker_req_tx,
             callback_tx,
+            termination_token,
         })
     }
 
+    pub async fn terminate(&self) {
+        self.termination_token.cancel_and_wait().await;
+    }
+
     pub async fn listen(&mut self) -> Result<(), Error> {
         let addr = SocketAddr::new(IpAddr::V4(self.ip), self.port);
         let listener = TcpListener::bind(&addr).await?;
+        let termination_token = self.termination_token.clone();
+
         debug!("edge-runtime is listening on {:?}", listener.local_addr()?);
 
         if let Some(callback) = self.callback_tx.clone() {
@@ -260,6 +272,12 @@ impl Server {
                         Err(e) => error!("socket error: {}", e)
                     }
                 }
+
+                _ = termination_token.outbound.cancelled() => {
+                    info!("termination token resolved");
+                    break;
+                }
+
                 // wait for shutdown signal...
                 _ = tokio::signal::ctrl_c() => {
                     info!("shutdown signal received");
diff --git a/crates/base/tests/custom_readable_response_tests.rs b/crates/base/tests/custom_readable_response_tests.rs
@@ -1,22 +1,18 @@
 extern crate core;
 
-use base::commands::start_server;
 use base::integration_test;
-use base::server::ServerCodes;
-use tokio::select;
-use tokio::sync::mpsc;
 
 #[tokio::test]
 async fn test_custom_readable_stream_response() {
     integration_test!(
         "./test_cases/main",
         8999,
         "readable-stream-resp",
-        |resp: Result<reqwest::Response, reqwest::Error>| async {
+        (|resp: Result<reqwest::Response, reqwest::Error>| async {
             assert_eq!(
                 resp.unwrap().text().await.unwrap(),
                 "Hello world from streams"
             );
-        }
+        })
     );
 }
diff --git a/crates/base/tests/linux_pku_sigsegv_test_cli.rs b/crates/base/tests/linux_pku_sigsegv_test_cli.rs
@@ -1,16 +1,14 @@
-use base::commands::start_server;
 use base::integration_test;
-use base::server::ServerCodes;
-use tokio::select;
-use tokio::sync::mpsc;
 
 #[cfg(target_os = "linux")]
 #[tokio::test]
 async fn test_not_trigger_pku_sigsegv_due_to_jit_compilation_cli() {
-    integration_test!("./test_cases/main", 8999, "slow_resp", |resp: Result<
-        reqwest::Response,
-        reqwest::Error,
-    >| async {
-        assert!(resp.unwrap().text().await.unwrap().starts_with("meow: "));
-    });
+    integration_test!(
+        "./test_cases/main",
+        8999,
+        "slow_resp",
+        (|resp: Result<reqwest::Response, reqwest::Error>| async {
+            assert!(resp.unwrap().text().await.unwrap().starts_with("meow: "));
+        })
+    );
 }
diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs
@@ -172,6 +172,7 @@ fn main() -> Result<(), anyhow::Error> {
                         main: maybe_main_entrypoint,
                         events: maybe_events_entrypoint,
                     },
+                    None,
                 )
                 .await?;
             }

Original file line number	Diff line number	Diff line change
`@@ -172,6 +172,7 @@ fn main() -> Result<(), anyhow::Error> {`
`172`	`172`	`main: maybe_main_entrypoint,`
`173`	`173`	`events: maybe_events_entrypoint,`
`174`	`174`	`},`
	`175`	`+ None,`
`175`	`176`	`)`
`176`	`177`	`.await?;`
`177`	`178`	`}`