blocktracker

lnbits/core/views/blockexplorer_api.py

@@ -1,6 +1,6 @@
 import asyncio
 from http import HTTPStatus
-from typing import Annotated, Any
+from typing import Annotated
 
 from fastapi import APIRouter, Depends, HTTPException, Request, WebSocket
 from loguru import logger
@@ -150,27 +150,39 @@ async def ws_blocks(websocket: WebSocket) -> None:
         await websocket.close(code=1008)
         return
     await websocket.accept()
+
+    queue: asyncio.Queue[BlockInfo] = asyncio.Queue()
+    task_manager.register_ws_block_queue(queue)
     try:
-        async with _client() as c:
+        await _ws_blocks_loop(websocket, queue)
-            tip = await c.subscribe_headers()
+    finally:
-            await websocket.send_json(parse_block_header(tip.hex, tip.height).dict())
+        task_manager.unregister_ws_block_queue(queue)
 
-            async def on_header(params: list[Any]) -> None:
-                h = params[0]
-                try:
-                    await websocket.send_json(
-                        parse_block_header(h["hex"], h["height"]).dict()
-                    )
-                except Exception as e:
-                    logger.debug(f"ws_blocks send error: {e}")
 
-            c.on("blockchain.headers.subscribe", on_header)
+async def _ws_blocks_loop(
-            while True:
+    websocket: WebSocket,
-                msg = await websocket.receive()
+    queue: asyncio.Queue[BlockInfo],
-                if msg["type"] == "websocket.disconnect":
+) -> 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:
+            if event_task in done:
-        logger.debug(f"ws_blocks error: {e}")
+                try:
+                    await websocket.send_json(event_task.result().dict())
+                except Exception as exc:
+                    logger.debug(f"ws_blocks send error: {exc}")
+                    break
+    except Exception as exc:
+        logger.debug(f"ws_blocks error: {exc}")
 
 
 def _address_event_to_response(event: OnchainAddressEvent) -> AddressResponse:

lnbits/utils/electrum.py

@@ -772,53 +772,108 @@ class OnchainAddressEvent(BaseModel):
 
 class AddressTracker:
     """
-    Subscribes to a Bitcoin address via Electrum and calls a callback on every
+    Subscribes to a set of Bitcoin addresses over a single shared Electrum
-    balance/history change.  Reconnects automatically on failure.
+    connection and calls a callback on every balance/history change.
+    Addresses can be added/removed at runtime via :meth:`add`/:meth:`remove`.
+    Reconnects automatically on failure.
 
     Args:
-        url:       Electrum server URL (e.g. ``ssl://electrum.blockstream.info:50002``).
+        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
+        self._addresses: set[str] = set()
+        self._updated = asyncio.Event()
 
-    async def track(
+    def add(self, address: str) -> None:
+        """Start tracking an address on the shared connection."""
+        if address not in self._addresses:
+            self._addresses.add(address)
+            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)
+            self._updated.set()
+
+    async def run(
         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._run_once(callback, is_active)
-
-                    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)
-                    await self._fetch_and_dispatch(
-                        client, address, scripthash, callback
-                    )
-                    while is_active():
-                        try:
-                            await asyncio.wait_for(client.closed.wait(), timeout=30)
-                            break  # connection closed; reconnect
-                        except asyncio.TimeoutError:
-                            pass
             except asyncio.CancelledError:
                 raise
             except Exception as exc:
                 if not is_active():
                     return
-                logger.warning(f"AddressTracker {address}: {exc!s}, retrying in 5s")
+                logger.warning(f"AddressTracker: {exc!s}, retrying in 5s")
                 await asyncio.sleep(5)
 
+    async def _run_once(
+        self,
+        callback: Callable[[OnchainAddressEvent], Coroutine[Any, Any, None]],
+        is_active: Callable[[], bool],
+    ) -> None:
+        async with ElectrumClient(self.url) as client:
+            subscribed: dict[str, str] = {}  # scripthash -> address
+
+            async def on_status_change(params: list[Any]) -> None:
+                if not params:
+                    return
+                address = subscribed.get(params[0])
+                if address:
+                    await self._fetch_and_dispatch(
+                        client, address, params[0], callback
+                    )
+
+            client.on("blockchain.scripthash.subscribe", on_status_change)
+
+            while is_active():
+                self._updated.clear()
+                await self._sync_subscriptions(client, subscribed, callback)
+                if await self._wait_for_change_or_close(client):
+                    break  # connection closed; reconnect
+
+    async def _sync_subscriptions(
+        self,
+        client: ElectrumClient,
+        subscribed: dict[str, str],
+        callback: Callable[[OnchainAddressEvent], Coroutine[Any, Any, None]],
+    ) -> None:
+        wanted = {a: scripthash_from_address(a) for a in self._addresses}
+        for address, scripthash in wanted.items():
+            if scripthash not in subscribed:
+                subscribed[scripthash] = address
+                await client.subscribe_scripthash(scripthash)
+                await self._fetch_and_dispatch(client, address, scripthash, callback)
+        still_wanted = set(wanted.values())
+        for scripthash, address in list(subscribed.items()):
+            if address not in still_wanted:
+                del subscribed[scripthash]
+                await client.unsubscribe_scripthash(scripthash)
+
+    async def _wait_for_change_or_close(self, client: ElectrumClient) -> bool:
+        """Waits until addresses change or the connection closes; returns True
+        if it was the connection that closed."""
+        wait_task = asyncio.create_task(self._updated.wait())
+        closed_task = asyncio.create_task(client.closed.wait())
+        try:
+            done, _ = await asyncio.wait(
+                [wait_task, closed_task],
+                timeout=30,
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+        finally:
+            for t in (wait_task, closed_task):
+                if not t.done():
+                    t.cancel()
+        return closed_task in done
+
     @staticmethod
     async def _fetch_and_dispatch(
         client: ElectrumClient,
@@ -979,3 +1034,58 @@ class TransactionTracker:
                 except ElectrumError:
                     pass
         return OnchainTxEvent(txid=txid, confirmed=False)
+
+
+# ---------------------------------------------------------------------------
+# Block tracking
+# ---------------------------------------------------------------------------
+
+
+class BlockTracker:
+    """
+    Subscribes to new block headers via Electrum and calls a callback on
+    every new block.  Reconnects automatically on failure.
+
+    Args:
+        url: Electrum server URL (e.g. ``ssl://electrum.blockstream.info:50002``).
+    """
+
+    def __init__(self, url: str) -> None:
+        self.url = url
+
+    async def run(
+        self,
+        callback: Callable[[BlockInfo], Coroutine[Any, Any, None]],
+        is_active: Callable[[], bool],
+    ) -> None:
+        while is_active():
+            try:
+                await self._run_once(callback, is_active)
+            except asyncio.CancelledError:
+                raise
+            except Exception as exc:
+                if not is_active():
+                    return
+                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 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"]))
+
+            tip = await client.subscribe_headers(on_header)
+            await callback(parse_block_header(tip.hex, tip.height))
+
+            while is_active():
+                try:
+                    await asyncio.wait_for(client.closed.wait(), timeout=30)
+                    break  # connection closed; reconnect
+                except asyncio.TimeoutError:
+                    pass