diff --git a/keep/api/api.py b/keep/api/api.py index 4ae2072457..ddbc4bd997 100644 --- a/keep/api/api.py +++ b/keep/api/api.py @@ -105,6 +105,20 @@ def no_redirect_request(self, method, url, **kwargs): requests.Session.request = no_redirect_request +# The event loop only keeps weak references to tasks, so fire-and-forget +# tasks must be strongly referenced somewhere to avoid being garbage-collected +# (and silently cancelled) mid-execution. +# https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task +service_tasks = set() + + +def create_service_task(coro) -> asyncio.Task: + task = asyncio.create_task(coro) + service_tasks.add(task) + task.add_done_callback(service_tasks.discard) + return task + + async def check_pending_tasks(background_tasks: set): while True: events_in_queue = len(background_tasks) @@ -180,7 +194,7 @@ async def startup(): except Exception: logger.exception("Failed to start the maintenance windows") else: - asyncio.create_task(process_watcher_task.async_process_watcher()) + create_service_task(process_watcher_task.async_process_watcher()) logger.info( "Added task", extra={ @@ -231,7 +245,7 @@ async def lifespan(app: FastAPI): # if debug tasks are enabled, create a task to check for pending tasks if KEEP_DEBUG_TASKS: logger.info("Starting background task to check for pending tasks") - asyncio.create_task(check_pending_tasks(background_tasks)) + create_service_task(check_pending_tasks(background_tasks)) # Startup await startup() diff --git a/tests/test_service_tasks.py b/tests/test_service_tasks.py new file mode 100644 index 0000000000..f3d5306846 --- /dev/null +++ b/tests/test_service_tasks.py @@ -0,0 +1,28 @@ +import asyncio + +import pytest + +from keep.api import api + + +@pytest.mark.asyncio +async def test_create_service_task_holds_strong_reference(): + started = asyncio.Event() + release = asyncio.Event() + + async def service(): + started.set() + await release.wait() + + task = api.create_service_task(service()) + await started.wait() + + # the task is strongly referenced while running, so the event loop's + # weak reference is not the only thing keeping it alive + assert task in api.service_tasks + + release.set() + await task + + # completed tasks are discarded so the set doesn't grow unboundedly + assert task not in api.service_tasks