🐛 Is638/fixes cancel on disconnect (ITISFoundation#3164)

Author: pcrespov

packages/service-library/requirements/_fastapi.in

@@ -9,7 +9,7 @@
 async-asgi-testclient # replacement for fastapi.testclient.TestClient [see b) below]
 fastapi
 fastapi_contrib[jaegertracing]
-
+uvicorn
 
 # NOTE: What test client to use for fastapi-based apps?
 #

packages/service-library/requirements/_fastapi.txt

@@ -12,12 +12,16 @@ certifi==2022.6.15
     # via requests
 charset-normalizer==2.0.12
     # via requests
+click==8.1.3
+    # via uvicorn
 fastapi==0.76.0
     # via
     #   -r requirements/_fastapi.in
     #   fastapi-contrib
 fastapi-contrib==0.2.11
     # via -r requirements/_fastapi.in
+h11==0.13.0
+    # via uvicorn
 idna==3.3
     # via
     #   anyio
@@ -59,3 +63,5 @@ urllib3==1.26.9
     # via
     #   -c requirements/./../../../requirements/constraints.txt
     #   requests
+uvicorn==0.18.2
+    # via -r requirements/_fastapi.in

packages/service-library/setup.cfg

@@ -12,3 +12,6 @@ universal = 1
 
 [aliases]
 test = pytest
+
+[tool:pytest]
+asyncio_mode = auto

packages/service-library/src/servicelib/fastapi/requests_decorators.py

@@ -1,110 +1,124 @@
 import asyncio
 import inspect
 import logging
-from asyncio import CancelledError
-from contextlib import suppress
 from functools import wraps
-from typing import Any, Callable, Coroutine, Optional
+from typing import Any, Protocol
 
-from fastapi import Request, Response
+from fastapi import Request, status
+from fastapi.exceptions import HTTPException
 
 logger = logging.getLogger(__name__)
 
 
-_DEFAULT_CHECK_INTERVAL_S: float = 0.5
+class _HandlerWithRequestArg(Protocol):
+    __name__: str
 
-HTTP_499_CLIENT_CLOSED_REQUEST = 499
-# A non-standard status code introduced by nginx for the case when a client
-# closes the connection while nginx is processing the request.
-# SEE https://www.webfx.com/web-development/glossary/http-status-codes/what-is-a-499-status-code/
+    async def __call__(self, request: Request, *args: Any) -> Any:
+        ...
 
-TASK_NAME_PREFIX = "cancellable_request"
 
-_FastAPIHandlerCallable = Callable[..., Coroutine[Any, Any, Optional[Any]]]
+def _validate_signature(handler: _HandlerWithRequestArg):
+    """Raises ValueError if handler does not have expected signature"""
+    try:
+        p = next(iter(inspect.signature(handler).parameters.values()))
+        if p.kind != inspect.Parameter.POSITIONAL_OR_KEYWORD or p.annotation != Request:
+            raise TypeError(
+                f"Invalid handler {handler.__name__} signature: first parameter must be a Request, got {p.annotation}"
+            )
+    except StopIteration as e:
+        raise TypeError(
+            f"Invalid handler {handler.__name__} signature: first parameter must be a Request, got none"
+        ) from e
 
 
-async def _cancel_task_if_client_disconnected(
-    request: Request, task: asyncio.Task, interval: float = _DEFAULT_CHECK_INTERVAL_S
-) -> None:
-    try:
-        while True:
-            if task.done():
-                logger.debug("task %s is done", task)
-                break
-            if await request.is_disconnected():
-                logger.warning(
-                    "client %s disconnected! Cancelling handler for %s",
-                    request.client,
-                    f"{request.url=}",
-                )
-                task.cancel()
-                break
-            await asyncio.sleep(interval)
-    except CancelledError:
-        logger.debug("task monitoring %s handler was cancelled", f"{request.url=}")
-        raise
-    finally:
-        logger.debug("task monitoring %s handler completed", f"{request.url}")
-
-
-def cancellable_request(handler_fun: _FastAPIHandlerCallable):
-    """This decorator periodically checks if the client disconnected and
-    then will cancel the request and return a HTTP_499_CLIENT_CLOSED_REQUEST code (a la nginx).
-
-    Usage: decorate the cancellable route and add request: Request as an argument
-
-        @cancellable_request
-        async def route(
-            _request: Request,
-            ...
-        )
+#
+# cancel_on_disconnect/disconnect_poller based
+# on https://github.com/RedRoserade/fastapi-disconnect-example/blob/main/app.py
+#
+_POLL_INTERVAL_S: float = 0.01
+
+
+async def disconnect_poller(request: Request, result: Any):
     """
-    # CHECK: Early check that will raise upon import
-    # IMPROVEMENT: inject this parameter to handler_fun here before it returned in the wrapper and consumed by fastapi.router?
-    found_required_arg = any(
-        parameter.name == "_request" and parameter.annotation == Request
-        for parameter in inspect.signature(handler_fun).parameters.values()
+    Poll for a disconnect.
+    If the request disconnects, stop polling and return.
+    """
+    while not await request.is_disconnected():
+        await asyncio.sleep(_POLL_INTERVAL_S)
+
+    logger.debug(
+        "client %s disconnected! Cancelling handler for request %s %s",
+        request.client,
+        request.method,
+        request.url,
     )
-    if not found_required_arg:
-        raise ValueError(
-            f"Invalid handler {handler_fun.__name__} signature: missing required parameter _request: Request"
-        )
+    return result
+
 
-    # WRAPPER ----
-    @wraps(handler_fun)
-    async def wrapper(*args, **kwargs) -> Optional[Any]:
-        request: Request = kwargs["_request"]
+def cancel_on_disconnect(handler: _HandlerWithRequestArg):
+    """
+    After client dicsonnects, handler gets cancelled in ~<3 secs
+    """
 
-        # Intercepts handler call and creates a task out of it
+    _validate_signature(handler)
+
+    @wraps(handler)
+    async def wrapper(request: Request, *args, **kwargs):
+        sentinel = object()
+
+        # Create two tasks:
+        # one to poll the request and check if the client disconnected
+        poller_task = asyncio.create_task(
+            disconnect_poller(request, sentinel),
+            name=f"cancel_on_disconnect/poller/{handler.__name__}/{id(sentinel)}",
+        )
+        # , and another which is the request handler
         handler_task = asyncio.create_task(
-            handler_fun(*args, **kwargs),
-            name=f"{TASK_NAME_PREFIX}/handler/{handler_fun.__name__}",
+            handler(request, *args, **kwargs),
+            name=f"cancel_on_disconnect/handler/{handler.__name__}/{id(sentinel)}",
         )
-        # An extra task to monitor when the client disconnects so it can
-        # cancel 'handler_task'
-        auto_cancel_task = asyncio.create_task(
-            _cancel_task_if_client_disconnected(request, handler_task),
-            name=f"{TASK_NAME_PREFIX}/auto_cancel/{handler_fun.__name__}",
+
+        done, pending = await asyncio.wait(
+            [poller_task, handler_task], return_when=asyncio.FIRST_COMPLETED
         )
 
-        try:
+        # One has completed, cancel the other
+        for t in pending:
+            t.cancel()
+            try:
+                await asyncio.wait_for(t, timeout=3)
+            except asyncio.CancelledError:
+                logger.debug("%s was cancelled", t)
+            except Exception as exc:  # pylint: disable=broad-except
+                if t is handler_task:
+                    logger.warning(
+                        "%s raised %s when being cancelled.", t, exc, exc_info=True
+                    )
+                    raise
+            finally:
+                assert t.done()  # nosec
+
+        # Return the result if the handler finished first
+        if handler_task in done:
+            assert poller_task.done()  # nosec
             return await handler_task
-        except CancelledError:
-            # TODO: check that 'auto_cancel_task' actually executed this cancellation
-            # E.g. app shutdown might cancel all pending tasks
-            logger.warning(
-                "Request %s was cancelled since client %s disconnected !",
-                f"{request.url}",
-                request.client,
-            )
-            return Response(
-                "Request cancelled because client disconnected",
-                status_code=HTTP_499_CLIENT_CLOSED_REQUEST,
-            )
-        finally:
-            # NOTE: This is ALSO called 'await handler_task' returns
-            auto_cancel_task.cancel()
-            with suppress(CancelledError):
-                await auto_cancel_task
+
+        # Otherwise, raise an exception. This is not exactly needed, but it will prevent
+        # validation errors if your request handler is supposed to return something.
+        logger.debug(
+            "Request %s %s cancelled:\n - %s\n - %s",
+            request.method,
+            request.url,
+            f"{poller_task=}",
+            f"{handler_task=}",
+        )
+        assert poller_task.done()  # nosec
+        assert handler_task.done()  # nosec
+
+        # NOTE: uvicorn server fails with 499
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail=f"client disconnected from {request=}",
+        )
 
     return wrapper