mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-15 21:11:05 +02:00
fix: resolve concurrency and DoS pitfalls from PR #272 review
- index/utils.py: fix an asyncio deadlock in _sync_llm_semaphore. Sync LLM calls (check_toc / process_no_toc / toc_transformer) run on the event-loop thread nested inside the async meta_processor, and its blocking ceiling acquire could wait forever for a permit held by async _llm_semaphore holders that can only release it once the (now-frozen) loop runs. Take the slot non-blocking when on a running loop; keep the blocking acquire off-loop. - index/utils.py: bound parse_pages ranges before materializing range() into the list, so a huge span like '1-2000000000' is rejected up front instead of exhausting memory before the 1000-page cap is ever checked (DoS). - storage/sqlite.py: bump a generation counter on close() so a thread that cached a connection in thread-local storage reconnects on its next call instead of reusing a closed handle (ProgrammingError). - backend/cloud.py: after connect, bail out of the SSE background thread if the consumer already abandoned the stream, instead of draining it in the background. Adds regression tests for each fix.
This commit is contained in:
parent
d97231b480
commit
9ad54122bb
6 changed files with 192 additions and 10 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue