Handle cleanup of ws tasks on shutdown (#803)

Author: gi0baro

src/asgi/callbacks.rs

@@ -171,6 +171,7 @@ pub(crate) fn call_http(
 pub(crate) fn call_ws(
     cb: ArcCBScheduler,
     rt: RuntimeRef,
+    disconnect_guard: Arc<Notify>,
     server_addr: SockAddr,
     client_addr: SockAddr,
     scheme: HTTPProto,
@@ -179,7 +180,7 @@ pub(crate) fn call_ws(
     upgrade: UpgradeData,
 ) -> oneshot::Receiver<WebsocketDetachedTransport> {
     let (tx, rx) = oneshot::channel();
-    let protocol = WebsocketProtocol::new(rt.clone(), tx, ws, upgrade);
+    let protocol = WebsocketProtocol::new(rt.clone(), tx, ws, upgrade, disconnect_guard);
 
     rt.spawn_blocking(move |py| {
         if let Ok(scope) = build_scope_ws(py, req, server_addr, client_addr, scheme)

src/asgi/http.rs

@@ -81,13 +81,15 @@ macro_rules! handle_request_with_ws {
                         let (restx, mut resrx) = mpsc::channel(1);
                         let (parts, _) = req.into_parts();
                         let rth = rt.clone();
+                        let cancel_sig = Arc::new(Notify::new());
 
-                        rt.spawn(async move {
+                        rt.spawn_cancellable(cancel_sig.clone(), async move {
                             let tx_ref = restx.clone();
 
                             match $handler_ws(
                                 callback,
                                 rth,
+                                cancel_sig,
                                 server_addr,
                                 client_addr,
                                 scheme,

src/asgi/io.rs

@@ -345,6 +345,7 @@ impl WebsocketDetachedTransport {
 pub(crate) struct ASGIWebsocketProtocol {
     rt: RuntimeRef,
     tx: Mutex<Option<oneshot::Sender<WebsocketDetachedTransport>>>,
+    disconnect_guard: Arc<Notify>,
     websocket: Mutex<Option<HyperWebsocket>>,
     upgrade: Mutex<Option<UpgradeData>>,
     response_intent: Mutex<Option<(u16, HeaderMap)>>,
@@ -362,10 +363,12 @@ impl ASGIWebsocketProtocol {
         tx: oneshot::Sender<WebsocketDetachedTransport>,
         websocket: HyperWebsocket,
         upgrade: UpgradeData,
+        disconnect_guard: Arc<Notify>,
     ) -> Self {
         Self {
             rt,
             tx: Mutex::new(Some(tx)),
+            disconnect_guard,
             websocket: Mutex::new(Some(websocket)),
             upgrade: Mutex::new(Some(upgrade)),
             response_intent: Mutex::new(None),
@@ -526,6 +529,7 @@ impl ASGIWebsocketProtocol {
         let accepted_ev = self.init_event.clone();
         let closed = self.closed.clone();
         let transport = self.ws_rx.clone();
+        let guard_disconnect = self.disconnect_guard.clone();
 
         future_into_py_futlike(self.rt.clone(), py, async move {
             if !accepted.load(atomic::Ordering::Acquire) {
@@ -534,7 +538,11 @@ impl ASGIWebsocketProtocol {
             }
 
             if let Some(ws) = &mut *(transport.lock().await) {
-                while let Some(recv) = ws.next().await {
+                while let Some(recv) = tokio::select! {
+                    biased;
+                    recv = ws.next() => recv,
+                    () = guard_disconnect.notified() => Some(Err(tokio_tungstenite::tungstenite::Error::ConnectionClosed)),
+                } {
                     match recv {
                         Ok(Message::Ping(_) | Message::Pong(_)) => {}
                         Ok(message @ Message::Close(_)) => {

src/rsgi/callbacks.rs

@@ -138,12 +138,13 @@ pub(crate) fn call_http(
 pub(crate) fn call_ws(
     cb: ArcCBScheduler,
     rt: RuntimeRef,
+    disconnect_guard: Arc<Notify>,
     ws: HyperWebsocket,
     upgrade: UpgradeData,
     scope: WebsocketScope,
 ) -> oneshot::Receiver<WebsocketDetachedTransport> {
     let (tx, rx) = oneshot::channel();
-    let protocol = WebsocketProtocol::new(rt.clone(), tx, ws, upgrade);
+    let protocol = WebsocketProtocol::new(rt.clone(), tx, ws, upgrade, disconnect_guard);
 
     rt.spawn_blocking(move |py| {
         if let Ok(watcher) = CallbackWatcherWebsocket::new(py, protocol, scope) {

src/rsgi/http.rs

@@ -82,11 +82,13 @@ macro_rules! handle_request_with_ws {
                         let scope = build_scope!(WebsocketScope, server_addr, client_addr, parts, scheme);
                         let (restx, mut resrx) = mpsc::channel(1);
                         let rth = rt.clone();
+                        let cancel_sig = Arc::new(Notify::new());
 
-                        rt.spawn(async move {
+                        rt.spawn_cancellable(cancel_sig.clone(), async move {
                             let tx_ref = restx.clone();
 
-                            match $handler_ws(callback, rth, ws, UpgradeData::new(res, restx), scope).await {
+                            match $handler_ws(callback, rth, cancel_sig, ws, UpgradeData::new(res, restx), scope).await
+                            {
                                 Ok((status, consumed, stream)) => match (consumed, stream) {
                                     (false, _) => {
                                         let _ = tx_ref

src/rsgi/io.rs

@@ -226,14 +226,16 @@ impl RSGIHTTPProtocol {
 #[pyclass(frozen, module = "granian._granian")]
 pub(crate) struct RSGIWebsocketTransport {
     rt: RuntimeRef,
+    dg: Arc<Notify>,
     tx: Arc<AsyncMutex<Option<WSTxStream>>>,
     rx: Arc<AsyncMutex<WSRxStream>>,
 }
 
 impl RSGIWebsocketTransport {
-    pub fn new(rt: RuntimeRef, tx: Arc<AsyncMutex<Option<WSTxStream>>>, rx: WSRxStream) -> Self {
+    pub fn new(rt: RuntimeRef, dg: Arc<Notify>, tx: Arc<AsyncMutex<Option<WSTxStream>>>, rx: WSRxStream) -> Self {
         Self {
             rt,
+            dg,
             tx,
             rx: Arc::new(AsyncMutex::new(rx)),
         }
@@ -244,9 +246,15 @@ impl RSGIWebsocketTransport {
 impl RSGIWebsocketTransport {
     fn receive<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
         let transport = self.rx.clone();
+        let dg = self.dg.clone();
+
         future_into_py_futlike(self.rt.clone(), py, async move {
             if let Ok(mut stream) = transport.try_lock() {
-                while let Some(recv) = stream.next().await {
+                while let Some(recv) = tokio::select! {
+                    biased;
+                    recv = stream.next() => recv,
+                    () = dg.notified() => Some(Err(tokio_tungstenite::tungstenite::Error::ConnectionClosed)),
+                } {
                     match recv {
                         Ok(Message::Ping(_) | Message::Pong(_)) => {}
                         Ok(message) => return FutureResultToPy::RSGIWSMessage(message),
@@ -297,6 +305,7 @@ impl RSGIWebsocketTransport {
 pub(crate) struct RSGIWebsocketProtocol {
     rt: RuntimeRef,
     tx: Mutex<Option<oneshot::Sender<WebsocketDetachedTransport>>>,
+    disconnect_guard: Arc<Notify>,
     websocket: Arc<AsyncMutex<HyperWebsocket>>,
     upgrade: RwLock<Option<UpgradeData>>,
     transport: Arc<AsyncMutex<Option<WSTxStream>>>,
@@ -308,10 +317,12 @@ impl RSGIWebsocketProtocol {
         tx: oneshot::Sender<WebsocketDetachedTransport>,
         websocket: HyperWebsocket,
         upgrade: UpgradeData,
+        disconnect_guard: Arc<Notify>,
     ) -> Self {
         Self {
             rt,
             tx: Mutex::new(Some(tx)),
+            disconnect_guard,
             websocket: Arc::new(AsyncMutex::new(websocket)),
             upgrade: RwLock::new(Some(upgrade)),
             transport: Arc::new(AsyncMutex::new(None)),
@@ -341,9 +352,11 @@ impl RSGIWebsocketProtocol {
 
     fn accept<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
         let rth = self.rt.clone();
+        let dg = self.disconnect_guard.clone();
         let mut upgrade = self.upgrade.write().unwrap().take().unwrap();
         let transport = self.websocket.clone();
         let itransport = self.transport.clone();
+
         future_into_py_futlike(self.rt.clone(), py, async move {
             let mut ws = transport.lock().await;
             match upgrade.send(None, None, None).await {
@@ -354,11 +367,7 @@ impl RSGIWebsocketProtocol {
                             let mut guard = itransport.lock().await;
                             *guard = Some(stx);
                         }
-                        FutureResultToPy::RSGIWSAccept(RSGIWebsocketTransport::new(
-                            rth.clone(),
-                            itransport.clone(),
-                            srx,
-                        ))
+                        FutureResultToPy::RSGIWSAccept(RSGIWebsocketTransport::new(rth, dg, itransport, srx))
                     }
                     _ => FutureResultToPy::Err(error_proto!()),
                 },

src/runtime.rs

@@ -34,6 +34,10 @@ pub trait Runtime: Send + 'static {
     fn spawn_blocking<F>(&self, task: F)
     where
         F: FnOnce(Python) + Send + 'static;
+
+    fn spawn_cancellable<F>(&self, on_cancel: Arc<tokio::sync::Notify>, fut: F) -> Self::JoinHandle
+    where
+        F: Future<Output = ()> + Send + 'static;
 }
 
 pub trait ContextExt: Runtime {
@@ -44,6 +48,7 @@ pub(crate) struct RuntimeWrapper {
     pub inner: tokio::runtime::Runtime,
     br: Arc<blocking::BlockingRunner>,
     pr: Arc<Py<PyAny>>,
+    sig: Arc<tokio::sync::Notify>,
 }
 
 impl RuntimeWrapper {
@@ -62,6 +67,7 @@ impl RuntimeWrapper {
             inner: default_runtime(blocking_threads),
             br: br.into(),
             pr: py_loop,
+            sig: tokio::sync::Notify::new().into(),
         }
     }
 
@@ -80,11 +86,17 @@ impl RuntimeWrapper {
             inner: rt,
             br: br.into(),
             pr: py_loop,
+            sig: tokio::sync::Notify::new().into(),
         }
     }
 
     pub fn handler(&self) -> RuntimeRef {
-        RuntimeRef::new(self.inner.handle().clone(), self.br.clone(), self.pr.clone())
+        RuntimeRef::new(
+            self.inner.handle().clone(),
+            self.br.clone(),
+            self.pr.clone(),
+            self.sig.clone(),
+        )
     }
 }
 
@@ -93,16 +105,27 @@ pub struct RuntimeRef {
     pub inner: tokio::runtime::Handle,
     innerb: Arc<blocking::BlockingRunner>,
     innerp: Arc<Py<PyAny>>,
+    sig: Arc<tokio::sync::Notify>,
 }
 
 impl RuntimeRef {
-    pub fn new(rt: tokio::runtime::Handle, br: Arc<blocking::BlockingRunner>, pyloop: Arc<Py<PyAny>>) -> Self {
+    pub fn new(
+        rt: tokio::runtime::Handle,
+        br: Arc<blocking::BlockingRunner>,
+        pyloop: Arc<Py<PyAny>>,
+        sig: Arc<tokio::sync::Notify>,
+    ) -> Self {
         Self {
             inner: rt,
             innerb: br,
             innerp: pyloop,
+            sig,
         }
     }
+
+    pub fn close(&self) {
+        self.sig.notify_waiters();
+    }
 }
 
 impl JoinError for tokio::task::JoinError {
@@ -129,6 +152,23 @@ impl Runtime for RuntimeRef {
     {
         _ = self.innerb.run(task);
     }
+
+    fn spawn_cancellable<F>(&self, on_cancel: Arc<tokio::sync::Notify>, fut: F) -> Self::JoinHandle
+    where
+        F: Future<Output = ()> + Send + 'static,
+    {
+        let sig = self.sig.clone();
+
+        self.inner.spawn(async move {
+            tokio::select! {
+                biased;
+                () = fut => {},
+                () = sig.notified() => {
+                    on_cancel.notify_one();
+                }
+            };
+        })
+    }
 }
 
 impl ContextExt for RuntimeRef {
@@ -293,7 +333,7 @@ where
 }
 
 #[allow(unused_must_use)]
-pub(crate) fn run_until_complete<F>(rt: RuntimeWrapper, event_loop: Bound<PyAny>, fut: F) -> PyResult<()>
+pub(crate) fn run_until_complete<F>(rt: &RuntimeWrapper, event_loop: Bound<PyAny>, fut: F) -> PyResult<()>
 where
     F: Future<Output = PyResult<()>> + Send + 'static,
 {

src/serve.rs

@@ -80,11 +80,12 @@ macro_rules! serve_fn {
             let wrk = crate::workers::Worker::new(ctx, acceptor, handler, rth, target, metrics.0);
             let tasks = wrk.tasks.clone();
 
-            let main_loop = crate::runtime::run_until_complete(rt, event_loop.clone(), async move {
+            let main_loop = crate::runtime::run_until_complete(&rt, event_loop.clone(), async move {
                 wrk.listen(srx, listener, backpressure).await;
 
                 log::info!("Stopping worker-{worker_id}");
 
+                wrk.rt.close();
                 tasks.close();
                 tasks.wait().await;
                 mc_notify.notified().await;
@@ -93,6 +94,8 @@ macro_rules! serve_fn {
                 Ok(())
             });
 
+            drop(rt);
+
             if let Err(err) = main_loop {
                 log::error!("{err}");
                 std::process::exit(1);
@@ -174,6 +177,7 @@ macro_rules! serve_fn {
 
                         log::info!("Stopping worker-{} runtime-{}", worker_id, thread_id + 1);
 
+                        wrk.rt.close();
                         tasks.close();
                         tasks.wait().await;
 
@@ -212,7 +216,7 @@ macro_rules! serve_fn {
                 mc_notify.notify_one();
             }
 
-            let main_loop = crate::runtime::run_until_complete(rtm, event_loop.clone(), async move {
+            let main_loop = crate::runtime::run_until_complete(&rtm, event_loop.clone(), async move {
                 let _ = pyrx.changed().await;
                 stx.send(true).unwrap();
                 log::info!("Stopping worker-{worker_id}");
@@ -223,6 +227,8 @@ macro_rules! serve_fn {
                 Ok(())
             });
 
+            drop(rtm);
+
             if let Err(err) = main_loop {
                 log::error!("{err}");
                 std::process::exit(1);
@@ -293,6 +299,7 @@ macro_rules! serve_fn {
 
                     log::info!("Stopping worker-{worker_id}");
 
+                    wrk.rt.close();
                     tasks.close();
                     tasks.wait().await;
 

src/workers.rs

@@ -291,7 +291,7 @@ pub(crate) struct Worker<C, A, H, F, M> {
     ctx: C,
     acceptor: A,
     handler: H,
-    rt: crate::runtime::RuntimeRef,
+    pub rt: crate::runtime::RuntimeRef,
     pub tasks: tokio_util::task::TaskTracker,
     target: F,
     metrics: M,