fix: recover expired Redis task reservations

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,
+    }

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

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,19 +100,18 @@ 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`` /
     ``beat_max_loop_interval``) ensures that only one Beat instance actually
     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)

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)