fix(personal): route upload and delete through the index job lock

/api/personal/upload and DELETE /api/personal/file mutated the same
vector and tracking state add/remove/reload serialize on, outside
_index_job_lock and inline on the event loop.

Both now stage async work on the loop, then run the complete transition
(vector writes, disk change, personal_docs_manager update) in one
offloaded critical section under the shared lock, acquired before the
offload so queued requests park on the loop rather than pinning a
threadpool worker.

Adds add-vs-upload and add-vs-file ordering regressions.
This commit is contained in:
StressTestor
2026-08-12 10:17:36 -06:00
parent e222e92153
commit 938251000b
2 changed files with 202 additions and 68 deletions
+42 -10
View File
@@ -328,10 +328,12 @@ def setup_personal_routes(personal_docs_manager, rag_manager, rag_available):
upload_dir = _personal_upload_dir_for_owner(user) upload_dir = _personal_upload_dir_for_owner(user)
total_indexed = 0
total_failed = 0 total_failed = 0
uploaded_files = []
# Read the request bodies on the event loop — that part is genuine async
# I/O — then stage them so every blocking step happens in one offloaded
# critical section below.
staged: List[Tuple[str, str, str, bytes]] = []
for upload in files: for upload in files:
try: try:
file_path, stored_name, safe_name = _unique_personal_upload_path(upload_dir, upload.filename) file_path, stored_name, safe_name = _unique_personal_upload_path(upload_dir, upload.filename)
@@ -340,6 +342,17 @@ def setup_personal_routes(personal_docs_manager, rag_manager, rag_available):
logger.warning(f"Rejected oversized personal upload: {upload.filename!r}") logger.warning(f"Rejected oversized personal upload: {upload.filename!r}")
total_failed += 1 total_failed += 1
continue continue
staged.append((file_path, stored_name, safe_name, content_bytes))
except Exception as e:
logger.error(f"Failed to read upload {upload.filename}: {e}")
total_failed += 1
def _index_uploads():
indexed = 0
failed = 0
names = []
for file_path, stored_name, safe_name, content_bytes in staged:
try:
with open(file_path, "wb") as f: with open(file_path, "wb") as f:
f.write(content_bytes) f.write(content_bytes)
@@ -351,7 +364,7 @@ def setup_personal_routes(personal_docs_manager, rag_manager, rag_available):
text = content_bytes.decode("utf-8", errors="replace") text = content_bytes.decode("utf-8", errors="replace")
if not text or not text.strip(): if not text or not text.strip():
total_failed += 1 failed += 1
continue continue
# Chunk and index # Chunk and index
@@ -368,18 +381,28 @@ def setup_personal_routes(personal_docs_manager, rag_manager, rag_available):
if user: if user:
metadata["owner"] = user metadata["owner"] = user
if rag.add_document(chunk, metadata): if rag.add_document(chunk, metadata):
total_indexed += 1 indexed += 1
else: else:
total_failed += 1 failed += 1
uploaded_files.append(safe_name) names.append(safe_name)
except Exception as e: except Exception as e:
logger.error(f"Failed to upload/index {upload.filename}: {e}") logger.error(f"Failed to upload/index {safe_name}: {e}")
total_failed += 1 failed += 1
# Track uploads directory # Same transition, same lock: the tracking update must not land
if uploaded_files and hasattr(personal_docs_manager, "add_directory"): # while another job is mid-write over the same state.
if names and hasattr(personal_docs_manager, "add_directory"):
personal_docs_manager.add_directory(upload_dir, index=False) personal_docs_manager.add_directory(upload_dir, index=False)
return indexed, failed, names
# Chunking, embedding and the tracking update are blocking work over the
# same vector/tracking state add_directory mutates (#5634). Take the
# shared job lock BEFORE offloading so a queued request parks on the loop
# instead of pinning a threadpool worker, matching add_directory.
async with _index_job_lock:
total_indexed, indexed_failed, uploaded_files = await run_in_threadpool(_index_uploads)
total_failed += indexed_failed
return { return {
"success": True, "success": True,
@@ -392,6 +415,7 @@ def setup_personal_routes(personal_docs_manager, rag_manager, rag_available):
async def delete_file_from_rag(filepath: str = Query(...), owner: str = Depends(require_user), _admin: None = Depends(require_admin)): async def delete_file_from_rag(filepath: str = Query(...), owner: str = Depends(require_user), _admin: None = Depends(require_admin)):
"""Delete a specific file from RAG index and optionally from disk.""" """Delete a specific file from RAG index and optionally from disk."""
try: try:
def _delete_file():
# Remove chunks from RAG vector store (best-effort) # Remove chunks from RAG vector store (best-effort)
removed = 0 removed = 0
rag = _rag() rag = _rag()
@@ -424,6 +448,14 @@ def setup_personal_routes(personal_docs_manager, rag_manager, rag_available):
# Exclude the file from the listing (persists across restarts) # Exclude the file from the listing (persists across restarts)
personal_docs_manager.exclude_file(filepath) personal_docs_manager.exclude_file(filepath)
return removed, deleted_from_disk
# Vector removal, the disk unlink and the exclusion write are one
# transition over the same state add_directory mutates (#5634), and
# all three block. Take the shared job lock BEFORE offloading, as
# add_directory does.
async with _index_job_lock:
removed, deleted_from_disk = await run_in_threadpool(_delete_file)
return { return {
"success": True, "success": True,
+102
View File
@@ -61,6 +61,17 @@ class _FakeRag:
self._record["index_thread"] = threading.get_ident() self._record["index_thread"] = threading.get_ident()
return {"success": True, "indexed_count": 3, "failed_count": 0} return {"success": True, "indexed_count": 3, "failed_count": 0}
def _split_into_chunks(self, text, chunk_size=500):
return [text]
def add_document(self, chunk, metadata):
self._record["add_document_thread"] = threading.get_ident()
return True
def delete_by_source(self, filepath):
self._record["delete_thread"] = threading.get_ident()
return 1
class _FakeDocsManager: class _FakeDocsManager:
def __init__(self, record): def __init__(self, record):
@@ -71,6 +82,9 @@ class _FakeDocsManager:
self._record["bookkeeping_thread"] = threading.get_ident() self._record["bookkeeping_thread"] = threading.get_ident()
self._record["bookkeeping_index_flag"] = index self._record["bookkeeping_index_flag"] = index
def exclude_file(self, filepath):
self._record["exclude_thread"] = threading.get_ident()
def _build_app(tmp_path, monkeypatch, record): def _build_app(tmp_path, monkeypatch, record):
monkeypatch.setattr(personal_routes, "PERSONAL_DIR", str(tmp_path)) monkeypatch.setattr(personal_routes, "PERSONAL_DIR", str(tmp_path))
@@ -219,6 +233,94 @@ async def test_add_and_remove_serialize(tmp_path, monkeypatch):
) )
async def test_add_and_upload_serialize(tmp_path, monkeypatch):
"""#5634 follow-up: POST /upload writes chunks into the vector store and then
calls personal_docs_manager.add_directory — the same vector/tracking state
add_directory mutates. It must hold the SAME job lock, or an upload landing
mid-add interleaves two writers over unsynchronized state."""
import time
state, enter, leave = _serialization_probe()
def _slow_index(self, directory, owner=None):
enter(); time.sleep(0.25); leave()
return {"success": True, "indexed_count": 1, "failed_count": 0}
def _slow_add_document(self, chunk, metadata):
self._record["add_document_thread"] = threading.get_ident()
enter(); time.sleep(0.25); leave()
return True
monkeypatch.setattr(_FakeRag, "index_personal_documents", _slow_index)
monkeypatch.setattr(_FakeRag, "add_document", _slow_add_document)
record = {}
app = _build_app(tmp_path, monkeypatch, record)
monkeypatch.setattr(personal_routes, "UPLOADS_DIR", str(tmp_path / "uploads"))
monkeypatch.setattr(personal_routes, "require_privilege", lambda request, key: "tester")
(tmp_path / "docs_a").mkdir()
async with _async_client(app) as ac:
results = await asyncio.gather(
ac.post("/api/personal/add_directory", json={"directory": str(tmp_path / "docs_a")}),
ac.post("/api/personal/upload", files={"files": ("a.txt", b"hello world", "text/plain")}),
)
assert all(r.status_code == 200 for r in results)
# The test coroutine runs on the event loop, so this IS the loop thread.
assert record["add_document_thread"] != threading.get_ident(), (
"rag.add_document ran on the event loop thread — chunk writes block "
"every other request for the duration of the upload"
)
assert state["max_active"] == 1, (
f"{state['max_active']} add/upload critical sections overlapped — "
"upload must hold the same index job lock as add"
)
async def test_add_and_delete_file_serialize(tmp_path, monkeypatch):
"""#5634 follow-up: DELETE /file removes chunks from the vector store and
calls personal_docs_manager.exclude_file. Both mutate state add_directory
also touches, so the delete must hold the SAME job lock as add."""
import time
state, enter, leave = _serialization_probe()
def _slow_index(self, directory, owner=None):
enter(); time.sleep(0.25); leave()
return {"success": True, "indexed_count": 1, "failed_count": 0}
def _slow_delete(self, filepath):
self._record["delete_thread"] = threading.get_ident()
enter(); time.sleep(0.25); leave()
return 1
monkeypatch.setattr(_FakeRag, "index_personal_documents", _slow_index)
monkeypatch.setattr(_FakeRag, "delete_by_source", _slow_delete)
record = {}
app = _build_app(tmp_path, monkeypatch, record)
monkeypatch.setattr(personal_routes, "UPLOADS_DIR", str(tmp_path / "uploads"))
(tmp_path / "docs_a").mkdir()
doomed = tmp_path / "doomed.txt"
doomed.write_text("bye")
async with _async_client(app) as ac:
results = await asyncio.gather(
ac.post("/api/personal/add_directory", json={"directory": str(tmp_path / "docs_a")}),
ac.delete("/api/personal/file", params={"filepath": str(doomed)}),
)
assert all(r.status_code == 200 for r in results)
assert record["delete_thread"] != threading.get_ident(), (
"rag.delete_by_source ran on the event loop thread"
)
assert state["max_active"] == 1, (
f"{state['max_active']} add/delete critical sections overlapped — "
"delete must hold the same index job lock as add"
)
async def test_reload_serializes_with_add(tmp_path, monkeypatch): async def test_reload_serializes_with_add(tmp_path, monkeypatch):
"""#5634: POST /reload rebuilds the index via refresh_index(); it must hold """#5634: POST /reload rebuilds the index via refresh_index(); it must hold
the same job lock so it cannot race an in-flight add job.""" the same job lock so it cannot race an in-flight add job."""