diff --git a/pageindex/index/utils.py b/pageindex/index/utils.py index cb4b9b5..7267fbc 100644 --- a/pageindex/index/utils.py +++ b/pageindex/index/utils.py @@ -103,11 +103,18 @@ async def _llm_semaphore(): node at once and exhausts the process file-descriptor limit (Errno 24). """ ceiling_sem = _process_ceiling_semaphore() - # threading.Semaphore.acquire() blocks the calling thread, so run it off - # the event loop thread — otherwise it would freeze every other coroutine - # on this loop while waiting for a slot. release() is non-blocking and - # safe to call directly from any thread. - await asyncio.to_thread(ceiling_sem.acquire) + # A blocking ceiling_sem.acquire() run via asyncio.to_thread() would be + # unsafe under cancellation: the worker thread can't be interrupted, so if + # this coroutine is cancelled (Ctrl-C, an outer timeout) while the thread + # is still parked inside acquire(), the thread can go on to actually + # acquire a permit *after* we've already unwound — leaking it forever, + # since the matching finally: release() below never runs for that attempt. + # Poll with the non-blocking form instead: each check returns immediately + # (no OS-level wait), so it's safe to call straight from the event loop + # thread and there's no window for a background acquire to succeed after + # we've already given up on it. + while not ceiling_sem.acquire(False): + await asyncio.sleep(0.05) try: effective = get_max_concurrency() ceiling = _process_wide_max_concurrency() diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py index f90244d..bf2e2b8 100644 --- a/tests/test_concurrency.py +++ b/tests/test_concurrency.py @@ -15,7 +15,7 @@ from pageindex.config import ( set_llm_params, set_max_concurrency, ) -from pageindex.index.utils import _llm_semaphore, llm_acompletion +from pageindex.index.utils import _llm_semaphore, _process_ceiling_semaphore, llm_acompletion @pytest.fixture(autouse=True) @@ -91,6 +91,40 @@ def test_llm_semaphore_is_a_true_process_wide_ceiling_across_threads(): assert state["peak"] == 3 +def test_llm_semaphore_cancellation_while_waiting_does_not_leak_a_permit(): + # A blocking ceiling_sem.acquire() run via asyncio.to_thread() would leak a + # permit under cancellation: the worker thread can't be interrupted, so if + # the awaiting coroutine is cancelled while the thread is still parked + # inside acquire(), the thread can go on to actually acquire the permit + # *after* the coroutine already unwound, and the matching finally: + # release() never runs for that attempt. Cancel a task waiting on an + # already-exhausted ceiling and confirm the permit count fully recovers. + set_max_concurrency(1) + + async def run(): + async def hold(): + async with _llm_semaphore(): + await asyncio.sleep(10) + + holder = asyncio.create_task(hold()) + await asyncio.sleep(0.1) # let it acquire the single permit + + waiter = asyncio.create_task(_llm_semaphore().__aenter__()) + await asyncio.sleep(0.1) # let it start waiting for the permit + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + + holder.cancel() + with pytest.raises(asyncio.CancelledError): + await holder + + await asyncio.sleep(0.2) # give any orphaned acquire a chance to land + assert _process_ceiling_semaphore()._value == 1 + + asyncio.run(run()) + + def test_llm_semaphore_uses_scoped_override(): # A per-index max_concurrency_scope active when the loop's semaphore is first # created must set its size, and must not mutate the process default.