Spaces:
Running
Running
| """下载准备、响应交付与应用关闭共用的有界归属。""" | |
| import asyncio | |
| from functools import wraps | |
| from fastapi import HTTPException | |
| from workflow_stream_response import OwnedDelivery | |
| def _spawn(function): | |
| async def execute(): | |
| try: | |
| return True, await function() | |
| except BaseException as error: | |
| return False, error | |
| coroutine = execute() | |
| try: | |
| return asyncio.create_task(coroutine) | |
| except BaseException: | |
| coroutine.close() | |
| raise | |
| async def _settle(task, interrupted=None): | |
| cancelled = None | |
| while not task.done(): | |
| try: | |
| await asyncio.shield(task) | |
| except asyncio.CancelledError as error: | |
| cancelled = error | |
| if interrupted is not None: | |
| interrupted(task) | |
| outcome = (False, asyncio.CancelledError()) if task.cancelled() else task.result() | |
| return outcome, cancelled | |
| class _Download: | |
| def __init__(self, group): | |
| self.group = group | |
| self.active = None | |
| self.response = None | |
| self.aborting = False | |
| self.released = False | |
| self.sending = False | |
| self.cancelled_tasks = set() | |
| self.cleanup = None | |
| self.finished = asyncio.Event() | |
| def cancel(self, task): | |
| self.aborting = True | |
| if not task.done() and task not in self.cancelled_tasks: | |
| self.cancelled_tasks.add(task) | |
| task.cancel() | |
| def release(self): | |
| self.released = True | |
| self.group._pending.discard(self) | |
| self.finished.set() | |
| async def prepare(self, function, args, kwargs): | |
| result = await function(*args, **kwargs) | |
| if isinstance(result, OwnedDelivery): | |
| result.bind_download(self) | |
| self.response = result | |
| return result | |
| async def abort(self): | |
| self.aborting = True | |
| self.cancel(self.active) | |
| await _settle(self.active, self.cancel) | |
| if self.response is None: | |
| self.release() | |
| return | |
| if self.sending: | |
| # 发送拥有者先完成迭代器和连接关闭，再归还名额。 | |
| await _settle(_spawn(self.finished.wait)) | |
| return | |
| if self.cleanup is None: | |
| async def close(): | |
| try: | |
| await self.response.aclose() | |
| finally: | |
| self.release() | |
| self.cleanup = _spawn(close) | |
| outcome, cancelled = await _settle(self.cleanup) | |
| succeeded, result = outcome | |
| if cancelled is not None: | |
| raise cancelled | |
| if not succeeded: | |
| raise result | |
| async def deliver(self, response, scope, receive, send): | |
| if self.released or self.aborting or self.sending or response is not self.response: | |
| raise asyncio.CancelledError() | |
| self.sending = True | |
| try: | |
| self.active = _spawn(lambda: response.send_owned(scope, receive, send)) | |
| outcome, cancelled = await _settle(self.active, self.cancel) | |
| succeeded, result = outcome | |
| if cancelled is not None: | |
| raise cancelled | |
| if not succeeded: | |
| raise result | |
| return result | |
| finally: | |
| try: | |
| # 发送任务在首次执行前被取消时，也必须关闭尚未消费的响应。 | |
| await response.aclose() | |
| finally: | |
| self.release() | |
| class ProxyDownloadWork: | |
| def __init__(self, capacity=8): | |
| if type(capacity) is not int or capacity < 1: | |
| raise ValueError('invalid download capacity') | |
| self.capacity = capacity | |
| self._pending = set() | |
| self._accepting = True | |
| self._closing = None | |
| self._closers = 0 | |
| def start(self): | |
| if self._pending or self._closers or self._closing is not None and not self._closing.done(): | |
| raise RuntimeError('download work is still pending') | |
| if self._closing is not None: | |
| succeeded, result = self._closing.result() | |
| if not succeeded: | |
| raise result | |
| self._closing = None | |
| self._accepting = True | |
| def stop(self): | |
| self._accepting = False | |
| def endpoint(self, function): | |
| async def invoke(*args, **kwargs): | |
| return await self.run(function, *args, **kwargs) | |
| return invoke | |
| async def run(self, function, *args, **kwargs): | |
| if not self._accepting or len(self._pending) >= self.capacity: | |
| raise HTTPException(status_code=503, detail='service_unavailable') | |
| owner = _Download(self) | |
| self._pending.add(owner) | |
| handed_off = False | |
| try: | |
| owner.active = _spawn(lambda: owner.prepare(function, args, kwargs)) | |
| outcome, cancelled = await _settle(owner.active, owner.cancel) | |
| succeeded, result = outcome | |
| if cancelled is not None or owner.aborting: | |
| raise cancelled or asyncio.CancelledError() | |
| if not succeeded: | |
| raise result | |
| handed_off = owner.response is not None | |
| return result | |
| finally: | |
| if owner.response is None: | |
| owner.release() | |
| elif not handed_off: | |
| await owner.abort() | |
| async def close(self): | |
| self.stop() | |
| self._closers += 1 | |
| try: | |
| if self._closing is None: | |
| async def drain(): | |
| tasks = [_spawn(owner.abort) for owner in tuple(self._pending)] | |
| results = await asyncio.gather(*tasks) | |
| for succeeded, result in results: | |
| if not succeeded: | |
| raise result | |
| self._closing = _spawn(drain) | |
| outcome, cancelled = await _settle(self._closing) | |
| succeeded, result = outcome | |
| if cancelled is not None: | |
| raise cancelled | |
| if not succeeded: | |
| raise result | |
| finally: | |
| self._closers -= 1 | |
| _downloads = ProxyDownloadWork() | |
| proxy_download_endpoint = _downloads.endpoint | |
| start_proxy_downloads = _downloads.start | |
| stop_proxy_downloads = _downloads.stop | |
| close_proxy_downloads = _downloads.close | |