From e4d3faefa0f968b990a8eec0c7911d520ebc1344 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?dni=20=E2=9A=A1?= <office@dnilabs.com>
Date: Mon, 18 Mar 2024 14:02:04 +0100
Subject: [PATCH] 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
---
 lnbits/tasks.py | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/lnbits/tasks.py b/lnbits/tasks.py
index e45328a3b..a4ec5caa5 100644
--- a/lnbits/tasks.py
+++ b/lnbits/tasks.py
@@ -4,7 +4,7 @@ import time
 import traceback
 import uuid
 from http import HTTPStatus
-from typing import Dict, List, Optional
+from typing import Coroutine, Dict, List, Optional
 
 from loguru import logger
 from py_vapid import Vapid
@@ -20,6 +20,7 @@ from lnbits.settings import settings
 from lnbits.wallets import get_wallet_class
 
 tasks: List[asyncio.Task] = []
+unique_tasks: Dict[str, asyncio.Task] = {}
 
 
 def create_task(coro):
@@ -32,12 +33,33 @@ def create_permanent_task(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():
     for task in tasks:
         try:
             task.cancel()
         except Exception as 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):