Work in progress

dedoc/api/cancellation.py

@@ -0,0 +1,37 @@
+import logging
+from contextlib import asynccontextmanager
+
+from anyio import create_task_group
+from fastapi import Request
+
+from dedoc.config import get_config
+
+
+@asynccontextmanager
+async def cancel_on_disconnect(request: Request) -> None:
+    """
+    Async context manager for async code that needs to be cancelled if client disconnects prematurely.
+    The client disconnect is monitored through the Request object.
+
+    Source: https://github.com/dorinclisu/runner-with-api
+    See discussion: https://github.com/fastapi/fastapi/discussions/8805
+    """
+    logger = get_config().get("logger", logging.getLogger())
+    async with create_task_group() as task_group:
+        async def watch_disconnect() -> None:
+            while True:
+                message = await request.receive()
+
+                if message["type"] == "http.disconnect":
+                    client = f"{request.client.host}:{request.client.port}" if request.client else "-:-"
+                    logger.info(f"{client} - `{request.method} {request.url.path}` 499 DISCONNECTED")
+
+                    task_group.cancel_scope.cancel()
+                    break
+
+        task_group.start_soon(watch_disconnect)
+
+        try:
+            yield
+        finally:
+            task_group.cancel_scope.cancel()

dedoc/api/dedoc_api.py

@@ -15,6 +15,7 @@
 import dedoc.version
 from dedoc.api.api_args import QueryParameters
 from dedoc.api.api_utils import json2collapsed_tree, json2html, json2tree, json2txt
+from dedoc.api.cancellation import cancel_on_disconnect
 from dedoc.api.schema.parsed_document import ParsedDocument
 from dedoc.common.exceptions.dedoc_error import DedocError
 from dedoc.common.exceptions.missing_file_error import MissingFileError
@@ -69,19 +70,20 @@ def __add_base64_info_to_attachments(document_tree: ParsedDocument, attachments_
 
 
 @app.post("/upload", response_model=ParsedDocument)
-async def upload(file: UploadFile = File(...), query_params: QueryParameters = Depends()) -> Response:
+async def upload(request: Request, file: UploadFile = File(...), query_params: QueryParameters = Depends()) -> Response:
     parameters = dataclasses.asdict(query_params)
     if not file or file.filename == "":
         raise MissingFileError("Error: Missing content in request_post file parameter", version=dedoc.version.__version__)
 
     return_format = str(parameters.get("return_format", "json")).lower()
 
-    with tempfile.TemporaryDirectory() as tmpdir:
-        file_path = save_upload_file(file, tmpdir)
-        document_tree = manager.parse(file_path, parameters={**dict(parameters), "attachments_dir": tmpdir})
+    async with cancel_on_disconnect(request):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            file_path = save_upload_file(file, tmpdir)
+            document_tree = manager.parse(file_path, parameters={**dict(parameters), "attachments_dir": tmpdir})
 
-        if return_format == "html":
-            __add_base64_info_to_attachments(document_tree, tmpdir)
+            if return_format == "html":
+                __add_base64_info_to_attachments(document_tree, tmpdir)
 
     if return_format == "html":
         html_content = json2html(