refactor: expose more precise server health messages to callback receiver

Author: nyannyacha

crates/base/src/commands.rs

@@ -1,6 +1,6 @@
 use crate::{
     rt_worker::{worker_ctx::TerminationToken, worker_pool::WorkerPoolPolicy},
-    server::{Server, ServerCodes, WorkerEntrypoints},
+    server::{Server, ServerHealth, WorkerEntrypoints},
 };
 use anyhow::Error;
 use tokio::sync::mpsc::Sender;
@@ -14,7 +14,7 @@ pub async fn start_server(
     user_worker_policy: Option<WorkerPoolPolicy>,
     import_map_path: Option<String>,
     no_module_cache: bool,
-    callback_tx: Option<Sender<ServerCodes>>,
+    callback_tx: Option<Sender<ServerHealth>>,
     entrypoints: WorkerEntrypoints,
     termination_token: Option<TerminationToken>,
 ) -> Result<(), Error> {

crates/base/src/macros/test_macros.rs

@@ -1,11 +1,11 @@
 #[macro_export]
 macro_rules! integration_test {
     ($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 (tx, mut rx) = tokio::sync::mpsc::channel::<base::server::ServerHealth>(1);
 
         let signal = tokio::spawn(async move {
-            while let Some(base::server::ServerCodes::Listening) = rx.recv().await {
-                integration_test!(@req $port, $url, ($($function)+));
+            while let Some(base::server::ServerHealth::Listening(event_rx)) = rx.recv().await {
+                integration_test!(@req event_rx, $port, $url, ($($function)+));
             }
             None
         });
@@ -47,11 +47,11 @@ macro_rules! integration_test {
         None
     };
 
-    (@req $port:expr, $url:expr, ($req:expr, $_:expr)) => {
-        return $req($port, $url).await;
+    (@req $event_rx:ident, $port:expr, $url:expr, ($req:expr, $_:expr)) => {
+        return $req($port, $url, $event_rx).await;
     };
 
-    (@req $port:expr, $url:expr, $_:expr) => {
+    (@req $_:ident, $port:expr, $url:expr, $__:expr) => {
         let req = reqwest::get(format!("http://localhost:{}/{}", $port, $url)).await;
         return Some(req);
     };

crates/base/src/server.rs

@@ -23,8 +23,12 @@ use tokio::sync::mpsc::Sender;
 use tokio::sync::{mpsc, oneshot, watch};
 use tokio_util::sync::CancellationToken;
 
-pub enum ServerCodes {
-    Listening,
+pub enum ServerEvent {
+    ConnectionError(hyper::Error),
+}
+
+pub enum ServerHealth {
+    Listening(mpsc::UnboundedReceiver<ServerEvent>),
     Failure,
 }
 
@@ -158,7 +162,7 @@ pub struct Server {
     ip: Ipv4Addr,
     port: u16,
     main_worker_req_tx: mpsc::UnboundedSender<WorkerRequestMsg>,
-    callback_tx: Option<Sender<ServerCodes>>,
+    callback_tx: Option<Sender<ServerHealth>>,
     termination_token: TerminationToken,
 }
 
@@ -172,7 +176,7 @@ impl Server {
         maybe_user_worker_policy: Option<WorkerPoolPolicy>,
         import_map_path: Option<String>,
         no_module_cache: bool,
-        callback_tx: Option<Sender<ServerCodes>>,
+        callback_tx: Option<Sender<ServerHealth>>,
         entrypoints: WorkerEntrypoints,
         termination_token: Option<TerminationToken>,
     ) -> Result<Self, Error> {
@@ -237,10 +241,14 @@ impl Server {
         let listener = TcpListener::bind(&addr).await?;
         let termination_token = self.termination_token.clone();
 
+        let mut can_receive_event = false;
+        let (event_tx, event_rx) = mpsc::unbounded_channel();
+
         debug!("edge-runtime is listening on {:?}", listener.local_addr()?);
 
         if let Some(callback) = self.callback_tx.clone() {
-            let _ = callback.send(ServerCodes::Listening).await;
+            can_receive_event = true;
+            let _ = callback.send(ServerHealth::Listening(event_rx)).await;
         }
 
         loop {
@@ -250,21 +258,28 @@ impl Server {
                 msg = listener.accept() => {
                     match msg {
                         Ok((conn, _)) => {
-                            tokio::task::spawn(async move {
-                                let (service, cancel) = WorkerService::new(main_worker_req_tx);
-                                let _guard = cancel.drop_guard();
-
-                                let conn_fut = Http::new()
-                                    .serve_connection(conn, service);
-
-                                if let Err(e) = conn_fut.await {
-                                    // Most common cause for these errors are
-                                    // when the client closes the connection
-                                    // before we could send a response
-                                    if e.is_incomplete_message() {
-                                        debug!("connection reset ({:?})", e);
-                                    } else {
-                                        error!("client connection error ({:?})", e);
+                            tokio::task::spawn({
+                                let event_tx = event_tx.clone();
+                                async move {
+                                    let (service, cancel) = WorkerService::new(main_worker_req_tx);
+                                    let _guard = cancel.drop_guard();
+
+                                    let conn_fut = Http::new()
+                                        .serve_connection(conn, service);
+
+                                    if let Err(e) = conn_fut.await {
+                                        // Most common cause for these errors are
+                                        // when the client closes the connection
+                                        // before we could send a response
+                                        if e.is_incomplete_message() {
+                                            debug!("connection reset ({:?})", e);
+                                        } else {
+                                            error!("client connection error ({:?})", e);
+                                        }
+
+                                        if can_receive_event {
+                                            let _ = event_tx.send(ServerEvent::ConnectionError(e));
+                                        }
                                     }
                                 }
                             });