From 7dedb10dca5960abb3a5393884d95160ca08e58f Mon Sep 17 00:00:00 2001
From: NanPan <111261006+poijygfdyy@users.noreply.github.com>
Date: Thu, 8 Oct 2026 09:26:31 +0800
Subject: [PATCH] fix(scheduler): serialize notification worker lifecycle
 (#6433)

* fix(scheduler): serialize notification worker lifecycle

* test(scheduler): bound lifecycle wait
---
 .../app/scheduler/notification_delivery.py    | 46 ++++++++-------
 ..._notification_delivery_worker_lifecycle.py | 56 +++++++++++++++++++
 2 files changed, 82 insertions(+), 20 deletions(-)
 create mode 100644 backend/tests/test_notification_delivery_worker_lifecycle.py

diff --git a/backend/app/scheduler/notification_delivery.py b/backend/app/scheduler/notification_delivery.py
index 89f094a5f2..20f143b755 100644
--- a/backend/app/scheduler/notification_delivery.py
+++ b/backend/app/scheduler/notification_delivery.py
@@ -84,31 +84,37 @@ class NotificationDeliveryWorker:
         self._default_locale = default_locale
         self._task: asyncio.Task | None = None
         self._stop = asyncio.Event()
+        self._lifecycle_lock = asyncio.Lock()
 
     async def start(self) -> None:
-        if self._task is not None:
-            return
-        self._stop.clear()
-        self._task = asyncio.create_task(self._run_loop())
+        async with self._lifecycle_lock:
+            if self._task is not None:
+                return
+            self._stop.clear()
+            self._task = asyncio.create_task(self._run_loop())
 
     async def stop(self) -> None:
-        if self._task is None:
-            return
-        self._stop.set()
-        task = self._task
-        self._task = None
-        try:
-            await asyncio.wait_for(task, timeout=self._stop_timeout_seconds)
-        except TimeoutError:
-            logger.warning(
-                "Notification delivery worker stop exceeded %.1fs; cancelling in-flight poll",
-                self._stop_timeout_seconds,
-            )
-            task.cancel()
+        async with self._lifecycle_lock:
+            task = self._task
+            if task is None:
+                return
+            self._stop.set()
             try:
-                await task
-            except asyncio.CancelledError:
-                pass
+                try:
+                    await asyncio.wait_for(task, timeout=self._stop_timeout_seconds)
+                except TimeoutError:
+                    logger.warning(
+                        "Notification delivery worker stop exceeded %.1fs; cancelling in-flight poll",
+                        self._stop_timeout_seconds,
+                    )
+                    task.cancel()
+                    try:
+                        await task
+                    except asyncio.CancelledError:
+                        pass
+            finally:
+                if self._task is task:
+                    self._task = None
 
     async def run_once(self, *, now: datetime) -> None:
         await self._recover_stale_sending(now)
diff --git a/backend/tests/test_notification_delivery_worker_lifecycle.py b/backend/tests/test_notification_delivery_worker_lifecycle.py
new file mode 100644
index 0000000000..3c4e231530
--- /dev/null
+++ b/backend/tests/test_notification_delivery_worker_lifecycle.py
@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from app.scheduler.notification_delivery import NotificationDeliveryWorker
+
+
+@pytest.mark.asyncio
+async def test_restart_waits_for_in_progress_stop_to_release_poller_ownership() -> None:
+    worker = NotificationDeliveryWorker(
+        delivery_repo=object(),
+        resolve_channel=lambda _provider: None,
+        poll_interval_seconds=60,
+        stop_timeout_seconds=5,
+    )
+    first_started = asyncio.Event()
+    allow_first_exit = asyncio.Event()
+    second_started = asyncio.Event()
+    generation = 0
+
+    async def controlled_loop() -> None:
+        nonlocal generation
+        generation += 1
+        if generation == 1:
+            first_started.set()
+            await allow_first_exit.wait()
+            return
+        second_started.set()
+        await worker._stop.wait()
+
+    worker._run_loop = controlled_loop
+
+    await worker.start()
+    await first_started.wait()
+    first_poller = worker._task
+    assert first_poller is not None
+
+    stop_task = asyncio.create_task(worker.stop())
+    await asyncio.wait_for(worker._stop.wait(), timeout=5)
+
+    assert worker._task is first_poller
+
+    restart_task = asyncio.create_task(worker.start())
+    await asyncio.sleep(0)
+    assert not restart_task.done()
+    assert worker._task is first_poller
+
+    allow_first_exit.set()
+    await stop_task
+    await restart_task
+    await second_started.wait()
+
+    assert worker._task is not first_poller
+    await worker.stop()