Spaces:
Running
Running
| """代理请求的所属响应/客户端收尾：每对象一次，取消等待同一任务。""" | |
| 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) | |