sync-v2 step 6: memory + connection bounds (§7, §4)

- 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 <noreply@anthropic.com>
This commit is contained in:
2026-07-19 07:13:03 +02:00
co-authored by Claude Opus 4.8
parent 19335116da
commit 9dbc5ed998
7 changed files with 254 additions and 34 deletions
+32
View File
@@ -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
+4
View File
@@ -50,6 +50,10 @@ class Settings(BaseSettings):
tier_session_grace_min: int = Field(default=30, alias="TIER_SESSION_GRACE_MIN") 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") 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") 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_id: str = Field(default="", alias="OAUTH_GOOGLE_CLIENT_ID")
oauth_google_client_secret: str = Field(default="", alias="OAUTH_GOOGLE_CLIENT_SECRET") oauth_google_client_secret: str = Field(default="", alias="OAUTH_GOOGLE_CLIENT_SECRET")
+19 -2
View File
@@ -17,8 +17,10 @@ from sqlalchemy import select
from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession 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 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.indexer.sanitize import sanitize_html
from app.models import Attachment, BodyPart, KnownRecipient, Message from app.models import Attachment, BodyPart, KnownRecipient, Message
from app.providers.base import RemoteMessage from app.providers.base import RemoteMessage
@@ -115,6 +117,21 @@ async def _resolve_message_id(
return parsed_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( async def persist_message(
session: AsyncSession, session: AsyncSession,
*, *,
@@ -132,7 +149,7 @@ async def persist_message(
return None return None
sha = blobstore.store(user_id, remote.raw) 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( message_id_hdr = await _resolve_message_id(
session, user_id=user_id, account_id=account_id, parsed_id=parsed.message_id, raw=remote.raw session, user_id=user_id, account_id=account_id, parsed_id=parsed.message_id, raw=remote.raw
+13 -1
View File
@@ -7,6 +7,7 @@ Header RFC2047 decoding is automatic under ``policy.default``.
from __future__ import annotations from __future__ import annotations
import email import email
import os
from dataclasses import dataclass, field from dataclasses import dataclass, field
from datetime import datetime from datetime import datetime
from email.headerregistry import AddressHeader from email.headerregistry import AddressHeader
@@ -104,10 +105,21 @@ def _extract_attachments(msg: EmailMessage) -> list[ParsedAttachment]:
def parse_message(raw: bytes) -> ParsedEmail: 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) msg = email.message_from_bytes(raw, policy=default_policy)
assert isinstance(msg, EmailMessage) 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")) message_id = normalize_message_id(msg.get("Message-ID"))
in_reply_to_ids = normalize_reference_list(msg.get("In-Reply-To")) 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 in_reply_to = in_reply_to_ids[0] if in_reply_to_ids else None
+68 -31
View File
@@ -17,6 +17,8 @@ from dataclasses import dataclass
import aiosmtplib import aiosmtplib
from aioimaplib import aioimaplib from aioimaplib import aioimaplib
from app.concurrency import loop_semaphore
from app.config import get_settings
from app.providers.base import ( from app.providers.base import (
FlagChange, FlagChange,
MailboxDelta, MailboxDelta,
@@ -37,6 +39,7 @@ from app.providers.imap_smtp.parsing import (
_FETCH_ITEMS = "(UID FLAGS INTERNALDATE RFC822.SIZE BODY.PEEK[])" _FETCH_ITEMS = "(UID FLAGS INTERNALDATE RFC822.SIZE BODY.PEEK[])"
_IDLE_TIMEOUT = 25 * 60 # SPEC §5.1: re-issue IDLE at least every 25 min _IDLE_TIMEOUT = 25 * 60 # SPEC §5.1: re-issue IDLE at least every 25 min
_CONNECT_TIMEOUT = 20 _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+)") _APPENDUID_RE = re.compile(rb"APPENDUID\s+\d+\s+(\d+)")
_COPYUID_RE = re.compile(rb"COPYUID\s+\d+\s+\S+\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._imap: aioimaplib.IMAP4 | None = None
self._selected: str | None = None self._selected: str | None = None
self._uidvalidity_seen: dict[str, int] = {} self._uidvalidity_seen: dict[str, int] = {}
self._conn_slots: list[asyncio.Semaphore] = []
# --- connection lifecycle ------------------------------------------------ # --- 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: async def _connect(self) -> aioimaplib.IMAP4:
if self._imap is not None: if self._imap is not None:
return self._imap return self._imap
cfg = self._cfg cfg = self._cfg
await self._acquire_conn_slots()
try: try:
if cfg.imap_tls == "ssl": try:
imap = aioimaplib.IMAP4_SSL( if cfg.imap_tls == "ssl":
host=cfg.imap_host, port=cfg.imap_port, timeout=_CONNECT_TIMEOUT imap = aioimaplib.IMAP4_SSL(
) host=cfg.imap_host, port=cfg.imap_port, timeout=_CONNECT_TIMEOUT
else: )
imap = aioimaplib.IMAP4( else:
host=cfg.imap_host, port=cfg.imap_port, timeout=_CONNECT_TIMEOUT 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.wait_hello_from_server()
await imap.starttls() if cfg.imap_tls == "starttls":
except (OSError, TimeoutError) as exc: await imap.starttls()
raise ProviderConnectError(f"cannot connect to {cfg.imap_host}: {exc}") from exc 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) resp = await imap.login(cfg.username, cfg.password)
if resp.result != "OK": if resp.result != "OK":
detail = ( detail = (
b" ".join(resp.lines).decode("utf-8", "ignore") if resp.lines else "login failed" b" ".join(resp.lines).decode("utf-8", "ignore")
) if resp.lines
raise ProviderAuthError(detail) else "login failed"
)
raise ProviderAuthError(detail)
except BaseException:
self._release_conn_slots() # never hold a slot without an open socket
raise
self._imap = imap self._imap = imap
return imap return imap
@@ -104,6 +138,7 @@ class ImapSmtpProvider(MailProvider):
await self._imap.logout() await self._imap.logout()
self._imap = None self._imap = None
self._selected = None self._selected = None
self._release_conn_slots() # free the per-account / per-host slots
# --- MailProvider protocol ---------------------------------------------- # --- MailProvider protocol ----------------------------------------------
@@ -293,19 +328,21 @@ class ImapSmtpProvider(MailProvider):
cfg = self._cfg cfg = self._cfg
use_tls = cfg.smtp_tls == "ssl" use_tls = cfg.smtp_tls == "ssl"
start_tls = cfg.smtp_tls == "starttls" start_tls = cfg.smtp_tls == "starttls"
host_slot = loop_semaphore("smtp_host", cfg.smtp_host, get_settings().per_host_conn_cap)
try: try:
await aiosmtplib.send( async with host_slot: # bound simultaneous sockets to the SMTP host (§4)
raw, await aiosmtplib.send(
sender=envelope_from, raw,
recipients=envelope_to, sender=envelope_from,
hostname=cfg.smtp_host, recipients=envelope_to,
port=cfg.smtp_port, hostname=cfg.smtp_host,
username=cfg.username, port=cfg.smtp_port,
password=cfg.password, username=cfg.username,
use_tls=use_tls, password=cfg.password,
start_tls=start_tls, use_tls=use_tls,
timeout=_CONNECT_TIMEOUT, start_tls=start_tls,
) timeout=_CONNECT_TIMEOUT,
)
except aiosmtplib.SMTPAuthenticationError as exc: except aiosmtplib.SMTPAuthenticationError as exc:
raise ProviderAuthError(f"SMTP auth failed: {exc}") from exc raise ProviderAuthError(f"SMTP auth failed: {exc}") from exc
except aiosmtplib.SMTPException as exc: except aiosmtplib.SMTPException as exc:
@@ -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"] = "<s1@example.com>"
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"]
@@ -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()