mirror of
https://github.com/domainaware/parsedmarc.git
synced 2026-08-02 05:32:17 +00:00
* Extend n_procs parallel parsing to mbox and mailbox sources (#147) n_procs previously applied only to report files passed directly as CLI arguments; messages from mbox files and mailbox connections (IMAP, Microsoft Graph, Gmail API, Maildir) were always parsed sequentially. A new parsedmarc.parallel module provides a shared bounded-window ProcessPoolExecutor helper (parallel_map) used by all three input paths. Only parsing fans out to a reused worker pool; message fetching, report deduplication, mailbox archiving/deletion, and output stay sequential in the main process. The submission window keeps at most ~2*n_procs messages in flight, so memory stays bounded even for huge mboxes, and the mailbox path fetches messages lazily on the connection-owning main thread. keep_alive never crosses the process boundary - the main process sends periodic IMAP keepalives while workers parse - and with n_procs > 1, invalid-message disposition happens after the parse phase, mirroring the existing deferred bulk archive moves. get_dmarc_reports_from_mbox, get_dmarc_reports_from_mailbox (including its tail-recursive re-check), and watch_inbox gain an n_procs keyword argument (default 1); sequential behavior at the default is unchanged. Replacing the CLI's hand-rolled Pipe/Process batching also fixes two defects in the direct-file path: a child process that died from a non-ParserError exception left the parent blocked forever on conn.recv(), and the hard batch barrier let one slow file idle every other worker slot. Workers are now a reused pool (no fresh interpreter per file), with worker logging reconstructed via a spawn-safe pool initializer instead of fork inheritance. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Address Copilot and code-quality review feedback on #849 - parallel_map now validates n_procs >= 1 itself with a clear error instead of surfacing ProcessPoolExecutor's max_workers error later. The check raises eagerly at the call (the generator body moved into an inner function) rather than on first iteration, with a regression test. - Aligned parallel_map's should_stop docstring with the implementation: queued-but-unstarted jobs are cancelled, while in-flight jobs are waited on and their results yielded, so the stop can block briefly but never discards completed work. - The parallel mailbox path keeps fetched message ids in a deque popped as each in-order result arrives, so the id queue stays bounded by the submission window instead of growing to message_limit. - Closed the three sample-file handles the new tests opened without a context manager. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Address second round of Copilot feedback on #849 - configure_logging no longer stacks duplicate FileHandlers when called again with the same log_file (compared by FileHandler.baseFilename, which stores the absolute path): a duplicate wrote every record twice and leaked a file descriptor per call, e.g. across SIGHUP config reloads. Latent in the pre-extraction cli._configure_logging too. Regression tests in the new tests/test_log.py. - Renamed the CHANGELOG's premature "10.4.0" heading to "Unreleased", matching the project convention where in-progress entries accumulate under Unreleased and the release PR renames the section and bumps parsedmarc/constants.py together (as in the 10.3.0 release). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Fable 5
parent
c95d4666dc
commit
48445c639e
@@ -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 <path>` 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
|
||||
|
||||
+15
-8
@@ -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.
|
||||
|
||||
+280
-102
@@ -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,
|
||||
|
||||
+45
-167
@@ -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__()}")
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -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()
|
||||
+105
-100
@@ -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)
|
||||
|
||||
+234
-8
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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 <path>" 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)
|
||||
Reference in New Issue
Block a user