utils/middleware/request_cancellation.py

import asyncio
from utils.logger import Logger

logger = Logger.get_logger(__name__)


class RequestCancellation:
    """
        RequestCancellation middleware handles request canceling
        * In case of API routes where very frequent/expensive requests are made.
    """

    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        queue = asyncio.Queue()

        async def message_poller(sentinel, handler_task):
            nonlocal queue
            while True:
                message = await receive()
                if message["type"] == "http.disconnect":
                    handler_task.cancel()
                    return sentinel
                await queue.put(message)

        sentinel = object()
        handler_task = asyncio.create_task(self.app(scope, queue.get, send))
        asyncio.create_task(message_poller(sentinel, handler_task))

        try:
            return await handler_task
        except asyncio.CancelledError:
            logger.info('Task Cancellatation Requested.')