"""等待所属同步工作真实终态;内部取消也不能遗留仍读写资源的线程。""" 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') @wraps(endpoint) 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]) @wraps(endpoint) 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