Improve coroutine handling

Qard · Qard · commit 232c8343f836 · 2025-07-08T23:06:35.000+08:00
diff --git a/src/asgi/mod.rs b/src/asgi/mod.rs
@@ -26,6 +26,27 @@ pub use websocket::{
   WebSocketConnectionScope, WebSocketReceiveMessage, WebSocketSendException, WebSocketSendMessage,
 };
 
+fn execute_coroutine(
+  py_func: PyObject,
+  scope: HttpConnectionScope,
+  rx_receiver: Receiver,
+  tx_sender: Sender,
+) -> PyResult<()> {
+  Python::with_gil(|py| {
+    let scope_py = scope.into_pyobject(py)?;
+    let coroutine = py_func.call1(py, (scope_py, rx_receiver, tx_sender))?;
+
+    let asyncio = py.import("asyncio")?;
+    let loop_ = asyncio.call_method0("new_event_loop")?;
+    asyncio.call_method1("set_event_loop", (&loop_,))?;
+
+    loop_.call_method1("run_until_complete", (coroutine,))?;
+    loop_.call_method0("close")?;
+
+    Ok(())
+  })
+}
+
 pub async fn execute_asgi_http_scope(
   py_func: PyObject,
   request: Request,
@@ -38,7 +59,7 @@ pub async fn execute_asgi_http_scope(
   let (tx_sender, mut tx) = Sender::http();
 
   // Channel to receive the response data
-  let (response_tx, response_rx) = oneshot::channel::<(u16, Vec<(String, String)>, Vec<u8>)>();
+  let (response_tx, response_rx) = oneshot::channel::<Result<(u16, Vec<(String, String)>, Vec<u8>), String>>();
 
   // Channel to signal Python execution completion
   let (python_done_tx, python_done_rx) = oneshot::channel::<Result<(), String>>();
@@ -50,115 +71,83 @@ pub async fn execute_asgi_http_scope(
     let mut body = Vec::new();
     let mut response_started = false;
 
-    while let Some(msg) = tx.recv().await {
-      match msg {
-        HttpSendMessage::HttpResponseStart {
-          status: s,
-          headers: h,
-          trailers: _
-        } => {
-          status = s;
-          headers = h;
-          response_started = true;
-        }
-        HttpSendMessage::HttpResponseBody {
-          body: b,
-          more_body
-        } => {
-          if response_started {
-            body.extend_from_slice(&b);
-            if !more_body {
-              // Response is complete
-              let _ = response_tx.send((status, headers, body));
-              break;
+    while let Some(msg_result) = tx.recv().await {
+      match msg_result {
+        Ok(msg) => {
+          match msg {
+            HttpSendMessage::HttpResponseStart {
+              status: s,
+              headers: h,
+              trailers: _
+            } => {
+              status = s;
+              headers = h;
+              response_started = true;
+            }
+            HttpSendMessage::HttpResponseBody {
+              body: b,
+              more_body
+            } => {
+              if response_started {
+                body.extend_from_slice(&b);
+                if !more_body {
+                  // Response is complete
+                  let _ = response_tx.send(Ok((status, headers, body)));
+                  break;
+                }
+              }
             }
           }
         }
+        Err(e) => {
+          // Error from Python side
+          let _ = response_tx.send(Err(e));
+          break;
+        }
       }
     }
   });
 
   // Send the request body to the ASGI app
   let request_body = request.body().clone();
+  let rx_for_request = rx.clone();
   tokio::spawn(async move {
     let request_message = HttpReceiveMessage::Request {
       body: request_body.to_vec(),
       more_body: false,
     };
 
-    if rx.send(request_message).is_err() {
-      eprintln!("Failed to send request message");
+    if rx_for_request.send(Ok(request_message)).is_err() {
+      // Channel closed, but we can't do much about it here
+      // The error will be handled by the main flow
     }
   });
 
   // Run the ASGI app in a Python thread
   tokio::task::spawn_blocking(move || {
-    Python::with_gil(|py| {
-      let scope_py = scope.into_pyobject(py).unwrap();
-
-      // Call the ASGI app to get a coroutine
-      let coroutine = match py_func.call1(py, (scope_py, rx_receiver, tx_sender)) {
-        Ok(coro) => coro,
-        Err(e) => {
-          eprintln!("Failed to call ASGI app: {e}");
-          let _ = python_done_tx.send(Err(format!("Failed to call ASGI app: {e}")));
-          return;
-        }
-      };
+    let result = execute_coroutine(py_func, scope, rx_receiver, tx_sender)
+      .map_err(|e| format!("Failed to execute ASGI app: {e}"));
 
-      // Run the coroutine using asyncio
-      let asyncio = match py.import("asyncio") {
-        Ok(module) => module,
-        Err(e) => {
-          eprintln!("Failed to import asyncio: {e}");
-          let _ = python_done_tx.send(Err(format!("Failed to import asyncio: {e}")));
-          return;
-        }
-      };
-
-      // Create a new event loop for this request
-      let loop_ = match asyncio.call_method0("new_event_loop") {
-        Ok(loop_) => loop_,
-        Err(e) => {
-          eprintln!("Failed to create event loop: {e}");
-          let _ = python_done_tx.send(Err(format!("Failed to create event loop: {e}")));
-          return;
-        }
-      };
-
-      // Set it as the current event loop
-      if let Err(e) = asyncio.call_method1("set_event_loop", (&loop_,)) {
-        eprintln!("Failed to set event loop: {e}");
-        let _ = python_done_tx.send(Err(format!("Failed to set event loop: {e}")));
-        return;
-      }
-
-      // Run the coroutine
-      let result = loop_.call_method1("run_until_complete", (coroutine,));
-
-      // Close the loop
-      let _ = loop_.call_method0("close");
-
-      // Send the result after a small delay to allow response to be processed
-      // This ensures the response handler has time to receive messages from the ASGI app
-      std::thread::sleep(std::time::Duration::from_millis(10));
-
-      let _ = python_done_tx.send(result.map(|_| ()).map_err(|e| {
-        format!("Failed to run ASGI coroutine: {e}")
-      }));
-    });
+    // If the coroutine execution was successful, we send a success signal
+    // Otherwise, we send the error
+    if let Err(err) = python_done_tx.send(result) {
+      eprintln!("Send error: {err:?}");
+    }
   });
 
   // Wait for either the response or Python completion
   let result = tokio::select! {
     response = response_rx => {
-      response.map_err(|_| pyo3::exceptions::PyRuntimeError::new_err("Failed to receive response"))
+      match response {
+        Ok(Ok(response_data)) => Ok(response_data),
+        Ok(Err(e)) => Err(pyo3::exceptions::PyRuntimeError::new_err(e)),
+        Err(_) => Err(pyo3::exceptions::PyRuntimeError::new_err("Failed to receive response"))
+      }
     }
     python_result = python_done_rx => {
       match python_result {
         Ok(Ok(())) => {
           // Python completed but no response was sent
-          eprintln!("ASGI app completed without sending response");
           Err(pyo3::exceptions::PyRuntimeError::new_err("ASGI app completed without sending response"))
         }
         Ok(Err(e)) => {
diff --git a/src/asgi/receiver.rs b/src/asgi/receiver.rs
@@ -11,9 +11,9 @@ use crate::asgi::{
 };
 
 enum ReceiverType {
-  Http(Arc<Mutex<mpsc::UnboundedReceiver<HttpReceiveMessage>>>),
-  WebSocket(Arc<Mutex<mpsc::UnboundedReceiver<WebSocketReceiveMessage>>>),
-  Lifespan(Arc<Mutex<mpsc::UnboundedReceiver<LifespanReceiveMessage>>>),
+  Http(Arc<Mutex<mpsc::UnboundedReceiver<Result<HttpReceiveMessage, String>>>>),
+  WebSocket(Arc<Mutex<mpsc::UnboundedReceiver<Result<WebSocketReceiveMessage, String>>>>),
+  Lifespan(Arc<Mutex<mpsc::UnboundedReceiver<Result<LifespanReceiveMessage, String>>>>),
 }
 
 /// Allows Python to receive messages from Rust.
@@ -22,22 +22,22 @@ pub struct Receiver(ReceiverType);
 
 impl Receiver {
   /// Create a new Receiver instance for http ASGI message types.
-  pub fn http() -> (Receiver, mpsc::UnboundedSender<HttpReceiveMessage>) {
-    let (tx, rx) = mpsc::unbounded_channel::<HttpReceiveMessage>();
+  pub fn http() -> (Receiver, mpsc::UnboundedSender<Result<HttpReceiveMessage, String>>) {
+    let (tx, rx) = mpsc::unbounded_channel::<Result<HttpReceiveMessage, String>>();
     let rx = Arc::new(Mutex::new(rx));
     (Receiver(ReceiverType::Http(rx)), tx)
   }
 
   /// Create a new Receiver instance for websocket ASGI message types.
-  pub fn websocket() -> (Receiver, mpsc::UnboundedSender<WebSocketReceiveMessage>) {
-    let (tx, rx) = mpsc::unbounded_channel::<WebSocketReceiveMessage>();
+  pub fn websocket() -> (Receiver, mpsc::UnboundedSender<Result<WebSocketReceiveMessage, String>>) {
+    let (tx, rx) = mpsc::unbounded_channel::<Result<WebSocketReceiveMessage, String>>();
     let rx = Arc::new(Mutex::new(rx));
     (Receiver(ReceiverType::WebSocket(rx)), tx)
   }
 
   /// Create a new Receiver instance for lifespan ASGI message types.
-  pub fn lifespan() -> (Receiver, mpsc::UnboundedSender<LifespanReceiveMessage>) {
-    let (tx, rx) = mpsc::unbounded_channel::<LifespanReceiveMessage>();
+  pub fn lifespan() -> (Receiver, mpsc::UnboundedSender<Result<LifespanReceiveMessage, String>>) {
+    let (tx, rx) = mpsc::unbounded_channel::<Result<LifespanReceiveMessage, String>>();
     let rx = Arc::new(Mutex::new(rx));
     (Receiver(ReceiverType::Lifespan(rx)), tx)
   }
@@ -49,26 +49,26 @@ impl Receiver {
     match &self.0 {
       ReceiverType::Http(rx) => {
         let message = rx.lock().await.recv().await;
-        if let Some(msg) = message {
-          Python::with_gil(|py| Ok(msg.into_pyobject(py)?.unbind()))
-        } else {
-          Err(PyValueError::new_err("No message received"))
+        match message {
+          Some(Ok(msg)) => Python::with_gil(|py| Ok(msg.into_pyobject(py)?.unbind())),
+          Some(Err(err)) => Err(PyValueError::new_err(err)),
+          None => Err(PyValueError::new_err("No message received"))
         }
       }
       ReceiverType::WebSocket(rx) => {
         let message = rx.lock().await.recv().await;
-        if let Some(msg) = message {
-          Python::with_gil(|py| Ok(msg.into_pyobject(py)?.unbind()))
-        } else {
-          Err(PyValueError::new_err("No message received"))
+        match message {
+          Some(Ok(msg)) => Python::with_gil(|py| Ok(msg.into_pyobject(py)?.unbind())),
+          Some(Err(err)) => Err(PyValueError::new_err(err)),
+          None => Err(PyValueError::new_err("No message received"))
         }
       }
       ReceiverType::Lifespan(rx) => {
         let message = rx.lock().await.recv().await;
-        if let Some(msg) = message {
-          Python::with_gil(|py| Ok(msg.into_pyobject(py)?.unbind()))
-        } else {
-          Err(PyValueError::new_err("No message received"))
+        match message {
+          Some(Ok(msg)) => Python::with_gil(|py| Ok(msg.into_pyobject(py)?.unbind())),
+          Some(Err(err)) => Err(PyValueError::new_err(err)),
+          None => Err(PyValueError::new_err("No message received"))
         }
       }
     }
diff --git a/src/asgi/sender.rs b/src/asgi/sender.rs
@@ -9,28 +9,28 @@ use crate::asgi::{
 };
 
 enum SenderType {
-  Http(mpsc::UnboundedSender<HttpSendMessage>),
-  WebSocket(mpsc::UnboundedSender<WebSocketSendMessage>),
-  Lifespan(mpsc::UnboundedSender<LifespanSendMessage>),
+  Http(mpsc::UnboundedSender<Result<HttpSendMessage, String>>),
+  WebSocket(mpsc::UnboundedSender<Result<WebSocketSendMessage, String>>),
+  Lifespan(mpsc::UnboundedSender<Result<LifespanSendMessage, String>>),
 }
 
 /// Allows Python to send messages to Rust.
 #[pyclass]
 pub struct Sender(SenderType);
 
 impl Sender {
-  pub fn http() -> (Sender, mpsc::UnboundedReceiver<HttpSendMessage>) {
-    let (tx, rx) = mpsc::unbounded_channel::<HttpSendMessage>();
+  pub fn http() -> (Sender, mpsc::UnboundedReceiver<Result<HttpSendMessage, String>>) {
+    let (tx, rx) = mpsc::unbounded_channel::<Result<HttpSendMessage, String>>();
     (Sender(SenderType::Http(tx)), rx)
   }
 
-  pub fn websocket() -> (Sender, mpsc::UnboundedReceiver<WebSocketSendMessage>) {
-    let (tx, rx) = mpsc::unbounded_channel::<WebSocketSendMessage>();
+  pub fn websocket() -> (Sender, mpsc::UnboundedReceiver<Result<WebSocketSendMessage, String>>) {
+    let (tx, rx) = mpsc::unbounded_channel::<Result<WebSocketSendMessage, String>>();
     (Sender(SenderType::WebSocket(tx)), rx)
   }
 
-  pub fn lifespan() -> (Sender, mpsc::UnboundedReceiver<LifespanSendMessage>) {
-    let (tx, rx) = mpsc::unbounded_channel::<LifespanSendMessage>();
+  pub fn lifespan() -> (Sender, mpsc::UnboundedReceiver<Result<LifespanSendMessage, String>>) {
+    let (tx, rx) = mpsc::unbounded_channel::<Result<LifespanSendMessage, String>>();
     (Sender(SenderType::Lifespan(tx)), rx)
   }
 }
@@ -42,17 +42,17 @@ impl Sender {
       match &self.0 {
         SenderType::Http(tx) => {
           let msg: HttpSendMessage = args.extract(py)?;
-          tx.send(msg)
+          tx.send(Ok(msg))
             .map_err(|_| PyValueError::new_err("connection closed"))?;
         }
         SenderType::WebSocket(tx) => {
           let msg: WebSocketSendMessage = args.extract(py)?;
-          tx.send(msg)
+          tx.send(Ok(msg))
             .map_err(|_| PyValueError::new_err("connection closed"))?;
         }
         SenderType::Lifespan(tx) => {
           let msg: LifespanSendMessage = args.extract(py)?;
-          tx.send(msg)
+          tx.send(Ok(msg))
             .map_err(|_| PyValueError::new_err("connection closed"))?;
         }
       };
