blocktracker
This commit is contained in:
dni
@@ -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:
|
||||
tip = await c.subscribe_headers()
|
||||
await websocket.send_json(parse_block_header(tip.hex, tip.height).dict())
|
||||
await _ws_blocks_loop(websocket, queue)
|
||||
finally:
|
||||
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)
|
||||
while True:
|
||||
msg = await websocket.receive()
|
||||
if msg["type"] == "websocket.disconnect":
|
||||
async def _ws_blocks_loop(
|
||||
websocket: WebSocket,
|
||||
queue: asyncio.Queue[BlockInfo],
|
||||
) -> 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_blocks error: {e}")
|
||||
if event_task in done:
|
||||
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:
|
||||
|
||||
+137
-27
@@ -772,53 +772,108 @@ class OnchainAddressEvent(BaseModel):
|
||||
|
||||
class AddressTracker:
|
||||
"""
|
||||
Subscribes to a Bitcoin address via Electrum and calls a callback on every
|
||||
balance/history change. Reconnects automatically on failure.
|
||||
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`.
|
||||
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.
|
||||
url: Electrum server URL (e.g. ``ssl://electrum.blockstream.info:50002``).
|
||||
"""
|
||||
|
||||
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:
|
||||
|
||||
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
|
||||
await self._run_once(callback, is_active)
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user