fix: ceiling semaphore permit leak under cancellation

asyncio.to_thread(ceiling_sem.acquire) blocks a worker thread that
can't be interrupted. If the awaiting coroutine is cancelled (Ctrl-C,
an outer timeout) while that thread is still parked inside acquire(),
the thread can go on to actually acquire the permit after the
coroutine has already unwound — leaking it forever, since the
matching finally: release() never runs for that attempt. Poll with
the non-blocking acquire(False) form instead, which returns instantly
and closes the leak window entirely.
This commit is contained in:
mountain 2026-07-09 16:33:09 +08:00
parent e12495dd5b
commit 214098f489
2 changed files with 47 additions and 6 deletions

View file

@ -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.