simbarag/blueprints/scheduled_messages/scheduler.py

import logging
from datetime import datetime, timezone

from dateutil.relativedelta import relativedelta
from tortoise.transactions import in_transaction

from . import delivery as transport
from .models import ScheduledMessage, MessageStatus, Recurrence, JobDelivery

logger = logging.getLogger(__name__)
RECURRENCE_DELTAS = {
    Recurrence.DAILY: relativedelta(days=1),
    Recurrence.WEEKLY: relativedelta(weeks=1),
    Recurrence.MONTHLY: relativedelta(months=1),
}


def run_id(msg):
    # Each recurring occurrence has its own row; edits keep the same identity.
    return str(msg.id)


async def _schedule_next_occurrence(msg, db=None):
    if db is None:
        async with in_transaction() as connection:
            return await _schedule_next_occurrence(msg, connection)
    delta = RECURRENCE_DELTAS.get(msg.recurrence)
    if not delta:
        return
    series = msg.series_id or msg.id
    # Serialize series creation/cancellation on the original occurrence.
    await (
        ScheduledMessage.filter(id=series)
        .using_db(db)
        .update(updated_at=transport.now())
    )
    if (
        await ScheduledMessage.filter(series_id=series, status=MessageStatus.CANCELLED)
        .using_db(db)
        .exists()
    ):
        return
    next_at = msg.scheduled_at + delta
    now = datetime.now(timezone.utc)
    while next_at <= now:
        next_at += delta
    await ScheduledMessage.get_or_create(
        previous_occurrence_id=msg.id,
        defaults=dict(
            series_id=series,
            recipient=msg.recipient,
            channel=msg.channel,
            content=msg.content,
            subject=msg.subject,
            scheduled_at=next_at,
            recurrence=msg.recurrence,
            use_agent=msg.use_agent,
            created_by_id=msg.created_by_id,
        ),
        using_db=db,
    )


async def process_message(msg, *, retry=False):
    async with in_transaction() as db:
        # Always lock the series before an occurrence, matching cancellation.
        await (
            ScheduledMessage.filter(id=msg.series_id or msg.id)
            .using_db(db)
            .update(updated_at=transport.now())
        )
        eligible = ScheduledMessage.filter(
            id=msg.id, status=MessageStatus.FAILED if retry else MessageStatus.PENDING
        )
        if not retry:
            eligible = eligible.filter(scheduled_at__lte=transport.now())
        won = await eligible.using_db(db).update(
            status=MessageStatus.PROCESSING, series_id=msg.series_id or msg.id
        )
        if not won:
            return "skipped"
        msg = await ScheduledMessage.get(id=msg.id).using_db(db)
        claimed = await transport.claim("scheduled", msg.id, run_id(msg), retry=retry)
        if not claimed:
            # Preserve the authoritative delivery state if the job was rescheduled
            # onto an occurrence already delivered.
            d = await JobDelivery.get(
                kind="scheduled", job_id=str(msg.id), run_id=run_id(msg)
            )
            status = (
                MessageStatus(d.status)
                if d.status in {s.value for s in MessageStatus}
                else MessageStatus.UNCERTAIN
            )
            await ScheduledMessage.filter(id=msg.id).using_db(db).update(status=status)
            return "skipped"
    sending = False
    try:
        content = msg.content
        if msg.use_agent:
            user = await msg.created_by
            content = await transport.agent_content(
                user, content, f"scheduled-{claimed.id}", "scheduled"
            )
        # A cancellation which wins before this transaction prevents sending.
        async with in_transaction() as db:
            await (
                ScheduledMessage.filter(id=msg.series_id or msg.id)
                .using_db(db)
                .update(updated_at=transport.now())
            )
            active = (
                await ScheduledMessage.filter(
                    id=msg.id, status=MessageStatus.PROCESSING
                )
                .using_db(db)
                .update(updated_at=transport.now())
            )
            if not active:
                await transport.transition(claimed, "cancelled")
                return "cancelled"
            if not await transport.transition(claimed, "sending"):
                return "skipped"
        sending = True
        await transport.send(msg.channel, msg.recipient, content, msg.subject)
        async with in_transaction() as db:
            await (
                ScheduledMessage.filter(id=msg.series_id or msg.id)
                .using_db(db)
                .update(updated_at=transport.now())
            )
            await transport.transition(claimed, "sent")
            active = (
                await ScheduledMessage.filter(
                    id=msg.id, status=MessageStatus.PROCESSING
                )
                .using_db(db)
                .update(status=MessageStatus.SENT, error_message=None)
            )
            if active:
                await _schedule_next_occurrence(msg, db)
        return "sent"
    except Exception as exc:
        status = (
            "uncertain"
            if sending and not isinstance(exc, transport.SafeDeliveryError)
            else "failed"
        )
        await transport.transition(claimed, status, str(exc))
        await ScheduledMessage.filter(
            id=msg.id, status=MessageStatus.PROCESSING
        ).update(status=MessageStatus(status), error_message=str(exc))
        logger.exception("Scheduled delivery %s ended as %s", msg.id, status)
        return status


async def process_due_messages() -> dict:
    await transport.expire_claims("scheduled")
    uncertain = await JobDelivery.filter(kind="scheduled", status="uncertain")
    for d in uncertain:
        await ScheduledMessage.filter(
            id=d.job_id, status=MessageStatus.PROCESSING
        ).update(status=MessageStatus.UNCERTAIN, error_message=d.error)
    # Recover recurrence creation if an earlier successful send was committed.
    # Only rows sent through delivery claims qualify: rows sent before migration
    # 20 have no successor links, and repairing them would fork each historical
    # occurrence into its own new series.
    orphans = await ScheduledMessage.filter(
        status=MessageStatus.SENT,
        deleted_at__isnull=True,
        next_occurrence__id__isnull=True,
    ).exclude(recurrence=Recurrence.NONE)
    if orphans:
        claimed_sends = set(
            await JobDelivery.filter(
                kind="scheduled",
                status="sent",
                job_id__in=[str(m.id) for m in orphans],
            ).values_list("job_id", flat=True)
        )
        for msg in orphans:
            if str(msg.id) in claimed_sends:
                await _schedule_next_occurrence(msg)
    due = await ScheduledMessage.filter(
        status=MessageStatus.PENDING, scheduled_at__lte=transport.now()
    )
    counts = {"sent": 0, "failed": 0, "uncertain": 0}
    for msg in due:
        result = await process_message(msg)
        if result in counts:
            counts[result] += 1
    return counts