From 9dbc5ed9983557a5aeb8c43eb772b37617ba9f55 Mon Sep 17 00:00:00 2001 From: Lukas Remis Date: Sun, 19 Jul 2026 07:13:03 +0200 Subject: [PATCH] =?UTF-8?q?sync-v2=20step=206:=20memory=20+=20connection?= =?UTF-8?q?=20bounds=20(=C2=A77,=20=C2=A74)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - concurrency.loop_semaphore: per-event-loop semaphores (safe across pytest loops and process restarts via a WeakKeyDictionary keyed by the running loop). - provider: holds a per-account (≤4) and per-host (PER_HOST_CONN_CAP) slot for the whole lifetime of an IMAP socket, released on close; SMTP send takes a per-host slot for its duration. A worker can no longer open hundreds of sockets to one host/account. - indexer: parsing goes through a PARSE_CONCURRENCY semaphore, and a message over STREAM_PARSE_MB is parsed from its blob file (parse_message_file) instead of an in-memory tree, so peak RAM ~= PARSE_CONCURRENCY x STREAM_PARSE_MB regardless of mailbox sizes. Backfill parallelism stays bounded by the bulk queue only. - config: PER_HOST_CONN_CAP, PARSE_CONCURRENCY, STREAM_PARSE_MB. - tests: per-account cap blocks the 5th connection / distinct accounts independent; file-vs-bytes parse equivalence + large-message streaming routing. Co-Authored-By: Claude Opus 4.8 --- backend/app/concurrency.py | 32 ++++++ backend/app/config.py | 4 + backend/app/indexer/indexer.py | 21 +++- backend/app/indexer/mime.py | 14 ++- backend/app/providers/imap_smtp/provider.py | 99 +++++++++++++------ backend/tests/indexer/test_streaming_parse.py | 60 +++++++++++ backend/tests/providers/test_conn_limits.py | 58 +++++++++++ 7 files changed, 254 insertions(+), 34 deletions(-) create mode 100644 backend/app/concurrency.py create mode 100644 backend/tests/indexer/test_streaming_parse.py create mode 100644 backend/tests/providers/test_conn_limits.py diff --git a/backend/app/concurrency.py b/backend/app/concurrency.py new file mode 100644 index 0000000..5296595 --- /dev/null +++ b/backend/app/concurrency.py @@ -0,0 +1,32 @@ +"""Per-event-loop semaphores for bounded concurrency (refactor §4, §7). + +Connection and parse caps are process-local but must be bound to the running +event loop: an :class:`asyncio.Semaphore` created on one loop cannot be awaited +on another (a real hazard under pytest, which uses a fresh loop per test). These +helpers hand back a semaphore keyed by ``(group, key)`` for the current loop, +creating it on first use and dropping the whole set when the loop is collected. +""" + +from __future__ import annotations + +import asyncio +import weakref + +# group -> {loop: {key: Semaphore}}. WeakKeyDictionary drops a loop's semaphores +# once the loop is gone, so no state leaks across tests or restarts. +_registry: dict[ + str, weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[str, asyncio.Semaphore]] +] = {} + + +def loop_semaphore(group: str, key: str, cap: int) -> asyncio.Semaphore: + """Return the current loop's semaphore for ``(group, key)`` (cap on first use).""" + loop = asyncio.get_running_loop() + per_group = _registry.setdefault(group, weakref.WeakKeyDictionary()) + per_loop = per_group.get(loop) + if per_loop is None: + per_loop = per_group[loop] = {} + sem = per_loop.get(key) + if sem is None: + sem = per_loop[key] = asyncio.Semaphore(cap) + return sem diff --git a/backend/app/config.py b/backend/app/config.py index 5e2075c..6617060 100644 --- a/backend/app/config.py +++ b/backend/app/config.py @@ -50,6 +50,10 @@ class Settings(BaseSettings): tier_session_grace_min: int = Field(default=30, alias="TIER_SESSION_GRACE_MIN") tier_cold_after_d: int = Field(default=3, alias="TIER_COLD_AFTER_D") hot_budget_per_worker: int = Field(default=250, alias="HOT_BUDGET_PER_WORKER") + # Memory/connection bounds (§7, §4): peak RAM ~= PARSE_CONCURRENCY x STREAM_PARSE_MB. + per_host_conn_cap: int = Field(default=50, alias="PER_HOST_CONN_CAP") + parse_concurrency: int = Field(default=4, alias="PARSE_CONCURRENCY") + stream_parse_mb: int = Field(default=8, alias="STREAM_PARSE_MB") oauth_google_client_id: str = Field(default="", alias="OAUTH_GOOGLE_CLIENT_ID") oauth_google_client_secret: str = Field(default="", alias="OAUTH_GOOGLE_CLIENT_SECRET") diff --git a/backend/app/indexer/indexer.py b/backend/app/indexer/indexer.py index 9698588..6bc1e3f 100644 --- a/backend/app/indexer/indexer.py +++ b/backend/app/indexer/indexer.py @@ -17,8 +17,10 @@ from sqlalchemy import select from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.ext.asyncio import AsyncSession +from app.concurrency import loop_semaphore +from app.config import get_settings from app.indexer import blobstore -from app.indexer.mime import ParsedEmail, parse_message +from app.indexer.mime import ParsedEmail, parse_message, parse_message_file from app.indexer.sanitize import sanitize_html from app.models import Attachment, BodyPart, KnownRecipient, Message from app.providers.base import RemoteMessage @@ -115,6 +117,21 @@ async def _resolve_message_id( return parsed_id +async def _parse_bounded(user_id: uuid.UUID, sha: str, raw: bytes) -> ParsedEmail: + """Parse under the PARSE_CONCURRENCY cap; stream large messages from disk (§7). + + Peak RAM stays ~PARSE_CONCURRENCY x STREAM_PARSE_MB: only a few parses run at + once, and a message above the threshold is parsed from its (already written) + blob file instead of being held as an in-memory tree beside its raw bytes. + """ + settings = get_settings() + threshold = settings.stream_parse_mb * 1024 * 1024 + async with loop_semaphore("mime_parse", "parse", settings.parse_concurrency): + if len(raw) > threshold: + return parse_message_file(blobstore.blob_path(user_id, sha)) + return parse_message(raw) + + async def persist_message( session: AsyncSession, *, @@ -132,7 +149,7 @@ async def persist_message( return None sha = blobstore.store(user_id, remote.raw) - parsed = parse_message(remote.raw) + parsed = await _parse_bounded(user_id, sha, remote.raw) message_id_hdr = await _resolve_message_id( session, user_id=user_id, account_id=account_id, parsed_id=parsed.message_id, raw=remote.raw diff --git a/backend/app/indexer/mime.py b/backend/app/indexer/mime.py index 63e0178..3451ddc 100644 --- a/backend/app/indexer/mime.py +++ b/backend/app/indexer/mime.py @@ -7,6 +7,7 @@ Header RFC2047 decoding is automatic under ``policy.default``. from __future__ import annotations import email +import os from dataclasses import dataclass, field from datetime import datetime from email.headerregistry import AddressHeader @@ -104,10 +105,21 @@ def _extract_attachments(msg: EmailMessage) -> list[ParsedAttachment]: def parse_message(raw: bytes) -> ParsedEmail: - """Parse a raw RFC822 message into a :class:`ParsedEmail`.""" + """Parse a raw RFC822 message (in memory) into a :class:`ParsedEmail`.""" msg = email.message_from_bytes(raw, policy=default_policy) assert isinstance(msg, EmailMessage) + return _parse_email(msg) + +def parse_message_file(path: str | os.PathLike[str]) -> ParsedEmail: + """Parse a message straight from a blob file handle (large-message path, §7).""" + with open(path, "rb") as handle: + msg = email.message_from_binary_file(handle, policy=default_policy) + assert isinstance(msg, EmailMessage) + return _parse_email(msg) + + +def _parse_email(msg: EmailMessage) -> ParsedEmail: message_id = normalize_message_id(msg.get("Message-ID")) in_reply_to_ids = normalize_reference_list(msg.get("In-Reply-To")) in_reply_to = in_reply_to_ids[0] if in_reply_to_ids else None diff --git a/backend/app/providers/imap_smtp/provider.py b/backend/app/providers/imap_smtp/provider.py index 856295f..0c937b3 100644 --- a/backend/app/providers/imap_smtp/provider.py +++ b/backend/app/providers/imap_smtp/provider.py @@ -17,6 +17,8 @@ from dataclasses import dataclass import aiosmtplib from aioimaplib import aioimaplib +from app.concurrency import loop_semaphore +from app.config import get_settings from app.providers.base import ( FlagChange, MailboxDelta, @@ -37,6 +39,7 @@ from app.providers.imap_smtp.parsing import ( _FETCH_ITEMS = "(UID FLAGS INTERNALDATE RFC822.SIZE BODY.PEEK[])" _IDLE_TIMEOUT = 25 * 60 # SPEC §5.1: re-issue IDLE at least every 25 min _CONNECT_TIMEOUT = 20 +_PER_ACCOUNT_CONN_CAP = 4 # SPEC §6.3 / §4: IDLE + transient + short-lived jobs _APPENDUID_RE = re.compile(rb"APPENDUID\s+\d+\s+(\d+)") _COPYUID_RE = re.compile(rb"COPYUID\s+\d+\s+\S+\s+(\d+)") @@ -59,34 +62,65 @@ class ImapSmtpProvider(MailProvider): self._imap: aioimaplib.IMAP4 | None = None self._selected: str | None = None self._uidvalidity_seen: dict[str, int] = {} + self._conn_slots: list[asyncio.Semaphore] = [] # --- connection lifecycle ------------------------------------------------ + async def _acquire_conn_slots(self) -> None: + """Bound simultaneous open sockets: ≤4 per account and ≤cap per host (§4). + + Held for the connection's whole lifetime and released in :meth:`close`. + """ + cfg = self._cfg + account = loop_semaphore( + "imap_account", f"{cfg.username}@{cfg.imap_host}", _PER_ACCOUNT_CONN_CAP + ) + host = loop_semaphore("imap_host", cfg.imap_host, get_settings().per_host_conn_cap) + await account.acquire() # consistent order (account, then host) avoids deadlock + try: + await host.acquire() + except BaseException: + account.release() + raise + self._conn_slots = [host, account] + + def _release_conn_slots(self) -> None: + for sem in self._conn_slots: + sem.release() + self._conn_slots = [] + async def _connect(self) -> aioimaplib.IMAP4: if self._imap is not None: return self._imap cfg = self._cfg + await self._acquire_conn_slots() try: - if cfg.imap_tls == "ssl": - imap = aioimaplib.IMAP4_SSL( - host=cfg.imap_host, port=cfg.imap_port, timeout=_CONNECT_TIMEOUT - ) - else: - imap = aioimaplib.IMAP4( - host=cfg.imap_host, port=cfg.imap_port, timeout=_CONNECT_TIMEOUT - ) - await imap.wait_hello_from_server() - if cfg.imap_tls == "starttls": - await imap.starttls() - except (OSError, TimeoutError) as exc: - raise ProviderConnectError(f"cannot connect to {cfg.imap_host}: {exc}") from exc + try: + if cfg.imap_tls == "ssl": + imap = aioimaplib.IMAP4_SSL( + host=cfg.imap_host, port=cfg.imap_port, timeout=_CONNECT_TIMEOUT + ) + else: + imap = aioimaplib.IMAP4( + host=cfg.imap_host, port=cfg.imap_port, timeout=_CONNECT_TIMEOUT + ) + await imap.wait_hello_from_server() + if cfg.imap_tls == "starttls": + await imap.starttls() + except (OSError, TimeoutError) as exc: + raise ProviderConnectError(f"cannot connect to {cfg.imap_host}: {exc}") from exc - resp = await imap.login(cfg.username, cfg.password) - if resp.result != "OK": - detail = ( - b" ".join(resp.lines).decode("utf-8", "ignore") if resp.lines else "login failed" - ) - raise ProviderAuthError(detail) + resp = await imap.login(cfg.username, cfg.password) + if resp.result != "OK": + detail = ( + b" ".join(resp.lines).decode("utf-8", "ignore") + if resp.lines + else "login failed" + ) + raise ProviderAuthError(detail) + except BaseException: + self._release_conn_slots() # never hold a slot without an open socket + raise self._imap = imap return imap @@ -104,6 +138,7 @@ class ImapSmtpProvider(MailProvider): await self._imap.logout() self._imap = None self._selected = None + self._release_conn_slots() # free the per-account / per-host slots # --- MailProvider protocol ---------------------------------------------- @@ -293,19 +328,21 @@ class ImapSmtpProvider(MailProvider): cfg = self._cfg use_tls = cfg.smtp_tls == "ssl" start_tls = cfg.smtp_tls == "starttls" + host_slot = loop_semaphore("smtp_host", cfg.smtp_host, get_settings().per_host_conn_cap) try: - await aiosmtplib.send( - raw, - sender=envelope_from, - recipients=envelope_to, - hostname=cfg.smtp_host, - port=cfg.smtp_port, - username=cfg.username, - password=cfg.password, - use_tls=use_tls, - start_tls=start_tls, - timeout=_CONNECT_TIMEOUT, - ) + async with host_slot: # bound simultaneous sockets to the SMTP host (§4) + await aiosmtplib.send( + raw, + sender=envelope_from, + recipients=envelope_to, + hostname=cfg.smtp_host, + port=cfg.smtp_port, + username=cfg.username, + password=cfg.password, + use_tls=use_tls, + start_tls=start_tls, + timeout=_CONNECT_TIMEOUT, + ) except aiosmtplib.SMTPAuthenticationError as exc: raise ProviderAuthError(f"SMTP auth failed: {exc}") from exc except aiosmtplib.SMTPException as exc: diff --git a/backend/tests/indexer/test_streaming_parse.py b/backend/tests/indexer/test_streaming_parse.py new file mode 100644 index 0000000..5fc2dc7 --- /dev/null +++ b/backend/tests/indexer/test_streaming_parse.py @@ -0,0 +1,60 @@ +"""Bounded / streaming MIME parse (refactor §7). + +Large messages are parsed from their blob file rather than an in-memory tree, and +all parsing passes through the PARSE_CONCURRENCY semaphore. +""" + +from __future__ import annotations + +import uuid +from email.message import EmailMessage + +from app.indexer import indexer +from app.indexer.mime import parse_message, parse_message_file + + +def _raw(body: str = "hello") -> bytes: + msg = EmailMessage() + msg["Subject"] = "Streamed" + msg["From"] = "a@example.com" + msg["To"] = "b@example.com" + msg["Message-ID"] = "" + msg.set_content(body) + msg.add_attachment( + b"x" * 1024, maintype="application", subtype="octet-stream", filename="a.bin" + ) + return msg.as_bytes() + + +def test_parse_file_matches_bytes(tmp_path) -> None: # type: ignore[no-untyped-def] + raw = _raw() + path = tmp_path / "msg.eml" + path.write_bytes(raw) + # Parsing from the file handle yields exactly the same ParsedEmail as from bytes. + assert parse_message_file(path) == parse_message(raw) + + +async def test_parse_bounded_streams_large_messages(monkeypatch) -> None: # type: ignore[no-untyped-def] + from app.config import get_settings + + threshold = get_settings().stream_parse_mb * 1024 * 1024 + calls: list[str] = [] + + def fake_file(path: object) -> str: + calls.append("file") + return "file" + + def fake_bytes(raw: bytes) -> str: + calls.append("bytes") + return "bytes" + + monkeypatch.setattr(indexer, "parse_message_file", fake_file) + monkeypatch.setattr(indexer, "parse_message", fake_bytes) + + user_id = uuid.uuid4() + small = await indexer._parse_bounded(user_id, "sha", b"x" * 16) + big = await indexer._parse_bounded(user_id, "sha", b"x" * (threshold + 1)) + + assert small == "bytes" # small stays in memory + assert big == "file" # over threshold => parsed from the blob file + assert calls == ["bytes", "file"] diff --git a/backend/tests/providers/test_conn_limits.py b/backend/tests/providers/test_conn_limits.py new file mode 100644 index 0000000..2e85c14 --- /dev/null +++ b/backend/tests/providers/test_conn_limits.py @@ -0,0 +1,58 @@ +"""Per-account / per-host connection caps (refactor §4). + +The provider holds a per-account (≤4) and per-host semaphore slot for the whole +lifetime of an open socket, so a worker can never open hundreds of simultaneous +connections to one account or host. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from app.providers.imap_smtp.provider import _PER_ACCOUNT_CONN_CAP, ImapSmtpProvider, ProviderConfig + + +def _cfg(username: str, host: str) -> ProviderConfig: + return ProviderConfig( + imap_host=host, + imap_port=993, + imap_tls="ssl", + smtp_host=host, + smtp_port=465, + smtp_tls="ssl", + username=username, + password="pw", + ) + + +async def test_per_account_connection_cap() -> None: + cfg = _cfg("user@example.com", "imap.example.com") + providers = [ImapSmtpProvider(cfg) for _ in range(_PER_ACCOUNT_CONN_CAP + 1)] + + # Fill the account's slots (host cap defaults to 50, so account is the limiter). + for provider in providers[:_PER_ACCOUNT_CONN_CAP]: + await provider._acquire_conn_slots() + + # The next acquire blocks until a slot frees. + with pytest.raises(TimeoutError): + await asyncio.wait_for(providers[-1]._acquire_conn_slots(), timeout=0.2) + + providers[0]._release_conn_slots() + await asyncio.wait_for(providers[-1]._acquire_conn_slots(), timeout=0.5) # now succeeds + + for provider in providers[1:]: + provider._release_conn_slots() + + +async def test_distinct_accounts_do_not_share_slots() -> None: + # Saturating account A must not block account B (independent semaphores). + a = [ImapSmtpProvider(_cfg("a@example.com", "imap.example.com")) for _ in range(4)] + for provider in a: + await provider._acquire_conn_slots() + b = ImapSmtpProvider(_cfg("b@example.com", "imap.example.com")) + await asyncio.wait_for(b._acquire_conn_slots(), timeout=0.5) + b._release_conn_slots() + for provider in a: + provider._release_conn_slots()