plm-lottery/app/tx/confirmation.py

import asyncio
import logging
from collections.abc import Awaitable, Callable

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from app.db.base import begin_immediate
from app.db.models import PendingTransaction
from app.electrum.client import ElectrumClient
from app.electrum.scripthash import address_to_scripthash
from app.rounds.events import broadcaster
from app.tx.broadcast import adopt_version
from app.tx.pending_address import own_address_for
from app.tx.versions import txids_by_row

logger = logging.getLogger(__name__)

_POLL_INTERVAL_SECONDS = 10

ConfirmationHandler = Callable[[AsyncSession, PendingTransaction], Awaitable[None]]
_handlers: dict[str, ConfirmationHandler] = {}


def register_handler(kind: str, handler: ConfirmationHandler) -> None:
    """Domain modules (bets, rounds, withdrawals) register here so this generic
    poller can notify them when one of their outgoing txs gets its 1st
    confirmation, without this module importing them directly."""
    _handlers[kind] = handler


async def poll_once(session_factory: async_sessionmaker, client: ElectrumClient) -> int:
    async with session_factory() as session:
        # Plain columns, not entities: nothing then outlives the session, so this
        # can't break if expire_on_commit is ever turned on (B-21).
        candidates = (
            await session.execute(
                select(
                    PendingTransaction.id,
                    PendingTransaction.current_txid,
                    PendingTransaction.kind,
                    PendingTransaction.user_id,
                    PendingTransaction.round_id,
                ).where(PendingTransaction.status == "pending")
            )
        ).all()

        # Resolved once per candidate while the session is still open, and cached
        # by scripthash below — the payout rows of one round share its pool
        # address (B-80), so this also avoids asking the server the same history twice.
        scripthash_by_id: dict[int, str] = {}
        # B-77 (TX-06, TX-19): every txid a row ever had, not only the current one — an
        # earlier version of a bumped transaction may be the one that confirms.
        txids_by_id = await txids_by_row(session, {pending_id: txid for pending_id, txid, *_ in candidates})
        for pending_id, _txid, kind, user_id, round_id in candidates:
            try:
                address = await own_address_for(session, kind, user_id, round_id)
                scripthash_by_id[pending_id] = address_to_scripthash(address)
            except Exception:
                logger.exception("could not derive the address for pending_transaction %s", pending_id)

    confirmed = 0
    history_cache: dict[str, list[dict]] = {}
    for pending_id, txid, kind, _user_id, _round_id in candidates:
        scripthash = scripthash_by_id.get(pending_id)
        if scripthash is None:
            continue  # address derivation failed above; already logged

        try:
            if scripthash not in history_cache:
                history_cache[scripthash] = await client.get_history(scripthash)
        except Exception:
            # One unresolvable scripthash must not stop the others: a tx the server
            # no longer knows about (dropped from the mempool, replaced) used to
            # abort the whole pass via a verbose blockchain.transaction.get call
            # that some servers reject outright (B-41), so nothing confirmed again
            # until an operator intervened (B-03). Abandoning such a row is
            # app/tx/reconcile.py's job, not ours.
            logger.warning("could not fetch history for pending_transaction %s (txid %s)", pending_id, txid, exc_info=True)
            continue

        known_txids = txids_by_id[pending_id]
        # height > 0 means confirmed at that height; 0 or absent means still in
        # the mempool (or the server doesn't know this txid at all yet) — either
        # way, not confirmed, so keep waiting.
        entry = next(
            (e for e in history_cache[scripthash] if e.get("tx_hash") in known_txids and e.get("height", 0) > 0),
            None,
        )
        if entry is None:
            continue

        async with session_factory() as session:
            await begin_immediate(session)  # B-88: read the row and write its outcome under one lock
            row = await session.get(PendingTransaction, pending_id)
            if row is None or row.status != "pending":
                continue
            if entry["tx_hash"] != row.current_txid:
                # The confirmed version is not the current one (B-77, DAT-14): realign
                # the row and everything that recorded its txid *before* the handler runs.
                await adopt_version(session, row, entry["tx_hash"])
            row.status = "confirmed"
            handler = _handlers.get(kind)
            if handler is not None:
                await handler(session, row)
            await session.commit()
        broadcaster.publish()  # a bet/withdrawal/payout just confirmed — balance and/or round state changed
        confirmed += 1

    return confirmed


class ConfirmationPoller:
    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:
        while True:
            client = self._get_client()
            if client is not None:
                try:
                    await poll_once(self._session_factory, client)
                except asyncio.CancelledError:
                    raise
                except Exception:
                    logger.exception("confirmation poll failed")
            await asyncio.sleep(_POLL_INTERVAL_SECONDS)