diff --git a/pageindex/backend/cloud.py b/pageindex/backend/cloud.py index ae6e8da..ade6983 100644 --- a/pageindex/backend/cloud.py +++ b/pageindex/backend/cloud.py @@ -396,6 +396,13 @@ class CloudBackend: timeout=120, ) resp_holder["resp"] = resp + # The consumer may have abandoned the stream while we were still + # blocked in requests.post() (its connect phase, before resp + # existed to close). Now that resp exists, bail immediately + # rather than reading/draining a stream nobody is listening to; + # the finally block closes resp and pushes the sentinel. + if stop.is_set(): + return if resp.status_code != 200: body = resp.text[:500] if resp.text else "" raise CloudAPIError( diff --git a/pageindex/index/utils.py b/pageindex/index/utils.py index ee541e0..1339708 100644 --- a/pageindex/index/utils.py +++ b/pageindex/index/utils.py @@ -116,11 +116,34 @@ def _sync_llm_semaphore(): It uses the same process-wide ceiling so sync and async LLM calls share one real cap. A scoped override can only narrow that cap for the active context. + + A *blocking* ceiling acquire is only safe OFF the event-loop thread. Several + sync LLM helpers (``check_toc`` → ``toc_detector_single_page``, + ``process_no_toc`` → ``generate_toc_init``, ``toc_transformer``, …) are + called synchronously from inside async coroutines (``meta_processor`` → + ``process_large_node_recursively``), i.e. ON the running loop. There, the + async ``_llm_semaphore`` holders own the ceiling permits and can only + release them by resuming on that same loop — so a blocking acquire here + would freeze the loop and *deadlock*: the permit it waits for can never be + freed. When we detect a running loop we therefore take a slot only if one is + immediately free (non-blocking) and otherwise proceed without it. That's + safe: a sync call monopolizes the loop thread while it runs, so it's already + serialized on this loop and can't multiply the in-flight count beyond one + extra per loop. """ + try: + asyncio.get_running_loop() + on_event_loop = True + except RuntimeError: + on_event_loop = False + 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. + # Blocking acquire() (off-loop) always returns True; acquire(False) (on-loop) + # may return False, meaning "no free permit — proceed without one" rather + # than block the loop into a deadlock. + held_ceiling = ceiling_sem.acquire(False) if on_event_loop else ceiling_sem.acquire() + # Only track a permit we actually hold (mirrors _llm_semaphore): guards + # against releasing one we never acquired. scoped_sem = None try: effective = get_max_concurrency() @@ -128,13 +151,19 @@ def _sync_llm_semaphore(): if effective < ceiling: candidate = _max_concurrency_scope_semaphore() if candidate is not None: - candidate.acquire() - scoped_sem = candidate + # Same rule for the scoped cap: never block the loop for it. + if on_event_loop: + if candidate.acquire(False): + scoped_sem = candidate + else: + candidate.acquire() + scoped_sem = candidate yield finally: if scoped_sem is not None: scoped_sem.release() - ceiling_sem.release() + if held_ceiling: + ceiling_sem.release() def llm_completion(model, prompt, chat_history=None, return_finish_reason=False): @@ -530,6 +559,9 @@ def remove_structure_text(data): # ── Functions migrated from retrieve.py ────────────────────────────────────── +_MAX_PAGES = 1000 + + def parse_pages(pages: str) -> list[int]: """Parse a pages string like '5-7', '3,8', or '12' into a sorted list of ints.""" result = [] @@ -539,13 +571,23 @@ def parse_pages(pages: str) -> list[int]: start, end = int(part.split('-', 1)[0].strip()), int(part.split('-', 1)[1].strip()) if start > end: raise ValueError(f"Invalid range '{part}': start must be <= end") + # Bound the span BEFORE materializing range() into the list. Checking + # len(result) only after `result.extend(range(...))` is too late: a + # single huge span like '1-2000000000' allocates billions of ints + # and exhausts memory before the cap is ever reached (DoS). page_nums + # is attacker/LLM-reachable via get_page_content. + span = end - start + 1 + if span > _MAX_PAGES or len(result) + span > _MAX_PAGES: + raise ValueError(f"Page range too large: max {_MAX_PAGES} pages") result.extend(range(start, end + 1)) else: + if len(result) + 1 > _MAX_PAGES: + raise ValueError(f"Page range too large: max {_MAX_PAGES} pages") result.append(int(part)) result = [p for p in result if p >= 1] result = sorted(set(result)) - if len(result) > 1000: - raise ValueError(f"Page range too large: {len(result)} pages (max 1000)") + if len(result) > _MAX_PAGES: + raise ValueError(f"Page range too large: {len(result)} pages (max {_MAX_PAGES})") return result diff --git a/pageindex/storage/sqlite.py b/pageindex/storage/sqlite.py index 20d5997..03e0dac 100644 --- a/pageindex/storage/sqlite.py +++ b/pageindex/storage/sqlite.py @@ -27,6 +27,14 @@ class SQLiteStorage: self._local = threading.local() self._connections: list[sqlite3.Connection] = [] self._conn_lock = threading.Lock() + # Bumped by close(). A thread caches its connection in thread-local + # storage, so after close() every OTHER thread's thread-local still + # points at a now-closed connection. Comparing the cached generation + # against this counter lets _get_conn detect that and reconnect, instead + # of handing back a closed connection (sqlite3.ProgrammingError). close() + # can only touch its OWN thread-local, so this is the only way to + # invalidate the others consistently. + self._generation = 0 # Serializes the (fast) write operations within this process so # concurrent indexing threads don't collide on WAL's single writer # ("database is locked"). Reads stay concurrent; the expensive LLM @@ -36,8 +44,13 @@ class SQLiteStorage: self._init_schema() def _get_conn(self) -> sqlite3.Connection: - """Return a thread-local SQLite connection.""" - if not hasattr(self._local, "conn"): + """Return a thread-local SQLite connection. + + Reconnects if this thread has no connection yet OR its cached connection + was invalidated by a close() on another thread (generation mismatch). + """ + if (not hasattr(self._local, "conn") + or getattr(self._local, "generation", None) != self._generation): # Each thread gets its own connection (threading.local), so # statements never race. check_same_thread=False exists solely so # close() can close every tracked connection from whichever thread @@ -55,6 +68,7 @@ class SQLiteStorage: conn.execute("PRAGMA foreign_keys=ON") conn.execute("PRAGMA busy_timeout=10000") self._local.conn = conn + self._local.generation = self._generation with self._conn_lock: self._connections.append(conn) return self._local.conn @@ -213,6 +227,10 @@ class SQLiteStorage: except Exception: pass self._connections.clear() + # Invalidate every thread's cached connection. close() can only + # del its OWN thread-local, so the bump is what makes _get_conn on + # any other thread reconnect instead of reusing a closed handle. + self._generation += 1 if hasattr(self._local, "conn"): del self._local.conn diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py index 3dd4f9d..ede813c 100644 --- a/tests/test_concurrency.py +++ b/tests/test_concurrency.py @@ -290,6 +290,56 @@ def test_llm_completion_holds_the_shared_semaphore(monkeypatch): assert state["peak"] == 1 +def test_sync_llm_completion_on_event_loop_does_not_deadlock(monkeypatch): + # Regression: _sync_llm_semaphore used a BLOCKING ceiling acquire. Sync LLM + # helpers (check_toc, process_no_toc, toc_transformer, …) run synchronously + # ON the event loop (nested inside the async meta_processor). If async + # llm_acompletion holders occupy every ceiling permit across their awaits, + # a blocking acquire froze the loop -> the holders could never resume to + # release their permits -> permanent deadlock. The sync path must never + # block the running loop. + set_max_concurrency(2) + + async def fake_acompletion(**kwargs): + await asyncio.sleep(0.3) # hold a ceiling permit across the await + return SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))] + ) + + def fake_completion(**kwargs): + return SimpleNamespace( + choices=[SimpleNamespace( + message=SimpleNamespace(content="sync-ok"), finish_reason="stop")] + ) + + monkeypatch.setattr("litellm.acompletion", fake_acompletion) + monkeypatch.setattr("litellm.completion", fake_completion) + + async def run(): + # Both ceiling permits taken by async holders, held across their await. + holders = [asyncio.create_task(llm_acompletion("m", f"p{i}")) for i in range(2)] + await asyncio.sleep(0.05) # let them acquire the permits + # Sync call on the loop thread: pre-fix this blocks forever waiting for a + # permit the holders own and can't release (loop is frozen). + result = llm_completion("m", "sync") + await asyncio.gather(*holders) + return result + + # Run in a thread with a join timeout so a regression FAILS instead of + # hanging CI: a real deadlock freezes the loop, so asyncio.wait_for can't + # cancel it (its timeout callback never runs on the frozen loop). + box = {} + + def target(): + box["result"] = asyncio.run(run()) + + t = threading.Thread(target=target, daemon=True) + t.start() + t.join(timeout=8) + assert not t.is_alive(), "deadlock: sync llm_completion blocked the event loop" + assert box["result"] == "sync-ok" + + def test_run_async_propagates_scope_into_worker_thread(): # When build_index runs inside an already-running loop, _run_async hops to a # worker thread. The max_concurrency_scope override must ride along (copied diff --git a/tests/test_page_content.py b/tests/test_page_content.py index c610180..23df66d 100644 --- a/tests/test_page_content.py +++ b/tests/test_page_content.py @@ -50,6 +50,31 @@ def test_retrieve_parse_pages_delegates_to_canonical_and_enforces_dos_cap(): _parse_pages("1-99999999") +def test_parse_pages_caps_a_huge_range_without_materializing_it(): + """Regression: the 1000-page cap was checked only AFTER + `result.extend(range(start, end + 1))`, so a single huge span like + '1-2000000000' allocated billions of ints and OOM'd before the check ran. + The span must be rejected up front, quickly, without building the list.""" + import time + import pytest + from pageindex.index.utils import parse_pages + + start = time.monotonic() + with pytest.raises(ValueError, match="too large"): + parse_pages("1-2000000000") + # Must be near-instant (no billion-element allocation). Generous bound to + # avoid flakiness while still failing loudly on a re-materializing regression. + assert time.monotonic() - start < 1.0 + + # Boundary: exactly 1000 pages is allowed; 1001 is rejected. + assert parse_pages("1-1000") == list(range(1, 1001)) + with pytest.raises(ValueError, match="too large"): + parse_pages("1-1001") + # A range that fits but whose accumulation across parts crosses the cap. + with pytest.raises(ValueError, match="too large"): + parse_pages("1-600,700-1400") + + def test_retrieve_get_pdf_page_content_falls_back_to_canonical(tmp_path, monkeypatch): """When no cached 'pages' are present, the file-read fallback must delegate to the canonical get_pdf_page_content instead of re-implementing diff --git a/tests/test_sqlite_storage.py b/tests/test_sqlite_storage.py index aa8f757..fe47cc9 100644 --- a/tests/test_sqlite_storage.py +++ b/tests/test_sqlite_storage.py @@ -139,6 +139,46 @@ def test_close_closes_connections_created_in_other_threads(storage): conns["worker"].execute("SELECT 1") +def test_worker_reconnects_via_get_conn_after_close(storage): + """Regression: after close(), a thread that had already cached a connection + in thread-local storage would get that now-CLOSED handle back from + _get_conn (close() can only del its own thread-local), raising + ProgrammingError instead of transparently reconnecting. A generation bump + on close() must make the SAME thread's next _get_conn hand back a fresh, + working connection.""" + import threading + + storage.create_collection("papers") + + cached = threading.Event() + closed = threading.Event() + result = {} + + def worker(): + # 1. cache a connection in this thread's thread-local + storage._get_conn().execute("SELECT 1") + cached.set() + # 2. wait until the main thread closed the storage (invalidating it) + closed.wait(timeout=5) + # 3. reuse from the SAME thread -> must reconnect, not reuse closed conn + try: + result["val"] = storage._get_conn().execute("SELECT 1").fetchone()[0] + result["list"] = storage.list_collections() + except Exception as e: # noqa: BLE001 - record for assertion + result["err"] = f"{type(e).__name__}: {e}" + + t = threading.Thread(target=worker) + t.start() + cached.wait(timeout=5) + storage.close() # closes + invalidates the worker's cached connection + closed.set() + t.join(timeout=5) + + assert "err" not in result, f"reconnect after close failed: {result.get('err')}" + assert result["val"] == 1 + assert result["list"] == ["papers"] + + def test_duplicate_file_hash_in_collection_raises(storage): """UNIQUE(collection_name, file_hash) guards the add-same-file race.""" import sqlite3