Improved SSE listeners (#865)

* logging listeners

* comments

* generate privkey upon init

* listener queue

* remove duplicate check

* make format

* reuse channel

* error handling in sse listener

* uuid for listeners

* register named invoices

* uuid for listeners and listener list

* fix poetry lock

* setuptools

* requirements asyncio timeout

* setuptool;s

* make format

* remove async-timeout

* async_timeout readd

* try lower setuptools version

* try lower lower setuptools version

* back to current version + fix, maybe

* fix worflows to use poetry 1.2.1

* remove uneeded setuptools from build-system

* fix up formatting workflow

* debug to trace

* more traces

* debug logs to trace

Co-authored-by: dni ⚡ <office@dnilabs.com>

lnbits/core/services.py

@@ -186,9 +186,9 @@ async def pay_invoice(
             )
 
         # notify receiver asynchronously
-
         from lnbits.tasks import internal_invoice_queue
 
+        logger.debug(f"enqueuing internal invoice {internal_checking_id}")
         await internal_invoice_queue.put(internal_checking_id)
     else:
         logger.debug(f"backend: sending payment {temp_id}")

lnbits/core/tasks.py

@@ -1,30 +1,43 @@
 import asyncio
-from typing import List
+from typing import Dict
 
 import httpx
 from loguru import logger
 
-from lnbits.tasks import register_invoice_listener
+from lnbits.helpers import get_current_extension_name
+from lnbits.tasks import SseListenersDict, register_invoice_listener
 
 from . import db
 from .crud import get_balance_notify
 from .models import Payment
 
-api_invoice_listeners: List[asyncio.Queue] = []
+api_invoice_listeners: Dict[str, asyncio.Queue] = SseListenersDict(
+    "api_invoice_listeners"
+)
 
 
 async def register_task_listeners():
+    """
+    Registers an invoice listener queue for the core tasks.
+    Incoming payaments in this queue will eventually trigger the signals sent to all other extensions
+    and fulfill other core tasks such as dispatching webhooks.
+    """
     invoice_paid_queue = asyncio.Queue(5)
-    register_invoice_listener(invoice_paid_queue)
+    # we register invoice_paid_queue to receive all incoming invoices
+    register_invoice_listener(invoice_paid_queue, "core/tasks.py")
+    # register a worker that will react to invoices
     asyncio.create_task(wait_for_paid_invoices(invoice_paid_queue))
 
 
 async def wait_for_paid_invoices(invoice_paid_queue: asyncio.Queue):
+    """
+    This worker dispatches events to all extensions, dispatches webhooks and balance notifys.
+    """
     while True:
         payment = await invoice_paid_queue.get()
-        logger.debug("received invoice paid event")
+        logger.trace("received invoice paid event")
         # send information to sse channel
-        await dispatch_invoice_listener(payment)
+        await dispatch_api_invoice_listeners(payment)
 
         # dispatch webhook
         if payment.webhook and not payment.webhook_status:
@@ -41,16 +54,23 @@ async def wait_for_paid_invoices(invoice_paid_queue: asyncio.Queue):
                     pass
 
 
-async def dispatch_invoice_listener(payment: Payment):
-    for send_channel in api_invoice_listeners:
+async def dispatch_api_invoice_listeners(payment: Payment):
+    """
+    Emits events to invoice listener subscribed from the API.
+    """
+    for chan_name, send_channel in api_invoice_listeners.items():
         try:
+            logger.debug(f"sending invoice paid event to {chan_name}")
             send_channel.put_nowait(payment)
         except asyncio.QueueFull:
-            logger.debug("removing sse listener", send_channel)
-            api_invoice_listeners.remove(send_channel)
+            logger.error(f"removing sse listener {send_channel}:{chan_name}")
+            api_invoice_listeners.pop(chan_name)
 
 
 async def dispatch_webhook(payment: Payment):
+    """
+    Dispatches the webhook to the webhook url.
+    """
     async with httpx.AsyncClient() as client:
         data = payment.dict()
         try:

lnbits/core/views/api.py

@@ -3,11 +3,13 @@ import binascii
 import hashlib
 import json
 import time
+import uuid
 from http import HTTPStatus
 from io import BytesIO
 from typing import Dict, List, Optional, Tuple, Union
 from urllib.parse import ParseResult, parse_qs, urlencode, urlparse, urlunparse
 
+import async_timeout
 import httpx
 import pyqrcode
 from fastapi import Depends, Header, Query, Request
@@ -16,7 +18,7 @@ from fastapi.params import Body
 from loguru import logger
 from pydantic import BaseModel
 from pydantic.fields import Field
-from sse_starlette.sse import EventSourceResponse
+from sse_starlette.sse import EventSourceResponse, ServerSentEvent
 from starlette.responses import HTMLResponse, StreamingResponse
 
 from lnbits import bolt11, lnurl
@@ -366,37 +368,48 @@ async def api_payments_pay_lnurl(
     }
 
 
-async def subscribe(request: Request, wallet: Wallet):
+async def subscribe_wallet_invoices(request: Request, wallet: Wallet):
+    """
+    Subscribe to new invoices for a wallet. Can be wrapped in EventSourceResponse.
+    Listenes invoming payments for a wallet and yields jsons with payment details.
+    """
     this_wallet_id = wallet.id
 
     payment_queue: asyncio.Queue[Payment] = asyncio.Queue(0)
 
-    logger.debug("adding sse listener", payment_queue)
-    api_invoice_listeners.append(payment_queue)
+    uid = f"{this_wallet_id}_{str(uuid.uuid4())[:8]}"
+    logger.debug(f"adding sse listener for wallet: {uid}")
+    api_invoice_listeners[uid] = payment_queue
 
     send_queue: asyncio.Queue[Tuple[str, Payment]] = asyncio.Queue(0)
 
     async def payment_received() -> None:
         while True:
-            payment: Payment = await payment_queue.get()
-            if payment.wallet_id == this_wallet_id:
-                logger.debug("payment received", payment)
-                await send_queue.put(("payment-received", payment))
+            try:
+                async with async_timeout.timeout(1):
+                    payment: Payment = await payment_queue.get()
+                    if payment.wallet_id == this_wallet_id:
+                        logger.debug("sse listener: payment receieved", payment)
+                        await send_queue.put(("payment-received", payment))
+            except asyncio.TimeoutError:
+                pass
 
-    asyncio.create_task(payment_received())
+    task = asyncio.create_task(payment_received())
 
     try:
         while True:
+            if await request.is_disconnected():
+                await request.close()
+                break
             typ, data = await send_queue.get()
-
             if data:
                 jdata = json.dumps(dict(data.dict(), pending=False))
 
-            # yield dict(id=1, event="this", data="1234")
-            # await asyncio.sleep(2)
             yield dict(data=jdata, event=typ)
-            # yield dict(data=jdata.encode("utf-8"), event=typ.encode("utf-8"))
-    except asyncio.CancelledError:
+    except asyncio.CancelledError as e:
+        logger.debug(f"CancelledError on listener {uid}: {e}")
+        api_invoice_listeners.pop(uid)
+        task.cancel()
         return
 
 
@@ -405,7 +418,9 @@ async def api_payments_sse(
     request: Request, wallet: WalletTypeInfo = Depends(get_key_type)
 ):
     return EventSourceResponse(
-        subscribe(request, wallet.wallet), ping=20, media_type="text/event-stream"
+        subscribe_wallet_invoices(request, wallet.wallet),
+        ping=20,
+        media_type="text/event-stream",
     )
 
 

lnbits/core/views/public_api.py

@@ -46,8 +46,8 @@ async def api_public_payment_longpolling(payment_hash):
 
     payment_queue = asyncio.Queue(0)
 
-    logger.debug("adding standalone invoice listener", payment_hash, payment_queue)
-    api_invoice_listeners.append(payment_queue)
+    logger.debug(f"adding standalone invoice listener for hash: {payment_hash}")
+    api_invoice_listeners[payment_hash] = payment_queue
 
     response = None