Spaces:
Running
Running
| """等待所属同步工作真实终态；内部取消也不能遗留仍读写资源的线程。""" | |
| import asyncio | |
| from functools import wraps | |
| from inspect import Parameter, signature, iscoroutinefunction, isasyncgenfunction, isgeneratorfunction | |
| import threading | |
| async def run_owned_thread(function, *args, **kwargs): | |
| ownership = threading.Lock() | |
| finished = threading.Event() | |
| started = False | |
| outcome = None | |
| def abandon_unstarted(result): | |
| nonlocal outcome | |
| with ownership: | |
| if not started and not finished.is_set(): | |
| outcome = result | |
| finished.set() | |
| def execute(): | |
| nonlocal started, outcome | |
| with ownership: | |
| if finished.is_set(): | |
| return | |
| started = True | |
| try: | |
| result = (True, function(*args, **kwargs)) | |
| except BaseException as error: | |
| # 同步异常作为数据交回；取消shield不暴露私有异常正文。 | |
| result = (False, error) | |
| with ownership: | |
| outcome = result | |
| finished.set() | |
| async def finish(): | |
| try: | |
| await asyncio.to_thread(execute) | |
| except asyncio.CancelledError: | |
| raise | |
| except BaseException as error: | |
| abandon_unstarted((False, error)) | |
| coroutine = finish() | |
| try: | |
| task = asyncio.create_task(coroutine) | |
| except BaseException: | |
| coroutine.close() | |
| raise | |
| def consume(completed): | |
| if completed.cancelled(): | |
| # 取消和线程启动竞争时，只有尚未开始的工作能被撤销。 | |
| abandon_unstarted((False, asyncio.CancelledError())) | |
| else: | |
| error = completed.exception() | |
| if error is not None: | |
| abandon_unstarted((False, error)) | |
| task.add_done_callback(consume) | |
| cancelled = None | |
| while not finished.is_set(): | |
| try: | |
| if not task.done(): | |
| await asyncio.shield(task) | |
| else: | |
| # 内部task已取消但线程仍在运行，不占线程池等待其终态。 | |
| await asyncio.sleep(0.01) | |
| except asyncio.CancelledError as error: | |
| cancelled = error | |
| with ownership: | |
| succeeded, result = outcome | |
| if cancelled is not None: | |
| raise cancelled | |
| if not succeeded: | |
| raise result | |
| return result | |
| class OwnedThreadGroup: | |
| def __init__(self, capacity, unavailable): | |
| if type(capacity) is not int or capacity < 1: | |
| raise ValueError('invalid owned thread capacity') | |
| self._capacity = threading.BoundedSemaphore(capacity) | |
| self._unavailable = unavailable | |
| self._gate = threading.Lock() | |
| self._accepting = True | |
| self._pending = 0 | |
| self._closers = 0 | |
| self._drained = threading.Event() | |
| self._drained.set() | |
| def start(self): | |
| with self._gate: | |
| if self._pending or self._closers: | |
| raise RuntimeError('owned work is still pending') | |
| self._accepting = True | |
| def stop(self): | |
| with self._gate: | |
| self._accepting = False | |
| async def close(self): | |
| with self._gate: | |
| self._accepting = False | |
| self._closers += 1 | |
| try: | |
| cancelled = False | |
| while not self._drained.is_set(): | |
| try: | |
| await asyncio.sleep(.01) | |
| except asyncio.CancelledError: | |
| cancelled = True | |
| if cancelled: | |
| raise asyncio.CancelledError | |
| finally: | |
| with self._gate: | |
| self._closers -= 1 | |
| async def run(self, function, *args, **kwargs): | |
| with self._gate: | |
| if not self._accepting or not self._capacity.acquire(blocking=False): | |
| raise self._unavailable() | |
| self._pending += 1 | |
| self._drained.clear() | |
| try: | |
| return await run_owned_thread(function, *args, **kwargs) | |
| finally: | |
| with self._gate: | |
| self._capacity.release() | |
| self._pending -= 1 | |
| if not self._pending: | |
| self._drained.set() | |
| def owned_sync_endpoint(group): | |
| def decorate(endpoint): | |
| if iscoroutinefunction(endpoint) or isasyncgenfunction(endpoint) or isgeneratorfunction(endpoint): | |
| raise TypeError('owned endpoint must be a synchronous function') | |
| async def invoke(*args, **kwargs): | |
| return await group.run(endpoint, *args, **kwargs) | |
| return invoke | |
| return decorate | |
| def owned_session_endpoint(group, session_factory, session_parameter): | |
| def decorate(endpoint): | |
| if iscoroutinefunction(endpoint) or isasyncgenfunction(endpoint) or isgeneratorfunction(endpoint): | |
| raise TypeError('owned session endpoint must be a synchronous function') | |
| original = signature(endpoint) | |
| if session_parameter not in original.parameters: | |
| raise ValueError('missing owned session parameter') | |
| if any(parameter.kind not in (Parameter.POSITIONAL_OR_KEYWORD, Parameter.KEYWORD_ONLY) | |
| for parameter in original.parameters.values()): | |
| raise TypeError('owned session endpoint requires named parameters') | |
| public = original.replace(parameters=[parameter for name, parameter in original.parameters.items() | |
| if name != session_parameter]) | |
| async def invoke(*args, **kwargs): | |
| arguments = public.bind(*args, **kwargs).arguments | |
| def execute(): | |
| with session_factory() as session: | |
| return endpoint(**arguments, **{session_parameter: session}) | |
| return await group.run(execute) | |
| invoke.__signature__ = public | |
| return invoke | |
| return decorate | |