plm-lottery/app/tx/reconcile.py

"""Resolves in-flight transactions against the chain.

Everything else in this codebase assumes a broadcast either confirms or gets
fee-bumped until it does. Neither is guaranteed: an RBF bump raises RbfError
whenever there's no change output big enough to absorb it (see tx/broadcast.py),
a node can drop a low-fee tx from its mempool, and the process can die between
building a transaction and broadcasting it. Without this module those cases were
permanent: `spent_txid` was set at build time and never cleared, so the coins
stayed spendable on-chain while the database considered them gone — the user's
balance simply lost them, with no path back short of editing the DB by hand
(B-04, and the "building" half of B-08).

What it does, per PendingTransaction that isn't already terminal:

  * status "building" — we crashed (or were killed) between writing the row and
    broadcasting. Ask the chain: if the tx is there after all, promote everything
    to its live state; if it isn't, release the UTXOs and undo the intent.
  * status "pending" — broadcast, still unconfirmed. Left alone until it has been
    unconfirmed for `_ABANDON_AFTER_SECONDS`, since absence from one server's
    mempool is not proof of death; only then is it abandoned like the above.

Deliberately conservative: it never touches a tx the chain knows about, and the
grace period is long (multiples of the RBF timeout) so a slow-but-alive tx is
bumped by RbfBumper rather than abandoned here.
"""

import asyncio
import logging
from collections.abc import Callable
from datetime import datetime, timedelta, timezone

from embit.transaction import Transaction
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from app.audit.log import write_audit_log
from app.db.models import PendingTransaction, Round, RoundParticipant, UtxoEvent, Withdrawal
from app.electrum.client import ElectrumClient
from app.rounds.events import broadcaster
from app.wallet.balance import recompute_balance

logger = logging.getLogger(__name__)

_POLL_INTERVAL_SECONDS = 120

# A "building" row means we never got confirmation that the broadcast happened, so
# it only needs long enough to rule out a request still in flight.
_BUILDING_GRACE_SECONDS = 120

# A "pending" row was accepted by a node once. Give it a wide margin — the RBF
# bumper gets several attempts inside this window — before concluding it's gone.
_ABANDON_AFTER_SECONDS = 6 * 60 * 60


class PendingTransactionReconciler:
    def __init__(self, session_factory: async_sessionmaker, get_client: Callable[[], ElectrumClient | None]):
        self._session_factory = session_factory
        self._get_client = get_client

    async def run(self) -> None:
        # Runs once promptly at startup: a crash mid-broadcast is exactly the case
        # that leaves a "building" row, and the restart is when we can clear it.
        while True:
            client = self._get_client()
            if client is not None:
                try:
                    await reconcile_once(self._session_factory, client)
                except asyncio.CancelledError:
                    raise
                except Exception:
                    logger.exception("pending-transaction reconciliation failed")
            await asyncio.sleep(_POLL_INTERVAL_SECONDS)


async def _tx_exists_on_chain(client: ElectrumClient, txid: str) -> bool:
    """True if the server knows this txid at all (mempool or mined). An error reply
    means "unknown", which is the answer we're looking for; a transport failure is
    *not* — that raises, and the caller leaves the row alone until next time."""
    try:
        tx = await client.get_transaction(txid, verbose=True)
    except Exception as exc:
        message = str(exc).lower()
        if "missing" in message or "not found" in message or "no such" in message or "unknown" in message:
            return False
        raise
    return bool(tx)


async def reconcile_once(session_factory: async_sessionmaker, client: ElectrumClient) -> int:
    """Returns how many rows were resolved (promoted or abandoned)."""
    now = datetime.now(timezone.utc)
    async with session_factory() as session:
        rows = (
            await session.scalars(
                select(PendingTransaction).where(PendingTransaction.status.in_(("building", "pending")))
            )
        ).all()
        candidates = [
            (row.id, row.status, row.current_txid)
            for row in rows
            if _is_due(row, now)
        ]

    resolved = 0
    for row_id, status, txid in candidates:
        try:
            exists = await _tx_exists_on_chain(client, txid)
        except Exception:
            # Transport/server problem — say nothing about this tx and try again on
            # the next pass rather than abandoning a tx that may be perfectly alive.
            logger.warning("could not check pending_transaction %s (txid %s) against the chain", row_id, txid)
            continue

        async with session_factory() as session:
            row = await session.get(PendingTransaction, row_id)
            if row is None or row.status != status:
                continue  # something else moved it while we were asking
            if exists:
                if row.status == "building":
                    await _promote(session, row)
                    resolved += 1
            else:
                await _abandon(session, row, "not found on chain")
                resolved += 1
            await session.commit()
        broadcaster.publish()

    return resolved


