use taskgroup for error handling

Author: bisgaard-itis

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

@@ -29,96 +29,58 @@ def _validate_signature(handler: _HandlerWithRequestArg):
 
 
 #
-# cancel_on_disconnect/disconnect_poller based
-# on https://github.com/RedRoserade/fastapi-disconnect-example/blob/main/app.py
+# cancel_on_disconnect based on TaskGroup
 #
 _POLL_INTERVAL_S: float = 0.01
 
 
-async def _disconnect_poller(close_event: asyncio.Event, request: Request, result: Any):
+class _ClientDisconnectedError(Exception):
+    """Internal exception raised by the poller task when the client disconnects."""
+
+
+async def _disconnect_poller_for_task_group(request: Request):
     """
-    Poll for a disconnect.
-    If the request disconnects, stop polling and return.
+    Polls for client disconnection and raises _ClientDisconnectedError if it occurs.
     """
     while not await request.is_disconnected():
         await asyncio.sleep(_POLL_INTERVAL_S)
-        if close_event.is_set():
-            break
-    return result
+    raise _ClientDisconnectedError()
 
 
 def cancel_on_disconnect(handler: _HandlerWithRequestArg):
     """
-    After client disconnects, handler gets cancelled in ~<3 secs
+    Decorator that cancels the request handler if the client disconnects.
+
+    Uses a TaskGroup to manage the handler and a poller task concurrently.
+    If the client disconnects, the poller raises an exception, which is
+    caught and translated into a 503 Service Unavailable response.
     """
 
     _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
-        # sometimes canceling a task doesn't cancel the task immediately. If the poller task is not "killed" immediately, the client doesn't
-        # get a response, and the request "hangs". For this reason, we use an event to signal the poller task to stop.
-        # See: https://github.com/ITISFoundation/osparc-issues/issues/1922
-        kill_poller_event = asyncio.Event()
-        poller_task = asyncio.create_task(
-            _disconnect_poller(kill_poller_event, 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(request, *args, **kwargs),
-            name=f"cancel_on_disconnect/handler/{handler.__name__}/{id(sentinel)}",
-        )
-
-        done, pending = await asyncio.wait(
-            [poller_task, handler_task], return_when=asyncio.FIRST_COMPLETED
-        )
-        kill_poller_event.set()
-
-        # One has completed, cancel the other
-        for t in pending:
-            t.cancel()
-
-            try:
-                await asyncio.wait_for(t, timeout=3)
-
-            except asyncio.CancelledError:
-                pass
-            except Exception:  # pylint: disable=broad-except
-                if t is handler_task:
-                    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
-
-        # 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.warning(
-            "Request %s %s cancelled since client %s disconnected:\n - %s\n - %s",
-            request.method,
-            request.url,
-            request.client,
-            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=}",
-        )
+        try:
+            async with asyncio.TaskGroup() as tg:
+                handler_task = tg.create_task(handler(request, *args, **kwargs))
+                tg.create_task(_disconnect_poller_for_task_group(request))
+
+            return handler_task.result()
+
+        except* _ClientDisconnectedError as eg:
+            logger.info(
+                "Request %s %s cancelled since client %s disconnected.",
+                request.method,
+                request.url,
+                request.client,
+            )
+            raise HTTPException(
+                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+                detail="Client disconnected",
+            ) from eg
+
+        except* Exception as eg:
+            raise eg.exceptions[0]
 
     return wrapper