mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
fix(scheduler): serialize notification worker lifecycle (#6433)
* fix(scheduler): serialize notification worker lifecycle * test(scheduler): bound lifecycle wait
This commit is contained in:
parent
220d62d79a
commit
7dedb10dca
@ -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)
|
||||
|
||||
56
backend/tests/test_notification_delivery_worker_lifecycle.py
Normal file
56
backend/tests/test_notification_delivery_worker_lifecycle.py
Normal 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()
|
||||
Loading…
x
Reference in New Issue
Block a user