fix: scoped concurrency semaphore over-release on cancellation

The sync-concurrency fix marked the scoped-override semaphore for release
before it was actually acquired, so a coroutine cancelled while polling for
a scoped permit (or a sync acquire interrupted mid-wait) ran the finally and
released a permit it never held — inflating the scoped cap for later calls
(the mirror of the ceiling-leak fix). Only bind the release guard after the
acquire succeeds, in both the async and sync semaphores. Regression test
included: cancelling a waiter leaves the scoped permit count at 1, not 2.
This commit is contained in:
mountain 2026-07-09 23:08:39 +08:00
parent 56590c63d5
commit d3ea9b9f2e
2 changed files with 49 additions and 9 deletions

View file

@ -90,18 +90,20 @@ async def _llm_semaphore():
# we've already given up on it.
while not ceiling_sem.acquire(False):
await asyncio.sleep(0.05)
# Only set once the permit is actually held, so a cancellation while polling
# for it doesn't make the finally release a permit we never acquired (which
# would inflate the scoped cap — the mirror of the ceiling leak fixed above).
scoped_sem = None
try:
effective = get_max_concurrency()
ceiling = _process_wide_max_concurrency()
if effective < ceiling:
scoped_sem = _max_concurrency_scope_semaphore()
if scoped_sem is not None:
while not scoped_sem.acquire(False):
candidate = _max_concurrency_scope_semaphore()
if candidate is not None:
while not candidate.acquire(False):
await asyncio.sleep(0.05)
yield
else:
yield
scoped_sem = candidate
yield
finally:
if scoped_sem is not None:
scoped_sem.release()
@ -117,14 +119,17 @@ def _sync_llm_semaphore():
"""
ceiling_sem = _process_ceiling_semaphore()
ceiling_sem.acquire()
# Only set once the permit is actually held (mirrors _llm_semaphore): guards
# against releasing a permit we never acquired if acquire() is interrupted.
scoped_sem = None
try:
effective = get_max_concurrency()
ceiling = _process_wide_max_concurrency()
if effective < ceiling:
scoped_sem = _max_concurrency_scope_semaphore()
if scoped_sem is not None:
scoped_sem.acquire()
candidate = _max_concurrency_scope_semaphore()
if candidate is not None:
candidate.acquire()
scoped_sem = candidate
yield
finally:
if scoped_sem is not None:

View file

@ -11,6 +11,7 @@ from pageindex.config import (
IndexConfig,
_env_llm_timeout_default,
_env_max_concurrency_default,
_max_concurrency_scope_semaphore,
get_llm_params,
get_max_concurrency,
llm_params_scope,
@ -133,6 +134,40 @@ def test_llm_semaphore_cancellation_while_waiting_does_not_leak_a_permit():
asyncio.run(run())
def test_scoped_semaphore_cancellation_while_waiting_does_not_over_release():
# Mirror of the ceiling test for the scoped override: a coroutine cancelled
# while polling for a scoped permit must NOT let the finally release a permit
# it never acquired, which would inflate the scoped cap for later calls.
set_max_concurrency(5) # high ceiling so the scope is the narrower cap
async def run():
with max_concurrency_scope(1):
sem = _max_concurrency_scope_semaphore()
assert sem._value == 1
async def hold():
async with _llm_semaphore():
await asyncio.sleep(10)
holder = asyncio.create_task(hold())
await asyncio.sleep(0.1) # holder takes the single scoped permit
waiter = asyncio.create_task(_llm_semaphore().__aenter__())
await asyncio.sleep(0.1) # waiter is now polling for the scoped permit
waiter.cancel()
with pytest.raises(asyncio.CancelledError):
await waiter
holder.cancel()
with pytest.raises(asyncio.CancelledError):
await holder
await asyncio.sleep(0.2)
assert sem._value == 1 # fully recovered, not inflated to 2
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.