77 lines
3.0 KiB
Python
77 lines
3.0 KiB
Python
"""Bounded file work, isolated from authentication's default executor."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextvars
|
|
import logging
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from functools import partial
|
|
from weakref import WeakKeyDictionary
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class FileIOPool:
|
|
def __init__(self, workers: int, admitted: int, name: str):
|
|
self._pool = ThreadPoolExecutor(max_workers=workers, thread_name_prefix=name)
|
|
self._slots: WeakKeyDictionary = WeakKeyDictionary()
|
|
self._admitted = admitted
|
|
|
|
async def __call__(self, function, /, *args, **kwargs):
|
|
return await self._run(function, args, kwargs, wait_on_cancel=False, on_cancel=None)
|
|
|
|
async def run_until_complete(self, function, /, *args, on_cancel=None, **kwargs):
|
|
"""Keep a submitted worker alive until it finishes, even after cancellation.
|
|
|
|
Callers can safely close inputs or remove temporary directories once this
|
|
returns or raises. Cancellation while waiting for admission does not
|
|
submit work. Once submitted, ``on_cancel`` may request cooperative
|
|
cancellation; admission remains held until the actual worker exits.
|
|
"""
|
|
return await self._run(function, args, kwargs, wait_on_cancel=True, on_cancel=on_cancel)
|
|
|
|
async def _run(self, function, args, kwargs, *, wait_on_cancel, on_cancel):
|
|
loop = asyncio.get_running_loop()
|
|
semaphore = self._slots.setdefault(loop, asyncio.Semaphore(self._admitted))
|
|
await semaphore.acquire()
|
|
context = contextvars.copy_context()
|
|
try:
|
|
future = loop.run_in_executor(self._pool, context.run, partial(function, *args, **kwargs))
|
|
except BaseException:
|
|
semaphore.release()
|
|
raise
|
|
|
|
def finished(done):
|
|
semaphore.release()
|
|
if not done.cancelled():
|
|
done.exception()
|
|
|
|
future.add_done_callback(finished)
|
|
# Cancellation cannot release admission before the real disk job ends.
|
|
try:
|
|
return await asyncio.shield(future)
|
|
except asyncio.CancelledError:
|
|
if not wait_on_cancel:
|
|
raise
|
|
if on_cancel is not None:
|
|
try:
|
|
on_cancel()
|
|
except Exception:
|
|
logger.exception("File worker cancellation callback failed")
|
|
# The request owns the worker's inputs. Do not let its finally
|
|
# block clean those up while a thread still reads or writes them.
|
|
# Repeated Task.cancel() calls must not break that guarantee.
|
|
while not future.done():
|
|
try:
|
|
await asyncio.shield(future)
|
|
except asyncio.CancelledError:
|
|
continue
|
|
except BaseException:
|
|
break # finished() consumes the worker's exception.
|
|
raise
|
|
|
|
|
|
file_io = FileIOPool(4, 8, "artifact-files")
|
|
event_io = FileIOPool(2, 4, "event-spool")
|