owned_thread.py

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