From d3ea9b9f2ec13ea5674f3f418ca3aea8d14d6882 Mon Sep 17 00:00:00 2001
From: mountain <kose2livs@gmail.com>
Date: Thu, 9 Jul 2026 23:08:39 +0800
Subject: [PATCH] fix: scoped concurrency semaphore over-release on
 cancellation
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 pageindex/index/utils.py  | 23 ++++++++++++++---------
 tests/test_concurrency.py | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/pageindex/index/utils.py b/pageindex/index/utils.py
index 1156086..11bd39a 100644
--- a/pageindex/index/utils.py
+++ b/pageindex/index/utils.py
@@ -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:
diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py
index aa697e2..58ee20b 100644
--- a/tests/test_concurrency.py
+++ b/tests/test_concurrency.py
@@ -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.