From aa6ca7eb1ad5c686c14d87b6a53d811c48dfd9f3 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 14 Aug 2026 01:20:37 +0800 Subject: [PATCH 1/3] fix: recover expired redis task reservations --- .../app/core/tasks/visibility_recovery.py | 40 ++++++++ apps/worker/app/core/worker_bootstrap.py | 1 + .../test_visibility_recovery_contract.py | 91 +++++++++++++++++++ .../test_worker_bootstrap_contract.py | 16 ++++ .../shared-python/shared/core/celery_app.py | 4 + 5 files changed, 152 insertions(+) create mode 100644 apps/worker/app/core/tasks/visibility_recovery.py create mode 100644 apps/worker/tests/contract/test_visibility_recovery_contract.py diff --git a/apps/worker/app/core/tasks/visibility_recovery.py b/apps/worker/app/core/tasks/visibility_recovery.py new file mode 100644 index 000000000..c2a6e9d21 --- /dev/null +++ b/apps/worker/app/core/tasks/visibility_recovery.py @@ -0,0 +1,40 @@ +"""Recover expired Redis broker reservations independently of task execution.""" + +from __future__ import annotations + +from loguru import logger + +from shared.core.celery_app import get_celery_app +from shared.services.redis.periodic_task_lock import periodic_task_lock + +_RECOVERY_PERIOD_SECONDS: int = 30 + +celery_app = get_celery_app() + + +@celery_app.task(name="app.core.tasks.visibility_recovery.restore_expired_reservations") +def restore_expired_reservations() -> dict[str, str]: + """Restore expired Redis reservations through Kombu's transport API.""" + with periodic_task_lock( + "app.core.tasks.visibility_recovery.restore_expired_reservations", + period_seconds=_RECOVERY_PERIOD_SECONDS, + buffer_seconds=5, + ) as acquired: + if not acquired: + return {"status": "skipped"} + + connection = celery_app.connection_for_read() + try: + connection.ensure_connection(max_retries=1) + channel = connection.channel() + try: + channel.qos.restore_visible(interval=1) + finally: + channel.close() + except Exception as exc: + logger.error(f"Expired Celery reservation recovery failed: {exc}") + return {"status": "error"} + finally: + connection.release() + + return {"status": "success"} diff --git a/apps/worker/app/core/worker_bootstrap.py b/apps/worker/app/core/worker_bootstrap.py index 5a0827a22..d62f5aa3f 100644 --- a/apps/worker/app/core/worker_bootstrap.py +++ b/apps/worker/app/core/worker_bootstrap.py @@ -17,6 +17,7 @@ def _register_task_modules() -> None: """Import task modules for Celery side-effect registration.""" import app.core.tasks.document_ingestion_tasks # noqa: F401 import app.core.tasks.stale_job_sweeper # noqa: F401 + import app.core.tasks.visibility_recovery # noqa: F401 import app.core.tasks.webhook_tasks # noqa: F401 diff --git a/apps/worker/tests/contract/test_visibility_recovery_contract.py b/apps/worker/tests/contract/test_visibility_recovery_contract.py new file mode 100644 index 000000000..c618fa358 --- /dev/null +++ b/apps/worker/tests/contract/test_visibility_recovery_contract.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +from pytest import MonkeyPatch + + +def test_should_restore_expired_reservations_with_a_fresh_kombu_connection( + worker_contract_environment: None, + monkeypatch: MonkeyPatch, +) -> None: + from app.core.tasks import visibility_recovery + + calls: list[tuple[str, int | None]] = [] + + class FakeQualityOfService: + def restore_visible(self, *, interval: int) -> None: + calls.append(("restore_visible", interval)) + + class FakeChannel: + qos = FakeQualityOfService() + + def close(self) -> None: + calls.append(("channel.close", None)) + + class FakeConnection: + def ensure_connection(self, *, max_retries: int) -> None: + calls.append(("ensure_connection", max_retries)) + + def channel(self) -> FakeChannel: + calls.append(("channel", None)) + return FakeChannel() + + def release(self) -> None: + calls.append(("connection.release", None)) + + connection = FakeConnection() + monkeypatch.setattr( + visibility_recovery.celery_app, + "connection_for_read", + lambda: connection, + ) + monkeypatch.setattr( + visibility_recovery, + "periodic_task_lock", + lambda *args, **kwargs: _acquired_lock(), + ) + + result = visibility_recovery.restore_expired_reservations() + + assert result == {"status": "success"} + assert calls == [ + ("ensure_connection", 1), + ("channel", None), + ("restore_visible", 1), + ("channel.close", None), + ("connection.release", None), + ] + + +def test_should_report_recovery_connection_errors_without_raising( + worker_contract_environment: None, + monkeypatch: MonkeyPatch, +) -> None: + from app.core.tasks import visibility_recovery + + class FailingConnection: + def ensure_connection(self, *, max_retries: int) -> None: + raise RuntimeError("broker unavailable") + + def release(self) -> None: + pass + + monkeypatch.setattr( + visibility_recovery.celery_app, + "connection_for_read", + lambda: FailingConnection(), + ) + monkeypatch.setattr( + visibility_recovery, + "periodic_task_lock", + lambda *args, **kwargs: _acquired_lock(), + ) + + result = visibility_recovery.restore_expired_reservations() + + assert result == {"status": "error"} + + +def _acquired_lock(): + from contextlib import nullcontext + + return nullcontext(True) diff --git a/apps/worker/tests/contract/test_worker_bootstrap_contract.py b/apps/worker/tests/contract/test_worker_bootstrap_contract.py index c3701b0c6..138bd3b27 100644 --- a/apps/worker/tests/contract/test_worker_bootstrap_contract.py +++ b/apps/worker/tests/contract/test_worker_bootstrap_contract.py @@ -17,11 +17,13 @@ def test_should_register_worker_task_modules_for_celery_consumers( "app.core.tasks.kb_tasks.upload_url_file_task", "app.core.tasks.kb_tasks.parse_task", "app.core.tasks.stale_job_sweeper.expire_stale_jobs", + "app.core.tasks.visibility_recovery.restore_expired_reservations", "app.core.tasks.webhook_tasks.recover_orphaned_webhooks", ) task_module_names: tuple[str, ...] = ( "app.core.tasks.document_ingestion_tasks", "app.core.tasks.stale_job_sweeper", + "app.core.tasks.visibility_recovery", "app.core.tasks.webhook_tasks", ) @@ -76,3 +78,17 @@ def record_worker_main(args: list[str]) -> None: "kb_medium", "kb_low", }.issubset(consumed_queues) + + +def test_should_register_periodic_expired_reservation_recovery_task( + worker_contract_environment: None, +) -> None: + from shared.core.celery_app import celery_app + + task_name = "app.core.tasks.visibility_recovery.restore_expired_reservations" + + assert task_name in celery_app.tasks + assert celery_app.conf.beat_schedule["restore-expired-celery-reservations"] == { + "task": task_name, + "schedule": 30.0, + } diff --git a/packages/shared-python/shared/core/celery_app.py b/packages/shared-python/shared/core/celery_app.py index d4b5e39b6..5b6d4c1cf 100644 --- a/packages/shared-python/shared/core/celery_app.py +++ b/packages/shared-python/shared/core/celery_app.py @@ -109,6 +109,10 @@ def get_unique_node_name() -> str: "task": "app.core.tasks.webhook_tasks.recover_orphaned_webhooks", "schedule": 1800.0, # Every 30 minutes }, + "restore-expired-celery-reservations": { + "task": "app.core.tasks.visibility_recovery.restore_expired_reservations", + "schedule": 30.0, + }, "expire-stale-jobs": { "task": "app.core.tasks.stale_job_sweeper.expire_stale_jobs", "schedule": 1800.0, # Every 30 minutes From 549b72fd8534ec5161b11e177405428ce42effde Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 14 Aug 2026 10:46:40 +0800 Subject: [PATCH 2/3] fix: isolate redis visibility recovery --- .../app/core/tasks/visibility_recovery.py | 40 ------ apps/worker/app/core/visibility_recovery.py | 93 +++++++++++++ .../app/core/visibility_recovery_watchdog.py | 65 +++++++++ apps/worker/app/core/worker_bootstrap.py | 54 ++++++-- .../test_visibility_recovery_contract.py | 87 +++++++++--- ...t_visibility_recovery_watchdog_contract.py | 70 ++++++++++ .../test_worker_bootstrap_contract.py | 131 +++++++++++++++--- .../contract/test_worker_health_contract.py | 83 +++++++++++ .../shared-python/shared/core/celery_app.py | 4 - .../shared/core/config/celery.py | 15 ++ .../shared/services/worker_health.py | 65 +++++++-- 11 files changed, 599 insertions(+), 108 deletions(-) delete mode 100644 apps/worker/app/core/tasks/visibility_recovery.py create mode 100644 apps/worker/app/core/visibility_recovery.py create mode 100644 apps/worker/app/core/visibility_recovery_watchdog.py create mode 100644 apps/worker/tests/contract/test_visibility_recovery_watchdog_contract.py create mode 100644 apps/worker/tests/contract/test_worker_health_contract.py diff --git a/apps/worker/app/core/tasks/visibility_recovery.py b/apps/worker/app/core/tasks/visibility_recovery.py deleted file mode 100644 index c2a6e9d21..000000000 --- a/apps/worker/app/core/tasks/visibility_recovery.py +++ /dev/null @@ -1,40 +0,0 @@ -"""Recover expired Redis broker reservations independently of task execution.""" - -from __future__ import annotations - -from loguru import logger - -from shared.core.celery_app import get_celery_app -from shared.services.redis.periodic_task_lock import periodic_task_lock - -_RECOVERY_PERIOD_SECONDS: int = 30 - -celery_app = get_celery_app() - - -@celery_app.task(name="app.core.tasks.visibility_recovery.restore_expired_reservations") -def restore_expired_reservations() -> dict[str, str]: - """Restore expired Redis reservations through Kombu's transport API.""" - with periodic_task_lock( - "app.core.tasks.visibility_recovery.restore_expired_reservations", - period_seconds=_RECOVERY_PERIOD_SECONDS, - buffer_seconds=5, - ) as acquired: - if not acquired: - return {"status": "skipped"} - - connection = celery_app.connection_for_read() - try: - connection.ensure_connection(max_retries=1) - channel = connection.channel() - try: - channel.qos.restore_visible(interval=1) - finally: - channel.close() - except Exception as exc: - logger.error(f"Expired Celery reservation recovery failed: {exc}") - return {"status": "error"} - finally: - connection.release() - - return {"status": "success"} diff --git a/apps/worker/app/core/visibility_recovery.py b/apps/worker/app/core/visibility_recovery.py new file mode 100644 index 000000000..b32b22dce --- /dev/null +++ b/apps/worker/app/core/visibility_recovery.py @@ -0,0 +1,93 @@ +"""Recover expired Redis broker reservations through a fresh Kombu channel.""" + +from __future__ import annotations + +from typing import Literal, Protocol, TypedDict, cast + +from celery import Celery +from kombu import Connection + +from shared.core.config import app_config +from shared.core.celery_app import get_celery_app +from shared.services.redis.periodic_task_lock import periodic_task_lock + +_RECOVERY_LOCK_NAME: str = "visibility-recovery-watchdog" +_RECOVERY_LOCK_BUFFER_SECONDS: int = 5 + +class VisibilityRecoveryAttemptedResult(TypedDict): + """Describe one bounded recovery sweep request.""" + + status: Literal["attempted"] + batch_count: int + batch_size: int + recovery_limit: int + + +class VisibilityRecoverySkippedResult(TypedDict): + """Describe a sweep skipped because another replica holds the lock.""" + + status: Literal["skipped"] + + +VisibilityRecoveryResult = ( + VisibilityRecoveryAttemptedResult | VisibilityRecoverySkippedResult +) + + +class VisibilityRecoveryQualityOfService(Protocol): + """Expose the Kombu QoS operation required by recovery.""" + + def restore_visible(self, *, num: int, interval: int) -> None: ... + + +class VisibilityRecoveryChannel(Protocol): + """Expose the minimal channel interface required by recovery.""" + + qos: VisibilityRecoveryQualityOfService + + def close(self) -> None: ... + + +celery_app: Celery = get_celery_app() + + +def restore_expired_reservations() -> VisibilityRecoveryResult: + """Attempt a bounded restoration sweep through Kombu's Redis transport.""" + period_seconds: int = app_config.VISIBILITY_RECOVERY_PERIOD_SECONDS + batch_size: int = app_config.VISIBILITY_RECOVERY_BATCH_SIZE + batch_count: int = app_config.VISIBILITY_RECOVERY_BATCH_COUNT + recovery_limit: int = batch_size * batch_count + + with periodic_task_lock( + _RECOVERY_LOCK_NAME, + period_seconds=period_seconds, + buffer_seconds=_RECOVERY_LOCK_BUFFER_SECONDS, + ) as acquired: + if not acquired: + return {"status": "skipped"} + + connection: Connection = celery_app.connection_for_read() + try: + connection.ensure_connection(max_retries=1) + channel: VisibilityRecoveryChannel = cast( + VisibilityRecoveryChannel, + connection.channel(), + ) + try: + _batch_index: int + for _batch_index in range(batch_count): + channel.qos.restore_visible( + num=batch_size, + interval=1, + ) + finally: + channel.close() + finally: + connection.release() + + return { + "status": "attempted", + "batch_count": batch_count, + "batch_size": batch_size, + "recovery_limit": recovery_limit, + } diff --git a/apps/worker/app/core/visibility_recovery_watchdog.py b/apps/worker/app/core/visibility_recovery_watchdog.py new file mode 100644 index 000000000..0d2d3cd5b --- /dev/null +++ b/apps/worker/app/core/visibility_recovery_watchdog.py @@ -0,0 +1,65 @@ +"""Run expired-reservation recovery outside the saturated Celery task pool.""" + +from __future__ import annotations + +import signal +from threading import Event +from types import FrameType + +from loguru import logger + +from app.core.visibility_recovery import ( + VisibilityRecoveryResult, + restore_expired_reservations, +) +from shared.core.config import app_config +from shared.core.logging import setup_logging +from shared.services.worker_health import ( + remove_visibility_recovery_heartbeat, + write_visibility_recovery_heartbeat, +) + +_stop_event: Event = Event() + + +def _request_stop(signal_number: int, frame: FrameType | None) -> None: + """Request watchdog shutdown after the current bounded recovery attempt.""" + logger.info(f"Visibility recovery watchdog stopping on signal {signal_number}") + _stop_event.set() + + +def run_visibility_recovery_watchdog() -> None: + """Attempt recovery periodically and remain healthy across broker errors.""" + setup_logging(service_name="knowhere-worker") + _stop_event.clear() + signal.signal(signal.SIGTERM, _request_stop) + signal.signal(signal.SIGINT, _request_stop) + period_seconds: float = float(app_config.VISIBILITY_RECOVERY_PERIOD_SECONDS) + + write_visibility_recovery_heartbeat() + logger.info( + "Visibility recovery watchdog started: " + f"period={period_seconds:.0f}s" + ) + + try: + while True: + try: + result: VisibilityRecoveryResult = restore_expired_reservations() + logger.bind(**result).info( + "Expired Celery reservation recovery sweep attempted" + ) + except Exception: + logger.exception("Expired Celery reservation recovery sweep failed") + finally: + write_visibility_recovery_heartbeat() + + if _stop_event.wait(timeout=period_seconds): + break + finally: + remove_visibility_recovery_heartbeat() + logger.info("Visibility recovery watchdog stopped") + + +if __name__ == "__main__": + run_visibility_recovery_watchdog() diff --git a/apps/worker/app/core/worker_bootstrap.py b/apps/worker/app/core/worker_bootstrap.py index d62f5aa3f..23d48ded8 100644 --- a/apps/worker/app/core/worker_bootstrap.py +++ b/apps/worker/app/core/worker_bootstrap.py @@ -17,10 +17,26 @@ def _register_task_modules() -> None: """Import task modules for Celery side-effect registration.""" import app.core.tasks.document_ingestion_tasks # noqa: F401 import app.core.tasks.stale_job_sweeper # noqa: F401 - import app.core.tasks.visibility_recovery # noqa: F401 import app.core.tasks.webhook_tasks # noqa: F401 +def _stop_child_process( + process: subprocess.Popen[bytes], + process_name: str, +) -> None: + """Stop a colocated worker child without exceeding the ECS stop window.""" + if process.poll() is not None: + return + + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + logger.warning(f"{process_name} did not stop after SIGTERM; killing it") + process.kill() + process.wait(timeout=5) + + @worker_init.connect def init_worker(**kwargs) -> None: """Initialize structured logging and sync Redis when worker process starts.""" @@ -84,7 +100,7 @@ def shutdown_worker(**kwargs) -> None: def run_worker() -> None: - """Start the gevent Celery worker and its colocated Beat process. + """Start Celery with colocated Beat and visibility-recovery processes. Every worker replica unconditionally spawns a Celery Beat subprocess. RedBeat's own distributed lock (``redbeat_lock_timeout`` / @@ -92,11 +108,10 @@ def run_worker() -> None: drives the scheduler tick loop — all other instances block on lock acquisition and remain idle. - Even if the RedBeat startup-burst window allows multiple Beat instances - to enqueue the same periodic task simultaneously, each task body is - guarded by a ``periodic_task_lock`` (Redis ``SET NX EX``) keyed on the - task name. Only the first invocation within each scheduling window - executes; all subsequent duplicates log a skip and return immediately. + Each replica also starts an independent visibility-recovery watchdog. + Recovery runs outside the Celery gevent pool so ingestion saturation cannot + starve it. The watchdogs coordinate through the application Redis periodic + lock, while Kombu's broker mutex protects the restoration transaction. """ from shared.core.config import settings @@ -141,8 +156,25 @@ def run_worker() -> None: "beat", f"--loglevel={log_level}", ] + visibility_recovery_cmd: list[str] = [ + sys.executable, + "-m", + "app.core.visibility_recovery_watchdog", + ] - logger.info("Starting Celery Beat subprocess") - subprocess.Popen(beat_cmd) - - celery_app.worker_main(celery_args) + child_processes: list[tuple[str, subprocess.Popen[bytes]]] = [] + try: + logger.info("Starting Celery Beat subprocess") + beat_process: subprocess.Popen[bytes] = subprocess.Popen(beat_cmd) + child_processes.append(("Celery Beat", beat_process)) + + logger.info("Starting visibility recovery watchdog subprocess") + recovery_process: subprocess.Popen[bytes] = subprocess.Popen( + visibility_recovery_cmd + ) + child_processes.append(("Visibility recovery watchdog", recovery_process)) + + celery_app.worker_main(celery_args) + finally: + for process_name, child_process in reversed(child_processes): + _stop_child_process(child_process, process_name) diff --git a/apps/worker/tests/contract/test_visibility_recovery_contract.py b/apps/worker/tests/contract/test_visibility_recovery_contract.py index c618fa358..ad54ab1e4 100644 --- a/apps/worker/tests/contract/test_visibility_recovery_contract.py +++ b/apps/worker/tests/contract/test_visibility_recovery_contract.py @@ -1,5 +1,8 @@ from __future__ import annotations +from contextlib import AbstractContextManager, nullcontext + +import pytest from pytest import MonkeyPatch @@ -7,30 +10,30 @@ def test_should_restore_expired_reservations_with_a_fresh_kombu_connection( worker_contract_environment: None, monkeypatch: MonkeyPatch, ) -> None: - from app.core.tasks import visibility_recovery + from app.core import visibility_recovery - calls: list[tuple[str, int | None]] = [] + calls: list[tuple[str, int | None, int | None]] = [] class FakeQualityOfService: - def restore_visible(self, *, interval: int) -> None: - calls.append(("restore_visible", interval)) + def restore_visible(self, *, num: int, interval: int) -> None: + calls.append(("restore_visible", num, interval)) class FakeChannel: qos = FakeQualityOfService() def close(self) -> None: - calls.append(("channel.close", None)) + calls.append(("channel.close", None, None)) class FakeConnection: def ensure_connection(self, *, max_retries: int) -> None: - calls.append(("ensure_connection", max_retries)) + calls.append(("ensure_connection", max_retries, None)) def channel(self) -> FakeChannel: - calls.append(("channel", None)) + calls.append(("channel", None, None)) return FakeChannel() def release(self) -> None: - calls.append(("connection.release", None)) + calls.append(("connection.release", None, None)) connection = FakeConnection() monkeypatch.setattr( @@ -46,28 +49,63 @@ def release(self) -> None: result = visibility_recovery.restore_expired_reservations() - assert result == {"status": "success"} - assert calls == [ - ("ensure_connection", 1), - ("channel", None), - ("restore_visible", 1), - ("channel.close", None), - ("connection.release", None), + assert result == { + "status": "attempted", + "batch_count": 10, + "batch_size": 100, + "recovery_limit": 1000, + } + assert calls[:2] == [ + ("ensure_connection", 1, None), + ("channel", None, None), + ] + assert calls[2:12] == [("restore_visible", 100, 1)] * 10 + assert calls[12:] == [ + ("channel.close", None, None), + ("connection.release", None, None), ] -def test_should_report_recovery_connection_errors_without_raising( +def test_should_skip_recovery_when_another_invocation_holds_the_periodic_lock( worker_contract_environment: None, monkeypatch: MonkeyPatch, ) -> None: - from app.core.tasks import visibility_recovery + from app.core import visibility_recovery + + def fail_if_connection_is_created() -> None: + raise AssertionError("a skipped recovery must not open a broker connection") + + monkeypatch.setattr( + visibility_recovery.celery_app, + "connection_for_read", + fail_if_connection_is_created, + ) + monkeypatch.setattr( + visibility_recovery, + "periodic_task_lock", + lambda *args, **kwargs: _skipped_lock(), + ) + + result = visibility_recovery.restore_expired_reservations() + + assert result == {"status": "skipped"} + + +def test_should_raise_recovery_connection_errors_for_celery_observability( + worker_contract_environment: None, + monkeypatch: MonkeyPatch, +) -> None: + from app.core import visibility_recovery + + connection_was_released: bool = False class FailingConnection: def ensure_connection(self, *, max_retries: int) -> None: raise RuntimeError("broker unavailable") def release(self) -> None: - pass + nonlocal connection_was_released + connection_was_released = True monkeypatch.setattr( visibility_recovery.celery_app, @@ -80,12 +118,15 @@ def release(self) -> None: lambda *args, **kwargs: _acquired_lock(), ) - result = visibility_recovery.restore_expired_reservations() - - assert result == {"status": "error"} + with pytest.raises(RuntimeError, match="broker unavailable"): + visibility_recovery.restore_expired_reservations() + assert connection_was_released is True -def _acquired_lock(): - from contextlib import nullcontext +def _acquired_lock() -> AbstractContextManager[bool]: return nullcontext(True) + + +def _skipped_lock() -> AbstractContextManager[bool]: + return nullcontext(False) diff --git a/apps/worker/tests/contract/test_visibility_recovery_watchdog_contract.py b/apps/worker/tests/contract/test_visibility_recovery_watchdog_contract.py new file mode 100644 index 000000000..a01eef115 --- /dev/null +++ b/apps/worker/tests/contract/test_visibility_recovery_watchdog_contract.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from pytest import MonkeyPatch + + +def test_should_keep_watchdog_alive_after_a_recovery_error( + worker_contract_environment: None, + monkeypatch: MonkeyPatch, +) -> None: + from app.core import visibility_recovery_watchdog + + calls: list[tuple[str, float | None]] = [] + + class StopAfterOneCycle: + def clear(self) -> None: + calls.append(("stop.clear", None)) + + def set(self) -> None: + calls.append(("stop.set", None)) + + def wait(self, timeout: float) -> bool: + calls.append(("stop.wait", timeout)) + return True + + def raise_recovery_error() -> dict[str, str]: + calls.append(("recover", None)) + raise RuntimeError("broker unavailable") + + monkeypatch.setattr( + visibility_recovery_watchdog, + "restore_expired_reservations", + raise_recovery_error, + ) + monkeypatch.setattr( + visibility_recovery_watchdog, + "write_visibility_recovery_heartbeat", + lambda: calls.append(("heartbeat.write", None)), + ) + monkeypatch.setattr( + visibility_recovery_watchdog, + "remove_visibility_recovery_heartbeat", + lambda: calls.append(("heartbeat.remove", None)), + ) + monkeypatch.setattr( + visibility_recovery_watchdog, + "setup_logging", + lambda *, service_name: calls.append(("logging", None)), + ) + monkeypatch.setattr( + visibility_recovery_watchdog.signal, + "signal", + lambda *args: None, + ) + monkeypatch.setattr( + visibility_recovery_watchdog, + "_stop_event", + StopAfterOneCycle(), + ) + + visibility_recovery_watchdog.run_visibility_recovery_watchdog() + + assert calls == [ + ("logging", None), + ("stop.clear", None), + ("heartbeat.write", None), + ("recover", None), + ("heartbeat.write", None), + ("stop.wait", 30.0), + ("heartbeat.remove", None), + ] diff --git a/apps/worker/tests/contract/test_worker_bootstrap_contract.py b/apps/worker/tests/contract/test_worker_bootstrap_contract.py index 138bd3b27..b997796a9 100644 --- a/apps/worker/tests/contract/test_worker_bootstrap_contract.py +++ b/apps/worker/tests/contract/test_worker_bootstrap_contract.py @@ -1,7 +1,9 @@ from __future__ import annotations +import subprocess import sys +import pytest from pytest import MonkeyPatch @@ -17,13 +19,11 @@ def test_should_register_worker_task_modules_for_celery_consumers( "app.core.tasks.kb_tasks.upload_url_file_task", "app.core.tasks.kb_tasks.parse_task", "app.core.tasks.stale_job_sweeper.expire_stale_jobs", - "app.core.tasks.visibility_recovery.restore_expired_reservations", "app.core.tasks.webhook_tasks.recover_orphaned_webhooks", ) task_module_names: tuple[str, ...] = ( "app.core.tasks.document_ingestion_tasks", "app.core.tasks.stale_job_sweeper", - "app.core.tasks.visibility_recovery", "app.core.tasks.webhook_tasks", ) @@ -46,26 +46,55 @@ def test_should_consume_current_and_legacy_ingestion_queues( from app.core import worker_bootstrap worker_main_calls: list[list[str]] = [] - beat_commands: list[list[str]] = [] + child_process_commands: list[list[str]] = [] + child_processes: list[FakeChildProcess] = [] - class FakeBeatProcess: - pass + class FakeChildProcess: + def __init__(self) -> None: + self.was_terminated: bool = False + self.was_killed: bool = False + self.wait_timeouts: list[float] = [] + + def poll(self) -> None: + return None + + def terminate(self) -> None: + self.was_terminated = True + + def kill(self) -> None: + self.was_killed = True - def record_beat_command(command: list[str]) -> FakeBeatProcess: - beat_commands.append(command) - return FakeBeatProcess() + def wait(self, *, timeout: float) -> int: + self.wait_timeouts.append(timeout) + return 0 + + def record_child_process(command: list[str]) -> FakeChildProcess: + child_process_commands.append(command) + process = FakeChildProcess() + child_processes.append(process) + return process def record_worker_main(args: list[str]) -> None: worker_main_calls.append(args) - monkeypatch.setattr(worker_bootstrap.subprocess, "Popen", record_beat_command) + monkeypatch.setattr(worker_bootstrap.subprocess, "Popen", record_child_process) monkeypatch.setattr(worker_bootstrap.celery_app, "worker_main", record_worker_main) worker_bootstrap.run_worker() - assert beat_commands != [] + assert len(child_process_commands) == 2 assert len(worker_main_calls) == 1 + assert any("beat" in command for command in child_process_commands) + assert [ + sys.executable, + "-m", + "app.core.visibility_recovery_watchdog", + ] in child_process_commands + assert all(process.was_terminated for process in child_processes) + assert all(process.wait_timeouts == [5] for process in child_processes) + assert all(not process.was_killed for process in child_processes) + worker_args = worker_main_calls[0] queue_arg = worker_args[worker_args.index("-Q") + 1] consumed_queues = set(queue_arg.split(",")) @@ -79,16 +108,82 @@ def record_worker_main(args: list[str]) -> None: "kb_low", }.issubset(consumed_queues) + assert "worker_maintenance" not in consumed_queues -def test_should_register_periodic_expired_reservation_recovery_task( + +def test_should_kill_a_child_that_ignores_sigterm( worker_contract_environment: None, ) -> None: - from shared.core.celery_app import celery_app + from app.core import worker_bootstrap + + class StubbornChildProcess: + def __init__(self) -> None: + self.was_terminated: bool = False + self.was_killed: bool = False + self.wait_timeouts: list[float] = [] + + def poll(self) -> None: + return None + + def terminate(self) -> None: + self.was_terminated = True + + def kill(self) -> None: + self.was_killed = True + + def wait(self, *, timeout: float) -> int: + self.wait_timeouts.append(timeout) + if len(self.wait_timeouts) == 1: + raise subprocess.TimeoutExpired("stubborn-child", timeout) + return 0 + + process: StubbornChildProcess = StubbornChildProcess() + + worker_bootstrap._stop_child_process(process, "Stubborn child") + + assert process.was_terminated is True + assert process.was_killed is True + assert process.wait_timeouts == [5, 5] + + +def test_should_stop_beat_when_the_watchdog_fails_to_start( + worker_contract_environment: None, + monkeypatch: MonkeyPatch, +) -> None: + from app.core import worker_bootstrap + + class FakeBeatProcess: + def __init__(self) -> None: + self.was_terminated: bool = False + self.wait_timeouts: list[float] = [] + + def poll(self) -> None: + return None + + def terminate(self) -> None: + self.was_terminated = True + + def kill(self) -> None: + raise AssertionError("responsive Beat must not be killed") + + def wait(self, *, timeout: float) -> int: + self.wait_timeouts.append(timeout) + return 0 + + beat_process: FakeBeatProcess = FakeBeatProcess() + spawn_count: int = 0 + + def fail_watchdog_spawn(command: list[str]) -> FakeBeatProcess: + nonlocal spawn_count + spawn_count += 1 + if spawn_count == 1: + return beat_process + raise OSError("watchdog spawn failed") + + monkeypatch.setattr(worker_bootstrap.subprocess, "Popen", fail_watchdog_spawn) - task_name = "app.core.tasks.visibility_recovery.restore_expired_reservations" + with pytest.raises(OSError, match="watchdog spawn failed"): + worker_bootstrap.run_worker() - assert task_name in celery_app.tasks - assert celery_app.conf.beat_schedule["restore-expired-celery-reservations"] == { - "task": task_name, - "schedule": 30.0, - } + assert beat_process.was_terminated is True + assert beat_process.wait_timeouts == [5] diff --git a/apps/worker/tests/contract/test_worker_health_contract.py b/apps/worker/tests/contract/test_worker_health_contract.py new file mode 100644 index 000000000..635347128 --- /dev/null +++ b/apps/worker/tests/contract/test_worker_health_contract.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest +from pytest import MonkeyPatch + + +def test_should_require_main_worker_and_visibility_recovery_heartbeats( + tmp_path: Path, + monkeypatch: MonkeyPatch, +) -> None: + from shared.services import worker_health + + worker_heartbeat_path = tmp_path / "worker.json" + recovery_heartbeat_path = tmp_path / "visibility-recovery.json" + monkeypatch.setattr(worker_health, "HEARTBEAT_PATH", worker_heartbeat_path) + monkeypatch.setattr( + worker_health, + "VISIBILITY_RECOVERY_HEARTBEAT_PATH", + recovery_heartbeat_path, + ) + worker_health.write_worker_heartbeat() + worker_health.write_visibility_recovery_heartbeat() + worker_health.assert_worker_healthy() + + recovery_heartbeat_path.unlink() + + with pytest.raises(SystemExit, match="Visibility recovery heartbeat file not found"): + worker_health.assert_worker_healthy() + + +def test_should_remove_visibility_recovery_heartbeat( + tmp_path: Path, + monkeypatch: MonkeyPatch, +) -> None: + from shared.services import worker_health + + recovery_heartbeat_path = tmp_path / "visibility-recovery.json" + monkeypatch.setattr( + worker_health, + "VISIBILITY_RECOVERY_HEARTBEAT_PATH", + recovery_heartbeat_path, + ) + + worker_health.write_visibility_recovery_heartbeat() + worker_health.remove_visibility_recovery_heartbeat() + + assert not recovery_heartbeat_path.exists() + + +def test_should_reject_a_stale_visibility_recovery_heartbeat( + tmp_path: Path, + monkeypatch: MonkeyPatch, +) -> None: + from shared.services import worker_health + + worker_heartbeat_path: Path = tmp_path / "worker.json" + recovery_heartbeat_path: Path = tmp_path / "visibility-recovery.json" + monkeypatch.setattr(worker_health, "HEARTBEAT_PATH", worker_heartbeat_path) + monkeypatch.setattr( + worker_health, + "VISIBILITY_RECOVERY_HEARTBEAT_PATH", + recovery_heartbeat_path, + ) + monkeypatch.setattr( + worker_health, + "HEARTBEAT_STALE_AFTER_SECONDS", + worker_health.VISIBILITY_RECOVERY_HEARTBEAT_STALE_AFTER_SECONDS + 2, + ) + + worker_health.write_worker_heartbeat() + worker_health.write_visibility_recovery_heartbeat() + recovery_mtime: float = recovery_heartbeat_path.stat().st_mtime + stale_time: float = ( + recovery_mtime + + worker_health.VISIBILITY_RECOVERY_HEARTBEAT_STALE_AFTER_SECONDS + + 1 + ) + monkeypatch.setattr(worker_health.time, "time", lambda: stale_time) + + with pytest.raises(SystemExit, match="Visibility recovery heartbeat stale"): + worker_health.assert_worker_healthy() diff --git a/packages/shared-python/shared/core/celery_app.py b/packages/shared-python/shared/core/celery_app.py index 5b6d4c1cf..d4b5e39b6 100644 --- a/packages/shared-python/shared/core/celery_app.py +++ b/packages/shared-python/shared/core/celery_app.py @@ -109,10 +109,6 @@ def get_unique_node_name() -> str: "task": "app.core.tasks.webhook_tasks.recover_orphaned_webhooks", "schedule": 1800.0, # Every 30 minutes }, - "restore-expired-celery-reservations": { - "task": "app.core.tasks.visibility_recovery.restore_expired_reservations", - "schedule": 30.0, - }, "expire-stale-jobs": { "task": "app.core.tasks.stale_job_sweeper.expire_stale_jobs", "schedule": 1800.0, # Every 30 minutes diff --git a/packages/shared-python/shared/core/config/celery.py b/packages/shared-python/shared/core/config/celery.py index 4d0c4a8fa..480a8207a 100644 --- a/packages/shared-python/shared/core/config/celery.py +++ b/packages/shared-python/shared/core/config/celery.py @@ -24,6 +24,21 @@ class CeleryConfig(BaseModel): BROKER_POOL_LIMIT: int = Field( default=10, description="Celery broker connection pool limit" ) + VISIBILITY_RECOVERY_PERIOD_SECONDS: int = Field( + default=30, + ge=10, + description="Interval between independent expired-reservation sweeps", + ) + VISIBILITY_RECOVERY_BATCH_SIZE: int = Field( + default=100, + ge=1, + description="Maximum reservations restored by one Kombu sweep batch", + ) + VISIBILITY_RECOVERY_BATCH_COUNT: int = Field( + default=10, + ge=1, + description="Maximum Kombu sweep batches attempted per interval", + ) # Task retry configuration DOCUMENT_INGESTION_TASK_MAX_RETRIES: int = Field( diff --git a/packages/shared-python/shared/services/worker_health.py b/packages/shared-python/shared/services/worker_health.py index 2d7a5d224..affadcc3c 100644 --- a/packages/shared-python/shared/services/worker_health.py +++ b/packages/shared-python/shared/services/worker_health.py @@ -14,6 +14,8 @@ from gevent.lock import Semaphore from loguru import logger +from shared.core.config import app_config + HEARTBEAT_PATH = Path( os.getenv("WORKER_HEARTBEAT_FILE", "/tmp/knowhere-worker-heartbeat.json") ) @@ -21,17 +23,38 @@ HEARTBEAT_STALE_AFTER_SECONDS = float( os.getenv("WORKER_HEARTBEAT_STALE_AFTER_SECONDS", "45") ) +VISIBILITY_RECOVERY_HEARTBEAT_PATH = Path( + os.getenv( + "VISIBILITY_RECOVERY_HEARTBEAT_FILE", + "/tmp/knowhere-visibility-recovery-heartbeat.json", + ) +) +VISIBILITY_RECOVERY_HEARTBEAT_STALE_AFTER_SECONDS: float = float( + max(app_config.VISIBILITY_RECOVERY_PERIOD_SECONDS * 3, 90) +) _heartbeat_greenlet: Optional[gevent.Greenlet] = None _heartbeat_stop_event = Event() _heartbeat_lock = Semaphore() -def write_worker_heartbeat() -> None: - HEARTBEAT_PATH.parent.mkdir(parents=True, exist_ok=True) - temp_path = HEARTBEAT_PATH.with_suffix(f"{HEARTBEAT_PATH.suffix}.tmp") +def _write_heartbeat(path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temp_path = path.with_suffix(f"{path.suffix}.tmp") temp_path.write_text(str(os.getpid()), encoding="utf-8") - os.replace(temp_path, HEARTBEAT_PATH) + os.replace(temp_path, path) + + +def write_worker_heartbeat() -> None: + _write_heartbeat(HEARTBEAT_PATH) + + +def write_visibility_recovery_heartbeat() -> None: + _write_heartbeat(VISIBILITY_RECOVERY_HEARTBEAT_PATH) + + +def remove_visibility_recovery_heartbeat() -> None: + VISIBILITY_RECOVERY_HEARTBEAT_PATH.unlink(missing_ok=True) def _heartbeat_loop() -> None: @@ -75,14 +98,32 @@ def stop_worker_heartbeat() -> None: logger.warning(f"Failed to remove worker heartbeat file: {exc}") -def assert_worker_healthy() -> None: - if not HEARTBEAT_PATH.exists(): - raise SystemExit(f"Worker heartbeat file not found: {HEARTBEAT_PATH}") +def _assert_heartbeat_is_fresh( + *, + name: str, + path: Path, + stale_after_seconds: float, +) -> None: + if not path.exists(): + raise SystemExit(f"{name} heartbeat file not found: {path}") - age_seconds = time.time() - HEARTBEAT_PATH.stat().st_mtime - if age_seconds > HEARTBEAT_STALE_AFTER_SECONDS: + age_seconds = time.time() - path.stat().st_mtime + if age_seconds > stale_after_seconds: raise SystemExit( - "Worker heartbeat stale: " - f"path={HEARTBEAT_PATH}, age={age_seconds:.1f}s, " - f"threshold={HEARTBEAT_STALE_AFTER_SECONDS:.1f}s" + f"{name} heartbeat stale: " + f"path={path}, age={age_seconds:.1f}s, " + f"threshold={stale_after_seconds:.1f}s" ) + + +def assert_worker_healthy() -> None: + _assert_heartbeat_is_fresh( + name="Worker", + path=HEARTBEAT_PATH, + stale_after_seconds=HEARTBEAT_STALE_AFTER_SECONDS, + ) + _assert_heartbeat_is_fresh( + name="Visibility recovery", + path=VISIBILITY_RECOVERY_HEARTBEAT_PATH, + stale_after_seconds=VISIBILITY_RECOVERY_HEARTBEAT_STALE_AFTER_SECONDS, + ) From 140f8201daef948bdfdd7a85ee12d01290845a28 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 14 Aug 2026 11:04:33 +0800 Subject: [PATCH 3/3] fix: satisfy visibility recovery codeql checks --- apps/worker/app/core/visibility_recovery.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/apps/worker/app/core/visibility_recovery.py b/apps/worker/app/core/visibility_recovery.py index b32b22dce..4c6a0c02a 100644 --- a/apps/worker/app/core/visibility_recovery.py +++ b/apps/worker/app/core/visibility_recovery.py @@ -37,7 +37,10 @@ class VisibilityRecoverySkippedResult(TypedDict): class VisibilityRecoveryQualityOfService(Protocol): """Expose the Kombu QoS operation required by recovery.""" - def restore_visible(self, *, num: int, interval: int) -> None: ... + def restore_visible(self, *, num: int, interval: int) -> None: + raise NotImplementedError( + "Visibility recovery QoS must restore visible reservations" + ) class VisibilityRecoveryChannel(Protocol): @@ -45,7 +48,8 @@ class VisibilityRecoveryChannel(Protocol): qos: VisibilityRecoveryQualityOfService - def close(self) -> None: ... + def close(self) -> None: + raise NotImplementedError("Visibility recovery channel must close") celery_app: Celery = get_celery_app()