450 lines
16 KiB
Python
450 lines
16 KiB
Python
import asyncio
|
|
import traceback
|
|
import uuid
|
|
from collections.abc import Callable, Coroutine
|
|
from datetime import datetime, timezone
|
|
from typing import TypeVar
|
|
|
|
from fastapi import WebSocket
|
|
from loguru import logger
|
|
from pydantic import BaseModel
|
|
|
|
from lnbits.core.models import Payment
|
|
from lnbits.settings import settings
|
|
from lnbits.utils.electrum import (
|
|
AddressTracker,
|
|
BlockInfo,
|
|
BlockTracker,
|
|
OnchainAddressEvent,
|
|
OnchainTxEvent,
|
|
TransactionTracker,
|
|
scripthash_from_address,
|
|
)
|
|
|
|
|
|
class PublicTask(BaseModel):
|
|
"""Public model used to expose task information via the API."""
|
|
|
|
name: str
|
|
created_at: datetime
|
|
|
|
|
|
class Task:
|
|
"""Model used on the backend to keep track of background tasks."""
|
|
|
|
coro: Coroutine
|
|
name: str
|
|
created_at: datetime
|
|
task: asyncio.Task
|
|
invoice_queue: asyncio.Queue[Payment] | None = None
|
|
onchain_queue: asyncio.Queue[OnchainAddressEvent] | None = None
|
|
|
|
def __init__(
|
|
self,
|
|
coro: Coroutine,
|
|
name: str | None = None,
|
|
invoice_queue: asyncio.Queue | None = None,
|
|
onchain_queue: asyncio.Queue | None = None,
|
|
) -> None:
|
|
self.coro = coro
|
|
self.name = name or f"task_{uuid.uuid4()}"
|
|
self.created_at = datetime.now(timezone.utc)
|
|
self.task = asyncio.create_task(self.coro, name=self.name)
|
|
self.invoice_queue = invoice_queue
|
|
self.onchain_queue = onchain_queue
|
|
|
|
|
|
class TaskManager:
|
|
"""Singleton class to manage background tasks."""
|
|
|
|
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]] = []
|
|
|
|
def init(self) -> None:
|
|
self.create_permanent_task(
|
|
func=self._heart_beat,
|
|
interval=settings.task_heart_beat_interval,
|
|
)
|
|
self.create_permanent_task(self._invoice_listener_consumer)
|
|
self.create_permanent_task(self._internal_invoice_listener_consumer)
|
|
|
|
def get_task(self, name: str) -> Task | None:
|
|
"""Get a running task by name."""
|
|
for task in self.tasks:
|
|
if task.name == name:
|
|
return task
|
|
return None
|
|
|
|
def get_public_tasks(self) -> list[PublicTask]:
|
|
"""Get a list of public tasks."""
|
|
return [PublicTask(name=t.name, created_at=t.created_at) for t in self.tasks]
|
|
|
|
def cancel_task(self, task: Task) -> None:
|
|
"""Cancel a running task."""
|
|
self.tasks.remove(task)
|
|
try:
|
|
task.task.cancel()
|
|
except Exception as exc:
|
|
logger.warning(f"error while cancelling task `{task.name}`: {exc!s}")
|
|
|
|
def cancel_all_tasks(self) -> None:
|
|
"""Cancel all running tasks."""
|
|
for task in list(self.tasks):
|
|
self.cancel_task(task)
|
|
|
|
def create_task(
|
|
self,
|
|
coro: Coroutine,
|
|
name: str | None = None,
|
|
invoice_queue: asyncio.Queue | None = None,
|
|
onchain_queue: asyncio.Queue | None = None,
|
|
) -> Task:
|
|
"""Create a task. If a task with the same name exists, it will be cancelled."""
|
|
if name:
|
|
task = self.get_task(name)
|
|
if task:
|
|
self.cancel_task(task)
|
|
task = Task(
|
|
coro=coro,
|
|
name=name,
|
|
invoice_queue=invoice_queue,
|
|
onchain_queue=onchain_queue,
|
|
)
|
|
self.tasks.append(task)
|
|
return task
|
|
|
|
def create_permanent_task(
|
|
self,
|
|
func: Callable[[], Coroutine],
|
|
invoice_queue: asyncio.Queue | None = None,
|
|
onchain_queue: asyncio.Queue | None = None,
|
|
name: str | None = None,
|
|
interval: int = 0,
|
|
) -> Task:
|
|
"""Create a task that runs forever and restarts on failure."""
|
|
|
|
async def wrapper():
|
|
while settings.lnbits_running:
|
|
await self._catch_everything_and_restart(func)
|
|
if interval > 0:
|
|
await asyncio.sleep(interval)
|
|
|
|
return self.create_task(
|
|
coro=wrapper(),
|
|
name=name or func.__name__,
|
|
invoice_queue=invoice_queue,
|
|
onchain_queue=onchain_queue,
|
|
)
|
|
|
|
def register_invoice_listener(
|
|
self,
|
|
func: Callable[[Payment], Coroutine],
|
|
name: str | None = None,
|
|
) -> Task:
|
|
"""
|
|
A method intended for extensions to call when they want to be notified about
|
|
incoming payments. Will call provided Coroutine with the updated payment.
|
|
"""
|
|
name = f"{name or uuid.uuid4()}_invoice_listener"
|
|
queue: asyncio.Queue[Payment] = asyncio.Queue()
|
|
return self.create_permanent_task(
|
|
self._invoice_listener_worker(func, queue),
|
|
name=name,
|
|
invoice_queue=queue,
|
|
)
|
|
|
|
def register_onchain_listener(
|
|
self,
|
|
func: Callable[[OnchainAddressEvent], Coroutine],
|
|
name: str | None = None,
|
|
) -> Task:
|
|
"""
|
|
Register a callback for onchain address events dispatched by track_address.
|
|
Will call the provided coroutine with an OnchainAddressEvent on each update.
|
|
"""
|
|
name = f"{name or uuid.uuid4()}_onchain_listener"
|
|
queue: asyncio.Queue[OnchainAddressEvent] = asyncio.Queue()
|
|
return self.create_permanent_task(
|
|
self._onchain_listener_worker(func, queue),
|
|
name=name,
|
|
onchain_queue=queue,
|
|
)
|
|
|
|
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)
|
|
|
|
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
|
|
|
|
def _get_address_tracker(self) -> "AddressTracker":
|
|
if self._address_tracker is None:
|
|
self._address_tracker = AddressTracker(
|
|
settings.lnbits_blockexplorer_electrum_url
|
|
)
|
|
if not self.get_task("address_tracker"):
|
|
self.create_task(
|
|
self._address_tracker.run(
|
|
self._dispatch_onchain_event,
|
|
lambda: settings.lnbits_running,
|
|
),
|
|
name="address_tracker",
|
|
)
|
|
return self._address_tracker
|
|
|
|
def register_ws_address_queue(
|
|
self, address: str, queue: asyncio.Queue[OnchainAddressEvent]
|
|
) -> None:
|
|
"""Register a per-connection queue for a watched address.
|
|
|
|
Raises ValueError if the address is invalid.
|
|
"""
|
|
scripthash_from_address(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 register_ws_block_queue(self, queue: asyncio.Queue[BlockInfo]) -> None:
|
|
"""Register a per-connection queue for new block events."""
|
|
self._ws_block_queues.append(queue)
|
|
if len(self._ws_block_queues) == 1:
|
|
if self._block_tracker is None:
|
|
self._block_tracker = BlockTracker(
|
|
settings.lnbits_blockexplorer_electrum_url
|
|
)
|
|
self.create_task(
|
|
self._block_tracker.run(
|
|
self._dispatch_block_event,
|
|
lambda: bool(self._ws_block_queues) and settings.lnbits_running,
|
|
),
|
|
name="block_tracker",
|
|
)
|
|
|
|
def unregister_ws_block_queue(self, queue: asyncio.Queue[BlockInfo]) -> None:
|
|
"""Deregister a per-connection queue; cancel tracker when last one leaves."""
|
|
if queue in self._ws_block_queues:
|
|
self._ws_block_queues.remove(queue)
|
|
if not self._ws_block_queues:
|
|
task = self.get_task("block_tracker")
|
|
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."""
|
|
for task in self.tasks:
|
|
state = task.task._state if task.task else "NOT RUNNING"
|
|
if settings.task_heart_beat_verbose:
|
|
logger.debug(
|
|
f"Task Manager: `{task.name}` state: `{state}` "
|
|
f"created: {task.created_at.strftime('%Y-%m-%d %H:%M:%S')}`"
|
|
)
|
|
if task.task and task.task.done():
|
|
logger.debug(f"Task Manager: task `{task.name}` is done.")
|
|
self.cancel_task(task)
|
|
invoice_listeners = sum(1 for task in self.tasks if task.invoice_queue)
|
|
onchain_listeners = sum(1 for task in self.tasks if task.onchain_queue)
|
|
other_tasks = len(self.tasks) - invoice_listeners - onchain_listeners
|
|
logger.debug(
|
|
f"Task Manager: {other_tasks} tasks, "
|
|
f"{invoice_listeners} invoice listeners, "
|
|
f"{onchain_listeners} onchain listeners."
|
|
)
|
|
|
|
async def _catch_everything_and_restart(
|
|
self,
|
|
func: Callable[[], Coroutine],
|
|
restart_interval: int = 5,
|
|
) -> None:
|
|
"""Catches all exceptions from a function and restarts it after 5 seconds."""
|
|
while settings.lnbits_running:
|
|
try:
|
|
return await func()
|
|
except asyncio.CancelledError:
|
|
raise # because we must pass this up
|
|
except Exception as exc:
|
|
if not settings.lnbits_running:
|
|
return
|
|
logger.error(f"exception in background task `{func.__name__}`:", exc)
|
|
logger.error(traceback.format_exc())
|
|
logger.info(
|
|
f"`{func.__name__}` restarts in {restart_interval} seconds."
|
|
)
|
|
await asyncio.sleep(restart_interval)
|
|
|
|
def _invoice_listener_worker(
|
|
self, func: Callable[[Payment], Coroutine], queue: asyncio.Queue[Payment]
|
|
) -> Callable:
|
|
async def wrapper() -> None:
|
|
payment: Payment = await queue.get()
|
|
await func(payment)
|
|
|
|
return wrapper
|
|
|
|
def _onchain_listener_worker(
|
|
self,
|
|
func: Callable[[OnchainAddressEvent], Coroutine],
|
|
queue: asyncio.Queue[OnchainAddressEvent],
|
|
) -> Callable:
|
|
async def wrapper() -> None:
|
|
event: OnchainAddressEvent = await queue.get()
|
|
await func(event)
|
|
|
|
return wrapper
|
|
|
|
def _invoice_dispatcher(self, payment: Payment) -> None:
|
|
"""Dispatches a payment to all registered invoice listeners."""
|
|
for task in self.tasks:
|
|
if not task.invoice_queue:
|
|
continue
|
|
logger.debug(f"Enqueing payment to task {task.name}")
|
|
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."""
|
|
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."""
|
|
for q in list(self._ws_tx_queues.get(event.txid, [])):
|
|
q.put_nowait(event)
|
|
|
|
async def _dispatch_block_event(self, event: BlockInfo) -> None:
|
|
"""Dispatches a new block event to all WS queues."""
|
|
for q in list(self._ws_block_queues):
|
|
q.put_nowait(event)
|
|
|
|
async def _invoice_listener_consumer(self) -> None:
|
|
payment = await self.invoice_queue.get()
|
|
logger.info(f"got a payment notification {payment.checking_id}")
|
|
self._invoice_dispatcher(payment)
|
|
|
|
async def _internal_invoice_listener_consumer(self) -> None:
|
|
payment = await self.internal_invoice_queue.get()
|
|
logger.info(f"got an internal payment notification {payment.checking_id}")
|
|
self._invoice_dispatcher(payment)
|
|
|
|
async def _transaction_tracker(
|
|
self, txid: str, callback: Callable[[OnchainTxEvent], Coroutine]
|
|
) -> None:
|
|
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,
|
|
)
|
|
|
|
|
|
T = TypeVar("T", bound=BaseModel)
|
|
|
|
|
|
async def relay_ws_queue(
|
|
websocket: WebSocket,
|
|
queue: "asyncio.Queue[T]",
|
|
serialize: Callable[[T], BaseModel] = lambda e: e,
|
|
stop_after: Callable[[T], bool] = lambda _: False,
|
|
) -> None:
|
|
"""
|
|
Pumps events from `queue` to `websocket` as JSON until the client
|
|
disconnects, sending fails, or `stop_after` returns True for an event.
|
|
Shared by the blockexplorer address/tx/block websocket endpoints.
|
|
"""
|
|
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()
|
|
disconnect = recv_task in done and (
|
|
recv_task.result().get("type") == "websocket.disconnect"
|
|
)
|
|
if disconnect:
|
|
break
|
|
if event_task in done:
|
|
event = event_task.result()
|
|
try:
|
|
await websocket.send_json(serialize(event).dict())
|
|
except Exception as exc:
|
|
logger.debug(f"ws relay send error: {exc}")
|
|
break
|
|
if stop_after(event):
|
|
break
|
|
except Exception as exc:
|
|
logger.debug(f"ws relay error: {exc}")
|
|
|
|
|
|
task_manager = TaskManager()
|