proxy_resources.py

"""代理请求的所属响应/客户端收尾：每对象一次，取消等待同一任务。"""
import asyncio
import logging

logger = logging.getLogger("ComfyUI-Ranking")


async def _await_close(task):
    cancelled = None
    while not task.done():
        try:
            await asyncio.shield(task)
        except asyncio.CancelledError as error:
            cancelled = error
    task.result()
    if cancelled is not None:
        raise cancelled


async def close_redirect_response(response):
    """中间跳转响应关闭一次；完整等待后取消优先，普通关闭失败仍向外传播。"""
    async def finish():
        try:
            await response.aclose()
        except Exception as error:
            logger.warning("PROXY_REDIRECT_CLOSE_FAIL | error_type=%s", type(error).__name__)
            return error
    task = asyncio.create_task(finish())
    await _await_close(task)
    error = task.result()
    if error is not None:
        raise error


class ProxyResources:
    def __init__(self, client):
        self.client = client
        self._responses = []
        self._response_tasks = {}
        self._close_task = None

    def track(self, response):
        if self._close_task is not None:
            raise RuntimeError("proxy resources already closing")
        if not any(previous is response for previous in self._responses):
            self._responses.append(response)
        return response

    async def close_response(self, response):
        key = id(response)
        if key not in self._response_tasks:
            async def finish():
                try:
                    await response.aclose()
                except Exception as error:
                    logger.warning("PROXY_RESPONSE_CLOSE_FAIL | error_type=%s", type(error).__name__)
            self._response_tasks[key] = asyncio.create_task(finish())
        await _await_close(self._response_tasks[key])

    async def close(self):
        if self._close_task is None:
            async def finish():
                cancelled = None
                try:
                    for response in self._responses:
                        try:
                            await self.close_response(response)
                        except asyncio.CancelledError as error:
                            cancelled = error
                finally:
                    try:
                        await self.client.aclose()
                    except Exception as error:
                        logger.warning("PROXY_CLIENT_CLOSE_FAIL | error_type=%s", type(error).__name__)
                if cancelled is not None:
                    raise cancelled
            self._close_task = asyncio.create_task(finish())
        await _await_close(self._close_task)