mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-15 21:11:05 +02:00
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:
parent
56590c63d5
commit
d3ea9b9f2e
2 changed files with 49 additions and 9 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue