Spaces:

sitcod3
/

modul-x

Running

modul-x / app /validator /retry_queue.py

Upload folder using huggingface_hub

0f6b827 verified 9 days ago

11 kB

"""Greylist retry queue.
Real mail servers often answer the first RCPT for an unknown sender with a
**4xx** (greylisting / "try again later"). That is not a verdict — it is an
instruction to come back. Without a retry, a perfectly valid mailbox on such a
domain stays permanently "risky", which is exactly the false-negative that makes
a validator feel unreliable.
Two rules keep this safe:
  * **Only 4xx is retried.** A 5xx is final, a 250 is done. A *transport*
    failure (connect timeout / unreachable) is systemic to the MX — re-queuing
    every address behind a dead host would flood the queue for no gain, so the
    MX-level circuit breaker handles that instead.
  * **Bounded attempts with exponential backoff.** A queue that retries forever
    is a queue that never drains; entries are dropped after `max_attempts`.
The store is pluggable (Redis in production, memory in tests) and every store
error degrades to "no queue" rather than breaking a validation run.
"""
from __future__ import annotations

import json
import time
from dataclasses import asdict, dataclass
from typing import Awaitable, Callable, Optional, Protocol

from ..models import SMTPResult


def is_greylist_result(smtp: Optional[SMTPResult]) -> bool:
    """True only for a genuine SMTP 4xx — the retryable case."""
    if smtp is None:
        return False
    code = smtp.smtp_code
    return code is not None and 400 <= code < 500


@dataclass
class RetryEntry:
    email: str
    mx_host: str
    attempts: int = 0
    due_at: float = 0.0
    last_code: Optional[int] = None

    def as_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict) -> "RetryEntry":
        return cls(
            email=data["email"],
            mx_host=data["mx_host"],
            attempts=int(data.get("attempts", 0)),
            due_at=float(data.get("due_at", 0.0)),
            last_code=data.get("last_code"),
        )


class RetryStore(Protocol):
    async def add(self, entry: RetryEntry) -> None: ...
    async def due(self, now: float, limit: int) -> list[RetryEntry]: ...
    async def remove(self, email: str) -> None: ...


class MemoryRetryStore:
    """In-process store. Adequate for a single worker or tests."""

    def __init__(self) -> None:
        self._items: dict[str, RetryEntry] = {}

    @staticmethod
    def _key(email: str) -> str:
        return email.strip().lower()

    async def add(self, entry: RetryEntry) -> None:
        self._items[self._key(entry.email)] = entry

    async def due(self, now: float, limit: int) -> list[RetryEntry]:
        ready = [e for e in self._items.values() if e.due_at <= now]
        ready.sort(key=lambda e: e.due_at)
        return ready[:limit]

    async def remove(self, email: str) -> None:
        self._items.pop(self._key(email), None)

    async def get(self, email: str) -> Optional[RetryEntry]:
        return self._items.get(self._key(email))

    async def count(self) -> int:
        return len(self._items)


class RedisRetryStore:
    """Durable, worker-shared store: a hash of entries plus a sorted-set index.
    The sorted set is scored by `due_at`, so "what is due now" is one
    ZRANGEBYSCORE rather than a full scan of every pending address.
    """

    HASH_KEY = "vx:retry:entries"
    DUE_KEY = "vx:retry:due"

    def __init__(self, client=None):
        self._client = client

    async def _redis(self):
        if self._client is not None:
            return self._client
        from ..utils.cache import get_redis

        self._client = await get_redis()
        return self._client

    @staticmethod
    def _key(email: str) -> str:
        return email.strip().lower()

    async def add(self, entry: RetryEntry) -> None:
        client = await self._redis()
        if client is None:
            raise RuntimeError("redis unavailable")
        key = self._key(entry.email)
        await client.hset(self.HASH_KEY, key, json.dumps(entry.as_dict()))
        await client.zadd(self.DUE_KEY, {key: entry.due_at})

    async def due(self, now: float, limit: int) -> list[RetryEntry]:
        client = await self._redis()
        if client is None:
            raise RuntimeError("redis unavailable")
        members = await client.zrangebyscore(self.DUE_KEY, float("-inf"), now, start=0, num=limit)
        entries: list[RetryEntry] = []
        for member in members:
            raw = await client.hget(self.HASH_KEY, member)
            if not raw:
                # Index without a payload (partial write): drop the orphan.
                await client.zrem(self.DUE_KEY, member)
                continue
            try:
                entries.append(RetryEntry.from_dict(json.loads(raw)))
            except Exception:
                # Corrupt payload must be skipped, never crash the drain.
                await client.zrem(self.DUE_KEY, member)
                await client.hdel(self.HASH_KEY, member)
        return entries

    async def remove(self, email: str) -> None:
        client = await self._redis()
        if client is None:
            raise RuntimeError("redis unavailable")
        key = self._key(email)
        await client.hdel(self.HASH_KEY, key)
        await client.zrem(self.DUE_KEY, key)

    async def get(self, email: str) -> Optional[RetryEntry]:
        client = await self._redis()
        if client is None:
            raise RuntimeError("redis unavailable")
        raw = await client.hget(self.HASH_KEY, self._key(email))
        if not raw:
            return None
        try:
            return RetryEntry.from_dict(json.loads(raw))
        except Exception:
            return None

    async def count(self) -> int:
        client = await self._redis()
        if client is None:
            raise RuntimeError("redis unavailable")
        return int(await client.hlen(self.HASH_KEY) or 0)


def build_retry_store():
    """Redis when configured/reachable, else the in-process store."""
    try:
        from ..utils.cache import _HAS_REDIS  # type: ignore[attr-defined]

        if _HAS_REDIS:
            return RedisRetryStore()
    except Exception:
        pass
    return MemoryRetryStore()


class RetryQueue:
    """Schedules greylisted addresses for a later attempt."""

    def __init__(
        self,
        store: Optional[RetryStore] = None,
        backoff_seconds: int = 900,      # 15 min — inside the usual greylist window
        max_attempts: int = 3,
        now: Optional[Callable[[], float]] = None,
    ):
        self.store = store or MemoryRetryStore()
        # 0 is legitimate (immediate retry — tests and hot-path recovery);
        # the bounded max_attempts prevents any runaway.
        self.backoff_seconds = max(0, backoff_seconds)
        self.max_attempts = max(1, max_attempts)
        self._now = now or time.time

    def _delay_for(self, attempts: int) -> float:
        # attempts=0 -> base, 1 -> 2x base, 2 -> 4x base (capped growth).
        return self.backoff_seconds * (2 ** min(attempts, 4))

    async def enqueue(self, email: str, mx_host: str, code: Optional[int] = None) -> bool:
        """Schedule `email` for a later attempt. Returns False if not queued."""
        try:
            existing = None
            if hasattr(self.store, "get"):
                existing = await self.store.get(email)  # type: ignore[attr-defined]
            if existing is not None:
                return False  # already queued; never create duplicates
            entry = RetryEntry(
                email=email,
                mx_host=mx_host,
                attempts=0,
                due_at=self._now() + self._delay_for(0),
                last_code=code,
            )
            await self.store.add(entry)
            return True
        except Exception:
            return False  # a queue outage must not break validation

    async def due(self, limit: int = 50) -> list[RetryEntry]:
        try:
            return await self.store.due(self._now(), limit)
        except Exception:
            return []

    async def mark_attempted(self, entry: RetryEntry) -> None:
        """Record an attempt; reschedule with backoff or drop when exhausted."""
        try:
            attempts = entry.attempts + 1
            if attempts >= self.max_attempts:
                await self.store.remove(entry.email)
                return
            entry.attempts = attempts
            entry.due_at = self._now() + self._delay_for(attempts)
            await self.store.add(entry)
        except Exception:
            return

    async def resolve(self, entry: RetryEntry) -> None:
        """The address got a definitive answer — stop retrying it."""
        try:
            await self.store.remove(entry.email)
        except Exception:
            return

    async def resolve_email(self, email: str) -> None:
        """Convenience: stop retrying by address (definitive verdict arrived)."""
        try:
            await self.store.remove(email)
        except Exception:
            return

    async def size(self) -> int:
        try:
            if hasattr(self.store, "count"):
                return await self.store.count()  # type: ignore[attr-defined]
            return len(await self.store.due(float("inf"), 10_000))
        except Exception:
            return 0


