feat: create_unique_task helper (#2292)

this is especially useful for extension.

with this an extension can only call `create_permanent_unique_task("my-ext", task)` and it will restarted when it exists and cancels when lnbits is stopped
This commit is contained in:
dni ⚡
2024-03-18 14:02:04 +01:00
committed by GitHub
parent 116ca7f011
commit e4d3faefa0

View File

@@ -4,7 +4,7 @@ import time
import traceback import traceback
import uuid import uuid
from http import HTTPStatus from http import HTTPStatus
from typing import Dict, List, Optional from typing import Coroutine, Dict, List, Optional
from loguru import logger from loguru import logger
from py_vapid import Vapid from py_vapid import Vapid
@@ -20,6 +20,7 @@ from lnbits.settings import settings
from lnbits.wallets import get_wallet_class from lnbits.wallets import get_wallet_class
tasks: List[asyncio.Task] = [] tasks: List[asyncio.Task] = []
unique_tasks: Dict[str, asyncio.Task] = {}
def create_task(coro): def create_task(coro):
@@ -32,12 +33,33 @@ def create_permanent_task(func):
return create_task(catch_everything_and_restart(func)) return create_task(catch_everything_and_restart(func))
def create_unique_task(name: str, coro: Coroutine):
if unique_tasks.get(name):
logger.warning(f"task `{name}` already exists, cancelling it")
try:
unique_tasks[name].cancel()
except Exception as exc:
logger.warning(f"error while cancelling task `{name}`: {str(exc)}")
task = asyncio.create_task(coro)
unique_tasks[name] = task
return task
def create_permanent_unique_task(name: str, coro: Coroutine):
return create_unique_task(name, catch_everything_and_restart(coro))
def cancel_all_tasks(): def cancel_all_tasks():
for task in tasks: for task in tasks:
try: try:
task.cancel() task.cancel()
except Exception as exc: except Exception as exc:
logger.warning(f"error while cancelling task: {str(exc)}") logger.warning(f"error while cancelling task: {str(exc)}")
for name, task in unique_tasks.items():
try:
task.cancel()
except Exception as exc:
logger.warning(f"error while cancelling task `{name}`: {str(exc)}")
async def catch_everything_and_restart(func): async def catch_everything_and_restart(func):