From 5514e35e18705be0d9500db2a0802b766e6c6779 Mon Sep 17 00:00:00 2001 From: dni Date: Thu, 9 Jul 2026 10:10:26 +0200 Subject: [PATCH] addresstracker --- lnbits/task_manager.py | 37 ++++------- lnbits/utils/electrum.py | 130 ++++++++++++++++++++++++++++----------- 2 files changed, 106 insertions(+), 61 deletions(-) diff --git a/lnbits/task_manager.py b/lnbits/task_manager.py index 43972aad5..dbf558daa 100644 --- a/lnbits/task_manager.py +++ b/lnbits/task_manager.py @@ -60,12 +60,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, int] = {} # address -> ref count _address_tracker: "AddressTracker | None" = None _block_tracker: "BlockTracker | None" = None - _ws_address_queues: dict[str, list[asyncio.Queue[OnchainAddressEvent]]] = {} - _ws_tx_queues: dict[str, list[asyncio.Queue[OnchainTxEvent]]] = {} - _ws_block_queues: list[asyncio.Queue[BlockInfo]] = [] + _tx_trackers: dict[str, "TransactionTracker"] = {} def init(self) -> None: self.create_permanent_task( @@ -179,20 +176,12 @@ class TaskManager: def track_address(self, address: str) -> None: """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._get_address_tracker().add(address) + self._get_address_tracker().add(address) def untrack_address(self, address: str) -> None: """Decrement ref count; remove from shared tracker when last caller leaves.""" - count = self._tracked_addresses.get(address, 0) - if count <= 1: - self._tracked_addresses.pop(address, None) - if self._address_tracker: - self._address_tracker.remove(address) - else: - self._tracked_addresses[address] = count - 1 + if self._address_tracker: + self._address_tracker.remove(address) def _get_address_tracker(self) -> "AddressTracker": if self._address_tracker is None: @@ -217,19 +206,14 @@ class TaskManager: Raises ValueError if the address is invalid. """ scripthash_from_address(address) - self._ws_address_queues.setdefault(address, []).append(queue) - self.track_address(address) + self._get_address_tracker().register_queue(address, queue) 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) + if self._address_tracker: + self._address_tracker.unregister_queue(address, queue) def register_ws_tx_queue( self, txid: str, queue: asyncio.Queue[OnchainTxEvent] @@ -361,12 +345,13 @@ class TaskManager: task.invoice_queue.put_nowait(payment) async def _dispatch_onchain_event(self, event: OnchainAddressEvent) -> None: - """Dispatches an onchain address event to listeners and WS queues.""" + """Dispatches an onchain address event to registered listeners. + + Per-address WS queue fan-out is handled by AddressTracker itself. + """ for task in self.tasks: 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.""" diff --git a/lnbits/utils/electrum.py b/lnbits/utils/electrum.py index 80092f3f5..34ed1e620 100644 --- a/lnbits/utils/electrum.py +++ b/lnbits/utils/electrum.py @@ -774,7 +774,9 @@ class AddressTracker: """ Subscribes to a set of Bitcoin addresses over a single shared Electrum connection and calls a callback on every balance/history change. - Addresses can be added/removed at runtime via :meth:`add`/:meth:`remove`. + Addresses can be added/removed at runtime via :meth:`add`/:meth:`remove`, + and per-connection queues can be attached via :meth:`register_queue` for + consumers (e.g. websockets) that want events for one specific address. Reconnects automatically on failure. Args: @@ -783,20 +785,43 @@ class AddressTracker: def __init__(self, url: str) -> None: self.url = url - self._addresses: set[str] = set() + self._ref_counts: dict[str, int] = {} + self._queues: dict[str, list[asyncio.Queue[OnchainAddressEvent]]] = {} self._updated = asyncio.Event() def add(self, address: str) -> None: - """Start tracking an address on the shared connection.""" - if address not in self._addresses: - self._addresses.add(address) + """Start tracking an address on the shared connection (ref-counted).""" + count = self._ref_counts.get(address, 0) + self._ref_counts[address] = count + 1 + if count == 0: self._updated.set() def remove(self, address: str) -> None: - """Stop tracking an address on the shared connection.""" - if address in self._addresses: - self._addresses.discard(address) + """Decrement ref count; drop the subscription once the last caller leaves.""" + count = self._ref_counts.get(address, 0) + if count <= 1: + self._ref_counts.pop(address, None) self._updated.set() + else: + self._ref_counts[address] = count - 1 + + def register_queue( + self, address: str, queue: "asyncio.Queue[OnchainAddressEvent]" + ) -> None: + """Register a per-connection queue to receive events for `address`.""" + self._queues.setdefault(address, []).append(queue) + self.add(address) + + def unregister_queue( + self, address: str, queue: "asyncio.Queue[OnchainAddressEvent]" + ) -> None: + """Deregister a per-connection queue for `address`.""" + queues = self._queues.get(address, []) + if queue in queues: + queues.remove(queue) + if not queues: + self._queues.pop(address, None) + self.remove(address) async def run( self, @@ -843,7 +868,7 @@ class AddressTracker: subscribed: dict[str, str], callback: Callable[[OnchainAddressEvent], Coroutine[Any, Any, None]], ) -> None: - wanted = {a: scripthash_from_address(a) for a in self._addresses} + wanted = {a: scripthash_from_address(a) for a in self._ref_counts} for address, scripthash in wanted.items(): if scripthash not in subscribed: subscribed[scripthash] = address @@ -872,8 +897,8 @@ class AddressTracker: t.cancel() return closed_task in done - @staticmethod async def _fetch_and_dispatch( + self, client: ElectrumClient, address: str, scripthash: str, @@ -898,15 +923,16 @@ class AddressTracker: for m in mempool_r: if m.tx_hash not in seen: history.append(HistoryEntry(tx_hash=m.tx_hash, height=0, fee=m.fee)) - await callback( - OnchainAddressEvent( - address=address, - confirmed=balance_r.confirmed, - unconfirmed=balance_r.unconfirmed, - history=history, - history_error=history_error, - ) + event = OnchainAddressEvent( + address=address, + confirmed=balance_r.confirmed, + unconfirmed=balance_r.unconfirmed, + history=history, + history_error=history_error, ) + for q in list(self._queues.get(address, [])): + q.put_nowait(event) + await callback(event) # --------------------------------------------------------------------------- @@ -935,6 +961,8 @@ 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``. + Per-connection queues can be attached via :meth:`register_queue` for + consumers (e.g. websockets) that want events for this transaction. Args: url: Electrum server URL (e.g. ``ssl://electrum.blockstream.info:50002``). @@ -942,6 +970,19 @@ class TransactionTracker: def __init__(self, url: str) -> None: self.url = url + self._queues: list[asyncio.Queue[OnchainTxEvent]] = [] + + def register_queue(self, queue: "asyncio.Queue[OnchainTxEvent]") -> None: + """Register a per-connection queue to receive events for this tx.""" + self._queues.append(queue) + + def unregister_queue(self, queue: "asyncio.Queue[OnchainTxEvent]") -> None: + """Deregister a per-connection queue.""" + if queue in self._queues: + self._queues.remove(queue) + + def has_queues(self) -> bool: + return bool(self._queues) async def track( self, @@ -989,7 +1030,7 @@ class TransactionTracker: ) -> None: if params and params[0] == _sh: ev = await self._fetch_status(client, txid, _sh) - await callback(ev) + await self._dispatch(ev, callback) if ev.confirmed: _done.set() @@ -997,7 +1038,7 @@ class TransactionTracker: await client.subscribe_scripthash(scripthash, on_change) event = await self._fetch_status(client, txid, scripthash) - await callback(event) + await self._dispatch(event, callback) if event.confirmed: return True @@ -1009,6 +1050,15 @@ class TransactionTracker: pass return confirmed_event.is_set() + async def _dispatch( + self, + event: OnchainTxEvent, + callback: Callable[[OnchainTxEvent], Coroutine[Any, Any, None]], + ) -> None: + for q in list(self._queues): + q.put_nowait(event) + await callback(event) + @staticmethod async def _fetch_status( client: ElectrumClient, txid: str, scripthash: str | None @@ -1041,8 +1091,9 @@ class TransactionTracker: class BlockTracker: """ - Subscribes to new block headers via Electrum and calls a callback on - every new block. Reconnects automatically on failure. + Subscribes to new block headers via Electrum and dispatches them to + registered queues. Per-connection queues can be attached via + :meth:`register_queue`. Reconnects automatically on failure. Args: url: Electrum server URL (e.g. ``ssl://electrum.blockstream.info:50002``). @@ -1050,15 +1101,24 @@ class BlockTracker: def __init__(self, url: str) -> None: self.url = url + self._queues: list[asyncio.Queue[BlockInfo]] = [] - async def run( - self, - callback: Callable[[BlockInfo], Coroutine[Any, Any, None]], - is_active: Callable[[], bool], - ) -> None: + def register_queue(self, queue: "asyncio.Queue[BlockInfo]") -> None: + """Register a per-connection queue to receive new block events.""" + self._queues.append(queue) + + def unregister_queue(self, queue: "asyncio.Queue[BlockInfo]") -> None: + """Deregister a per-connection queue.""" + if queue in self._queues: + self._queues.remove(queue) + + def has_queues(self) -> bool: + return bool(self._queues) + + async def run(self, is_active: Callable[[], bool]) -> None: while is_active(): try: - await self._run_once(callback, is_active) + await self._run_once(is_active) except asyncio.CancelledError: raise except Exception as exc: @@ -1067,19 +1127,15 @@ class BlockTracker: logger.warning(f"BlockTracker: {exc!s}, retrying in 5s") await asyncio.sleep(5) - async def _run_once( - self, - callback: Callable[[BlockInfo], Coroutine[Any, Any, None]], - is_active: Callable[[], bool], - ) -> None: + async def _run_once(self, is_active: Callable[[], bool]) -> None: async with ElectrumClient(self.url) as client: async def on_header(params: list[Any]) -> None: h = params[0] - await callback(parse_block_header(h["hex"], h["height"])) + self._dispatch(parse_block_header(h["hex"], h["height"])) tip = await client.subscribe_headers(on_header) - await callback(parse_block_header(tip.hex, tip.height)) + self._dispatch(parse_block_header(tip.hex, tip.height)) while is_active(): try: @@ -1087,3 +1143,7 @@ class BlockTracker: break # connection closed; reconnect except asyncio.TimeoutError: pass + + def _dispatch(self, event: BlockInfo) -> None: + for q in list(self._queues): + q.put_nowait(event)