async def drain_retry_queue(
    prober: Callable[[str, str], Awaitable[SMTPResult]],
    queue: Optional[RetryQueue] = None,
    limit: int = 50,
) -> dict:
    """Re-probe addresses that were greylisted earlier.
    A queue that only fills is a leak, so this is the other half of the feature.
    For each due entry:
      * definitive answer (250 or 5xx) -> drop it from the queue
      * still 4xx                       -> re-queue with exponential backoff
      * probe raised                    -> leave queued for a later attempt
    Returns a stats dict including the raw results so the caller can persist them.
    """
    q = queue or retry_queue
    stats: dict = {"checked": 0, "resolved": 0, "requeued": 0, "errors": 0, "results": []}

    for entry in await q.due(limit=limit):
        stats["checked"] += 1
        try:
            result = await prober(entry.email, entry.mx_host)
        except Exception:
            # Transport-level blow-up: keep it queued, try again next drain.
            stats["errors"] += 1
            await q.mark_attempted(entry)
            continue

        stats["results"].append((entry.email, result))

        if is_greylist_result(result):
            stats["requeued"] += 1
            await q.mark_attempted(entry)
        else:
            # 250 or 5xx — a real answer, so stop retrying this address.
            stats["resolved"] += 1
            await q.resolve(entry)

    return stats