def _is_due(row: PendingTransaction, now: datetime) -> bool:
    grace = _BUILDING_GRACE_SECONDS if row.status == "building" else _ABANDON_AFTER_SECONDS
    return now >= row.broadcast_at.replace(tzinfo=timezone.utc) + timedelta(seconds=grace)


async def _promote(session: AsyncSession, row: PendingTransaction) -> None:
    """The tx did make it onto the chain before we died — finish what phase 2 of
    place_bet/request_withdrawal would have done."""
    row.status = "pending"
    if row.kind == "bet":
        participant = await session.scalar(
            select(RoundParticipant).where(
                RoundParticipant.round_id == row.round_id, RoundParticipant.user_id == row.user_id
            )
        )
        if participant is not None and participant.status == "building":
            participant.status = "broadcast"
    elif row.kind == "withdrawal" and row.withdrawal_id is not None:
        withdrawal = await session.get(Withdrawal, row.withdrawal_id)
        if withdrawal is not None and withdrawal.status == "building":
            withdrawal.status = "broadcast"
    await write_audit_log(
        session,
        "pending_tx_recovered",
        {"pending_transaction_id": row.id, "kind": row.kind, "txid": row.current_txid},
        user_id=row.user_id,
        round_id=row.round_id,
    )
    logger.info("recovered %s pending_transaction %s: tx %s is on-chain", row.kind, row.id, row.current_txid)


async def _abandon(session: AsyncSession, row: PendingTransaction, reason: str) -> None:
    """The tx is gone for good. Release whatever it reserved so the funds come back,
    and roll the domain row back to something truthful."""
    row.status = "failed"
    row.failure_reason = reason[:128]

    released = await _release_inputs(session, row)

    if row.kind == "bet":
        participant = await session.scalar(
            select(RoundParticipant).where(
                RoundParticipant.round_id == row.round_id, RoundParticipant.user_id == row.user_id
            )
        )
        if participant is not None and participant.status in ("building", "broadcast"):
            # The bet never happened, so the user is not in this round. Removing the
            # row also unblocks the scheduler, which waits for every non-confirmed
            # participant before closing the round.
            await session.delete(participant)
    elif row.kind == "withdrawal" and row.withdrawal_id is not None:
        withdrawal = await session.get(Withdrawal, row.withdrawal_id)
        if withdrawal is not None and withdrawal.status in ("building", "broadcast"):
            withdrawal.status = "failed"
            withdrawal.txid = None
    elif row.kind == "payout":
        # Payout funds come from the pool address, which isn't tracked in
        # utxo_events, so there's nothing to release. Clearing payout_txid leaves the
        # round in "paying_out" with no tx attached, which is the state an operator
        # (or a future payout-retry routine — still an open gap) can act on.
        round_ = await session.get(Round, row.round_id) if row.round_id else None
        if round_ is not None and round_.status == "paying_out":
            round_.payout_txid = None
            logger.error(
                "round %s payout tx %s vanished — round needs operator attention", round_.id, row.current_txid
            )

    if row.user_id is not None:
        await recompute_balance(session, row.user_id)

    await write_audit_log(
        session,
        "pending_tx_abandoned",
        {
            "pending_transaction_id": row.id,
            "kind": row.kind,
            "txid": row.current_txid,
            "reason": reason,
            "utxos_released": released,
        },
        user_id=row.user_id,
        round_id=row.round_id,
    )
    logger.warning(
        "abandoned %s pending_transaction %s (txid %s): %s — released %s UTXO(s)",
        row.kind,
        row.id,
        row.current_txid,
        reason,
        released,
    )


async def _release_inputs(session: AsyncSession, row: PendingTransaction) -> int:
    """Clear spent_txid on every UTXO this transaction consumed, so the balance
    counts them again. The inputs come from the stored raw tx, which is kept current
    across RBF bumps, so this works for a bumped tx too."""
    try:
        tx = Transaction.parse(bytes.fromhex(row.raw_tx_hex))
    except Exception:
        logger.exception("could not parse raw tx of pending_transaction %s; inputs not released", row.id)
        return 0

    released = 0
    for vin in tx.vin:
        utxo = await session.scalar(
            select(UtxoEvent).where(UtxoEvent.txid == vin.txid.hex(), UtxoEvent.vout == vin.vout)
        )
        # Only release what this tx actually reserved: if another tx has since spent
        # the same UTXO, its claim is the live one and must not be cleared.
        if utxo is not None and utxo.spent_txid == row.current_txid:
            utxo.spent_txid = None
            released += 1
    return released