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:
@@ -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
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user