fix(scheduler): serialize notification worker lifecycle (#6433)

* fix(scheduler): serialize notification worker lifecycle

* test(scheduler): bound lifecycle wait
This commit is contained in:
NanPan 2026-10-08 09:26:31 +08:00 committed by GitHub
parent 220d62d79a
commit 7dedb10dca
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 82 additions and 20 deletions

View File

@ -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)

View File

@ -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()