plm-lottery/app/tx/confirmation.py

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

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


class ConfirmationProver(Protocol):
    """The part of ElectrumListener a confirmation needs (B-83): the quorum's verdict on
    a merkle proof of `txid` at `height`."""

    async def confirmation_proven(self, txid: str, height: int) -> bool: ...


async def poll_once(session_factory: async_sessionmaker, client: ElectrumClient, network: ConfirmationProver | None) -> int:
    """One pass over the `pending` rows. The active server's history only nominates a
    candidate: a row is confirmed when the quorum proves it with a merkle proof
    (B-83, TX-22). With no `network` to ask nothing can be confirmed — the active
    server's word alone is exactly what a hostile one would use to declare a payout that
    was never broadcast confirmed."""
    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

        # B-83 (TX-22, ELC-26): the height is a claim of the active server. The quorum has
        # to return a valid merkle proof of the txid against that block's header; without
        # one the row stays pending and is looked at again on the next pass.
        if network is None or not await network.confirmation_proven(entry["tx_hash"], entry["height"]):
            logger.warning(
                "pending_transaction %s: %s reported at height %s but the quorum did not prove it — still pending",
                pending_id,
                entry["tx_hash"],
                entry["height"],
            )
            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],
        network: ConfirmationProver | None = None,
    ):
        self._session_factory = session_factory
        self._get_client = get_client
        self._network = network  # the ElectrumListener, which holds the headers and asks the quorum

    async def run(self) -> None:
        while True:
            client = self._get_client()
            if client is not None:
                try:
                    await poll_once(self._session_factory, client, self._network)
                except asyncio.CancelledError:
                    raise
                except Exception:
                    logger.exception("confirmation poll failed")
            await asyncio.sleep(_POLL_INTERVAL_SECONDS)