# Shared queue used by the bulk pipeline. Prefers Redis so the queue survives a
# restart and is shared across workers; falls back to memory otherwise.
retry_queue = RetryQueue(store=build_retry_store())

	"""Greylist retry queue.

	Real mail servers often answer the first RCPT for an unknown sender with a
	4xx (greylisting / "try again later"). That is not a verdict — it is an
	instruction to come back. Without a retry, a perfectly valid mailbox on such a
	domain stays permanently "risky", which is exactly the false-negative that makes
	a validator feel unreliable.

	Two rules keep this safe:

	* Only 4xx is retried. A 5xx is final, a 250 is done. A transport
	failure (connect timeout / unreachable) is systemic to the MX — re-queuing
	every address behind a dead host would flood the queue for no gain, so the
	MX-level circuit breaker handles that instead.
	* Bounded attempts with exponential backoff. A queue that retries forever
	is a queue that never drains; entries are dropped after `max_attempts`.

	The store is pluggable (Redis in production, memory in tests) and every store
	error degrades to "no queue" rather than breaking a validation run.
	"""
	from __future__ import annotations

	import json
	import time
	from dataclasses import asdict, dataclass
	from typing import Awaitable, Callable, Optional, Protocol

	from ..models import SMTPResult


	def is_greylist_result(smtp: Optional[SMTPResult]) -> bool:
	"""True only for a genuine SMTP 4xx — the retryable case."""
	if smtp is None:
	return False
	code = smtp.smtp_code
	return code is not None and 400 <= code < 500


	@dataclass
	class RetryEntry:
	email: str
	mx_host: str
	attempts: int = 0
	due_at: float = 0.0
	last_code: Optional[int] = None

	def as_dict(self) -> dict:
	return asdict(self)

	@classmethod
	def from_dict(cls, data: dict) -> "RetryEntry":
	return cls(
	email=data["email"],
	mx_host=data["mx_host"],
	attempts=int(data.get("attempts", 0)),
	due_at=float(data.get("due_at", 0.0)),
	last_code=data.get("last_code"),
	)


	class RetryStore(Protocol):
	async def add(self, entry: RetryEntry) -> None: ...
	async def due(self, now: float, limit: int) -> list[RetryEntry]: ...
	async def remove(self, email: str) -> None: ...


	class MemoryRetryStore:
	"""In-process store. Adequate for a single worker or tests."""

	def __init__(self) -> None:
	self._items: dict[str, RetryEntry] = {}

	@staticmethod
	def _key(email: str) -> str:
	return email.strip().lower()

	async def add(self, entry: RetryEntry) -> None:
	self._items[self._key(entry.email)] = entry

	async def due(self, now: float, limit: int) -> list[RetryEntry]:
	ready = [e for e in self._items.values() if e.due_at <= now]
	ready.sort(key=lambda e: e.due_at)
	return ready[:limit]

	async def remove(self, email: str) -> None:
	self._items.pop(self._key(email), None)

	async def get(self, email: str) -> Optional[RetryEntry]:
	return self._items.get(self._key(email))

	async def count(self) -> int:
	return len(self._items)


	class RedisRetryStore:
	"""Durable, worker-shared store: a hash of entries plus a sorted-set index.

	The sorted set is scored by `due_at`, so "what is due now" is one
	ZRANGEBYSCORE rather than a full scan of every pending address.
	"""

	HASH_KEY = "vx:retry:entries"
	DUE_KEY = "vx:retry:due"

	def __init__(self, client=None):
	self._client = client

	async def _redis(self):
	if self._client is not None:
	return self._client
	from ..utils.cache import get_redis

	self._client = await get_redis()
	return self._client

	@staticmethod
	def _key(email: str) -> str:
	return email.strip().lower()

	async def add(self, entry: RetryEntry) -> None:
	client = await self._redis()
	if client is None:
	raise RuntimeError("redis unavailable")
	key = self._key(entry.email)
	await client.hset(self.HASH_KEY, key, json.dumps(entry.as_dict()))
	await client.zadd(self.DUE_KEY, {key: entry.due_at})

	async def due(self, now: float, limit: int) -> list[RetryEntry]:
	client = await self._redis()
	if client is None:
	raise RuntimeError("redis unavailable")
	members = await client.zrangebyscore(self.DUE_KEY, float("-inf"), now, start=0, num=limit)
	entries: list[RetryEntry] = []
	for member in members:
	raw = await client.hget(self.HASH_KEY, member)
	if not raw:
	# Index without a payload (partial write): drop the orphan.
	await client.zrem(self.DUE_KEY, member)
	continue
	try:
	entries.append(RetryEntry.from_dict(json.loads(raw)))
	except Exception:
	# Corrupt payload must be skipped, never crash the drain.
	await client.zrem(self.DUE_KEY, member)
	await client.hdel(self.HASH_KEY, member)
	return entries

	async def remove(self, email: str) -> None:
	client = await self._redis()
	if client is None:
	raise RuntimeError("redis unavailable")
	key = self._key(email)
	await client.hdel(self.HASH_KEY, key)
	await client.zrem(self.DUE_KEY, key)

	async def get(self, email: str) -> Optional[RetryEntry]:
	client = await self._redis()
	if client is None:
	raise RuntimeError("redis unavailable")
	raw = await client.hget(self.HASH_KEY, self._key(email))
	if not raw:
	return None
	try:
	return RetryEntry.from_dict(json.loads(raw))
	except Exception:
	return None

	async def count(self) -> int:
	client = await self._redis()
	if client is None:
	raise RuntimeError("redis unavailable")
	return int(await client.hlen(self.HASH_KEY) or 0)


	def build_retry_store():
	"""Redis when configured/reachable, else the in-process store."""
	try:
	from ..utils.cache import _HAS_REDIS # type: ignore[attr-defined]

	if _HAS_REDIS:
	return RedisRetryStore()
	except Exception:
	pass
	return MemoryRetryStore()


	class RetryQueue:
	"""Schedules greylisted addresses for a later attempt."""

	def __init__(
	self,
	store: Optional[RetryStore] = None,
	backoff_seconds: int = 900, # 15 min — inside the usual greylist window
	max_attempts: int = 3,
	now: Optional[Callable[[], float]] = None,
	):
	self.store = store or MemoryRetryStore()
	# 0 is legitimate (immediate retry — tests and hot-path recovery);
	# the bounded max_attempts prevents any runaway.
	self.backoff_seconds = max(0, backoff_seconds)
	self.max_attempts = max(1, max_attempts)
	self._now = now or time.time

	def _delay_for(self, attempts: int) -> float:
	# attempts=0 -> base, 1 -> 2x base, 2 -> 4x base (capped growth).
	return self.backoff_seconds * (2 ** min(attempts, 4))

	async def enqueue(self, email: str, mx_host: str, code: Optional[int] = None) -> bool:
	"""Schedule `email` for a later attempt. Returns False if not queued."""
	try:
	existing = None
	if hasattr(self.store, "get"):
	existing = await self.store.get(email) # type: ignore[attr-defined]
	if existing is not None:
	return False # already queued; never create duplicates
	entry = RetryEntry(
	email=email,
	mx_host=mx_host,
	attempts=0,
	due_at=self._now() + self._delay_for(0),
	last_code=code,
	)
	await self.store.add(entry)
	return True
	except Exception:
	return False # a queue outage must not break validation

	async def due(self, limit: int = 50) -> list[RetryEntry]:
	try:
	return await self.store.due(self._now(), limit)
	except Exception:
	return []

	async def mark_attempted(self, entry: RetryEntry) -> None:
	"""Record an attempt; reschedule with backoff or drop when exhausted."""
	try:
	attempts = entry.attempts + 1
	if attempts >= self.max_attempts:
	await self.store.remove(entry.email)
	return
	entry.attempts = attempts
	entry.due_at = self._now() + self._delay_for(attempts)
	await self.store.add(entry)
	except Exception:
	return

	async def resolve(self, entry: RetryEntry) -> None:
	"""The address got a definitive answer — stop retrying it."""
	try:
	await self.store.remove(entry.email)
	except Exception:
	return

	async def resolve_email(self, email: str) -> None:
	"""Convenience: stop retrying by address (definitive verdict arrived)."""
	try:
	await self.store.remove(email)
	except Exception:
	return

	async def size(self) -> int:
	try:
	if hasattr(self.store, "count"):
	return await self.store.count() # type: ignore[attr-defined]
	return len(await self.store.due(float("inf"), 10_000))
	except Exception:
	return 0


	async def drain_retry_queue(
	prober: Callable[[str, str], Awaitable[SMTPResult]],
	queue: Optional[RetryQueue] = None,
	limit: int = 50,
	) -> dict:
	"""Re-probe addresses that were greylisted earlier.

	A queue that only fills is a leak, so this is the other half of the feature.
	For each due entry:

	* definitive answer (250 or 5xx) -> drop it from the queue
	* still 4xx -> re-queue with exponential backoff
	* probe raised -> leave queued for a later attempt

	Returns a stats dict including the raw results so the caller can persist them.
	"""
	q = queue or retry_queue
	stats: dict = {"checked": 0, "resolved": 0, "requeued": 0, "errors": 0, "results": []}

	for entry in await q.due(limit=limit):
	stats["checked"] += 1
	try:
	result = await prober(entry.email, entry.mx_host)
	except Exception:
	# Transport-level blow-up: keep it queued, try again next drain.
	stats["errors"] += 1
	await q.mark_attempted(entry)
	continue

	stats["results"].append((entry.email, result))

	if is_greylist_result(result):
	stats["requeued"] += 1
	await q.mark_attempted(entry)
	else:
	# 250 or 5xx — a real answer, so stop retrying this address.
	stats["resolved"] += 1
	await q.resolve(entry)

	return stats


	# Shared queue used by the bulk pipeline. Prefers Redis so the queue survives a
	# restart and is shared across workers; falls back to memory otherwise.
	retry_queue = RetryQueue(store=build_retry_store())