proxy_download_work.py

"""下载准备、响应交付与应用关闭共用的有界归属。"""
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):
        @wraps(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