diff --git a/apps/worker/app/core/visibility_recovery.py b/apps/worker/app/core/visibility_recovery.py new file mode 100644 index 000000000..4c6a0c02a --- /dev/null +++ b/apps/worker/app/core/visibility_recovery.py @@ -0,0 +1,97 @@ +"""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: + raise NotImplementedError( + "Visibility recovery QoS must restore visible reservations" + ) + + +class VisibilityRecoveryChannel(Protocol): + """Expose the minimal channel interface required by recovery.""" + + qos: VisibilityRecoveryQualityOfService + + def close(self) -> None: + raise NotImplementedError("Visibility recovery channel must close") + + +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 5a0827a22..23d48ded8 100644 --- a/apps/worker/app/core/worker_bootstrap.py +++ b/apps/worker/app/core/worker_bootstrap.py @@ -20,6 +20,23 @@ def _register_task_modules() -> None: 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.""" @@ -83,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`` / @@ -91,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 @@ -140,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 new file mode 100644 index 000000000..ad54ab1e4 --- /dev/null +++ b/apps/worker/tests/contract/test_visibility_recovery_contract.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +from contextlib import AbstractContextManager, nullcontext + +import pytest +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 import visibility_recovery + + calls: list[tuple[str, int | None, int | None]] = [] + + class FakeQualityOfService: + 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, None)) + + class FakeConnection: + def ensure_connection(self, *, max_retries: int) -> None: + calls.append(("ensure_connection", max_retries, None)) + + def channel(self) -> FakeChannel: + calls.append(("channel", None, None)) + return FakeChannel() + + def release(self) -> None: + calls.append(("connection.release", None, 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": "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_skip_recovery_when_another_invocation_holds_the_periodic_lock( + worker_contract_environment: None, + monkeypatch: MonkeyPatch, +) -> None: + 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: + nonlocal connection_was_released + connection_was_released = True + + monkeypatch.setattr( + visibility_recovery.celery_app, + "connection_for_read", + lambda: FailingConnection(), + ) + monkeypatch.setattr( + visibility_recovery, + "periodic_task_lock", + lambda *args, **kwargs: _acquired_lock(), + ) + + with pytest.raises(RuntimeError, match="broker unavailable"): + visibility_recovery.restore_expired_reservations() + + assert connection_was_released is True + + +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 c3701b0c6..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 @@ -44,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 wait(self, *, timeout: float) -> int: + self.wait_timeouts.append(timeout) + return 0 - def record_beat_command(command: list[str]) -> FakeBeatProcess: - beat_commands.append(command) - return FakeBeatProcess() + 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(",")) @@ -76,3 +107,83 @@ def record_worker_main(args: list[str]) -> None: "kb_medium", "kb_low", }.issubset(consumed_queues) + + assert "worker_maintenance" not in consumed_queues + + +def test_should_kill_a_child_that_ignores_sigterm( + worker_contract_environment: None, +) -> None: + 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) + + with pytest.raises(OSError, match="watchdog spawn failed"): + worker_bootstrap.run_worker() + + 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/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, + )