diff --git a/CHANGELOG.md b/CHANGELOG.md index 6fde2c29..dcc0e8b1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,16 @@ # Changelog +## Unreleased + +### New features + +- **`n_procs` parallel parsing now covers messages from mbox files and mailbox connections** ([#147](https://github.com/domainaware/parsedmarc/issues/147)), not just report files passed directly as CLI arguments: IMAP, Microsoft Graph, Gmail API, and Maildir connections, including watch mode. `get_dmarc_reports_from_mbox()`, `get_dmarc_reports_from_mailbox()`, and `watch_inbox()` all gained an `n_procs` keyword argument. Parallel workers are now a reused process pool for the whole run, with a bounded submission window that keeps at most roughly `2 * n_procs` messages in flight at a time so memory stays bounded even for huge mboxes; the main process sends periodic IMAP keepalives while workers parse. + +### Bug fixes + +- **A report file whose parsing raised an unexpected non-parser exception no longer hangs the CLI forever.** The old direct-file parallel implementation spawned a fresh child process per file and blocked the parent on a pipe read; a child that crashed with anything other than a `ParserError` never wrote to that pipe, so the parent waited indefinitely. The new implementation returns the error as a value and logs `Failed to parse ` instead. +- **Direct-file parallel parsing no longer stalls a whole batch on one slow file, and no longer spawns a fresh interpreter per file.** The old implementation processed files in hard batches of `n_procs`, so a single slow file delayed every other file in its batch; the new pooled-worker implementation streams files through a reused pool instead. + ## 10.3.0 ### New features diff --git a/docs/source/usage.md b/docs/source/usage.md index c6a93ea2..41a1d7d0 100644 --- a/docs/source/usage.md +++ b/docs/source/usage.md @@ -170,19 +170,22 @@ The full set of configuration options are: reports (Default: `False`) - `log_file` - str: Write log messages to a file at this path - `n_procs` - int: Number of processes to run in parallel when - parsing report files passed directly as CLI arguments + parsing report files passed directly as CLI arguments, messages + in mbox files, and messages from mailbox connections (IMAP, + Microsoft Graph, Gmail API, Maildir), including watch mode (Default: `1`) :::{note} Setting this to a number larger than one can improve - performance when processing thousands of files + performance when processing thousands of files or messages ::: :::{note} - `n_procs` only applies to report files passed directly on the - command line. Messages from mbox files and from mailbox - connections (IMAP, Microsoft Graph, Gmail API, Maildir) are - always processed sequentially. + Only parsing is parallelized across worker processes. Fetching + messages, deduplicating reports, archiving/deleting mailbox + messages, and saving/publishing to outputs all stay sequential + in the main process. Each worker process keeps its own DNS/GeoIP + cache. ::: - `mailbox` @@ -951,8 +954,12 @@ imports more predictable: - Reduce `mailbox.batch_size` to smaller values such as `100-500` instead of processing a very large message set at once. Smaller batches trade throughput for lower peak memory use and less sink pressure. -- Keep `n_procs` low for mailbox-heavy runs. In practice, `1-2` workers is often - a safer starting point for large backfills than aggressive parallelism. +- `n_procs` now parallelizes parsing for mailbox and mbox runs too, not just + report files passed directly as CLI arguments. It pays off most when a run + is bound by DNS/GeoIP enrichment rather than fetching or output. The + trade-off is memory and DNS load: at most roughly `2 * n_procs` messages + are held in flight at once, each worker process keeps its own DNS/GeoIP + cache, and DNS query volume can multiply by up to `n_procs`. - Use `mailbox.since` to process reports in smaller time windows such as `1d`, `7d`, or another interval that fits the backlog. This makes it easier to catch up incrementally instead of loading an entire mailbox history in one run. diff --git a/parsedmarc/__init__.py b/parsedmarc/__init__.py index 18c25c71..bc53f4d1 100644 --- a/parsedmarc/__init__.py +++ b/parsedmarc/__init__.py @@ -7,6 +7,7 @@ from __future__ import annotations import binascii import email import email.utils +import functools import json import logging import mailbox @@ -22,6 +23,7 @@ from base64 import b64decode from csv import DictWriter from datetime import date, datetime, timedelta, timezone, tzinfo from io import BytesIO, StringIO +from collections import deque from collections.abc import Callable, Sequence from typing import ( Any, @@ -55,6 +57,7 @@ from parsedmarc.types import ( ForensicReport as ForensicReport, ParsedReport, ParsingResults, + ReportType, SMTPTLSReport, ) from parsedmarc.utils import ( @@ -2068,6 +2071,100 @@ def parse_report_file( return results +def _classify_parsed_email( + parsed_email: ParsedReport, + aggregate_reports: list[AggregateReport], + failure_reports: list[FailureReport], + smtp_tls_reports: list[SMTPTLSReport], +) -> ReportType: + """Classify a parsed report email, appending it to the matching list. + + Owns the ``SEEN_AGGREGATE_REPORT_IDS`` dedup check: an aggregate report + already seen (keyed on ``{org_name}_{report_id}``) is logged and + dropped instead of appended. Shared, unmodified, by the sequential and + parallel branches of ``get_dmarc_reports_from_mbox`` and + ``get_dmarc_reports_from_mailbox`` so both dedup identically. + + Returns the report type so mailbox callers know which UID list to + append the source message's UID to. + """ + report_type = parsed_email["report_type"] + # Compare against parsed_email["report_type"] directly (not the + # report_type local above) in each branch so pyright's TypedDict + # discriminated-union narrowing applies to parsed_email["report"]. + if parsed_email["report_type"] == "aggregate": + report_org = parsed_email["report"]["report_metadata"]["org_name"] + report_id = parsed_email["report"]["report_metadata"]["report_id"] + report_key = f"{report_org}_{report_id}" + if report_key not in SEEN_AGGREGATE_REPORT_IDS: + SEEN_AGGREGATE_REPORT_IDS[report_key] = True + aggregate_reports.append(parsed_email["report"]) + else: + logger.debug( + f"Skipping duplicate aggregate report from {report_org} " + f"with ID: {report_id}" + ) + elif parsed_email["report_type"] == "failure": + failure_reports.append(parsed_email["report"]) + elif parsed_email["report_type"] == "smtp_tls": + smtp_tls_reports.append(parsed_email["report"]) + return report_type + + +def _fetch_mailbox_message( + connection: MailboxConnection, msg_uid: Any, test: bool +) -> tuple[int | str, str]: + """Fetch one message from ``connection`` by UID, casting the UID to the + type each backend's ``fetch_message`` expects. + + Shared, unmodified, by the sequential and parallel branches of + ``get_dmarc_reports_from_mailbox`` so both fetch identically. + + Returns ``(message_id, msg_content)``; ``message_id`` is the + backend-appropriate id to use for later move/delete calls. + """ + message_id: int | str + if isinstance(connection, IMAPConnection): + message_id = int(msg_uid) + msg_content = connection.fetch_message(message_id) + elif isinstance(connection, MSGraphConnection): + message_id = str(msg_uid) + msg_content = connection.fetch_message(message_id, mark_read=not test) + elif isinstance(connection, MaildirConnection): + message_id = str(msg_uid) if not isinstance(msg_uid, str) else msg_uid + msg_content = connection.fetch_message(message_id, mark_read=not test) + else: + message_id = str(msg_uid) if not isinstance(msg_uid, str) else msg_uid + msg_content = connection.fetch_message(message_id) + return message_id, msg_content + + +def _dispose_invalid_message( + connection: MailboxConnection, + message_id: int | str, + delete: bool, + invalid_reports_folder: str, +) -> None: + """Delete or move an unparseable message, per ``delete``. + + Shared, unmodified, by the sequential and parallel branches of + ``get_dmarc_reports_from_mailbox`` so both dispose of invalid messages + identically. + """ + if delete: + logger.debug(f"Deleting message UID {message_id}") + if isinstance(connection, IMAPConnection): + connection.delete_message(int(message_id)) + else: + connection.delete_message(str(message_id)) + else: + logger.debug(f"Moving message UID {message_id} to {invalid_reports_folder}") + if isinstance(connection, IMAPConnection): + connection.move_message(int(message_id), invalid_reports_folder) + else: + connection.move_message(str(message_id), invalid_reports_folder) + + def get_dmarc_reports_from_mbox( input_: str, *, @@ -2081,6 +2178,7 @@ def get_dmarc_reports_from_mbox( reverse_dns_map_url: str | None = None, offline: bool = False, normalize_timespan_threshold_hours: float = 24.0, + n_procs: int = 1, ) -> ParsingResults: """Parses a mailbox in mbox format containing e-mails with attached DMARC reports @@ -2100,6 +2198,9 @@ def get_dmarc_reports_from_mbox( ip_db_path (str): Path to a MMDB file from IPinfo, MaxMind, or DBIP offline (bool): Do not make online queries for geolocation or DNS normalize_timespan_threshold_hours (float): Normalize timespans beyond this + n_procs (int): Number of processes to use for parsing messages in + parallel. Message reading, deduplication, and result assembly + stay in the calling process; only parsing is parallelized. Returns: dict: Lists of ``aggregate_reports``, ``failure_reports``, and ``smtp_tls_reports`` @@ -2113,43 +2214,73 @@ def get_dmarc_reports_from_mbox( message_keys = mbox.keys() total_messages = len(message_keys) logger.debug(f"Found {total_messages} messages in {input_}") - for i in tqdm(range(total_messages), disable=None): - message_key = message_keys[i] - logger.info(f"Processing message {i + 1} of {total_messages}") - msg_content = mbox.get_string(message_key) - try: - sa = strip_attachment_payloads - parsed_email = parse_report_email( - msg_content, - ip_db_path=ip_db_path, - always_use_local_files=always_use_local_files, - reverse_dns_map_path=reverse_dns_map_path, - reverse_dns_map_url=reverse_dns_map_url, - offline=offline, - nameservers=nameservers, - dns_timeout=dns_timeout, - dns_retries=dns_retries, - strip_attachment_payloads=sa, - normalize_timespan_threshold_hours=normalize_timespan_threshold_hours, - ) - if parsed_email["report_type"] == "aggregate": - report_org = parsed_email["report"]["report_metadata"]["org_name"] - report_id = parsed_email["report"]["report_metadata"]["report_id"] - report_key = f"{report_org}_{report_id}" - if report_key not in SEEN_AGGREGATE_REPORT_IDS: - SEEN_AGGREGATE_REPORT_IDS[report_key] = True - aggregate_reports.append(parsed_email["report"]) - else: - logger.debug( - "Skipping duplicate aggregate report " - f"from {report_org} with ID: {report_id}" - ) - elif parsed_email["report_type"] == "failure": - failure_reports.append(parsed_email["report"]) - elif parsed_email["report_type"] == "smtp_tls": - smtp_tls_reports.append(parsed_email["report"]) - except InvalidDMARCReport as error: - logger.warning(error.__str__()) + + if n_procs > 1 and total_messages > 1: + from parsedmarc.parallel import _parse_report_email_job, parallel_map + + parse_kwargs = { + "ip_db_path": ip_db_path, + "always_use_local_files": always_use_local_files, + "reverse_dns_map_path": reverse_dns_map_path, + "reverse_dns_map_url": reverse_dns_map_url, + "offline": offline, + "nameservers": nameservers, + "dns_timeout": dns_timeout, + "dns_retries": dns_retries, + "strip_attachment_payloads": strip_attachment_payloads, + "normalize_timespan_threshold_hours": ( + normalize_timespan_threshold_hours + ), + } + func = functools.partial(_parse_report_email_job, kwargs=parse_kwargs) + + def _jobs(): + for i in range(total_messages): + message_key = message_keys[i] + logger.info(f"Processing message {i + 1} of {total_messages}") + yield mbox.get_string(message_key) + + for result in tqdm( + parallel_map(func, _jobs(), n_procs), + total=total_messages, + disable=None, + ): + if isinstance(result, InvalidDMARCReport): + logger.warning(str(result)) + elif isinstance(result, ParserError): + raise result + else: + _classify_parsed_email( + result, aggregate_reports, failure_reports, smtp_tls_reports + ) + else: + for i in tqdm(range(total_messages), disable=None): + message_key = message_keys[i] + logger.info(f"Processing message {i + 1} of {total_messages}") + msg_content = mbox.get_string(message_key) + try: + sa = strip_attachment_payloads + parsed_email = parse_report_email( + msg_content, + ip_db_path=ip_db_path, + always_use_local_files=always_use_local_files, + reverse_dns_map_path=reverse_dns_map_path, + reverse_dns_map_url=reverse_dns_map_url, + offline=offline, + nameservers=nameservers, + dns_timeout=dns_timeout, + dns_retries=dns_retries, + strip_attachment_payloads=sa, + normalize_timespan_threshold_hours=normalize_timespan_threshold_hours, + ) + _classify_parsed_email( + parsed_email, + aggregate_reports, + failure_reports, + smtp_tls_reports, + ) + except InvalidDMARCReport as error: + logger.warning(error.__str__()) except mailbox.NoSuchMailboxError: raise InvalidDMARCReport(f"Mailbox {input_} does not exist") return { @@ -2218,6 +2349,7 @@ def get_dmarc_reports_from_mailbox( since: datetime | date | str | None = None, create_folders: bool = True, normalize_timespan_threshold_hours: float = 24, + n_procs: int = 1, ) -> ParsingResults: """ Fetches and parses DMARC reports from a mailbox @@ -2247,6 +2379,12 @@ def get_dmarc_reports_from_mailbox( create_folders (bool): Whether to create the destination folders (not used in watch) normalize_timespan_threshold_hours (float): Normalize timespans beyond this + n_procs (int): Number of processes to use for parsing messages in + parallel. Fetching, archiving, and deduplication remain + sequential in the calling process; only parsing is + parallelized. With ``n_procs > 1``, invalid-message disposition + happens after the parsing phase completes, rather than + interleaved message-by-message as it is when ``n_procs`` is 1. Returns: dict: Lists of ``aggregate_reports``, ``failure_reports``, and ``smtp_tls_reports`` @@ -2338,73 +2476,107 @@ def get_dmarc_reports_from_mailbox( logger.debug(f"Processing {message_limit} messages") - for i in range(message_limit): - msg_uid = messages[i] - logger.debug(f"Processing message {i + 1} of {message_limit}: UID {msg_uid}") - message_id: int | str - if isinstance(connection, IMAPConnection): - message_id = int(msg_uid) - msg_content = connection.fetch_message(message_id) - elif isinstance(connection, MSGraphConnection): - message_id = str(msg_uid) - msg_content = connection.fetch_message(message_id, mark_read=not test) - elif isinstance(connection, MaildirConnection): - message_id = str(msg_uid) if not isinstance(msg_uid, str) else msg_uid - msg_content = connection.fetch_message(message_id, mark_read=not test) - else: - message_id = str(msg_uid) if not isinstance(msg_uid, str) else msg_uid - msg_content = connection.fetch_message(message_id) - try: - sa = strip_attachment_payloads - parsed_email = parse_report_email( - msg_content, - nameservers=nameservers, - dns_timeout=dns_timeout, - dns_retries=dns_retries, - ip_db_path=ip_db_path, - always_use_local_files=always_use_local_files, - reverse_dns_map_path=reverse_dns_map_path, - reverse_dns_map_url=reverse_dns_map_url, - offline=offline, - strip_attachment_payloads=sa, - keep_alive=connection.keepalive, - normalize_timespan_threshold_hours=normalize_timespan_threshold_hours, + if n_procs > 1 and message_limit > 1: + from parsedmarc.parallel import _parse_report_email_job, parallel_map + + # keep_alive is a bound method of the live connection object and is + # not picklable, so it must never cross the process boundary; the + # heartbeat passed to parallel_map below keeps the connection alive + # instead. + parse_kwargs = { + "nameservers": nameservers, + "dns_timeout": dns_timeout, + "dns_retries": dns_retries, + "ip_db_path": ip_db_path, + "always_use_local_files": always_use_local_files, + "reverse_dns_map_path": reverse_dns_map_path, + "reverse_dns_map_url": reverse_dns_map_url, + "offline": offline, + "strip_attachment_payloads": strip_attachment_payloads, + "normalize_timespan_threshold_hours": (normalize_timespan_threshold_hours), + } + func = functools.partial(_parse_report_email_job, kwargs=parse_kwargs) + + # parallel_map yields results in submission order, so the oldest + # queued id always belongs to the next yielded result; popping as + # results arrive keeps this queue no larger than the in-flight + # submission window. + fetched_ids: deque[int | str] = deque() + invalid_msg_ids: list[int | str] = [] + + def _jobs(): + for i in range(message_limit): + msg_uid = messages[i] + logger.debug( + f"Processing message {i + 1} of {message_limit}: UID {msg_uid}" + ) + message_id, msg_content = _fetch_mailbox_message( + connection, msg_uid, test + ) + fetched_ids.append(message_id) + yield msg_content + + for result in parallel_map( + func, _jobs(), n_procs, heartbeat=connection.keepalive + ): + message_id = fetched_ids.popleft() + if isinstance(result, ParserError): + logger.warning(str(result)) + invalid_msg_ids.append(message_id) + else: + report_type = _classify_parsed_email( + result, aggregate_reports, failure_reports, smtp_tls_reports + ) + if report_type == "aggregate": + aggregate_report_msg_uids.append(message_id) + elif report_type == "failure": + failure_report_msg_uids.append(message_id) + elif report_type == "smtp_tls": + smtp_tls_msg_uids.append(message_id) + + if not test: + for invalid_message_id in invalid_msg_ids: + _dispose_invalid_message( + connection, invalid_message_id, delete, invalid_reports_folder + ) + else: + for i in range(message_limit): + msg_uid = messages[i] + logger.debug( + f"Processing message {i + 1} of {message_limit}: UID {msg_uid}" ) - if parsed_email["report_type"] == "aggregate": - report_org = parsed_email["report"]["report_metadata"]["org_name"] - report_id = parsed_email["report"]["report_metadata"]["report_id"] - report_key = f"{report_org}_{report_id}" - if report_key not in SEEN_AGGREGATE_REPORT_IDS: - SEEN_AGGREGATE_REPORT_IDS[report_key] = True - aggregate_reports.append(parsed_email["report"]) - else: - logger.debug( - f"Skipping duplicate aggregate report with ID: {report_id}" + message_id, msg_content = _fetch_mailbox_message(connection, msg_uid, test) + try: + sa = strip_attachment_payloads + parsed_email = parse_report_email( + msg_content, + nameservers=nameservers, + dns_timeout=dns_timeout, + dns_retries=dns_retries, + ip_db_path=ip_db_path, + always_use_local_files=always_use_local_files, + reverse_dns_map_path=reverse_dns_map_path, + reverse_dns_map_url=reverse_dns_map_url, + offline=offline, + strip_attachment_payloads=sa, + keep_alive=connection.keepalive, + normalize_timespan_threshold_hours=normalize_timespan_threshold_hours, + ) + report_type = _classify_parsed_email( + parsed_email, aggregate_reports, failure_reports, smtp_tls_reports + ) + if report_type == "aggregate": + aggregate_report_msg_uids.append(message_id) + elif report_type == "failure": + failure_report_msg_uids.append(message_id) + elif report_type == "smtp_tls": + smtp_tls_msg_uids.append(message_id) + except ParserError as error: + logger.warning(error.__str__()) + if not test: + _dispose_invalid_message( + connection, message_id, delete, invalid_reports_folder ) - aggregate_report_msg_uids.append(message_id) - elif parsed_email["report_type"] == "failure": - failure_reports.append(parsed_email["report"]) - failure_report_msg_uids.append(message_id) - elif parsed_email["report_type"] == "smtp_tls": - smtp_tls_reports.append(parsed_email["report"]) - smtp_tls_msg_uids.append(message_id) - except ParserError as error: - logger.warning(error.__str__()) - if not test: - if delete: - logger.debug(f"Deleting message UID {msg_uid}") - if isinstance(connection, IMAPConnection): - connection.delete_message(int(message_id)) - else: - connection.delete_message(str(message_id)) - else: - logger.debug( - f"Moving message UID {msg_uid} to {invalid_reports_folder}" - ) - if isinstance(connection, IMAPConnection): - connection.move_message(int(message_id), invalid_reports_folder) - else: - connection.move_message(str(message_id), invalid_reports_folder) if not test: if delete: @@ -2509,6 +2681,7 @@ def get_dmarc_reports_from_mailbox( offline=offline, since=current_time, normalize_timespan_threshold_hours=normalize_timespan_threshold_hours, + n_procs=n_procs, ) return results @@ -2536,6 +2709,7 @@ def watch_inbox( since: datetime | date | str | None = None, normalize_timespan_threshold_hours: float = 24, config_reloading: Callable | None = None, + n_procs: int = 1, ): """ Watches the mailbox for new messages and @@ -2569,6 +2743,9 @@ def watch_inbox( reload (or shutdown) has been requested (e.g. via SIGHUP/SIGTERM). Polled by the mailbox backend between checks, including the IMAP IDLE loop, so the watcher exits cleanly at a safe boundary. + n_procs (int): Number of processes to use for parsing messages in + parallel. Passed through to ``get_dmarc_reports_from_mailbox`` + on each check. """ def check_callback(connection): @@ -2585,6 +2762,7 @@ def watch_inbox( offline=offline, nameservers=nameservers, dns_timeout=dns_timeout, + n_procs=n_procs, dns_retries=dns_retries, strip_attachment_payloads=strip_attachment_payloads, batch_size=batch_size, diff --git a/parsedmarc/cli.py b/parsedmarc/cli.py index be5ed614..1133a38d 100644 --- a/parsedmarc/cli.py +++ b/parsedmarc/cli.py @@ -4,6 +4,7 @@ """A CLI for parsing DMARC reports""" import atexit +import functools import http.client import json import logging @@ -14,7 +15,6 @@ import time from argparse import ArgumentParser, Namespace from configparser import ConfigParser from glob import escape as glob_escape, glob -from multiprocessing import Pipe, Process from ssl import CERT_NONE, create_default_context import httpx @@ -38,7 +38,6 @@ from parsedmarc import ( kafkaclient, loganalytics, opensearch, - parse_report_file, postgres, s3, save_output, @@ -55,6 +54,7 @@ from parsedmarc.mail import ( MaildirConnection, MSGraphConnection, ) +from parsedmarc.parallel import _parse_report_file_job, parallel_map from parsedmarc.types import ParsingResults from parsedmarc.utils import ( InvalidIPinfoAPIKey, @@ -369,37 +369,9 @@ def _configure_logging(log_level, log_file=None): log_level: The logging level (e.g., logging.DEBUG, logging.WARNING) log_file: Optional path to log file """ - # Get the logger - from parsedmarc.log import logger + from parsedmarc.log import configure_logging - # Set the log level - logger.setLevel(log_level) - - # Add StreamHandler with formatter if not already present - # Check if we already have a StreamHandler to avoid duplicates - # Use exact type check to distinguish from FileHandler subclass - has_stream_handler = any(type(h) is logging.StreamHandler for h in logger.handlers) - - if not has_stream_handler: - formatter = logging.Formatter( - fmt="%(levelname)8s:%(filename)s:%(lineno)d:%(message)s", - datefmt="%Y-%m-%d:%H:%M:%S", - ) - handler = logging.StreamHandler() - handler.setFormatter(formatter) - logger.addHandler(handler) - - # Add FileHandler if log_file is specified - if log_file: - try: - fh = logging.FileHandler(log_file, "a") - formatter = logging.Formatter( - "%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s" - ) - fh.setFormatter(formatter) - logger.addHandler(fh) - except (IOError, OSError, PermissionError) as error: - logger.warning(f"Unable to write to log file: {error}") + configure_logging(log_level, log_file) # Loggers of the libraries that implement the mailbox and Microsoft Graph @@ -447,64 +419,6 @@ def _configure_dependency_logging(level: int) -> None: dep_logger.addHandler(wanted) -def cli_parse( - file_path, - sa, - nameservers, - dns_timeout, - dns_retries, - ip_db_path, - offline, - always_use_local_files, - reverse_dns_map_path, - reverse_dns_map_url, - normalize_timespan_threshold_hours, - conn, - log_level=logging.ERROR, - log_file=None, -): - """Separated this function for multiprocessing - - Args: - file_path: Path to the report file - sa: Strip attachment payloads flag - nameservers: List of nameservers - dns_timeout: DNS timeout - dns_retries: Number of DNS retries on transient errors - ip_db_path: Path to IP database - offline: Offline mode flag - always_use_local_files: Always use local files flag - reverse_dns_map_path: Path to reverse DNS map - reverse_dns_map_url: URL to reverse DNS map - normalize_timespan_threshold_hours: Timespan threshold - conn: Pipe connection for IPC - log_level: Logging level for this process - log_file: Optional path to log file - """ - # Configure logging in this child process - _configure_logging(log_level, log_file) - - try: - file_results = parse_report_file( - file_path, - ip_db_path=ip_db_path, - offline=offline, - always_use_local_files=always_use_local_files, - reverse_dns_map_path=reverse_dns_map_path, - reverse_dns_map_url=reverse_dns_map_url, - nameservers=nameservers, - dns_timeout=dns_timeout, - dns_retries=dns_retries, - strip_attachment_payloads=sa, - normalize_timespan_threshold_hours=normalize_timespan_threshold_hours, - ) - conn.send([file_results, file_path]) - except ParserError as error: - conn.send([error, file_path]) - finally: - conn.close() - - def _load_config(config_file: str | None = None) -> ConfigParser: """Load configuration from an INI file and/or environment variables. @@ -2367,10 +2281,6 @@ def _main(): for mbox_path in mbox_paths: file_paths.remove(mbox_path) - counter = 0 - - results = [] - pbar = None if sys.stderr.isatty() and len(file_paths) > 0: pbar = tqdm(total=len(file_paths)) @@ -2379,88 +2289,53 @@ def _main(): if n_procs < 1: n_procs = 1 - # Capture the current log level to pass to child processes - current_log_level = logger.level - current_log_file = opts.log_file - - for batch_index in range((len(file_paths) + n_procs - 1) // n_procs): - # Honor a shutdown request between batches before spawning the - # next pool. Anything already parsed is still in `results` and - # will go through process_reports() in the cleanup path so we - # don't lose work the operator already paid for. - if _shutdown_requested: - logger.info( - "Shutdown requested, stopping file processing after %d batch(es)", - batch_index, - ) - break - - processes = [] - connections = [] - - for proc_index in range(n_procs * batch_index, n_procs * (batch_index + 1)): - if proc_index >= len(file_paths): - break - - parent_conn, child_conn = Pipe() - connections.append(parent_conn) - - process = Process( - target=cli_parse, - args=( - file_paths[proc_index], - opts.strip_attachment_payloads, - opts.nameservers, - opts.dns_timeout, - opts.dns_retries, - opts.ip_db_path, - opts.offline, - opts.always_use_local_files, - opts.reverse_dns_map_path, - opts.reverse_dns_map_url, - opts.normalize_timespan_threshold_hours, - child_conn, - current_log_level, - current_log_file, - ), - ) - processes.append(process) - - for proc in processes: - proc.start() - - for conn in connections: - results.append(conn.recv()) - - for proc in processes: - proc.join() - if pbar is not None: - counter += 1 - pbar.update(1) - - if pbar is not None: - pbar.close() - - for result in results: - if isinstance(result[0], ParserError) or result[0] is None: - logger.error(f"Failed to parse {result[1]} - {result[0]}") + parse_kwargs = dict( + offline=opts.offline, + ip_db_path=opts.ip_db_path, + always_use_local_files=opts.always_use_local_files, + reverse_dns_map_path=opts.reverse_dns_map_path, + reverse_dns_map_url=opts.reverse_dns_map_url, + nameservers=opts.nameservers, + dns_timeout=opts.dns_timeout, + dns_retries=opts.dns_retries, + strip_attachment_payloads=opts.strip_attachment_payloads, + normalize_timespan_threshold_hours=opts.normalize_timespan_threshold_hours, + ) + func = functools.partial(_parse_report_file_job, kwargs=parse_kwargs) + for file_path, result in parallel_map( + func, file_paths, n_procs, should_stop=lambda: _shutdown_requested + ): + if pbar is not None: + pbar.update(1) + if isinstance(result, Exception): + logger.error(f"Failed to parse {file_path} - {result}") else: - if result[0]["report_type"] == "aggregate": - report_org = result[0]["report"]["report_metadata"]["org_name"] - report_id = result[0]["report"]["report_metadata"]["report_id"] + if result["report_type"] == "aggregate": + report_org = result["report"]["report_metadata"]["org_name"] + report_id = result["report"]["report_metadata"]["report_id"] report_key = f"{report_org}_{report_id}" if report_key not in SEEN_AGGREGATE_REPORT_IDS: SEEN_AGGREGATE_REPORT_IDS[report_key] = True - aggregate_reports.append(result[0]["report"]) + aggregate_reports.append(result["report"]) else: logger.debug( "Skipping duplicate aggregate report " f"from {report_org} with ID: {report_id}" ) - elif result[0]["report_type"] == "failure": - failure_reports.append(result[0]["report"]) - elif result[0]["report_type"] == "smtp_tls": - smtp_tls_reports.append(result[0]["report"]) + elif result["report_type"] == "failure": + failure_reports.append(result["report"]) + elif result["report_type"] == "smtp_tls": + smtp_tls_reports.append(result["report"]) + + if pbar is not None: + pbar.close() + + if _shutdown_requested: + # Anything already parsed is still in aggregate_reports / + # failure_reports / smtp_tls_reports and will go through + # process_reports() in the cleanup path so we don't lose work + # the operator already paid for. + logger.info("Shutdown requested, stopping file processing early") for mbox_path in mbox_paths: if _shutdown_requested: @@ -2484,6 +2359,7 @@ def _main(): reverse_dns_map_url=opts.reverse_dns_map_url, offline=opts.offline, normalize_timespan_threshold_hours=normalize_timespan_threshold_hours_value, + n_procs=n_procs, ) aggregate_reports += reports["aggregate_reports"] failure_reports += reports["failure_reports"] @@ -2673,6 +2549,7 @@ def _main(): since=opts.mailbox_since, dns_retries=opts.dns_retries, normalize_timespan_threshold_hours=normalize_timespan_threshold_hours_value, + n_procs=n_procs, ) aggregate_reports += reports["aggregate_reports"] @@ -2803,6 +2680,7 @@ def _main(): offline=opts.offline, normalize_timespan_threshold_hours=normalize_timespan_threshold_hours_value, config_reloading=lambda: _reload_requested or _shutdown_requested, + n_procs=n_procs, ) except FileExistsError as error: logger.error(f"{error.__str__()}") diff --git a/parsedmarc/log.py b/parsedmarc/log.py index c10988db..8ba69b7f 100644 --- a/parsedmarc/log.py +++ b/parsedmarc/log.py @@ -1,4 +1,59 @@ +from __future__ import annotations + import logging +import os logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) + + +def configure_logging(log_level: int, log_file: str | None = None) -> None: + """Configure the parsedmarc logger's handlers. + + This is needed for child processes (e.g. parallel report parsing + workers) to properly log messages, since a spawned/forkserver process + does not inherit the parent's configured logging handlers. + + Args: + log_level: The logging level (e.g., logging.DEBUG, logging.WARNING) + log_file: Optional path to log file + """ + # Set the log level + logger.setLevel(log_level) + + # Add StreamHandler with formatter if not already present + # Check if we already have a StreamHandler to avoid duplicates + # Use exact type check to distinguish from FileHandler subclass + has_stream_handler = any(type(h) is logging.StreamHandler for h in logger.handlers) + + if not has_stream_handler: + formatter = logging.Formatter( + fmt="%(levelname)8s:%(filename)s:%(lineno)d:%(message)s", + datefmt="%Y-%m-%d:%H:%M:%S", + ) + handler = logging.StreamHandler() + handler.setFormatter(formatter) + logger.addHandler(handler) + + # Add FileHandler if log_file is specified and no handler for that + # file is attached yet. FileHandler stores its target as an absolute + # path in baseFilename, so compare against the absolute path; without + # this check, repeated calls (e.g. a SIGHUP config reload) would stack + # duplicate handlers that write every record twice and leak file + # descriptors. + if log_file: + try: + log_file_path = os.path.abspath(log_file) + has_file_handler = any( + isinstance(h, logging.FileHandler) and h.baseFilename == log_file_path + for h in logger.handlers + ) + if not has_file_handler: + fh = logging.FileHandler(log_file, "a") + formatter = logging.Formatter( + "%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s" + ) + fh.setFormatter(formatter) + logger.addHandler(fh) + except (IOError, OSError, PermissionError) as error: + logger.warning(f"Unable to write to log file: {error}") diff --git a/parsedmarc/parallel.py b/parsedmarc/parallel.py new file mode 100644 index 00000000..5e007afe --- /dev/null +++ b/parsedmarc/parallel.py @@ -0,0 +1,192 @@ +"""Bounded-window multiprocessing helpers for parallel report parsing.""" + +from __future__ import annotations + +import logging +from collections import deque +from collections.abc import Callable, Iterable, Iterator +from concurrent.futures import Future, ProcessPoolExecutor, wait +from typing import TYPE_CHECKING, Any, TypeVar + +if TYPE_CHECKING: + from parsedmarc import ParserError + from parsedmarc.types import ParsedReport + +_J = TypeVar("_J") +_R = TypeVar("_R") + + +def _init_worker_logging(log_level: int, log_files: list[str]) -> None: + """Pool initializer that reconstructs the parent's logging handlers. + + A spawned or forkserver worker process does not inherit the parent + process's already-configured logging handlers, so without this the + child's log records (including any that would surface a parsing bug) + are silently dropped. Called once per worker process by + ``ProcessPoolExecutor``'s ``initializer``/``initargs``. + """ + from parsedmarc.log import configure_logging + + if not log_files: + configure_logging(log_level, None) + return + for log_file in log_files: + configure_logging(log_level, log_file) + + +def _parse_report_email_job( + msg_content: bytes | str, *, kwargs: dict[str, Any] +) -> ParsedReport | ParserError: + """Worker job that parses a single report email. + + Module-level and picklable so it can run in a ``ProcessPoolExecutor`` + worker. ``kwargs`` is forwarded to ``parse_report_email`` and must + never contain ``keep_alive`` - it is a bound method of the live + mailbox connection in the parent process and is not picklable. + + Returns the parsed report on success. A ``ParserError`` raised by + ``parse_report_email`` is caught and returned as a value (never + re-raised) so that a single invalid message doesn't need special + handling on the submission side; any other exception propagates and + surfaces as the future's exception. + """ + from parsedmarc import ParserError, parse_report_email + + try: + return parse_report_email(msg_content, **kwargs) + except ParserError as e: + return e + + +def _parse_report_file_job( + file_path: str, *, kwargs: dict[str, Any] +) -> tuple[str, ParsedReport | Exception]: + """Worker job that parses a single report file. + + Module-level and picklable so it can run in a ``ProcessPoolExecutor`` + worker. ``kwargs`` is forwarded to ``parse_report_file``. + + Catches any ``Exception`` (not just ``ParserError``) and returns it + paired with ``file_path`` rather than letting it propagate. This is a + deliberate behavior change from the old hand-rolled + ``Pipe``/``Process`` CLI worker, where a non-``ParserError`` exception + in the child crashed the child without sending anything back over the + pipe, leaving the parent blocked forever on ``conn.recv()``. Returning + the exception as a value lets the caller log it and move on. + """ + from parsedmarc import parse_report_file + + try: + return file_path, parse_report_file(file_path, **kwargs) + except Exception as e: + return file_path, e + + +def parallel_map( + func: Callable[[_J], _R], + jobs: Iterable[_J], + n_procs: int, + *, + window_factor: int = 2, + heartbeat: Callable[[], None] | None = None, + heartbeat_interval: float = 30.0, + should_stop: Callable[[], bool] | None = None, +) -> Iterator[_R]: + """Map ``func`` over ``jobs`` using a bounded-window process pool. + + ``jobs`` is iterated lazily and is never materialized into a list, so + this is safe to call with a generator over a very large input (e.g. a + 20,000 message mbox). At most ``window_factor * n_procs`` jobs are + submitted ahead of the harvesting point at any time, which bounds + memory use and lets the caller's job-producing generator itself stay + lazy (e.g. fetching mailbox messages just-in-time). + + Results are yielded in submission order (i.e. the same order as + ``jobs``), not completion order, so callers can rely on index-based + bookkeeping against the original input. + + If ``jobs`` is empty, returns without spawning a process pool. + + If ``heartbeat`` is given, it is called after every + ``heartbeat_interval`` seconds spent waiting on the oldest + outstanding future (e.g. to keep an IMAP connection alive while a + large report is being parsed). If ``heartbeat`` is ``None``, results + are awaited with a plain blocking ``future.result()`` call. + + If ``should_stop`` is given, it is checked after each yielded result; + if it returns ``True``, the pool is shut down with + ``cancel_futures=True``: jobs still queued in the submission window + that have not started are cancelled and their results dropped, while + jobs already running (or finished) are waited on and their results + yielded before returning - no completed work is discarded, but the + stop may block briefly until in-flight jobs finish. + + Raises: + ValueError: If ``n_procs`` is less than 1. Raised eagerly at the + call itself (not on first iteration of the returned iterator). + """ + if n_procs < 1: + raise ValueError(f"n_procs must be at least 1, got {n_procs}") + + def _generate() -> Iterator[_R]: + jobs_iter = iter(jobs) + try: + first_job = next(jobs_iter) + except StopIteration: + return + + from parsedmarc.log import logger + + log_level = logger.getEffectiveLevel() + log_files = [ + h.baseFilename + for h in logger.handlers + if isinstance(h, logging.FileHandler) + ] + + window_size = max(window_factor * n_procs, 1) + + with ProcessPoolExecutor( + max_workers=n_procs, + initializer=_init_worker_logging, + initargs=(log_level, log_files), + ) as executor: + window: deque[Future[_R]] = deque() + window.append(executor.submit(func, first_job)) + + # Prime the submission window up to its bound before harvesting. + while len(window) < window_size: + try: + job = next(jobs_iter) + except StopIteration: + break + window.append(executor.submit(func, job)) + + while window: + oldest = window.popleft() + if heartbeat is None: + result = oldest.result() + else: + while True: + done, _ = wait([oldest], timeout=heartbeat_interval) + if done: + result = oldest.result() + break + heartbeat() + + yield result + + if should_stop is not None and should_stop(): + executor.shutdown(cancel_futures=True) + for remaining in window: + if not remaining.cancelled(): + yield remaining.result() + return + + try: + job = next(jobs_iter) + except StopIteration: + continue + window.append(executor.submit(func, job)) + + return _generate() diff --git a/tests/test_cli.py b/tests/test_cli.py index f2bce156..e5afd995 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -992,13 +992,16 @@ class TestDirectoryFilePaths(unittest.TestCase): dst.write(src.read()) return reports_dir - def _write_config(self, tmp_dir, output_dirname): + def _write_config(self, tmp_dir, output_dirname, n_procs=None): cfg_path = os.path.join(tmp_dir, "parsedmarc.ini") output_dir = os.path.join(tmp_dir, output_dirname) + config_text = ( + f"[general]\noffline = True\nsilent = True\noutput = {output_dir}\n" + ) + if n_procs is not None: + config_text += f"n_procs = {n_procs}\n" with open(cfg_path, "w") as f: - f.write( - f"[general]\noffline = True\nsilent = True\noutput = {output_dir}\n" - ) + f.write(config_text) return cfg_path, output_dir def _report_ids(self, output_dir): @@ -1053,6 +1056,35 @@ class TestDirectoryFilePaths(unittest.TestCase): self.assertIn(top_level_id, report_ids) self.assertIn(nested_id, report_ids) + def test_directory_file_path_recursive_includes_nested_n_procs_2(self): + """The same recursive-directory scenario as + ``test_directory_file_path_recursive_includes_nested``, but with + ``n_procs = 2`` in the config file so the direct-file parsing path + runs through ``parallel_map``'s process pool with more than one + worker. The report-id sets in the output JSON must match the + ``n_procs = 1`` (default) run exactly — parallelism must not change + which reports get parsed or deduplicated. + """ + with tempfile.TemporaryDirectory() as tmp_dir: + reports_dir = self._build_reports_dir(tmp_dir) + cfg_path, output_dir = self._write_config( + tmp_dir, "output_recursive_n_procs_2", n_procs=2 + ) + + with patch.object( + sys, "argv", ["parsedmarc", "-c", cfg_path, "-r", reports_dir] + ): + parsedmarc.cli._main() + + report_ids = self._report_ids(output_dir) + + top_level_id = self._sample_report_id(self.TOP_LEVEL_SAMPLE) + nested_id = self._sample_report_id(self.NESTED_SAMPLE) + + self.assertIn(top_level_id, report_ids) + self.assertIn(nested_id, report_ids) + self.assertEqual(report_ids, {top_level_id, nested_id}) + class TestGmailAuthModes(unittest.TestCase): @patch("parsedmarc.cli.get_dmarc_reports_from_mailbox") @@ -3164,68 +3196,100 @@ watch = true @patch("parsedmarc.cli.get_dmarc_reports_from_mbox") @patch("parsedmarc.cli.is_mbox", side_effect=lambda p: p.endswith(".mbox")) @patch("parsedmarc.cli._init_output_clients") - @patch("parsedmarc.cli.Process") + @patch("parsedmarc.parallel.ProcessPoolExecutor") @patch("parsedmarc.cli.glob") - def testSigtermDuringOneShotStopsBetweenBatchesAndMbox( + def testSigtermDuringOneShotStopsEarlyAndSkipsMbox( self, mock_glob, - mock_process_cls, + mock_pool_cls, mock_init_clients, mock_is_mbox, mock_get_mbox, ): - """SIGTERM during one-shot processing: the in-flight child is - joined normally (no work lost), the file-batch loop stops before - spawning the next batch, and the subsequent mbox loop breaks on - its first iteration (the flag is already set). Output clients are - still closed. + """SIGTERM during one-shot processing: ``parallel_map``'s + ``should_stop`` check (polled after each yielded result, see + ``parsedmarc/parallel.py``) stops submitting new jobs once the flag + is set, and the subsequent mbox loop breaks on its first iteration + (the flag is already set). Output clients are still closed. - Two ``.xml`` files give the batch loop a second iteration to hit - its break; one ``.mbox`` file routes into ``mbox_paths`` so the - mbox break is exercised too. ``is_mbox`` is keyed by suffix so the - fake filenames don't trigger ``mailbox.mbox(path, create=True)``.""" - mock_glob.return_value = ["a.xml", "b.xml", "c.mbox"] + ``ProcessPoolExecutor`` is patched at the stdlib boundary (mirrors + the old ``cli.Process`` patch) with a fake whose ``submit(fn, arg)`` + runs ``fn(arg)`` inline and returns an already-completed + ``Future``-like object, so no real subprocess is spawned and no + pickling occurs. SIGTERM is raised on the very first ``submit`` + call, mirroring the old test's trigger on the first child's + ``start()``. + + ``parallel_map`` primes its submission window up to + ``2 * n_procs`` jobs *before* the first result is harvested and + ``should_stop`` is ever checked (see ``parallel_map``'s docstring), + so with the default ``n_procs = 1`` the first two files are + submitted and run before the stop takes effect — the ``should_stop`` + check only prevents a third submission. Four ``.xml`` files (more + than the window size) are supplied so that the stop is still + observable: exactly 2 of the 4 are submitted, not all of them. One + ``.mbox`` file routes into ``mbox_paths`` so the mbox break is + exercised too. ``is_mbox`` is keyed by suffix so the fake filenames + don't trigger ``mailbox.mbox(path, create=True)``.""" + mock_glob.return_value = ["a.xml", "b.xml", "c.xml", "d.xml", "e.mbox"] kafka_client = MagicMock(spec=["close"]) mock_init_clients.return_value = {"kafka": kafka_client} - starts = [] + submitted = [] - class FakeProc: - """Stand-in child that finishes its file and sends a result - even though SIGTERM arrived mid-batch.""" + class FakeFuture: + def __init__(self, value): + self._value = value - def __init__(self, target=None, args=()): - self._args = args + def result(self, timeout=None): + return self._value - def start(self): - starts.append(self._args[0]) - if len(starts) == 1: + def cancelled(self): + return False + + class FakeExecutor: + def __init__(self, *args, **kwargs): + pass + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + def submit(self, fn, arg): + if not submitted: + # Mirrors the old test's trigger on the first child's + # start(): SIGTERM arrives after the first job is + # dispatched but while the submission window is still + # being primed. os.kill(os.getpid(), signal.SIGTERM) - # Child still completes and reports back over the pipe. - self._args[-3].send([None, self._args[0]]) + submitted.append(arg) + return FakeFuture(fn(arg)) - def join(self, timeout=None): + def shutdown(self, cancel_futures=True): return None - mock_process_cls.side_effect = FakeProc + mock_pool_cls.side_effect = FakeExecutor - with patch.object(sys, "argv", ["parsedmarc", "a.xml", "b.xml", "c.mbox"]): + with patch.object( + sys, "argv", ["parsedmarc", "a.xml", "b.xml", "c.xml", "d.xml", "e.mbox"] + ): parsedmarc.cli._main() - # Only the first xml batch ran before the batch loop broke, and the - # mbox loop broke before processing its file. - self.assertEqual(len(starts), 1) + # The submission window (2 * n_procs, n_procs=1) primes 2 jobs + # before should_stop is first checked; the stop takes effect before + # a 3rd or 4th file is ever submitted. + self.assertEqual(len(submitted), 2) mock_get_mbox.assert_not_called() kafka_client.close.assert_called() @patch("parsedmarc.cli._init_output_clients") - @patch("parsedmarc.cli.cli_parse") @patch("parsedmarc.cli.glob") def testNormalOneShotExitClosesOutputClients( self, mock_glob, - mock_cli_parse, mock_init_clients, ): """A successful one-shot run with no signal still closes its @@ -4687,9 +4751,11 @@ class TestParseConfigWebhook(unittest.TestCase): class TestConfigureLogging(unittest.TestCase): - """_configure_logging is called in every child process for parallel - parsing — if it stops attaching a handler, log output goes dark in - multiprocessing mode.""" + """cli._configure_logging is a thin wrapper around + parsedmarc.log.configure_logging (the parallel-parsing worker + initializer in parsedmarc/parallel.py calls configure_logging + directly) — if this wrapper stops attaching a handler, any remaining + caller's log output goes dark.""" def setUp(self): from parsedmarc.log import logger as plog @@ -4766,66 +4832,5 @@ class TestConfigureLogging(unittest.TestCase): self.assertTrue(any("Unable to write to log file" in m for m in cm.output)) -class TestCliParse(unittest.TestCase): - """cli_parse is the multiprocessing worker — it shells out to - parse_report_file, then sends the result (or error) back over a - pipe. Both branches matter: a regression would silently drop - results in parallel mode.""" - - def test_cli_parse_sends_results_on_success(self): - from multiprocessing import Pipe - from unittest.mock import patch - from parsedmarc.cli import cli_parse - - parent_conn, child_conn = Pipe() - with patch("parsedmarc.cli.parse_report_file") as mock_parse: - mock_parse.return_value = {"report_type": "aggregate", "report": {}} - cli_parse( - "/path/to/report.xml", - False, - None, - 2.0, - 0, - None, - True, - True, - None, - None, - 24.0, - child_conn, - ) - sent = parent_conn.recv() - self.assertEqual(sent[0], {"report_type": "aggregate", "report": {}}) - self.assertEqual(sent[1], "/path/to/report.xml") - - def test_cli_parse_sends_error_on_parser_error(self): - from multiprocessing import Pipe - from unittest.mock import patch - from parsedmarc.cli import cli_parse - from parsedmarc import ParserError - - parent_conn, child_conn = Pipe() - with patch("parsedmarc.cli.parse_report_file") as mock_parse: - err = ParserError("bad report") - mock_parse.side_effect = err - cli_parse( - "/bad.xml", - False, - None, - 2.0, - 0, - None, - True, - True, - None, - None, - 24.0, - child_conn, - ) - sent = parent_conn.recv() - self.assertIsInstance(sent[0], ParserError) - self.assertEqual(sent[1], "/bad.xml") - - if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/tests/test_init.py b/tests/test_init.py index bfa83e09..ed8ec7a8 100644 --- a/tests/test_init.py +++ b/tests/test_init.py @@ -2628,6 +2628,87 @@ class TestGetDmarcReportsFromMbox(unittest.TestCase): os.remove(path) +class TestGetDmarcReportsFromMboxParallel(unittest.TestCase): + """n_procs=1 vs n_procs=2 parity for get_dmarc_reports_from_mbox, driven + against a real mbox file built from real sample emails (offline parsing, + no mocks of parsedmarc internals): one aggregate, one failure, one + SMTP TLS report, a duplicate copy of the aggregate, and a junk message + that isn't a report at all. + + Confirms the parallel branch classifies and dedups identically to the + sequential branch, and that the junk message produces a warning log + instead of raising -- both branches must only catch InvalidDMARCReport + (a ParserError subclass) around a single message and continue, since a + bare ParserError is deliberately re-raised (see the n_procs > 1 branch + of get_dmarc_reports_from_mbox, which mirrors the sequential branch's + `except InvalidDMARCReport` scope). + """ + + AGGREGATE = "samples/aggregate/twilight.eml" + FAILURE = "samples/failure/dmarc_ruf_report_linkedin.eml" + SMTP_TLS = "samples/smtp_tls/google.com_smtp_tls_report.eml" + JUNK = b"From: noise@example.com\nSubject: not a report\n\nplain text\n" + + def setUp(self): + self._tmp = mkdtemp() + self.addCleanup(rmtree, self._tmp, ignore_errors=True) + parsedmarc.SEEN_AGGREGATE_REPORT_IDS.clear() + self._path = os.path.join(self._tmp, "reports.mbox") + box = mailbox.mbox(self._path) + box.lock() + try: + # AGGREGATE appears twice: dedup must collapse it to one report. + for source in (self.AGGREGATE, self.FAILURE, self.SMTP_TLS, self.AGGREGATE): + with open(source, "rb") as source_file: + box.add(mailbox.mboxMessage(source_file.read())) + box.add(mailbox.mboxMessage(self.JUNK)) + box.flush() + finally: + box.unlock() + box.close() + + @staticmethod + def _aggregate_report_ids(results): + return {r["report_metadata"]["report_id"] for r in results["aggregate_reports"]} + + def test_parallel_matches_sequential(self): + parsedmarc.SEEN_AGGREGATE_REPORT_IDS.clear() + with self.assertLogs("parsedmarc.log", level="WARNING") as sequential_logs: + sequential = parsedmarc.get_dmarc_reports_from_mbox( + self._path, offline=True, n_procs=1 + ) + + parsedmarc.SEEN_AGGREGATE_REPORT_IDS.clear() + with self.assertLogs("parsedmarc.log", level="WARNING") as parallel_logs: + parallel = parsedmarc.get_dmarc_reports_from_mbox( + self._path, offline=True, n_procs=2 + ) + + # Dedup collapsed the duplicate aggregate to a single report, and + # both branches picked the same report. + self.assertEqual(len(sequential["aggregate_reports"]), 1) + self.assertEqual(len(parallel["aggregate_reports"]), 1) + self.assertEqual( + self._aggregate_report_ids(sequential), self._aggregate_report_ids(parallel) + ) + + # Same failure/smtp_tls counts in both branches. + self.assertEqual(len(sequential["failure_reports"]), 1) + self.assertEqual(len(parallel["failure_reports"]), 1) + self.assertEqual(len(sequential["smtp_tls_reports"]), 1) + self.assertEqual(len(parallel["smtp_tls_reports"]), 1) + + # The junk message warned in both branches without raising. + self.assertTrue( + any("not a valid report" in line for line in sequential_logs.output), + sequential_logs.output, + ) + self.assertTrue( + any("not a valid report" in line for line in parallel_logs.output), + parallel_logs.output, + ) + + class TestGetDmarcReportsFromMailboxValidation(unittest.TestCase): """Input validation on get_dmarc_reports_from_mailbox. @@ -2795,6 +2876,23 @@ class TestGetDmarcReportsFromMailboxMaildir(unittest.TestCase): ) return conn, result + def _assert_each_report_type_routed(self, conn, result): + """Shared assertions for one report of each type plus an + unparseable message: each is filed under the correct subfolder + (Aggregate / Failure / SMTP-TLS / Invalid) and the INBOX is + drained. Used by both the sequential and n_procs=2 variants below + so the two tests share their assertions instead of duplicating + them.""" + self.assertEqual(len(result["aggregate_reports"]), 1) + self.assertEqual(len(result["failure_reports"]), 1) + self.assertEqual(len(result["smtp_tls_reports"]), 1) + + self.assertEqual(conn.fetch_messages("INBOX"), []) + self.assertEqual(len(conn.fetch_messages("Archive/Aggregate")), 1) + self.assertEqual(len(conn.fetch_messages("Archive/Failure")), 1) + self.assertEqual(len(conn.fetch_messages("Archive/SMTP-TLS")), 1) + self.assertEqual(len(conn.fetch_messages("Archive/Invalid")), 1) + def test_each_report_type_routed_to_its_archive_subfolder(self): """One report of each type plus an unparseable message: each is filed under the correct subfolder (Aggregate / Failure / SMTP-TLS / Invalid) @@ -2806,15 +2904,20 @@ class TestGetDmarcReportsFromMailboxMaildir(unittest.TestCase): conn, result = self._run() - self.assertEqual(len(result["aggregate_reports"]), 1) - self.assertEqual(len(result["failure_reports"]), 1) - self.assertEqual(len(result["smtp_tls_reports"]), 1) + self._assert_each_report_type_routed(conn, result) - self.assertEqual(conn.fetch_messages("INBOX"), []) - self.assertEqual(len(conn.fetch_messages("Archive/Aggregate")), 1) - self.assertEqual(len(conn.fetch_messages("Archive/Failure")), 1) - self.assertEqual(len(conn.fetch_messages("Archive/SMTP-TLS")), 1) - self.assertEqual(len(conn.fetch_messages("Archive/Invalid")), 1) + def test_each_report_type_routed_to_its_archive_subfolder_parallel(self): + """Same scenario as above, with n_procs=2: fetching and archiving + stay sequential in the parent, but parsing runs in worker + processes. The routing outcome must be identical.""" + self._deliver(self.AGGREGATE) + self._deliver(self.FAILURE) + self._deliver(self.SMTP_TLS) + self._deliver(self.JUNK) + + conn, result = self._run(n_procs=2) + + self._assert_each_report_type_routed(conn, result) def test_delete_mode_removes_processed_messages(self): """delete=True: a parsed message is removed from the INBOX rather than @@ -2828,6 +2931,23 @@ class TestGetDmarcReportsFromMailboxMaildir(unittest.TestCase): # The Failure folder is created but nothing is filed there — deleted. self.assertEqual(conn.fetch_messages("Archive/Failure"), []) + def test_delete_mode_removes_processed_messages_parallel(self): + """Same as above, with n_procs=2 and enough messages (>1) to take + the parallel branch: invalid-message disposition happens after the + parse phase for n_procs > 1 (see get_dmarc_reports_from_mailbox's + docstring), but the end result -- both messages gone from the + INBOX and nothing archived -- must match delete mode exactly.""" + self._deliver(self.FAILURE) + self._deliver(self.AGGREGATE) + + conn, result = self._run(delete=True, n_procs=2) + + self.assertEqual(len(result["failure_reports"]), 1) + self.assertEqual(len(result["aggregate_reports"]), 1) + self.assertEqual(conn.fetch_messages("INBOX"), []) + self.assertEqual(conn.fetch_messages("Archive/Failure"), []) + self.assertEqual(conn.fetch_messages("Archive/Aggregate"), []) + def test_test_mode_parses_without_moving_or_creating_folders(self): """test=True: the report is parsed and returned, but the message stays in the INBOX and no archive folders are created/touched.""" @@ -2839,6 +2959,112 @@ class TestGetDmarcReportsFromMailboxMaildir(unittest.TestCase): self.assertEqual(len(conn.fetch_messages("INBOX")), 1) self.assertFalse(conn.folder_exists("Archive/Failure")) + def test_test_mode_parses_without_moving_or_creating_folders_parallel(self): + """Same as above, with n_procs=2 and enough messages (>1) to take + the parallel branch: test mode disposes of nothing regardless of + n_procs, so both messages stay put and no archive folders appear.""" + self._deliver(self.FAILURE) + self._deliver(self.AGGREGATE) + + conn, result = self._run(test=True, n_procs=2) + + self.assertEqual(len(result["failure_reports"]), 1) + self.assertEqual(len(result["aggregate_reports"]), 1) + self.assertEqual(len(conn.fetch_messages("INBOX")), 2) + self.assertFalse(conn.folder_exists("Archive/Failure")) + self.assertFalse(conn.folder_exists("Archive/Aggregate")) + + def test_duplicate_aggregate_parallel_archives_both_messages(self): + """Delivering the same aggregate sample twice: dedup means the + parsed *results* contain only one aggregate report, but the + sequential caller appends a message's UID to the aggregate archive + list unconditionally -- including for the duplicate -- so BOTH + source messages still get archived to Aggregate. The n_procs=2 + branch must match: the helper (_classify_parsed_email) owns the + dedup, but the caller appends the UID regardless of report_type + matching "aggregate", exactly as the sequential branch does.""" + self._deliver(self.AGGREGATE) + self._deliver(self.AGGREGATE) + + conn, result = self._run(n_procs=2) + + self.assertEqual(len(result["aggregate_reports"]), 1) + self.assertEqual(len(conn.fetch_messages("Archive/Aggregate")), 2) + self.assertEqual(conn.fetch_messages("INBOX"), []) + + +class _MidRunArrivalMaildirConnection(MaildirConnection): + """A MaildirConnection that delivers one extra message into the INBOX + right after its first fetch_messages() call returns, simulating mail + arriving while the first batch is being processed. + + Used only to make get_dmarc_reports_from_mailbox's batch_size=0 + tail-recursion branch (`if not test and not batch_size:`) actually + recurse once for real -- no parsedmarc internals are mocked, only this + real mailsuite MaildirConnection subclass's timing -- so the n_procs + pass-through in that recursive call is genuinely exercised rather than + merely inspected. The recursive call's own re-check then finds nothing + further (the extra message has already been archived), so recursion + terminates after one extra level -- cheap and deterministic. + """ + + def __init__(self, *args, extra_source, **kwargs): + super().__init__(*args, **kwargs) + self._extra_source = extra_source + self._delivered_extra = False + + def fetch_messages(self, reports_folder, **kwargs): + result = super().fetch_messages(reports_folder, **kwargs) + if not self._delivered_extra: + self._delivered_extra = True + with open(self._extra_source, "rb") as extra_file: + raw = extra_file.read() + box = mailbox.Maildir(self._maildir_path, create=False) + box.add(mailbox.MaildirMessage(raw)) + box.flush() + return result + + +class TestGetDmarcReportsFromMailboxMaildirBatchSizeZeroRecursion(unittest.TestCase): + """batch_size=0 makes get_dmarc_reports_from_mailbox re-check the + mailbox after its main pass and recurse if new messages showed up + (__init__.py's `if not test and not batch_size:` block). This exercises + that recursive call with n_procs=2 threaded through it for real.""" + + AGGREGATE = "samples/aggregate/twilight.eml" + FAILURE = "samples/failure/dmarc_ruf_report_linkedin.eml" + SMTP_TLS = "samples/smtp_tls/google.com_smtp_tls_report.eml" + + def setUp(self): + self._tmp = mkdtemp() + self.addCleanup(rmtree, self._tmp, ignore_errors=True) + parsedmarc.SEEN_AGGREGATE_REPORT_IDS.clear() + self._maildir = os.path.join(self._tmp, "Maildir") + inbox = mailbox.Maildir(self._maildir, create=True) + for source in (self.AGGREGATE, self.FAILURE): + with open(source, "rb") as source_file: + inbox.add(mailbox.MaildirMessage(source_file.read())) + inbox.flush() + + def test_batch_size_zero_recursion_threads_n_procs(self): + conn = _MidRunArrivalMaildirConnection( + self._maildir, maildir_create=True, extra_source=self.SMTP_TLS + ) + + result = parsedmarc.get_dmarc_reports_from_mailbox( + connection=conn, offline=True, n_procs=2, batch_size=0 + ) + + # The main pass parses AGGREGATE + FAILURE; the message that "arrives" + # after the main pass's fetch is only found and parsed via the + # recursive re-check call, which would be missing from the results + # entirely if the n_procs kwarg (or anything else) were dropped from + # that recursive call. + self.assertEqual(len(result["aggregate_reports"]), 1) + self.assertEqual(len(result["failure_reports"]), 1) + self.assertEqual(len(result["smtp_tls_reports"]), 1) + self.assertEqual(conn.fetch_messages("INBOX"), []) + class TestEmailResultsErrorBranches(unittest.TestCase): """email_results requires mail_to to be a list — this is enforced diff --git a/tests/test_log.py b/tests/test_log.py new file mode 100644 index 00000000..7c745fd0 --- /dev/null +++ b/tests/test_log.py @@ -0,0 +1,74 @@ +"""Tests for parsedmarc.log""" + +import logging +import os +import tempfile +import unittest + +from parsedmarc.log import configure_logging, logger + + +class TestConfigureLoggingFileHandlerDedup(unittest.TestCase): + """Repeated configure_logging calls with the same log_file (e.g. a + SIGHUP config reload re-running the CLI's logging setup) must not + stack duplicate FileHandlers - a duplicate would write every record + twice and leak a file descriptor per call.""" + + def setUp(self): + self._saved_handlers = list(logger.handlers) + self._saved_level = logger.level + + def tearDown(self): + for handler in list(logger.handlers): + if handler not in self._saved_handlers: + logger.removeHandler(handler) + if isinstance(handler, logging.FileHandler): + handler.close() + logger.handlers[:] = self._saved_handlers + logger.setLevel(self._saved_level) + + def _temp_log_path(self): + with tempfile.NamedTemporaryFile(suffix=".log", delete=False) as tf: + path = tf.name + self.addCleanup(lambda: os.path.exists(path) and os.remove(path)) + return path + + def test_same_log_file_twice_attaches_one_handler_and_logs_once(self): + log_path = self._temp_log_path() + + configure_logging(logging.INFO, log_path) + configure_logging(logging.INFO, log_path) + + file_handlers = [ + h + for h in logger.handlers + if isinstance(h, logging.FileHandler) + and h.baseFilename == os.path.abspath(log_path) + ] + self.assertEqual(len(file_handlers), 1) + + logger.info("dedup-check line") + for handler in file_handlers: + handler.flush() + with open(log_path) as f: + contents = f.read() + self.assertEqual(contents.count("dedup-check line"), 1) + + def test_different_log_files_attach_one_handler_each(self): + first_path = self._temp_log_path() + second_path = self._temp_log_path() + + configure_logging(logging.INFO, first_path) + configure_logging(logging.INFO, second_path) + + file_paths = [ + h.baseFilename + for h in logger.handlers + if isinstance(h, logging.FileHandler) + ] + self.assertIn(os.path.abspath(first_path), file_paths) + self.assertIn(os.path.abspath(second_path), file_paths) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_parallel.py b/tests/test_parallel.py new file mode 100644 index 00000000..fdedb48f --- /dev/null +++ b/tests/test_parallel.py @@ -0,0 +1,284 @@ +"""Tests for parsedmarc.parallel""" + +import functools +import logging +import os +import tempfile +import unittest +from unittest.mock import patch + +import parsedmarc +from parsedmarc.parallel import ( + _init_worker_logging, + _parse_report_email_job, + _parse_report_file_job, + parallel_map, +) + +# Stable sample files reused from tests/test_cli.py's TestDirectoryFilePaths, +# plus a third plain aggregate sample, so parity/order tests exercise more +# than one worker submission. +SAMPLE_PATHS = [ + "samples/aggregate/!example.com!1538204542!1538463818.xml", + "samples/aggregate/!large-example.com!1711897200!1711983600.xml", + "samples/aggregate/example.net!example.com!1529366400!1529452799.xml", +] + + +def _echo_job(x): + """Trivial module-level (spawn-picklable) worker used to exercise + parallel_map's scheduling behavior without involving real parsing.""" + return x + + +class _CountingIterable: + """Wraps a range so tests can observe how many items parallel_map has + pulled from a lazily-iterated jobs source at any point during + iteration, without materializing the whole sequence up front.""" + + def __init__(self, n): + self.n = n + self.pulled = 0 + + def __iter__(self): + for i in range(self.n): + self.pulled += 1 + yield i + + +class _ParallelTestCase(unittest.TestCase): + """Common env setup shared by parallel.py tests: offline mode, no DNS.""" + + def setUp(self): + self._env_patcher = patch.dict( + os.environ, {"GITHUB_ACTIONS": "true"}, clear=False + ) + self._env_patcher.start() + self.addCleanup(self._env_patcher.stop) + + +class TestParallelMapParseReportFile(_ParallelTestCase): + """parallel_map + _parse_report_file_job over real sample files must + behave identically to calling parse_report_file sequentially, and + must preserve submission order in its results.""" + + def test_results_match_sequential_parsing_in_order(self): + expected = [ + (path, parsedmarc.parse_report_file(path, offline=True)) + for path in SAMPLE_PATHS + ] + + job = functools.partial(_parse_report_file_job, kwargs=dict(offline=True)) + results = list(parallel_map(job, SAMPLE_PATHS, n_procs=2)) + + self.assertEqual(len(results), len(SAMPLE_PATHS)) + # Order must match submission order (SAMPLE_PATHS), not completion + # order. + self.assertEqual([path for path, _ in results], SAMPLE_PATHS) + for (path, report), (expected_path, expected_report) in zip(results, expected): + self.assertEqual(path, expected_path) + self.assertNotIsInstance(report, Exception) + self.assertEqual(report, expected_report) + + +class TestParallelMapJunkFile(_ParallelTestCase): + """A worker crash on an unparseable file must surface as an Exception + *value* in the result tuple, not hang the whole run. This is a + regression guard for the old Pipe/Process CLI worker, which left the + parent blocked forever on conn.recv() when a child died from a + non-ParserError exception.""" + + def test_junk_file_yields_exception_value_and_completes(self): + with tempfile.NamedTemporaryFile(suffix=".xml", delete=False, mode="wb") as tf: + tf.write(b"not a report") + junk_path = tf.name + self.addCleanup(os.remove, junk_path) + + job = functools.partial(_parse_report_file_job, kwargs=dict(offline=True)) + results = list(parallel_map(job, [junk_path, junk_path], n_procs=2)) + + self.assertEqual(len(results), 2) + for path, result in results: + self.assertEqual(path, junk_path) + self.assertIsInstance(result, Exception) + + +class TestParallelMapBoundedLaziness(unittest.TestCase): + """The jobs iterable must never be materialized up front. At any point + during iteration, the number of items pulled from the source should + stay within window_factor * n_procs of the number of results already + yielded, bounding memory use for very large inputs (e.g. a 20,000 + message mbox).""" + + def test_consumption_stays_bounded_and_results_are_complete_and_ordered(self): + n = 20 + window_factor = 2 + n_procs = 2 + window_size = window_factor * n_procs + + jobs = _CountingIterable(n) + results = [] + gen = parallel_map( + _echo_job, jobs, n_procs=n_procs, window_factor=window_factor + ) + for result in gen: + results.append(result) + # +1 buffer: the generator may pull one extra job before it + # can submit-then-harvest on a given step. + self.assertLessEqual(jobs.pulled, window_size + len(results) + 1) + + self.assertEqual(results, list(range(n))) + + +class TestParallelMapShouldStop(unittest.TestCase): + """should_stop lets a caller (e.g. a CLI handling SIGTERM) end a run + early without raising, while still yielding results already + completed in the submission window.""" + + def test_should_stop_ends_iteration_early(self): + n = 20 + jobs = _CountingIterable(n) + + results = list( + parallel_map(_echo_job, jobs, n_procs=2, should_stop=lambda: True) + ) + + self.assertLess(len(results), n) + self.assertLess(jobs.pulled, n) + # Results still seen so far must be a prefix of submission order. + self.assertEqual(results, list(range(len(results)))) + + +class TestParallelMapValidation(unittest.TestCase): + """parallel_map is a reusable helper, so it validates n_procs itself + with a clear message instead of surfacing ProcessPoolExecutor's + max_workers error later - and it must do so eagerly at the call, not + on first iteration of the returned iterator (callers that pass the + iterator elsewhere before consuming it would otherwise see the error + far from the bad argument).""" + + def test_n_procs_below_one_raises_value_error_eagerly(self): + with self.assertRaises(ValueError): + parallel_map(_echo_job, [1, 2], n_procs=0) + + +class TestParallelMapEmptyJobs(unittest.TestCase): + """An empty jobs iterable must return immediately without spawning a + process pool.""" + + def test_empty_jobs_yields_nothing_and_spawns_no_pool(self): + with patch("parsedmarc.parallel.ProcessPoolExecutor") as mock_executor: + results = list(parallel_map(_echo_job, [], n_procs=2)) + self.assertEqual(results, []) + mock_executor.assert_not_called() + + +class TestWorkerLogging(_ParallelTestCase): + """Worker processes must reconstruct the parent parsedmarc logger's + level and FileHandler(s) so records emitted during parsing (e.g. + parse_report_file's "Parsing " debug line) aren't silently + dropped just because they happened in a child process.""" + + def setUp(self): + super().setUp() + from parsedmarc.log import logger as plog + + self._saved_handlers = list(plog.handlers) + self._saved_level = plog.level + + def tearDown(self): + from parsedmarc.log import logger as plog + + for handler in list(plog.handlers): + if handler not in self._saved_handlers: + plog.removeHandler(handler) + if isinstance(handler, logging.FileHandler): + handler.close() + plog.handlers[:] = self._saved_handlers + plog.setLevel(self._saved_level) + super().tearDown() + + def test_worker_debug_log_reaches_parent_file_handler(self): + from parsedmarc.log import configure_logging + + with tempfile.NamedTemporaryFile(suffix=".log", delete=False) as tf: + log_path = tf.name + self.addCleanup(lambda: os.path.exists(log_path) and os.remove(log_path)) + + configure_logging(logging.DEBUG, log_path) + + sample = SAMPLE_PATHS[0] + job = functools.partial(_parse_report_file_job, kwargs=dict(offline=True)) + results = list(parallel_map(job, [sample], n_procs=2)) + + self.assertEqual(len(results), 1) + path, report = results[0] + self.assertEqual(path, sample) + self.assertNotIsInstance(report, Exception) + + with open(log_path) as f: + contents = f.read() + # parse_report_file logs `Parsing {file_path}` at DEBUG + # (parsedmarc/__init__.py) -- this line only appears if the + # worker process's reconstructed logger actually wrote to the + # parent's log file. + self.assertIn(f"Parsing {sample}", contents) + + +class TestInitWorkerLogging(unittest.TestCase): + """_init_worker_logging must be usable directly as a pool initializer: + with no log files it just sets the level (adds a console handler), + and with log files it attaches a FileHandler per path.""" + + def setUp(self): + from parsedmarc.log import logger as plog + + self._saved_handlers = list(plog.handlers) + self._saved_level = plog.level + + def tearDown(self): + from parsedmarc.log import logger as plog + + for handler in list(plog.handlers): + if handler not in self._saved_handlers: + plog.removeHandler(handler) + if isinstance(handler, logging.FileHandler): + handler.close() + plog.handlers[:] = self._saved_handlers + plog.setLevel(self._saved_level) + + def test_no_log_files_sets_level_only(self): + from parsedmarc.log import logger as plog + + _init_worker_logging(logging.WARNING, []) + self.assertEqual(plog.level, logging.WARNING) + + def test_log_files_attach_file_handlers(self): + from parsedmarc.log import logger as plog + + with tempfile.NamedTemporaryFile(suffix=".log", delete=False) as tf: + log_path = tf.name + self.addCleanup(lambda: os.path.exists(log_path) and os.remove(log_path)) + + _init_worker_logging(logging.DEBUG, [log_path]) + + self.assertEqual(plog.level, logging.DEBUG) + file_handlers = [h for h in plog.handlers if isinstance(h, logging.FileHandler)] + self.assertTrue(any(h.baseFilename == log_path for h in file_handlers)) + for h in file_handlers: + h.close() + + +class TestParseReportEmailJob(_ParallelTestCase): + """_parse_report_email_job must return a ParserError as a value (never + raise it) on an invalid message.""" + + def test_invalid_email_returns_parser_error_value(self): + result = _parse_report_email_job( + b"not a valid email", kwargs=dict(offline=True) + ) + self.assertIsInstance(result, parsedmarc.ParserError) + + +if __name__ == "__main__": + unittest.main(verbosity=2)