diff --git a/lnbits/app.py b/lnbits/app.py index c4c2d5651..bf3617c43 100644 --- a/lnbits/app.py +++ b/lnbits/app.py @@ -477,7 +477,6 @@ async def check_and_register_extensions(app: FastAPI) -> None: def register_async_tasks() -> None: - task_manager.init() # listen to all incoming payments and dispatch payment notifications diff --git a/lnbits/core/views/blockexplorer_api.py b/lnbits/core/views/blockexplorer_api.py index 838c66354..92ed728dc 100644 --- a/lnbits/core/views/blockexplorer_api.py +++ b/lnbits/core/views/blockexplorer_api.py @@ -8,19 +8,19 @@ from pydantic.types import UUID4 from lnbits.decorators import check_access_token, check_user_exists from lnbits.settings import settings +from lnbits.task_manager import OnchainAddressEvent, OnchainTxEvent, task_manager from lnbits.utils.electrum import ( AddressResponse, + Balance, BlockHeader, BlockInfo, ElectrumClient, ElectrumError, FeeResponse, - HistoryEntry, Transaction, parse_block_header, parse_raw_tx, scripthash_from_address, - scripthash_from_scriptpubkey, ) blockexplorer_router = APIRouter( @@ -51,46 +51,6 @@ def _client() -> ElectrumClient: return ElectrumClient(settings.lnbits_blockexplorer_electrum_url) -def _watch_scripthash(tx: Transaction) -> str | None: - """Return scripthash of first spendable output to watch for tx confirmation.""" - for out in tx.vout: - if out.scriptPubKey.type != "nulldata": - return scripthash_from_scriptpubkey(bytes.fromhex(out.scriptPubKey.hex)) - return None - - -async def _tx_status( - c: ElectrumClient, txid: str, scripthash: str | None -) -> dict[str, Any]: - if scripthash: - try: - history: list[HistoryEntry] = await c.get_history(scripthash) - for entry in history: - if entry.tx_hash == txid: - return { - "txid": txid, - "confirmed": entry.height > 0, - "height": entry.height if entry.height > 0 else None, - "fee": entry.fee, - } - except ElectrumError: - # History too large; check mempool to determine confirmation status. - try: - mempool = await c.get_mempool(scripthash) - for mem in mempool: - if mem.tx_hash == txid: - return { - "txid": txid, - "confirmed": False, - "height": None, - "fee": mem.fee, - } - return {"txid": txid, "confirmed": True, "height": None, "fee": None} - except ElectrumError: - pass - return {"txid": txid, "confirmed": False, "height": None, "fee": None} - - # ---- REST ---- @@ -213,63 +173,61 @@ async def ws_blocks(websocket: WebSocket) -> None: logger.debug(f"ws_blocks error: {e}") +def _address_event_to_response(event: OnchainAddressEvent) -> AddressResponse: + return AddressResponse( + balance=Balance(confirmed=event.confirmed, unconfirmed=event.unconfirmed), + history=event.history, + history_error=event.history_error, + ) + + @blockexplorer_router.websocket("/ws/address/{address}") async def ws_address(websocket: WebSocket, address: str) -> None: if not settings.lnbits_blockexplorer_enabled: await websocket.close(code=1008) return try: - scripthash = scripthash_from_address(address) + scripthash_from_address(address) # validate except ValueError as e: await websocket.close(code=1008, reason=str(e)) return await websocket.accept() + + queue: asyncio.Queue[OnchainAddressEvent] = asyncio.Queue() + task_manager.register_ws_address_queue(address, queue) try: - async with _client() as c: - await c.subscribe_scripthash(scripthash) - balance_res, history_res = await asyncio.gather( - c.get_balance(scripthash), - c.get_history(scripthash), - return_exceptions=True, - ) - if isinstance(balance_res, BaseException): - await websocket.send_json({"error": str(balance_res)}) - return - history = [] if isinstance(history_res, BaseException) else history_res - history_error = ( - str(history_res) if isinstance(history_res, BaseException) else None - ) - resp = AddressResponse( - balance=balance_res, history=history, history_error=history_error - ) - await websocket.send_json(resp.dict()) + await _ws_address_loop(websocket, address, queue) + finally: + task_manager.unregister_ws_address_queue(address, queue) - async def on_address_change(params: list[Any]) -> None: - try: - bal_r, hist_r = await asyncio.gather( - c.get_balance(scripthash), - c.get_history(scripthash), - return_exceptions=True, - ) - if isinstance(bal_r, BaseException): - return - hist = [] if isinstance(hist_r, BaseException) else hist_r - h_err = str(hist_r) if isinstance(hist_r, BaseException) else None - await websocket.send_json( - AddressResponse( - balance=bal_r, history=hist, history_error=h_err - ).dict() - ) - except Exception as e: - logger.debug(f"ws_address send error: {e}") - c.on("blockchain.scripthash.subscribe", on_address_change) - while True: - msg = await websocket.receive() - if msg["type"] == "websocket.disconnect": +async def _ws_address_loop( + websocket: WebSocket, + address: str, + queue: asyncio.Queue[OnchainAddressEvent], +) -> None: + try: + while True: + recv_task = asyncio.create_task(websocket.receive()) + event_task = asyncio.create_task(queue.get()) + done, pending = await asyncio.wait( + [recv_task, event_task], return_when=asyncio.FIRST_COMPLETED + ) + for t in pending: + t.cancel() + if recv_task in done: + if recv_task.result().get("type") == "websocket.disconnect": break - except Exception as e: - logger.debug(f"ws_address error: {e}") + if event_task in done: + try: + await websocket.send_json( + _address_event_to_response(event_task.result()).dict() + ) + except Exception as exc: + logger.debug(f"ws_address send error: {exc}") + break + except Exception as exc: + logger.debug(f"ws_address error: {exc}") @blockexplorer_router.websocket("/ws/tx/{txid}") @@ -278,29 +236,39 @@ async def ws_tx(websocket: WebSocket, txid: str) -> None: await websocket.close(code=1008) return await websocket.accept() + + queue: asyncio.Queue[OnchainTxEvent] = asyncio.Queue() + task_manager.register_ws_tx_queue(txid, queue) try: - async with _client() as c: - try: - raw = await c.get_transaction(txid) - except ElectrumError as e: - await websocket.send_json({"error": str(e)}) - return - tx = parse_raw_tx(raw) - watch = _watch_scripthash(tx) - await websocket.send_json(await _tx_status(c, txid, watch)) - if watch: - await c.subscribe_scripthash(watch) + await _ws_tx_loop(websocket, queue) + finally: + task_manager.unregister_ws_tx_queue(txid, queue) - async def on_tx_change(params: list[Any]) -> None: - try: - await websocket.send_json(await _tx_status(c, txid, watch)) - except Exception as e: - logger.debug(f"ws_tx send error: {e}") - c.on("blockchain.scripthash.subscribe", on_tx_change) - while True: - msg = await websocket.receive() - if msg["type"] == "websocket.disconnect": +async def _ws_tx_loop( + websocket: WebSocket, + queue: asyncio.Queue[OnchainTxEvent], +) -> None: + try: + while True: + recv_task = asyncio.create_task(websocket.receive()) + event_task = asyncio.create_task(queue.get()) + done, pending = await asyncio.wait( + [recv_task, event_task], return_when=asyncio.FIRST_COMPLETED + ) + for t in pending: + t.cancel() + if recv_task in done: + if recv_task.result().get("type") == "websocket.disconnect": break - except Exception as e: - logger.debug(f"ws_tx error: {e}") + if event_task in done: + event: OnchainTxEvent = event_task.result() + try: + await websocket.send_json(event.dict()) + except Exception as exc: + logger.debug(f"ws_tx send error: {exc}") + break + if event.confirmed: + break + except Exception as exc: + logger.debug(f"ws_tx error: {exc}") diff --git a/lnbits/task_manager.py b/lnbits/task_manager.py index 35f2c6b28..73c58eda2 100644 --- a/lnbits/task_manager.py +++ b/lnbits/task_manager.py @@ -3,14 +3,18 @@ import traceback import uuid from collections.abc import Callable, Coroutine from datetime import datetime, timezone -from typing import Any from loguru import logger from pydantic import BaseModel from lnbits.core.models import Payment from lnbits.settings import settings -from lnbits.utils.electrum import ElectrumClient, ElectrumError, scripthash_from_address +from lnbits.utils.electrum import ( + AddressTracker, + OnchainAddressEvent, + OnchainTxEvent, + TransactionTracker, +) class PublicTask(BaseModel): @@ -20,13 +24,6 @@ class PublicTask(BaseModel): created_at: datetime -class OnchainAddressEvent(BaseModel): - address: str - confirmed: int # satoshis - unconfirmed: int # satoshis - txids: list[str] - - class Task: """Model used on the backend to keep track of background tasks.""" @@ -58,7 +55,9 @@ class TaskManager: tasks: list[Task] = [] invoice_queue: asyncio.Queue[Payment] = asyncio.Queue() internal_invoice_queue: asyncio.Queue[Payment] = asyncio.Queue() - _tracked_addresses: dict[str, str] = {} # address -> task_name + _tracked_addresses: dict[str, int] = {} # address -> ref count + _ws_address_queues: dict[str, list[asyncio.Queue[OnchainAddressEvent]]] = {} + _ws_tx_queues: dict[str, list[asyncio.Queue[OnchainTxEvent]]] = {} def init(self) -> None: self.create_permanent_task( @@ -171,20 +170,76 @@ class TaskManager: ) def track_address(self, address: str) -> None: - """Start tracking a Bitcoin address via Electrum.""" - if address in self._tracked_addresses: - return - task_name = f"onchain_address_{address}" - self._tracked_addresses[address] = task_name - self.create_task(self._address_tracker(address), name=task_name) + """Start tracking a Bitcoin address via Electrum (ref-counted).""" + count = self._tracked_addresses.get(address, 0) + self._tracked_addresses[address] = count + 1 + if count == 0: + self.create_task( + self._address_tracker(address), name=f"onchain_address_{address}" + ) def untrack_address(self, address: str) -> None: - """Stop tracking a Bitcoin address.""" - task_name = self._tracked_addresses.pop(address, None) - if task_name: - task = self.get_task(task_name) + """Decrement ref count; cancel the tracker when the last caller unregisters.""" + count = self._tracked_addresses.get(address, 0) + if count <= 1: + self._tracked_addresses.pop(address, None) + task = self.get_task(f"onchain_address_{address}") if task: self.cancel_task(task) + else: + self._tracked_addresses[address] = count - 1 + + def register_ws_address_queue( + self, address: str, queue: asyncio.Queue[OnchainAddressEvent] + ) -> None: + """Register a per-connection queue for a watched address.""" + self._ws_address_queues.setdefault(address, []).append(queue) + self.track_address(address) + + def unregister_ws_address_queue( + self, address: str, queue: asyncio.Queue[OnchainAddressEvent] + ) -> None: + """Deregister a per-connection queue and decrement the address ref count.""" + queues = self._ws_address_queues.get(address, []) + if queue in queues: + queues.remove(queue) + if not queues: + self._ws_address_queues.pop(address, None) + self.untrack_address(address) + + def register_ws_tx_queue( + self, txid: str, queue: asyncio.Queue[OnchainTxEvent] + ) -> None: + """Register a per-connection queue for a watched transaction.""" + self._ws_tx_queues.setdefault(txid, []).append(queue) + if len(self._ws_tx_queues[txid]) == 1: + self.create_task( + self._transaction_tracker_dispatch(txid), name=f"ws_tx_{txid}" + ) + + def unregister_ws_tx_queue( + self, txid: str, queue: asyncio.Queue[OnchainTxEvent] + ) -> None: + """Deregister a per-connection queue; cancel tracker when last one leaves.""" + queues = self._ws_tx_queues.get(txid, []) + if queue in queues: + queues.remove(queue) + if not queues: + self._ws_tx_queues.pop(txid, None) + task = self.get_task(f"ws_tx_{txid}") + if task: + self.cancel_task(task) + + def track_transaction( + self, + txid: str, + callback: Callable[[OnchainTxEvent], Coroutine], + ) -> Task: + """Track a transaction until confirmed, calling callback on each change.""" + return self.create_task( + self._transaction_tracker(txid, callback), + name=f"onchain_tx_{txid}", + ) async def _heart_beat(self) -> None: """A heartbeat that removes done tasks logs the number of tasks.""" @@ -256,12 +311,18 @@ class TaskManager: logger.debug(f"Enqueing payment to task {task.name}") task.invoice_queue.put_nowait(payment) - def _dispatch_onchain_event(self, event: OnchainAddressEvent) -> None: - """Dispatches an onchain address event to all registered onchain listeners.""" + async def _dispatch_onchain_event(self, event: OnchainAddressEvent) -> None: + """Dispatches an onchain address event to listeners and WS queues.""" for task in self.tasks: - if not task.onchain_queue: - continue - task.onchain_queue.put_nowait(event) + if task.onchain_queue: + task.onchain_queue.put_nowait(event) + for q in list(self._ws_address_queues.get(event.address, [])): + q.put_nowait(event) + + async def _dispatch_onchain_tx_event(self, event: OnchainTxEvent) -> None: + """Dispatches a tx event to all WS queues watching that txid.""" + for q in list(self._ws_tx_queues.get(event.txid, [])): + q.put_nowait(event) async def _invoice_listener_consumer(self) -> None: payment = await self.invoice_queue.get() @@ -274,60 +335,26 @@ class TaskManager: self._invoice_dispatcher(payment) async def _address_tracker(self, address: str) -> None: - """Track an address via Electrum subscription; dispatch OnchainAddressEvents.""" - electrum_url = settings.lnbits_blockexplorer_electrum_url - scripthash = scripthash_from_address(address) + await AddressTracker(settings.lnbits_blockexplorer_electrum_url).track( + address, + self._dispatch_onchain_event, + lambda: address in self._tracked_addresses and settings.lnbits_running, + ) - while address in self._tracked_addresses and settings.lnbits_running: - try: - async with ElectrumClient(electrum_url) as client: - await self._fetch_and_dispatch_address(client, address, scripthash) - - async def on_status_change(params: list[Any]) -> None: - if params and params[0] == scripthash: - await self._fetch_and_dispatch_address( - client, address, scripthash - ) - - await client.subscribe_scripthash(scripthash, on_status_change) - - while ( - address in self._tracked_addresses and settings.lnbits_running - ): - await asyncio.sleep(30) - except asyncio.CancelledError: - raise - except Exception as exc: - if not settings.lnbits_running: - return - logger.warning(f"Address tracker {address}: {exc!s}, retrying in 5s") - await asyncio.sleep(5) - - async def _fetch_and_dispatch_address( - self, client: Any, address: str, scripthash: str + async def _transaction_tracker( + self, txid: str, callback: Callable[[OnchainTxEvent], Coroutine] ) -> None: - """Fetch balance + history for a scripthash and dispatch an onchain event.""" - balance = await client.get_balance(scripthash) - txids: list[str] = [] - try: - history = await client.get_history(scripthash) - txids = [e.tx_hash for e in history] - except ElectrumError: - pass - try: - mempool = await client.get_mempool(scripthash) - for e in mempool: - if e.tx_hash not in txids: - txids.append(e.tx_hash) - except ElectrumError: - pass - self._dispatch_onchain_event( - OnchainAddressEvent( - address=address, - confirmed=balance.confirmed, - unconfirmed=balance.unconfirmed, - txids=txids, - ) + await TransactionTracker(settings.lnbits_blockexplorer_electrum_url).track( + txid, + callback, + lambda: settings.lnbits_running, + ) + + async def _transaction_tracker_dispatch(self, txid: str) -> None: + await TransactionTracker(settings.lnbits_blockexplorer_electrum_url).track( + txid, + self._dispatch_onchain_tx_event, + lambda: txid in self._ws_tx_queues and settings.lnbits_running, ) diff --git a/lnbits/utils/electrum.py b/lnbits/utils/electrum.py index acc8bdaef..d05be8a5b 100644 --- a/lnbits/utils/electrum.py +++ b/lnbits/utils/electrum.py @@ -12,7 +12,7 @@ import itertools import json import ssl import struct -from collections.abc import Callable +from collections.abc import Callable, Coroutine from typing import Any from urllib.parse import urlparse @@ -712,3 +712,218 @@ class ElectrumClient: """Returns mempool fee histogram as FeeHistogramEntry(fee_rate, vsize) list.""" data = await self._call("mempool.get_fee_histogram") return [FeeHistogramEntry(fee_rate=r[0], vsize=r[1]) for r in data] + + +# --------------------------------------------------------------------------- +# Address tracking +# --------------------------------------------------------------------------- + + +class OnchainAddressEvent(BaseModel): + address: str + confirmed: int # satoshis + unconfirmed: int # satoshis + history: list[HistoryEntry] = [] + history_error: str | None = None + + @property + def txids(self) -> list[str]: + return [e.tx_hash for e in self.history] + + +class AddressTracker: + """ + Subscribes to a Bitcoin address via Electrum and calls a callback on every + balance/history change. Reconnects automatically on failure. + + Args: + url: Electrum server URL (e.g. ``ssl://electrum.blockstream.info:50002``). + callback: Async function receiving an :class:`OnchainAddressEvent`. + is_active: Callable returning ``False`` when tracking should stop. + """ + + def __init__(self, url: str) -> None: + self.url = url + + async def track( + self, + address: str, + callback: Callable[[OnchainAddressEvent], Coroutine[Any, Any, None]], + is_active: Callable[[], bool], + ) -> None: + scripthash = scripthash_from_address(address) + while is_active(): + try: + async with ElectrumClient(self.url) as client: + await self._fetch_and_dispatch( + client, address, scripthash, callback + ) + + async def on_status_change(params: list[Any]) -> None: + if params and params[0] == scripthash: + await self._fetch_and_dispatch( + client, address, scripthash, callback + ) + + await client.subscribe_scripthash(scripthash, on_status_change) + while is_active(): + await asyncio.sleep(30) + except asyncio.CancelledError: + raise + except Exception as exc: + if not is_active(): + return + logger.warning(f"AddressTracker {address}: {exc!s}, retrying in 5s") + await asyncio.sleep(5) + + @staticmethod + async def _fetch_and_dispatch( + client: ElectrumClient, + address: str, + scripthash: str, + callback: Callable[[OnchainAddressEvent], Coroutine[Any, Any, None]], + ) -> None: + balance = await client.get_balance(scripthash) + history: list[HistoryEntry] = [] + history_error: str | None = None + try: + history = await client.get_history(scripthash) + except ElectrumError as exc: + history_error = str(exc) + try: + seen = {e.tx_hash for e in history} + for m in await client.get_mempool(scripthash): + if m.tx_hash not in seen: + history.append(HistoryEntry(tx_hash=m.tx_hash, height=0, fee=m.fee)) + except ElectrumError: + pass + await callback( + OnchainAddressEvent( + address=address, + confirmed=balance.confirmed, + unconfirmed=balance.unconfirmed, + history=history, + history_error=history_error, + ) + ) + + +# --------------------------------------------------------------------------- +# Transaction tracking +# --------------------------------------------------------------------------- + + +class OnchainTxEvent(BaseModel): + txid: str + confirmed: bool + height: int | None = None + fee: int | None = None + + +def tx_watch_scripthash(tx: Transaction) -> str | None: + """Return the scripthash of the first spendable output, used to subscribe + for confirmation notifications.""" + for out in tx.vout: + if out.scriptPubKey.type != "nulldata": + return scripthash_from_scriptpubkey(bytes.fromhex(out.scriptPubKey.hex)) + return None + + +class TransactionTracker: + """ + Subscribes to a Bitcoin transaction via Electrum and calls a callback on + each status change (unconfirmed → confirmed). Stops automatically once + the transaction is confirmed or ``is_active()`` returns ``False``. + + Args: + url: Electrum server URL (e.g. ``ssl://electrum.blockstream.info:50002``). + """ + + def __init__(self, url: str) -> None: + self.url = url + + async def track( + self, + txid: str, + callback: Callable[[OnchainTxEvent], Coroutine[Any, Any, None]], + is_active: Callable[[], bool], + ) -> None: + while is_active(): + try: + confirmed = await self._track_once(txid, callback, is_active) + if confirmed: + return + except asyncio.CancelledError: + raise + except Exception as exc: + if not is_active(): + return + logger.warning( + f"TransactionTracker {txid[:8]}: {exc!s}, retrying in 5s" + ) + await asyncio.sleep(5) + + async def _track_once( + self, + txid: str, + callback: Callable[[OnchainTxEvent], Coroutine[Any, Any, None]], + is_active: Callable[[], bool], + ) -> bool: + """One connection attempt; returns True if the tx is confirmed.""" + async with ElectrumClient(self.url) as client: + try: + raw = await client.get_transaction(txid) + except ElectrumError as exc: + logger.warning(f"TransactionTracker {txid[:8]}: {exc!s}") + await asyncio.sleep(10) + return False + + scripthash = tx_watch_scripthash(parse_raw_tx(raw)) + event = await self._fetch_status(client, txid, scripthash) + await callback(event) + if event.confirmed: + return True + + confirmed_event = asyncio.Event() + + async def on_change( + params: list[Any], + _sh: str | None = scripthash, + _done: asyncio.Event = confirmed_event, + ) -> None: + if params and params[0] == _sh: + ev = await self._fetch_status(client, txid, _sh) + await callback(ev) + if ev.confirmed: + _done.set() + + if scripthash: + await client.subscribe_scripthash(scripthash, on_change) + + while is_active() and not confirmed_event.is_set(): + await asyncio.sleep(5) + return confirmed_event.is_set() + + @staticmethod + async def _fetch_status( + client: ElectrumClient, txid: str, scripthash: str | None + ) -> OnchainTxEvent: + if scripthash: + try: + for entry in await client.get_history(scripthash): + if entry.tx_hash == txid: + return OnchainTxEvent( + txid=txid, + confirmed=entry.height > 0, + height=entry.height if entry.height > 0 else None, + fee=entry.fee, + ) + except ElectrumError: + try: + for m in await client.get_mempool(scripthash): + if m.tx_hash == txid: + return OnchainTxEvent(txid=txid, confirmed=False, fee=m.fee) + return OnchainTxEvent(txid=txid, confirmed=True) + except ElectrumError: + pass + return OnchainTxEvent(txid=txid, confirmed=False)