408 lines
17 KiB
Python
408 lines
17 KiB
Python
"""Uploads must not monopolize the gateway's event loop or worker capacity."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import io
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import httpx
|
|
import pytest
|
|
from fastapi import FastAPI, File, HTTPException, Request, UploadFile
|
|
|
|
from app.gateway.routers import uploads as uploads_router
|
|
from deerflow.persistence.upload_policy import MemoryUploadPolicyStore
|
|
from deerflow.runtime.user_context import get_current_user, require_current_user, reset_current_user, set_current_user
|
|
from deerflow.uploads.policy import default_upload_policy
|
|
|
|
|
|
async def _upload_app(monkeypatch, uploads_dir: Path, *, provider=None, separate_threads=False):
|
|
uploads_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
def ensure_uploads_dir(thread_id):
|
|
directory = uploads_dir / thread_id if separate_threads else uploads_dir
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
return directory
|
|
|
|
monkeypatch.setattr(uploads_router, "ensure_uploads_dir", ensure_uploads_dir)
|
|
monkeypatch.setattr(
|
|
uploads_router,
|
|
"get_paths",
|
|
lambda: SimpleNamespace(sandbox_uploads_dir=lambda thread_id, user_id=None: uploads_dir / thread_id if separate_threads else uploads_dir),
|
|
)
|
|
|
|
async def path_user(_thread_id):
|
|
return "user-1"
|
|
|
|
monkeypatch.setattr(uploads_router, "aresolve_path_user_id", path_user)
|
|
provider = provider or SimpleNamespace(uses_thread_data_mounts=True)
|
|
monkeypatch.setattr(uploads_router, "get_sandbox_provider", lambda: provider)
|
|
app = FastAPI()
|
|
store = MemoryUploadPolicyStore()
|
|
await store.ensure_seed(default_upload_policy())
|
|
app.state.upload_policy_store = store
|
|
app.state.upload_finished = False
|
|
|
|
@app.post("/upload/{thread_id}")
|
|
async def upload(thread_id: str, request: Request, files: list[UploadFile] = File(...)):
|
|
try:
|
|
return await uploads_router.upload_files.__wrapped__(thread_id, request, files, SimpleNamespace(uploads={"auto_convert_documents": False}))
|
|
finally:
|
|
app.state.upload_finished = True
|
|
|
|
@app.get("/probe")
|
|
async def probe():
|
|
return {"upload_finished": app.state.upload_finished}
|
|
|
|
return app
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("stage", ["write", "validate", "acquire", "sync", "commit", "cleanup"])
|
|
async def test_multifile_upload_keeps_other_requests_responsive(monkeypatch, tmp_path: Path, stage: str):
|
|
"""Exercise actual multipart requests alongside deterministic slow upload work.
|
|
|
|
Busy work simulates CPU-heavy validation; the other stages stand in for
|
|
slow storage and a remote sandbox. The ordering assertion avoids a fragile
|
|
machine-specific latency threshold while still failing on the old router.
|
|
"""
|
|
loop = asyncio.get_running_loop()
|
|
started = asyncio.Event()
|
|
started_at: list[float] = []
|
|
|
|
def slow_work():
|
|
if not started_at:
|
|
started_at.append(time.perf_counter())
|
|
loop.call_soon_threadsafe(started.set)
|
|
deadline = time.perf_counter() + 0.15
|
|
while time.perf_counter() < deadline:
|
|
pass
|
|
|
|
class Sandbox:
|
|
def update_file(self, _path, _content):
|
|
if stage == "sync":
|
|
slow_work()
|
|
|
|
def acquire(_thread_id):
|
|
if stage == "acquire":
|
|
slow_work()
|
|
return "sandbox-1"
|
|
|
|
provider = SimpleNamespace(uses_thread_data_mounts=False, acquire=acquire, get=lambda _id: Sandbox())
|
|
app = await _upload_app(monkeypatch, tmp_path / "uploads", provider=provider)
|
|
|
|
if stage == "write":
|
|
original_open = uploads_router.open_upload_file_no_symlink
|
|
|
|
class SlowWriter:
|
|
def __init__(self, handle):
|
|
self.handle = handle
|
|
|
|
def write(self, data):
|
|
slow_work()
|
|
return self.handle.write(data)
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self.handle, name)
|
|
|
|
def open_slow(*args, **kwargs):
|
|
path, handle = original_open(*args, **kwargs)
|
|
return path, SlowWriter(handle)
|
|
|
|
monkeypatch.setattr(uploads_router, "open_upload_file_no_symlink", open_slow)
|
|
elif stage in {"validate", "commit", "cleanup"}:
|
|
owner, name = {
|
|
"validate": (uploads_router, "validate_upload_content"),
|
|
"commit": (uploads_router, "_commit_staged_files"),
|
|
"cleanup": (uploads_router.shutil, "rmtree"),
|
|
}[stage]
|
|
original = getattr(owner, name)
|
|
|
|
def slow_call(*args, **kwargs):
|
|
slow_work()
|
|
return original(*args, **kwargs)
|
|
|
|
monkeypatch.setattr(owner, name, slow_call)
|
|
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
|
|
|
async def probe_during_upload():
|
|
await asyncio.wait_for(started.wait(), timeout=5)
|
|
response = await client.get("/probe")
|
|
return response, time.perf_counter() - started_at[0]
|
|
|
|
upload, (probe, latency) = await asyncio.gather(
|
|
client.post(
|
|
"/upload/thread-1",
|
|
files=[("files", (f"sample-{index}.txt", b"test content", "text/plain")) for index in range(3)],
|
|
),
|
|
probe_during_upload(),
|
|
)
|
|
|
|
assert upload.status_code == 200, upload.text
|
|
assert len(upload.json()["files"]) == 3
|
|
assert probe.status_code == 200
|
|
print(f"{stage}: concurrent probe delay = {latency:.3f}s")
|
|
assert probe.json()["upload_finished"] is False, f"{stage} blocked other requests for {latency:.3f}s"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_processing_has_a_global_concurrency_limit(monkeypatch, tmp_path: Path):
|
|
app = await _upload_app(monkeypatch, tmp_path / "uploads", separate_threads=True)
|
|
loop = asyncio.get_running_loop()
|
|
two_workers_started = asyncio.Event()
|
|
release = threading.Event()
|
|
lock = threading.Lock()
|
|
active = peak = entered = 0
|
|
original_validate = uploads_router.validate_upload_content
|
|
|
|
def validate(path, extension):
|
|
nonlocal active, peak, entered
|
|
with lock:
|
|
active += 1
|
|
entered += 1
|
|
peak = max(peak, active)
|
|
if active == 2:
|
|
loop.call_soon_threadsafe(two_workers_started.set)
|
|
try:
|
|
assert release.wait(timeout=5), "Upload workers were not released"
|
|
original_validate(path, extension)
|
|
finally:
|
|
with lock:
|
|
active -= 1
|
|
|
|
monkeypatch.setattr(uploads_router, "validate_upload_content", validate)
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
|
uploads = [asyncio.create_task(client.post(f"/upload/thread-{index}", files={"files": ("file.txt", b"content")})) for index in range(6)]
|
|
try:
|
|
await asyncio.wait_for(two_workers_started.wait(), timeout=3)
|
|
await asyncio.sleep(0.05)
|
|
with lock:
|
|
assert entered == 2, "Excess uploads entered expensive work instead of waiting"
|
|
assert (await client.get("/probe")).status_code == 200
|
|
finally:
|
|
release.set()
|
|
results = await asyncio.gather(*uploads, return_exceptions=True)
|
|
|
|
assert all(isinstance(result, httpx.Response) and result.status_code == 200 for result in results), results
|
|
assert peak == 2
|
|
assert entered == 6
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("stage", ["validate", "sync"])
|
|
async def test_repeated_cancellation_waits_for_worker_and_restores_old_files(monkeypatch, tmp_path: Path, stage: str):
|
|
uploads_dir = tmp_path / "uploads"
|
|
loop = asyncio.get_running_loop()
|
|
started = asyncio.Event()
|
|
release = threading.Event()
|
|
original_validate = uploads_router.validate_upload_content
|
|
|
|
def pause_worker():
|
|
loop.call_soon_threadsafe(started.set)
|
|
assert release.wait(timeout=5), "Cancelled worker was not released"
|
|
|
|
class Sandbox:
|
|
def update_file(self, _path, _content):
|
|
if stage == "sync":
|
|
pause_worker()
|
|
|
|
provider = SimpleNamespace(uses_thread_data_mounts=False, acquire=lambda _id: "sandbox-1", get=lambda _id: Sandbox())
|
|
app = await _upload_app(monkeypatch, uploads_dir, provider=provider)
|
|
(uploads_dir / "same.txt").write_bytes(b"old content")
|
|
|
|
def validate(path, extension):
|
|
if stage == "validate":
|
|
pause_worker()
|
|
original_validate(path, extension)
|
|
|
|
monkeypatch.setattr(uploads_router, "validate_upload_content", validate)
|
|
upload_file = UploadFile(filename="same.txt", file=io.BytesIO(b"new content"))
|
|
task = asyncio.create_task(uploads_router.upload_files.__wrapped__("thread-1", SimpleNamespace(app=app), [upload_file], SimpleNamespace(uploads={"auto_convert_documents": False})))
|
|
try:
|
|
await asyncio.wait_for(started.wait(), timeout=3)
|
|
task.cancel()
|
|
await asyncio.sleep(0)
|
|
assert not task.done(), "Cancellation detached a still-running upload worker"
|
|
task.cancel()
|
|
await asyncio.sleep(0)
|
|
assert not task.done(), "A second cancellation detached the upload worker"
|
|
assert not upload_file.file.closed
|
|
finally:
|
|
release.set()
|
|
results = await asyncio.gather(task, return_exceptions=True)
|
|
await upload_file.close()
|
|
|
|
assert isinstance(results[0], asyncio.CancelledError)
|
|
assert (uploads_dir / "same.txt").read_bytes() == b"old content"
|
|
assert sorted(path.name for path in uploads_dir.iterdir()) == ["same.txt"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_worker_preserves_each_requests_user_context(monkeypatch, tmp_path: Path):
|
|
app = await _upload_app(monkeypatch, tmp_path / "uploads", separate_threads=True)
|
|
original_ensure = uploads_router.ensure_uploads_dir
|
|
original_validate = uploads_router.validate_upload_content
|
|
observed = []
|
|
caller_thread = threading.get_ident()
|
|
previous_user = get_current_user()
|
|
|
|
def ensure_directory(thread_id):
|
|
assert require_current_user().id == thread_id
|
|
assert threading.get_ident() != caller_thread
|
|
return original_ensure(thread_id)
|
|
|
|
def validate(path, extension):
|
|
assert require_current_user().id == path.stem
|
|
observed.append(require_current_user().id)
|
|
original_validate(path, extension)
|
|
|
|
monkeypatch.setattr(uploads_router, "ensure_uploads_dir", ensure_directory)
|
|
monkeypatch.setattr(uploads_router, "validate_upload_content", validate)
|
|
|
|
async def upload_as(user_id):
|
|
token = set_current_user(SimpleNamespace(id=user_id))
|
|
upload_file = UploadFile(filename=f"{user_id}.txt", file=io.BytesIO(b"content"))
|
|
try:
|
|
return await uploads_router.upload_files.__wrapped__(user_id, SimpleNamespace(app=app), [upload_file], SimpleNamespace(uploads={"auto_convert_documents": False}))
|
|
finally:
|
|
await upload_file.close()
|
|
reset_current_user(token)
|
|
|
|
results = await asyncio.gather(upload_as("user-alpha"), upload_as("user-beta"))
|
|
assert all(result.success for result in results)
|
|
assert sorted(observed) == ["user-alpha", "user-beta"]
|
|
assert get_current_user() is previous_user
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("collision", [False, True])
|
|
async def test_worker_conversion_sidecars_and_collision_rollback(monkeypatch, tmp_path: Path, collision: bool):
|
|
uploads_dir = tmp_path / "uploads"
|
|
app = await _upload_app(monkeypatch, uploads_dir)
|
|
(uploads_dir / "report.pdf").write_bytes(b"old pdf")
|
|
(uploads_dir / "report.pdf.md").write_bytes(b"old markdown")
|
|
caller_thread = threading.get_ident()
|
|
|
|
def convert(path):
|
|
assert threading.get_ident() != caller_thread
|
|
sidecar = path.with_name(f"{path.name}.md")
|
|
sidecar.write_text("# converted", encoding="utf-8")
|
|
return sidecar
|
|
|
|
monkeypatch.setattr(uploads_router, "convert_file_to_markdown_sync", convert)
|
|
files = [UploadFile(filename="report.pdf", file=io.BytesIO(b"%PDF-1.7\n"))]
|
|
if collision:
|
|
files.append(UploadFile(filename="report.pdf.md", file=io.BytesIO(b"user markdown")))
|
|
try:
|
|
operation = uploads_router.upload_files.__wrapped__("thread-1", SimpleNamespace(app=app), files, SimpleNamespace(uploads={"auto_convert_documents": True}))
|
|
if collision:
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await operation
|
|
assert exc_info.value.status_code == 409
|
|
assert (uploads_dir / "report.pdf").read_bytes() == b"old pdf"
|
|
assert (uploads_dir / "report.pdf.md").read_bytes() == b"old markdown"
|
|
else:
|
|
result = await operation
|
|
assert result.success
|
|
assert result.files[0]["markdown_file"] == "report.pdf.md"
|
|
assert (uploads_dir / "report.pdf.md").read_text(encoding="utf-8") == "# converted"
|
|
assert (uploads_dir / "report.pdf").read_bytes() == b"%PDF-1.7\n"
|
|
finally:
|
|
for upload_file in files:
|
|
await upload_file.close()
|
|
assert sorted(path.name for path in uploads_dir.iterdir()) == ["report.pdf", "report.pdf.md"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_same_thread_uploads_cannot_undo_each_others_files(monkeypatch, tmp_path: Path):
|
|
uploads_dir = tmp_path / "uploads"
|
|
loop = asyncio.get_running_loop()
|
|
first_sync_started = asyncio.Event()
|
|
release = threading.Event()
|
|
|
|
class Sandbox:
|
|
def update_file(self, _path, content):
|
|
if content == b"first":
|
|
loop.call_soon_threadsafe(first_sync_started.set)
|
|
assert release.wait(timeout=5), "First upload was not released"
|
|
raise OSError("First sandbox sync failed")
|
|
|
|
provider = SimpleNamespace(uses_thread_data_mounts=False, acquire=lambda _id: "sandbox-1", get=lambda _id: Sandbox())
|
|
app = await _upload_app(monkeypatch, uploads_dir, provider=provider)
|
|
(uploads_dir / "same.txt").write_bytes(b"original")
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
|
first = asyncio.create_task(client.post("/upload/thread-1", files={"files": ("same.txt", b"first")}))
|
|
tasks = [first]
|
|
try:
|
|
await asyncio.wait_for(first_sync_started.wait(), timeout=3)
|
|
second = asyncio.create_task(client.post("/upload/thread-1", files={"files": ("same.txt", b"second")}))
|
|
tasks.append(second)
|
|
await asyncio.sleep(0.05)
|
|
assert not second.done(), "Overlapping upload modified a batch that can still roll back"
|
|
finally:
|
|
release.set()
|
|
results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
assert [result.status_code for result in results] == [500, 200]
|
|
assert (uploads_dir / "same.txt").read_bytes() == b"second"
|
|
assert sorted(path.name for path in uploads_dir.iterdir()) == ["same.txt"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("operation", ["list", "delete"])
|
|
async def test_list_and_delete_wait_for_upload_rollback(monkeypatch, tmp_path: Path, operation: str):
|
|
"""No reader sees an uncommitted batch, and rollback never revives a deletion."""
|
|
uploads_dir = tmp_path / "uploads"
|
|
loop = asyncio.get_running_loop()
|
|
sync_started = asyncio.Event()
|
|
release = threading.Event()
|
|
|
|
class Sandbox:
|
|
def update_file(self, _path, _content):
|
|
loop.call_soon_threadsafe(sync_started.set)
|
|
assert release.wait(timeout=5), "Upload sync was not released"
|
|
raise OSError("Sandbox sync failed after the host commit")
|
|
|
|
provider = SimpleNamespace(uses_thread_data_mounts=False, acquire=lambda _id: "sandbox-1", get=lambda _id: Sandbox())
|
|
app = await _upload_app(monkeypatch, uploads_dir, provider=provider)
|
|
monkeypatch.setattr(uploads_router, "get_uploads_dir", lambda _thread_id: uploads_dir)
|
|
(uploads_dir / "same.txt").write_bytes(b"old")
|
|
request = SimpleNamespace(app=app)
|
|
files = [
|
|
UploadFile(filename="same.txt", file=io.BytesIO(b"new content")),
|
|
UploadFile(filename="added.txt", file=io.BytesIO(b"new file")),
|
|
]
|
|
upload = asyncio.create_task(uploads_router.upload_files.__wrapped__("thread-1", request, files, SimpleNamespace(uploads={"auto_convert_documents": False})))
|
|
tasks = [upload]
|
|
try:
|
|
await asyncio.wait_for(sync_started.wait(), timeout=3)
|
|
if operation == "list":
|
|
followup = asyncio.create_task(uploads_router.list_uploaded_files.__wrapped__("thread-1", request))
|
|
else:
|
|
followup = asyncio.create_task(uploads_router.delete_uploaded_file.__wrapped__("thread-1", "same.txt", request))
|
|
tasks.append(followup)
|
|
await asyncio.sleep(0.05)
|
|
assert not followup.done(), f"{operation} accessed files before the upload could finish rolling back"
|
|
finally:
|
|
release.set()
|
|
results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
for upload_file in files:
|
|
await upload_file.close()
|
|
|
|
assert isinstance(results[0], HTTPException)
|
|
assert results[0].status_code == 500
|
|
result = results[1]
|
|
assert isinstance(result, dict)
|
|
if operation == "list":
|
|
assert result["count"] == 1
|
|
assert [(file["filename"], file["size"]) for file in result["files"]] == [("same.txt", "3")]
|
|
assert (uploads_dir / "same.txt").read_bytes() == b"old"
|
|
else:
|
|
assert result["success"]
|
|
assert list(uploads_dir.iterdir()) == []
|
|
assert not (uploads_dir / "added.txt").exists()
|