Spaces:
Running
Running
| """资料文件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 | |