profile_io.py

"""资料文件I/O在后台线程执行；准入在排队前，已开始工作取消不提前释放容量。"""
import asyncio
import threading
from fastapi import HTTPException

MAX_PROFILE_IO = 8

class ProfileIO:
    def __init__(self, capacity=MAX_PROFILE_IO):
        self.capacity = threading.BoundedSemaphore(capacity)
        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("profile work is still pending")
            self._accepting = True

    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(0.01)
                except asyncio.CancelledError:
                    cancelled = True
            if cancelled:
                raise asyncio.CancelledError
        finally:
            with self._gate:
                self._closers -= 1

    def _finish(self):
        with self._gate:
            self.capacity.release()
            self._pending -= 1
            if not self._pending:
                self._drained.set()

    async def run(self, function, *args, **kwargs):
        with self._gate:
            if not self._accepting or not self.capacity.acquire(blocking=False):
                raise HTTPException(status_code=503, detail="用户数据暂不可用")
            self._pending += 1
            self._drained.clear()

        # 内部任务也可能在宿主关闭时被取消。锁区分未开始与已开始的工作，
        # 保证取消和线程启动竞争时只释放一次，且取消后不再开始业务写入。
        ownership = threading.Lock()
        started = False
        released = False

        def release_unstarted():
            nonlocal released
            with ownership:
                if not started and not released:
                    released = True
                    self._finish()

        def execute():
            nonlocal started, released
            with ownership:
                if released:
                    return None
                started = True
            try:
                return function(*args, **kwargs)
            finally:
                with ownership:
                    released = True
                    self._finish()

        coroutine = asyncio.to_thread(execute)
        try:
            task = asyncio.create_task(coroutine)
        except BaseException:
            coroutine.close(); release_unstarted()
            raise

        def consume(completed):
            if completed.cancelled():
                release_unstarted()
            else:
                # 提交默认线程池本身也可能失败，尚未开始的准入必须收回。
                release_unstarted()
                completed.exception()
        task.add_done_callback(consume)
        # 外层请求取消不取消所属任务；已开始的工作只在线程终态释放容量。
        return await asyncio.shield(task)

_profile_io = ProfileIO()
run_profile_io = _profile_io.run

start_profile_io = _profile_io.start
close_profile_io = _profile_io.close