diff --git a/CHANGELOG.md b/CHANGELOG.md index b0c255829..40e2cd15b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -83,6 +83,7 @@ All notable changes to this project are documented in this file. The format is b ### Fixed +- Offload S3 server-upload spool I/O, wait for in-flight I/O during cancellation, and reuse one client per multipart upload. Index the default tenant file browse order (migration `f6c2d4e8b109`; local performance review, 2026-11-05). - Users password creation, forgot-password fingerprints, reset verification and password updates offload hashing to worker threads, keeping unrelated async requests responsive during registration and bulk invites. Password policy, diff --git a/docs/database/migrations.md b/docs/database/migrations.md index 8602d6aa9..d08a1d651 100644 --- a/docs/database/migrations.md +++ b/docs/database/migrations.md @@ -71,7 +71,7 @@ The history is **not** a single linear chain. It has several roots and several h | Head | Owning module | Label Alembic reports | |---|---|---| -| `c5d8e2f1a903` | background_tasks (newest-first listing index) | `api_tokens`, `background_tasks`, `file_storage`, `users` | +| `f6c2d4e8b109` | file_storage (tenant browse index) | `api_tokens`, `background_tasks`, `file_storage`, `users` | | `b8d4f2a6c901` | audit_log | `audit_log` | | `f1a9c3d7b2e4` | tenants | `tenants` | | `d263ea727663` | framework (`sm_lease`) | `leases` | diff --git a/docs/modules/file_storage.md b/docs/modules/file_storage.md index e61262ebe..c00de307f 100644 --- a/docs/modules/file_storage.md +++ b/docs/modules/file_storage.md @@ -332,7 +332,11 @@ worker loop itself. S3 server-side `put` switches to multipart above `direct_upload_threshold_bytes`, even for streams whose size is unknown. It spools only the threshold prefix and one part at a time (each spills to disk -above 10 MiB). Part size comes from `multipart_part_size_bytes`, grown for known +above 10 MiB). Spool I/O runs in worker threads; cancellation waits for +in-flight I/O before closing its file. Each server-side multipart upload reuses +one S3 client through completion or abort, allowing connection-pool reuse +without sharing a client across requests or event loops. Part size comes from +`multipart_part_size_bytes`, grown for known sizes to fit 10,000 parts; unknown-size streams abort if they exceed that bound. Declared nonzero sizes must match the bytes read. Multipart is also used above S3's single-PUT limit, even if the threshold was configured higher. Stream, @@ -343,6 +347,11 @@ without masking the failure. Worker uploads still honor `max_file_size_bytes` 32-bit `INTEGER` on PostgreSQL, which refused files of 2 GiB or more), and the S3 backend's `put` goes multipart, so objects beyond 5 GiB work end to end. +Migration `f6c2d4e8b109` adds an index on `(tenant_id, is_deleted, +created_at DESC, id)` for the default live-file browse order. It extends the +shared `file_storage@head` trunk; apply with `upgrade heads`. Other sort orders +and filtered searches may still require sorting. + A process killed mid-upload cannot run compensation. Its incomplete multipart upload is an orphan like an abandoned browser one, and the upload reaper aborts it once it is older than twice `upload_session_ttl_seconds` (48 hours by diff --git a/docs/perf/2026-11-05-file-storage.md b/docs/perf/2026-11-05-file-storage.md new file mode 100644 index 000000000..1179bd69e --- /dev/null +++ b/docs/perf/2026-11-05-file-storage.md @@ -0,0 +1,53 @@ +# File-storage performance fixes, 2026-11-05 + +Local review against `4920a13a`: synchronous server-upload spool I/O, per-part +S3 clients, and a missing ordered tenant browse index. + +## Local measurements + +No S3, HTTP server, browser, or PostgreSQL service was used. These measurements +are evidence for the local changes, not production latency predictions. + +### Upload loop lag + +Two concurrent uploads (32 MiB and 64 MiB), 1 MiB source chunks, real temporary +files with the existing 10 MiB spool limit, and the existing `S3Wire` fake. +The wire consumed file bodies in executor threads, like aiohttp's file payload, +without simulating network latency. A 1 ms asyncio timer sampled scheduling +lag. Known sizes used the actual length; unknown sizes used zero. Each cell is +the median of three runs. Three implementations were compared in one process: +`4920a13a` (origin/main), the first offloading change where every spool call +went through the executor ("all-offload"), and the final change, where spool +operations that stay in memory (up to 2 MiB) run inline and only operations that +touch disk are offloaded. The final change also replays the prefix in 1 MiB +blocks instead of 64 KiB and tracks the spool position itself instead of +calling `tell()` through the executor. + +| Threshold | Size known | Loop lag p99, main / all-offload / final | Elapsed, main / all-offload / final | Client opens main → final | +| --- | --- | --- | --- | --- | +| 16 MiB (force multipart) | yes | 28.9 / 7.2 / 13.4 ms | 0.234 / 0.306 / 0.297 s | 10 → 2 | +| 16 MiB (force multipart) | no | 54.8 / 1.5 / 11.4 ms | 0.251 / 0.794 / 0.340 s | 10 → 2 | +| 64 MiB (default; single PUT) | yes | 37.1 / 10.4 / 13.8 ms | 0.261 / 0.314 / 0.301 s | 2 → 2 | +| 64 MiB (default; single PUT) | no | 36.8 / 11.5 / 12.8 ms | 0.257 / 0.340 / 0.290 s | 2 → 2 | + +The final change keeps most of the loop-lag gain (p99 2-4x lower than main) and +removes most of the throughput regression: elapsed is within 10-35% of main +(unknown-size multipart was 3.2x slower with full offloading, now 1.35x), +at the cost of a somewhat higher p99 than full offloading (inline memory copies +and the remaining executor hops). This fake-only run cannot establish +end-to-end throughput or connection reuse at the network level. + +`/health` p99 during real S3 uploads was not measured. SDK-internal signing/checksum work was not profiled by this fake. + +### Browse query + +A one-off script loaded 100,000 rows (the committed test uses 1,000) for one live tenant into +SQLite and checks the first 25 IDs, before and after migration `f6c2d4e8b109`. +Ten executions averaged **38.455 → 0.076 ms**. Before the change SQLite used a +temporary B-tree for ordering; afterward the ordered tenant index served the +page without a sort. Both model metadata and the upgrade/downgrade are tested. + +The test uses SQLite as a local proxy. A PostgreSQL `EXPLAIN (ANALYZE, BUFFERS)` +on 100k+ representative tenant files remains necessary to establish that +platform's plan and timings. The index targets the default live-file order; +counts, alternate sort orders and arbitrary filters may still scan or sort. diff --git a/host/migrations/versions/f6c2d4e8b109_file_storage_browse_index.py b/host/migrations/versions/f6c2d4e8b109_file_storage_browse_index.py new file mode 100644 index 000000000..d67c83a91 --- /dev/null +++ b/host/migrations/versions/f6c2d4e8b109_file_storage_browse_index.py @@ -0,0 +1,24 @@ +"""Index tenant file browsing in page order. + +Revision ID: f6c2d4e8b109 +Revises: c5d8e2f1a903 (current file_storage@head on the shared trunk) +""" + +import sqlalchemy as sa +from alembic import op + +revision = "f6c2d4e8b109" +down_revision = "c5d8e2f1a903" +branch_labels = None +depends_on = None + +_TABLE = "file_storage_stored_file" +_INDEX = "ix_file_storage_stored_file_tenant_browse" + + +def upgrade() -> None: + op.create_index(_INDEX, _TABLE, ["tenant_id", "is_deleted", sa.text("created_at DESC"), "id"]) + + +def downgrade() -> None: + op.drop_index(_INDEX, table_name=_TABLE) diff --git a/host/tests/test_migration_graph.py b/host/tests/test_migration_graph.py index 46da61036..02178e747 100644 --- a/host/tests/test_migration_graph.py +++ b/host/tests/test_migration_graph.py @@ -46,6 +46,7 @@ "c5d8e2f1a903": "background_tasks", # newest-first listing index "e731a1b2c3d4": "users", # unique lower(email) (#731) "e805a1b2c3d4": "users", # browser session retirement + "f6c2d4e8b109": "file_storage", # tenant browse index "51b2058d8afd": "api_tokens", "168a2882f443": "keycloak", "f6a3b9d2e8c4": "keycloak", @@ -72,7 +73,7 @@ # head revision -> module owning the tables it creates or alters. _HEADS = { "b8d4f2a6c901": "audit_log", - "c5d8e2f1a903": "background_tasks", # also carries ``api_tokens``, ``file_storage``, ``users`` + "f6c2d4e8b109": "file_storage", # also carries ``api_tokens``, ``background_tasks``, ``users`` "f1a9c3d7b2e4": "tenants", "f6a3b9d2e8c4": "keycloak", "c9d2e4f6a810": "oidc", @@ -122,6 +123,7 @@ "e731a1b2c3d4": "api_tokens", "e805a1b2c3d4": "api_tokens", "c5d8e2f1a903": "api_tokens", + "f6c2d4e8b109": "api_tokens", "a626b1c8d023": "api_tokens", # also inherits ``background_tasks`` "51b2058d8afd": "api_tokens", # ancestors of a labelled revision inherit its label too } diff --git a/modules/file_storage/file_storage/backends/s3.py b/modules/file_storage/file_storage/backends/s3.py index 8a7f15af5..4f8c088f2 100644 --- a/modules/file_storage/file_storage/backends/s3.py +++ b/modules/file_storage/file_storage/backends/s3.py @@ -30,9 +30,9 @@ class S3Backend(S3MultipartMixin): """S3-compatible object storage. - Each operation opens a fresh ``aioboto3`` client context (cheap; the - underlying ``aiobotocore`` session is reused via ``self._session``). We - deliberately do not hold one long-lived client across requests because + Each operation opens an ``aioboto3`` client context; a server-side + multipart upload shares that client through completion or abort. We + do not hold one long-lived client across requests because aioboto3 client contexts are per-event-loop and tests run multiple loops. """ diff --git a/modules/file_storage/file_storage/backends/s3_multipart.py b/modules/file_storage/file_storage/backends/s3_multipart.py index 854104bfb..e0ed165e0 100644 --- a/modules/file_storage/file_storage/backends/s3_multipart.py +++ b/modules/file_storage/file_storage/backends/s3_multipart.py @@ -44,9 +44,17 @@ def _presign_client(self): async def create_multipart(self, key: str, content_type: str) -> str: try: async with self._client() as client: - resp = await client.create_multipart_upload( - Bucket=self.bucket, Key=key, ContentType=content_type - ) + return await self._create_multipart_on(client, key, content_type) + except StorageBackendError: + raise + except Exception as exc: + raise StorageBackendError(f"S3 create_multipart_upload failed: {exc}") from exc + + async def _create_multipart_on(self, client, key: str, content_type: str) -> str: + try: + resp = await client.create_multipart_upload( + Bucket=self.bucket, Key=key, ContentType=content_type + ) return resp["UploadId"] except Exception as exc: raise StorageBackendError(f"S3 create_multipart_upload failed: {exc}") from exc @@ -108,24 +116,42 @@ async def list_parts(self, key: str, upload_id: str) -> list[Part]: async def complete_multipart(self, key: str, upload_id: str, parts: list[Part]) -> None: try: async with self._client() as client: - await client.complete_multipart_upload( - Bucket=self.bucket, - Key=key, - UploadId=upload_id, - MultipartUpload={ - "Parts": [ - {"PartNumber": p.no, "ETag": p.etag} - for p in sorted(parts, key=lambda p: p.no) - ] - }, - ) + await self._complete_multipart_on(client, key, upload_id, parts) + except StorageBackendError: + raise + except Exception as exc: + raise StorageBackendError(f"S3 complete_multipart_upload failed: {exc}") from exc + + async def _complete_multipart_on( + self, client, key: str, upload_id: str, parts: list[Part] + ) -> None: + try: + await client.complete_multipart_upload( + Bucket=self.bucket, + Key=key, + UploadId=upload_id, + MultipartUpload={ + "Parts": [ + {"PartNumber": p.no, "ETag": p.etag} + for p in sorted(parts, key=lambda p: p.no) + ] + }, + ) except Exception as exc: raise StorageBackendError(f"S3 complete_multipart_upload failed: {exc}") from exc async def abort_multipart(self, key: str, upload_id: str) -> None: try: async with self._client() as client: - await client.abort_multipart_upload(Bucket=self.bucket, Key=key, UploadId=upload_id) + await self._abort_multipart_on(client, key, upload_id) + except StorageBackendError: + raise + except Exception as exc: + raise StorageBackendError(f"S3 abort_multipart_upload failed: {exc}") from exc + + async def _abort_multipart_on(self, client, key: str, upload_id: str) -> None: + try: + await client.abort_multipart_upload(Bucket=self.bucket, Key=key, UploadId=upload_id) except Exception as exc: if _is_no_such_upload(exc): return diff --git a/modules/file_storage/file_storage/backends/s3_put.py b/modules/file_storage/file_storage/backends/s3_put.py index 330f6c184..3c1b7ed4e 100644 --- a/modules/file_storage/file_storage/backends/s3_put.py +++ b/modules/file_storage/file_storage/backends/s3_put.py @@ -2,10 +2,11 @@ from __future__ import annotations +import asyncio import logging -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Callable from tempfile import SpooledTemporaryFile -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from file_storage import constants, upload_support from file_storage.contracts.service import Part, StorageBackendError @@ -17,6 +18,67 @@ logger = logging.getLogger(__name__) +# Replay blocks are read from disk once per executor hop; keep them well above the +# request-chunk size so a rolled-over spool costs few hops, yet bounded. +_REPLAY_BLOCK = 1024 * 1024 +# Larger in-memory writes grow the BytesIO enough to stall the loop; offload them. +_INLINE_WRITE_LIMIT = 2 * 1024 * 1024 + + +class _Buf: + """A SpooledTemporaryFile whose in-memory operations run inline. + + Executor hops cost more than a memcpy, so only operations that touch disk (the + spool already rolled over, or this write would roll it) are offloaded. + """ + + def __init__(self) -> None: + self.max_size = constants.SPOOL_MAX_SIZE_BYTES + self.file = SpooledTemporaryFile(max_size=self.max_size) # noqa: SIM115 — closed in aclose + self.pos = 0 + + async def __aenter__(self) -> _Buf: + return self + + async def __aexit__(self, *exc: object) -> None: + await self.aclose() + + def _in_memory(self) -> bool: + # `_rolled` has no public API; missing -> treat as on disk. + return not getattr(self.file, "_rolled", True) + + async def _call(self, fn: Callable[..., Any], *args: Any) -> Any: + if self._in_memory(): + return fn(*args) + return await _io(fn, *args) + + async def write(self, data: bytes | memoryview) -> None: + # Growing a large BytesIO reallocates and faults pages: keep that off the loop too. + if self._in_memory() and self.pos + len(data) <= min(self.max_size, _INLINE_WRITE_LIMIT): + self.file.write(data) + else: + await _io(self.file.write, data) + self.pos += len(data) + + async def read(self, size: int) -> bytes: + return await self._call(self.file.read, size) + + async def rewind(self) -> None: + self.pos = 0 + await self._call(self.file.seek, 0) + + async def reset(self) -> None: + self.pos = 0 + await self._call(self._truncate_from_start) + + def _truncate_from_start(self) -> None: + self.file.seek(0) + self.file.truncate() + + async def aclose(self) -> None: + await self._call(self.file.close) + + async def put_stream( backend: S3Backend, key: str, stream: AsyncIterator[bytes], content_type: str, size: int ) -> None: @@ -26,39 +88,38 @@ async def put_stream( threshold = min(backend.multipart_threshold_bytes, constants.MAX_S3_PART_SIZE_BYTES) iterator = aiter(stream) tail: memoryview | None = None - with SpooledTemporaryFile(max_size=constants.SPOOL_MAX_SIZE_BYTES) as prefix: + async with _Buf() as prefix: if size <= threshold: async for chunk in iterator: - take = min(len(chunk), threshold + 1 - prefix.tell()) - prefix.write(memoryview(chunk)[:take]) - if prefix.tell() > threshold: + take = min(len(chunk), threshold + 1 - prefix.pos) + await prefix.write(memoryview(chunk)[:take]) + if prefix.pos > threshold: tail = memoryview(chunk)[take:] break else: - length = prefix.tell() + length = prefix.pos if size and length != size: raise SizeMismatchError(size, length) - prefix.seek(0) + await prefix.rewind() try: async with backend._client() as client: await client.put_object( Bucket=backend.bucket, Key=key, - Body=prefix, + Body=prefix.file, ContentType=content_type, ContentLength=length, ) except Exception as exc: raise StorageBackendError(f"S3 put_object failed: {exc}") from exc return - prefix.seek(0) + await prefix.rewind() async def body() -> AsyncIterator[bytes]: - while chunk := prefix.read(constants.DEFAULT_CHUNK_SIZE): + while chunk := await prefix.read(_REPLAY_BLOCK): yield chunk - if tail is not None: - for offset in range(0, len(tail), constants.DEFAULT_CHUNK_SIZE): - yield bytes(tail[offset : offset + constants.DEFAULT_CHUNK_SIZE]) + if tail: + yield bytes(tail) async for chunk in iterator: yield chunk @@ -69,53 +130,79 @@ async def _multipart(backend: S3Backend, key, stream, content_type, size) -> Non part_size = upload_support.choose_part_size(size, backend.multipart_part_size_bytes) if part_size > constants.MAX_S3_PART_SIZE_BYTES: raise ValueError("S3 part size exceeds 5 GiB") - upload_id = await backend.create_multipart(key, content_type) + async with backend._client() as client: + await _upload(client, backend, key, stream, content_type, size, part_size) + + +async def _upload(client, backend, key, stream, content_type, size, part_size) -> None: + upload_id = await backend._create_multipart_on(client, key, content_type) parts: list[Part] = [] seen = 0 try: - with SpooledTemporaryFile(max_size=constants.SPOOL_MAX_SIZE_BYTES) as buf: + async with _Buf() as buf: async for chunk in stream: seen += len(chunk) if size and seen > size: raise SizeMismatchError(size, seen) remaining = memoryview(chunk) while remaining: - take = min(len(remaining), part_size - buf.tell()) - buf.write(remaining[:take]) + take = min(len(remaining), part_size - buf.pos) + await buf.write(remaining[:take]) remaining = remaining[take:] - if buf.tell() == part_size: - parts.append(await _send(backend, key, upload_id, parts, buf)) - buf.seek(0) - buf.truncate() + if buf.pos == part_size: + parts.append(await _send(client, backend, key, upload_id, parts, buf)) + await buf.reset() if size and seen != size: raise SizeMismatchError(size, seen) - if buf.tell() or not parts: - parts.append(await _send(backend, key, upload_id, parts, buf)) - await backend.complete_multipart(key, upload_id, parts) + if buf.pos or not parts: + parts.append(await _send(client, backend, key, upload_id, parts, buf)) + await backend._complete_multipart_on(client, key, upload_id, parts) except BaseException: try: - await backend.abort_multipart(key, upload_id) + await backend._abort_multipart_on(client, key, upload_id) except BaseException: logger.exception("file_storage.server_upload_abort_failed key=%s", key) raise -async def _send(backend: S3Backend, key, upload_id, parts, buf) -> Part: +async def _send(client, backend: S3Backend, key, upload_id, parts, buf: _Buf) -> Part: number = len(parts) + 1 if number > constants.MAX_UPLOAD_PARTS: raise ValueError("S3 upload exceeds maximum part count") - length = buf.tell() - buf.seek(0) + length = buf.pos + await buf.rewind() try: - async with backend._client() as client: - result = await client.upload_part( - Bucket=backend.bucket, - Key=key, - UploadId=upload_id, - PartNumber=number, - Body=buf, - ContentLength=length, - ) + result = await client.upload_part( + Bucket=backend.bucket, + Key=key, + UploadId=upload_id, + PartNumber=number, + Body=buf.file, + ContentLength=length, + ) except Exception as exc: raise StorageBackendError(f"S3 upload_part failed: {exc}") from exc return Part(no=number, etag=result["ETag"], size=length) + + +async def _io(fn: Callable[..., Any], *args: Any) -> Any: + # Finish in-flight I/O before cancellation can close or reuse its spool. + task = asyncio.create_task(asyncio.to_thread(fn, *args)) + cancelled = None + while True: + try: + result = await asyncio.shield(task) + break + except asyncio.CancelledError as exc: + if task.cancelled(): + # The I/O task itself was cancelled (loop shutdown cancels every task): + # nothing is in flight, and shielding it again would spin forever. + raise + cancelled = exc + except BaseException: + if cancelled is not None: + raise cancelled from None + raise + if cancelled is not None: + raise cancelled + return result diff --git a/modules/file_storage/file_storage/models.py b/modules/file_storage/file_storage/models.py index ef7531f58..23a7aaa0f 100644 --- a/modules/file_storage/file_storage/models.py +++ b/modules/file_storage/file_storage/models.py @@ -62,6 +62,13 @@ class StoredFile(Base, AuditMixin, SoftDeleteMixin, MultiTenantMixin, table=True Index(f"ix_{constants.TABLE_STORED_FILE}_tenant_key", "tenant_id", "key", unique=True), Index(f"ix_{constants.TABLE_STORED_FILE}_created_by", "created_by"), Index(f"ix_{constants.TABLE_STORED_FILE}_is_deleted", "is_deleted"), + Index( + f"ix_{constants.TABLE_STORED_FILE}_tenant_browse", + "tenant_id", + "is_deleted", + sa.text("created_at DESC"), + "id", + ), ) diff --git a/modules/file_storage/tests/conftest.py b/modules/file_storage/tests/conftest.py index 53839ecd2..fe39fd01a 100644 --- a/modules/file_storage/tests/conftest.py +++ b/modules/file_storage/tests/conftest.py @@ -11,9 +11,12 @@ import contextlib from collections.abc import Generator +from contextlib import asynccontextmanager import pytest +from file_storage.backends import build_backend from file_storage.backends.filesystem import FilesystemBackend +from file_storage.settings import FileStorageSettings from file_storage_uploads_helpers import PART from file_storage_worker_helpers import worker_db as worker_db from sqlalchemy import event @@ -114,3 +117,75 @@ def storage(app, tmp_path): services.backend = FilesystemBackend(root=tmp_path) services.settings.multipart_part_size_bytes = PART return services + + +class S3Wire: + def __init__(self): + self.parts = [] + self.single = None + self.completed = None + self.aborted = False + self.failure = None + self.abort_failure = False + self.opened = 0 + self.closed = 0 + + async def put_object(self, **kw): + self.single = kw["Body"].read() + assert kw["ContentLength"] == len(self.single) + + async def create_multipart_upload(self, **kw): + return {"UploadId": "upload"} + + async def upload_part(self, **kw): + assert kw["PartNumber"] == len(self.parts) + 1 + body = kw["Body"] + self.parts.append(body if isinstance(body, bytes) else body.read()) + return {"ETag": str(kw["PartNumber"])} + + async def complete_multipart_upload(self, **kw): + if self.failure: + raise self.failure + self.completed = kw["MultipartUpload"]["Parts"] + + async def abort_multipart_upload(self, **kw): + self.aborted = True + if self.abort_failure: + raise RuntimeError("abort failed") + + +@pytest.fixture +def wire_backend(monkeypatch): + wire = S3Wire() + cfg = FileStorageSettings( + backend="s3", + s3_bucket="bucket", + s3_region="us-east-1", + direct_upload_threshold_bytes=6, + ) + backend = build_backend(cfg) + + @asynccontextmanager + async def client(): + wire.opened += 1 + try: + yield wire + finally: + wire.closed += 1 + + monkeypatch.setattr(backend, "_client", client) + # Tiny, hand-verifiable parts after testing settings propagation separately. + backend.multipart_part_size_bytes = 5 + monkeypatch.setattr( + "file_storage.upload_support.choose_part_size", lambda size, preferred: preferred + ) + return wire, backend + + +@pytest.fixture +def chunks(): + async def _chunks(*values): + for value in values: + yield value + + return _chunks diff --git a/modules/file_storage/tests/test_file_storage_browse_index.py b/modules/file_storage/tests/test_file_storage_browse_index.py new file mode 100644 index 000000000..256842226 --- /dev/null +++ b/modules/file_storage/tests/test_file_storage_browse_index.py @@ -0,0 +1,64 @@ +"""The default tenant browse page should not sort the tenant's whole bucket.""" + +import importlib.util +from pathlib import Path + +import sqlalchemy as sa +from alembic.migration import MigrationContext +from alembic.operations import Operations +from file_storage.models import StoredFile +from file_storage.queries import order_clauses + +_INDEX = "ix_file_storage_stored_file_tenant_browse" + + +def test_default_browse_uses_ordered_index(): + engine = sa.create_engine("sqlite://") + table = StoredFile.__table__ + query = ( + sa.select(table.c.id) + .where(table.c.tenant_id == "tenant", table.c.is_deleted.is_(False)) + .order_by(*order_clauses("-created_at")) + .limit(25) + ) + sql = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + with engine.begin() as conn: + table.create(conn) + plan = conn.exec_driver_sql("EXPLAIN QUERY PLAN " + sql).all() + assert _INDEX in str(plan) + assert "TEMP B-TREE" not in str(plan) + + +def test_browse_index_migration_roundtrip(): + path = ( + Path(__file__).resolve().parents[3] + / "host/migrations/versions/f6c2d4e8b109_file_storage_browse_index.py" + ) + spec = importlib.util.spec_from_file_location("browse_migration", path) + migration = importlib.util.module_from_spec(spec) + spec.loader.exec_module(migration) + engine = sa.create_engine("sqlite://") + with engine.begin() as conn: + conn.exec_driver_sql( + "CREATE TABLE file_storage_stored_file " + "(id INTEGER PRIMARY KEY, tenant_id TEXT, is_deleted BOOLEAN, created_at TEXT)" + ) + conn.exec_driver_sql( + "WITH RECURSIVE n(x) AS (VALUES(1) UNION ALL SELECT x+1 FROM n WHERE x<1000) " + "INSERT INTO file_storage_stored_file SELECT x, 'tenant', 0, printf('%06d', x) FROM n" + ) + sql = ( + "SELECT id FROM file_storage_stored_file WHERE tenant_id='tenant' AND is_deleted=0 " + "ORDER BY created_at DESC, id ASC LIMIT 25" + ) + + with Operations.context(MigrationContext.configure(conn)): + migration.upgrade() + rows = conn.exec_driver_sql(sql).scalars().all() + assert rows == list(range(1000, 975, -1)) + plan = conn.exec_driver_sql("EXPLAIN QUERY PLAN " + sql).all() + assert _INDEX in str(plan) and "TEMP B-TREE" not in str(plan) + migration.downgrade() + assert _INDEX not in { + i["name"] for i in sa.inspect(conn).get_indexes("file_storage_stored_file") + } diff --git a/modules/file_storage/tests/test_file_storage_server_put.py b/modules/file_storage/tests/test_file_storage_server_put.py index 73d8a448a..71ab7559b 100644 --- a/modules/file_storage/tests/test_file_storage_server_put.py +++ b/modules/file_storage/tests/test_file_storage_server_put.py @@ -1,7 +1,6 @@ """Server-side multipart boundaries without a network service (#708).""" import asyncio -from contextlib import asynccontextmanager import pytest from file_storage import constants @@ -10,74 +9,12 @@ from file_storage.upload_support import SizeMismatchError -class S3Wire: - def __init__(self): - self.parts = [] - self.single = None - self.completed = None - self.aborted = False - self.failure = None - self.abort_failure = False - - async def put_object(self, **kw): - self.single = kw["Body"].read() - assert kw["ContentLength"] == len(self.single) - - async def create_multipart_upload(self, **kw): - return {"UploadId": "upload"} - - async def upload_part(self, **kw): - assert kw["PartNumber"] == len(self.parts) + 1 - body = kw["Body"] - self.parts.append(body if isinstance(body, bytes) else body.read()) - return {"ETag": str(kw["PartNumber"])} - - async def complete_multipart_upload(self, **kw): - if self.failure: - raise self.failure - self.completed = kw["MultipartUpload"]["Parts"] - - async def abort_multipart_upload(self, **kw): - self.aborted = True - if self.abort_failure: - raise RuntimeError("abort failed") - - -@pytest.fixture -def wire_backend(monkeypatch): - wire = S3Wire() - cfg = FileStorageSettings( - backend="s3", - s3_bucket="bucket", - s3_region="us-east-1", - direct_upload_threshold_bytes=6, - ) - backend = build_backend(cfg) - - @asynccontextmanager - async def client(): - yield wire - - monkeypatch.setattr(backend, "_client", client) - # Tiny, hand-verifiable parts after testing settings propagation separately. - backend.multipart_part_size_bytes = 5 - monkeypatch.setattr( - "file_storage.upload_support.choose_part_size", lambda size, preferred: preferred - ) - return wire, backend - - -async def chunks(*values): - for value in values: - yield value - - @pytest.mark.parametrize( "payload,multipart", [(b"", False), (b"123456", False), (b"1234567", True), (b"1234567890", True)], ) @pytest.mark.parametrize("known", [False, True]) -async def test_threshold_and_exact_final_part(wire_backend, payload, multipart, known): +async def test_threshold_and_exact_final_part(wire_backend, payload, multipart, known, chunks): wire, backend = wire_backend await backend.put( "key", @@ -111,7 +48,7 @@ async def broken(): assert wire.completed is None -async def test_completion_failure_aborts(wire_backend): +async def test_completion_failure_aborts(wire_backend, chunks): wire, backend = wire_backend wire.failure = RuntimeError("complete failed") wire.abort_failure = True @@ -121,7 +58,7 @@ async def test_completion_failure_aborts(wire_backend): @pytest.mark.parametrize("declared", [2, 8, 12]) -async def test_declared_size_mismatch_is_not_published(wire_backend, declared): +async def test_declared_size_mismatch_is_not_published(wire_backend, declared, chunks): wire, backend = wire_backend with pytest.raises(SizeMismatchError): await backend.put("key", chunks(b"1234567890"), content_type="text/plain", size=declared) @@ -129,7 +66,7 @@ async def test_declared_size_mismatch_is_not_published(wire_backend, declared): assert wire.aborted -async def test_part_count_bound_aborts(wire_backend, monkeypatch): +async def test_part_count_bound_aborts(wire_backend, monkeypatch, chunks): wire, backend = wire_backend monkeypatch.setattr(constants, "MAX_UPLOAD_PARTS", 1) with pytest.raises(Exception, match="part"): @@ -164,7 +101,7 @@ async def source(): assert b"".join(wire.parts) == b"1234567890" -async def test_large_source_chunk_keeps_spools_bounded(wire_backend, monkeypatch): +async def test_large_source_chunk_keeps_spools_bounded(wire_backend, monkeypatch, chunks): from tempfile import SpooledTemporaryFile from file_storage.backends import s3_put @@ -173,19 +110,13 @@ async def test_large_source_chunk_keeps_spools_bounded(wire_backend, monkeypatch class TrackedBuffer: def __init__(self, **kwargs): - self.file = SpooledTemporaryFile(**kwargs) # noqa: SIM115 — closed by wrapper context + self.file = SpooledTemporaryFile(**kwargs) # noqa: SIM115 self.index = len(maximums) maximums.append(0) def __getattr__(self, name): return getattr(self.file, name) - def __enter__(self): - return self - - def __exit__(self, *args): - self.file.close() - def write(self, data): self.file.write(data) maximums[self.index] = max(maximums[self.index], self.file.tell()) @@ -197,7 +128,7 @@ def write(self, data): assert maximums == [7, 5] -async def test_failed_part_upload_aborts(wire_backend, monkeypatch): +async def test_failed_part_upload_aborts(wire_backend, monkeypatch, chunks): wire, backend = wire_backend async def fail(**kwargs): diff --git a/modules/file_storage/tests/test_file_storage_upload_io.py b/modules/file_storage/tests/test_file_storage_upload_io.py new file mode 100644 index 000000000..867e624b6 --- /dev/null +++ b/modules/file_storage/tests/test_file_storage_upload_io.py @@ -0,0 +1,140 @@ +"""Upload I/O must leave the event loop free and close resources safely.""" + +import asyncio +import threading +from tempfile import SpooledTemporaryFile + +import pytest +from file_storage.backends import s3_put + + +@pytest.mark.parametrize("size", [0, 6, 10]) +async def test_spool_io_runs_off_loop(wire_backend, monkeypatch, size, chunks): + loop_thread = threading.get_ident() + calls = [] + + class CheckedSpool(SpooledTemporaryFile): + def _note(self, extra=0): + disk = self._rolled or self.tell() + extra > self._max_size + calls.append((threading.get_ident(), disk)) + + def write(self, s): + self._note(len(s)) + return super().write(s) + + def read(self, n=-1, /): + self._note() + return super().read(n) + + def seek(self, offset, whence=0): + self._note() + return super().seek(offset, whence) + + def truncate(self, size=None): + self._note() + return super().truncate(size) + + def close(self): + self._note() + return super().close() + + monkeypatch.setattr(s3_put, "SpooledTemporaryFile", CheckedSpool) + monkeypatch.setattr(s3_put.constants, "SPOOL_MAX_SIZE_BYTES", 2) + wire, backend = wire_backend + + async def put_object(**kw): + wire.single = await asyncio.to_thread(kw["Body"].read) + + async def upload_part(**kw): + wire.parts.append(await asyncio.to_thread(kw["Body"].read)) + return {"ETag": str(kw["PartNumber"])} + + monkeypatch.setattr(wire, "put_object", put_object) + monkeypatch.setattr(wire, "upload_part", upload_part) + payload = b"123456" if size == 6 else b"1234567890" + await backend.put("key", chunks(payload), content_type="text/plain", size=size) + assert (wire.single if size == 6 else b"".join(wire.parts)) == payload + # In-memory operations run inline by design; anything on disk must be off-loop. + assert any(thread != loop_thread for thread, _ in calls) + assert all(thread != loop_thread for thread, disk in calls if disk) + + +@pytest.mark.parametrize("fail_completion", [False, True]) +async def test_multipart_reuses_one_client(wire_backend, monkeypatch, fail_completion, chunks): + wire, backend = wire_backend + if fail_completion: + wire.failure = RuntimeError("complete failed") + with pytest.raises(Exception, match="complete failed"): + await backend.put("key", chunks(b"12345678901"), content_type="text/plain", size=0) + assert wire.aborted + else: + await backend.put("key", chunks(b"12345678901"), content_type="text/plain", size=0) + assert b"".join(wire.parts) == b"12345678901" + assert wire.opened == wire.closed == 1 + + +@pytest.mark.parametrize("cancel_count", [1, 2]) +@pytest.mark.parametrize("write_fails", [False, True]) +async def test_cancellation_waits_for_write_before_closing( + wire_backend, monkeypatch, cancel_count, write_fails, chunks +): + entered = threading.Event() + release = threading.Event() + files = [] + + class SlowSpool(SpooledTemporaryFile): + def __init__(self, **kw): + super().__init__(**kw) + files.append(self) + + def write(self, s): + entered.set() + assert release.wait(3) + assert not self.closed + if write_fails: + raise OSError("disk failed after cancellation") + return super().write(s) + + monkeypatch.setattr(s3_put, "SpooledTemporaryFile", SlowSpool) + # Writes that stay in memory run inline; force the disk path this test models. + monkeypatch.setattr(s3_put.constants, "SPOOL_MAX_SIZE_BYTES", 1) + wire, backend = wire_backend + backend.multipart_threshold_bytes = 1 + task = asyncio.create_task( + backend.put("key", chunks(b"1234567"), content_type="text/plain", size=7) + ) + try: + assert await asyncio.to_thread(entered.wait, 3) + for _ in range(cancel_count): + task.cancel() + await asyncio.sleep( + 0.01 + ) # intentional-sleep: expose premature close while I/O is blocked + assert not task.done() + finally: + release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert files and all(f.closed for f in files) + assert wire.aborted and wire.completed is None + + +async def test_io_task_cancelled_by_loop_shutdown_does_not_spin(): + """Cancelling the inner executor task (as loop shutdown does) must propagate.""" + gate = threading.Event() + waiter = asyncio.create_task(s3_put._io(gate.wait, 3)) + + me = asyncio.current_task() + inner = None + for _ in range(1000): # yield until _io has scheduled its executor task + await asyncio.sleep(0) + inner = next((t for t in asyncio.all_tasks() if t not in (waiter, me)), None) + if inner is not None: + break + assert inner is not None + inner.cancel() + try: + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(waiter, 2) + finally: + gate.set()