Add client disconnection handling in ASGI and RSGI (#524)

* Handle client disconnection in ASGI

* Add `RSGIHTTPProtocol.client_disconnect`

Author: gi0baro

docs/spec/RSGI.md

@@ -1,6 +1,6 @@
 # RSGI Specification
 
-**Version:** 1.4
+**Version:** 1.5
 
 ## Abstract
 
@@ -165,10 +165,11 @@ And here are descriptions for the upper attributes:
 
 #### HTTP protocol interface
 
-HTTP protocol object implements two awaitable methods to receive the request body, and five different methods to send data, in particular:
+HTTP protocol object implements two awaitable methods to receive the request body, five different methods to send data, and one awaitable method to wait for client disconnection, in particular:
 
 - `__call__` to receive the entire body in `bytes` format
 - `__aiter__` to receive the body in `bytes` chunks
+- `client_disconnect` to watch for client disconnection
 - `response_empty` to send back an empty response
 - `response_str` to send back a response with a `str` body
 - `response_bytes` to send back a response with `bytes` body
@@ -180,6 +181,7 @@ All the upper-mentioned response methods accepts an integer `status` parameter,
 ```
 coroutine __call__() -> body
 asynciterator __aiter__() -> body chunks
+coroutine client_disconnect()
 function response_empty(status, headers)
 function response_str(status, headers, body)
 function response_bytes(status, headers, body)
@@ -197,6 +199,10 @@ coroutine send_bytes(bytes)
 coroutine send_str(str)
 ```
 
+The `client_disconnect` method will return a future that resolve ones the client has disconnected.
+
+> **Note:** as HTTP supports keep-alived connections, the lifecycle of the client connection might not be the same of the single request. This is why the RSGI specification doesn't imply `client_disconnect` should resolve in case a client sends multiple requests within the same connection, and thus the protocol delegates to the application the responsibility to cancel the disconnection watcher once the response is sent.
+
 ### Websocket protocol
 
 WebSockets share some HTTP details - they have a path and headers - but also have more state. Again, most of that state is in the scope, which will live as long as the socket does.

granian/_granian.pyi

@@ -21,6 +21,7 @@ class RSGIHTTPStreamTransport:
 class RSGIHTTPProtocol:
     async def __call__(self) -> bytes: ...
     def __aiter__(self) -> Any: ...
+    async def client_disconnect(self) -> None: ...
     def response_empty(self, status: int, headers: List[Tuple[str, str]]): ...
     def response_str(self, status: int, headers: List[Tuple[str, str]], body: str): ...
     def response_bytes(self, status: int, headers: List[Tuple[str, str]], body: bytes): ...

src/asgi/callbacks.rs

@@ -1,7 +1,10 @@
 use pyo3::prelude::*;
 use pyo3::types::PyDict;
-use std::{net::SocketAddr, sync::OnceLock};
-use tokio::sync::oneshot;
+use std::{
+    net::SocketAddr,
+    sync::{Arc, OnceLock},
+};
+use tokio::sync::{oneshot, Notify};
 
 use super::{
     io::{ASGIHTTPProtocol as HTTPProtocol, ASGIWebsocketProtocol as WebsocketProtocol, WebsocketDetachedTransport},
@@ -145,14 +148,15 @@ impl CallbackWatcherWebsocket {
 pub(crate) fn call_http(
     cb: ArcCBScheduler,
     rt: RuntimeRef,
+    disconnect_guard: Arc<Notify>,
     server_addr: SocketAddr,
     client_addr: SocketAddr,
     scheme: &str,
     req: hyper::http::request::Parts,
     body: hyper::body::Incoming,
 ) -> oneshot::Receiver<HTTPResponse> {
     let (tx, rx) = oneshot::channel();
-    let protocol = HTTPProtocol::new(rt.clone(), body, tx);
+    let protocol = HTTPProtocol::new(rt.clone(), body, tx, disconnect_guard);
     let scheme: Box<str> = scheme.into();
 
     rt.spawn_blocking(move |py| {

src/asgi/http.rs

@@ -1,7 +1,7 @@
 use http_body_util::BodyExt;
 use hyper::{header::SERVER as HK_SERVER, http::response::Builder as ResponseBuilder, StatusCode};
-use std::net::SocketAddr;
-use tokio::sync::mpsc;
+use std::{net::SocketAddr, sync::Arc};
+use tokio::sync::{mpsc, Notify};
 
 use super::callbacks::{call_http, call_ws};
 use crate::{
@@ -16,8 +16,19 @@ const SCHEME_WS: &str = "ws";
 const SCHEME_WSS: &str = "wss";
 
 macro_rules! handle_http_response {
-    ($handler:expr, $rt:expr, $callback:expr, $server_addr:expr, $client_addr:expr, $scheme:expr, $req:expr, $body:expr) => {
-        match $handler($callback, $rt, $server_addr, $client_addr, $req, $scheme, $body).await {
+    ($handler:expr, $rt:expr, $disconnect_guard:expr, $callback:expr, $server_addr:expr, $client_addr:expr, $scheme:expr, $req:expr, $body:expr) => {
+        match $handler(
+            $callback,
+            $rt,
+            $disconnect_guard,
+            $server_addr,
+            $client_addr,
+            $req,
+            $scheme,
+            $body,
+        )
+        .await
+        {
             Ok(res) => res,
             _ => {
                 log::error!("ASGI protocol failure");
@@ -32,6 +43,7 @@ macro_rules! handle_request {
         #[inline]
         pub(crate) async fn $func_name(
             rt: RuntimeRef,
+            disconnect_guard: Arc<Notify>,
             callback: ArcCBScheduler,
             server_addr: SocketAddr,
             client_addr: SocketAddr,
@@ -42,6 +54,7 @@ macro_rules! handle_request {
             handle_http_response!(
                 $handler,
                 rt,
+                disconnect_guard,
                 callback,
                 server_addr,
                 client_addr,
@@ -58,6 +71,7 @@ macro_rules! handle_request_with_ws {
         #[inline]
         pub(crate) async fn $func_name(
             rt: RuntimeRef,
+            disconnect_guard: Arc<Notify>,
             callback: ArcCBScheduler,
             server_addr: SocketAddr,
             client_addr: SocketAddr,
@@ -142,6 +156,7 @@ macro_rules! handle_request_with_ws {
             handle_http_response!(
                 $handler_req,
                 rt,
+                disconnect_guard,
                 callback,
                 server_addr,
                 client_addr,

src/asgi/io.rs

@@ -13,7 +13,7 @@ use std::{
 };
 use tokio::{
     fs::File,
-    sync::{mpsc, oneshot, Mutex as AsyncMutex},
+    sync::{mpsc, oneshot, Mutex as AsyncMutex, Notify},
 };
 use tokio_tungstenite::tungstenite::Message;
 use tokio_util::io::ReaderStream;
@@ -25,7 +25,9 @@ use super::{
 use crate::{
     conversion::FutureResultToPy,
     http::{response_404, HTTPResponse, HTTPResponseBody, HV_SERVER},
-    runtime::{empty_future_into_py, err_future_into_py, future_into_py_futlike, Runtime, RuntimeRef},
+    runtime::{
+        done_future_into_py, empty_future_into_py, err_future_into_py, future_into_py_futlike, Runtime, RuntimeRef,
+    },
     ws::{HyperWebsocket, UpgradeData, WSRxStream, WSTxStream},
 };
 
@@ -37,27 +39,36 @@ static WS_SUBPROTO_HNAME: &str = "Sec-WebSocket-Protocol";
 pub(crate) struct ASGIHTTPProtocol {
     rt: RuntimeRef,
     tx: Mutex<Option<oneshot::Sender<HTTPResponse>>>,
+    disconnect_guard: Arc<Notify>,
     request_body: Arc<AsyncMutex<http_body_util::BodyStream<body::Incoming>>>,
     response_started: atomic::AtomicBool,
     response_chunked: atomic::AtomicBool,
     response_intent: Mutex<Option<(u16, HeaderMap)>>,
     body_tx: Mutex<Option<mpsc::Sender<Result<body::Bytes, anyhow::Error>>>>,
     flow_rx_exhausted: Arc<atomic::AtomicBool>,
+    flow_rx_closed: Arc<atomic::AtomicBool>,
     flow_tx_waiter: Arc<tokio::sync::Notify>,
     sent_response_code: Arc<atomic::AtomicU16>,
 }
 
 impl ASGIHTTPProtocol {
-    pub fn new(rt: RuntimeRef, body: hyper::body::Incoming, tx: oneshot::Sender<HTTPResponse>) -> Self {
+    pub fn new(
+        rt: RuntimeRef,
+        body: hyper::body::Incoming,
+        tx: oneshot::Sender<HTTPResponse>,
+        disconnect_guard: Arc<Notify>,
+    ) -> Self {
         Self {
             rt,
             tx: Mutex::new(Some(tx)),
+            disconnect_guard,
             request_body: Arc::new(AsyncMutex::new(http_body_util::BodyStream::new(body))),
             response_started: false.into(),
             response_chunked: false.into(),
             response_intent: Mutex::new(None),
             body_tx: Mutex::new(None),
             flow_rx_exhausted: Arc::new(atomic::AtomicBool::new(false)),
+            flow_rx_closed: Arc::new(atomic::AtomicBool::new(false)),
             flow_tx_waiter: Arc::new(tokio::sync::Notify::new()),
             sent_response_code: Arc::new(atomic::AtomicU16::new(500)),
         }
@@ -108,35 +119,56 @@ impl ASGIHTTPProtocol {
 #[pymethods]
 impl ASGIHTTPProtocol {
     fn receive<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
+        if self.flow_rx_closed.load(atomic::Ordering::Acquire) {
+            return done_future_into_py(
+                py,
+                super::conversion::message_into_py(py, ASGIMessageType::HTTPDisconnect).map(Bound::unbind),
+            );
+        }
+
         if self.flow_rx_exhausted.load(atomic::Ordering::Acquire) {
             let flow_hld = self.flow_tx_waiter.clone();
+            let flow_dgr = self.disconnect_guard.clone();
+            let flow_dsr = self.flow_rx_closed.clone();
             return future_into_py_futlike(self.rt.clone(), py, async move {
-                let () = flow_hld.notified().await;
+                tokio::select! {
+                    () = flow_hld.notified() => {},
+                    () = flow_dgr.notified() => flow_dsr.store(true, atomic::Ordering::Release),
+                }
                 FutureResultToPy::ASGIMessage(ASGIMessageType::HTTPDisconnect)
             });
         }
 
         let body_ref = self.request_body.clone();
         let flow_ref = self.flow_rx_exhausted.clone();
         let flow_hld = self.flow_tx_waiter.clone();
+        let flow_dgr = self.disconnect_guard.clone();
+        let flow_dsr = self.flow_rx_closed.clone();
         future_into_py_futlike(self.rt.clone(), py, async move {
             let mut bodym = body_ref.lock().await;
             let body = &mut *bodym;
             let mut more_body = false;
-            let chunk = match body.next().await {
-                Some(Ok(buf)) => {
-                    more_body = true;
-                    Ok(buf.into_data().unwrap_or_default())
+
+            let chunk = tokio::select! {
+                frame = body.next() => match frame {
+                    Some(Ok(buf)) => {
+                        more_body = true;
+                        Some(buf.into_data().unwrap_or_default())
+                    }
+                    Some(Err(_)) => None,
+                    _ => Some(body::Bytes::new()),
+                },
+                () = flow_dgr.notified() => {
+                    flow_dsr.store(true, atomic::Ordering::Release);
+                    None
                 }
-                Some(Err(err)) => Err(err),
-                _ => Ok(body::Bytes::new()),
             };
             if !more_body {
                 flow_ref.store(true, atomic::Ordering::Release);
             }
 
             match chunk {
-                Ok(data) => FutureResultToPy::ASGIMessage(ASGIMessageType::HTTPRequestBody((data, more_body))),
+                Some(data) => FutureResultToPy::ASGIMessage(ASGIMessageType::HTTPRequestBody((data, more_body))),
                 _ => {
                     flow_hld.notify_one();
                     FutureResultToPy::ASGIMessage(ASGIMessageType::HTTPDisconnect)

src/callbacks.rs

@@ -333,6 +333,35 @@ impl PyEmptyAwaitable {
     }
 }
 
+#[pyclass(frozen, module = "granian._granian")]
+pub(crate) struct PyDoneAwaitable {
+    result: PyResult<PyObject>,
+}
+
+impl PyDoneAwaitable {
+    pub(crate) fn new(result: PyResult<PyObject>) -> Self {
+        Self { result }
+    }
+}
+
+#[pymethods]
+impl PyDoneAwaitable {
+    fn __await__(pyself: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        pyself
+    }
+
+    fn __iter__(pyself: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        pyself
+    }
+
+    fn __next__(&self, py: Python) -> PyResult<PyObject> {
+        self.result
+            .as_ref()
+            .map(|v| v.clone_ref(py))
+            .map_err(|v| v.clone_ref(py))
+    }
+}
+
 #[pyclass(frozen, module = "granian._granian")]
 pub(crate) struct PyErrAwaitable {
     result: PyResult<()>,

src/rsgi/callbacks.rs

@@ -1,6 +1,6 @@
 use pyo3::prelude::*;
-use std::sync::OnceLock;
-use tokio::sync::oneshot;
+use std::sync::{Arc, OnceLock};
+use tokio::sync::{oneshot, Notify};
 
 use super::{
     io::{RSGIHTTPProtocol as HTTPProtocol, RSGIWebsocketProtocol as WebsocketProtocol, WebsocketDetachedTransport},
@@ -118,11 +118,12 @@ impl CallbackWatcherWebsocket {
 pub(crate) fn call_http(
     cb: ArcCBScheduler,
     rt: RuntimeRef,
+    disconnect_guard: Arc<Notify>,
     body: hyper::body::Incoming,
     scope: HTTPScope,
 ) -> oneshot::Receiver<PyResponse> {
     let (tx, rx) = oneshot::channel();
-    let protocol = HTTPProtocol::new(rt.clone(), tx, body);
+    let protocol = HTTPProtocol::new(rt.clone(), tx, body, disconnect_guard);
 
     rt.spawn_blocking(move |py| {
         if let Ok(watcher) = CallbackWatcherHTTP::new(py, protocol, scope) {

src/rsgi/http.rs

@@ -1,8 +1,8 @@
 use futures::sink::SinkExt;
 use http_body_util::BodyExt;
 use hyper::{header::SERVER as HK_SERVER, http::response::Builder as ResponseBuilder, StatusCode};
-use std::net::SocketAddr;
-use tokio::sync::mpsc;
+use std::{net::SocketAddr, sync::Arc};
+use tokio::sync::{mpsc, Notify};
 
 use super::{
     callbacks::{call_http, call_ws},
@@ -30,8 +30,8 @@ macro_rules! build_scope {
 }
 
 macro_rules! handle_http_response {
-    ($handler:expr, $rt:expr, $callback:expr, $body:expr, $scope:expr) => {
-        match $handler($callback, $rt, $body, $scope).await {
+    ($handler:expr, $rt:expr, $disconnect_guard:expr, $callback:expr, $body:expr, $scope:expr) => {
+        match $handler($callback, $rt, $disconnect_guard, $body, $scope).await {
             Ok(PyResponse::Body(pyres)) => pyres.to_response(),
             Ok(PyResponse::File(pyres)) => pyres.to_response().await,
             _ => {
@@ -47,6 +47,7 @@ macro_rules! handle_request {
         #[inline]
         pub(crate) async fn $func_name(
             rt: RuntimeRef,
+            disconnect_guard: Arc<Notify>,
             callback: ArcCBScheduler,
             server_addr: SocketAddr,
             client_addr: SocketAddr,
@@ -55,7 +56,7 @@ macro_rules! handle_request {
         ) -> HTTPResponse {
             let (parts, body) = req.into_parts();
             let scope = build_scope!(HTTPScope, server_addr, client_addr, parts, scheme);
-            handle_http_response!($handler, rt, callback, body, scope)
+            handle_http_response!($handler, rt, disconnect_guard, callback, body, scope)
         }
     };
 }
@@ -65,6 +66,7 @@ macro_rules! handle_request_with_ws {
         #[inline]
         pub(crate) async fn $func_name(
             rt: RuntimeRef,
+            disconnect_guard: Arc<Notify>,
             callback: ArcCBScheduler,
             server_addr: SocketAddr,
             client_addr: SocketAddr,
@@ -131,7 +133,7 @@ macro_rules! handle_request_with_ws {
 
             let (parts, body) = req.into_parts();
             let scope = build_scope!(HTTPScope, server_addr, client_addr, parts, scheme);
-            handle_http_response!($handler_req, rt, callback, body, scope)
+            handle_http_response!($handler_req, rt, disconnect_guard, callback, body, scope)
         }
     };
 }