diff --git a/apps/worker/app/core/gevent_worker_shutdown.py b/apps/worker/app/core/gevent_worker_shutdown.py new file mode 100644 index 000000000..02567c912 --- /dev/null +++ b/apps/worker/app/core/gevent_worker_shutdown.py @@ -0,0 +1,180 @@ +"""Bound normal Celery warm shutdown for the gevent worker pool.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import cast + +import gevent +from celery.concurrency.gevent import TaskPool as GeventTaskPool +from celery.worker import WorkController, state +from celery.worker.request import Request +from gevent import Greenlet +from loguru import logger + +_TerminateJob = Callable[[GeventTaskPool, int, int | None], None] + + +class _ShutdownCancellationPool: + """Expose Celery's pool cancellation interface only for shutdown.""" + + def __init__( + self, + pool: GeventTaskPool, + terminate_job: _TerminateJob, + ) -> None: + self._pool: GeventTaskPool = pool + self._terminate_job: _TerminateJob = terminate_job + + def terminate_job(self, pid: int, signal: int | None = None) -> None: + self._terminate_job(self._pool, pid, signal) + + +def _ignore_reconnect_cancellation( + self: GeventTaskPool, + pid: int, + signal: int | None = None, +) -> None: + """Keep the existing broker-reconnect behavior for gevent tasks.""" + logger.warning( + "Ignoring gevent task cancellation outside bounded worker shutdown " + f"(pid={pid}); relying on RedisJobLock for redelivery deduplication" + ) + + +class GeventWorkerShutdownController: + """Cancel unacknowledged active tasks after a bounded warm shutdown. + + Celery's ``worker_soft_shutdown_timeout`` does not bound normal SIGTERM. + Celery applies that setting only after entering its cold/SIGQUIT path. That + path is unsafe for this worker because Celery calls its patched ``sleep`` + directly from the gevent signal callback, which raises + ``BlockingSwitchOutError``. This controller keeps normal warm SIGTERM and + schedules its own non-blocking timer instead. + + The worker also intentionally ignores Celery cancellation requests caused + by broker reconnects. The original gevent cancellation method is retained + privately and exposed only to this explicit shutdown path, so reconnects + cannot accidentally kill useful work while Redis redelivery is settling. + """ + + def __init__( + self, + worker: WorkController, + timeout_seconds: float, + ) -> None: + self._worker: WorkController = worker + self._timeout_seconds: float = timeout_seconds + self._shutdown_timer: Greenlet | None = None + self._has_scheduled_shutdown: bool = False + self._original_terminate_job: _TerminateJob = ( + self._patch_reconnect_cancellation() + ) + + def schedule(self) -> None: + """Schedule the one bounded cancellation pass for warm SIGTERM.""" + if self._has_scheduled_shutdown: + logger.info("Bounded worker shutdown is already scheduled") + return + + self._has_scheduled_shutdown = True + + # ``worker_shutting_down`` is emitted inside Celery's signal handler. + # Waiting or killing a greenlet there would try to switch out of the + # gevent hub callback and reproduce the staging BlockingSwitchOutError. + # ``spawn_later`` only arms a timer here; its callback runs in a normal + # greenlet where cooperative cancellation is safe. + self._shutdown_timer = gevent.spawn_later( + self._timeout_seconds, + self._cancel_unacknowledged_active_tasks, + ) + logger.warning( + "Scheduled bounded warm shutdown cancellation in " + f"{self._timeout_seconds:g} seconds" + ) + + def close(self) -> None: + """Disarm a pending timer after the worker finishes naturally.""" + shutdown_timer: Greenlet | None = self._shutdown_timer + if shutdown_timer is not None and not shutdown_timer.dead: + shutdown_timer.kill(block=False) + self._shutdown_timer = None + + @staticmethod + def _patch_reconnect_cancellation() -> _TerminateJob: + original_terminate_job: _TerminateJob = cast( + _TerminateJob, + getattr( + GeventTaskPool, + "_original_terminate_job", + GeventTaskPool.terminate_job, + ), + ) + + if not hasattr(GeventTaskPool, "_original_terminate_job"): + setattr( + GeventTaskPool, + "_original_terminate_job", + original_terminate_job, + ) + + # Celery uses ``pool.terminate_job`` both for broker reconnect recovery + # and for deliberate task cancellation. Replacing the class method + # only for a moment would race with reconnect handling, so the public + # pool behavior stays a no-op and shutdown uses the private adapter + # above to reach the saved original method. + GeventTaskPool.terminate_job = _ignore_reconnect_cancellation + logger.info( + "Patched gevent TaskPool.terminate_job for reconnect-safe recovery" + ) + return original_terminate_job + + @staticmethod + def _should_cancel(request: Request) -> bool: + if not request.task.acks_late: + return True + return not request.acknowledged + + def _cancel_unacknowledged_active_tasks(self) -> None: + requests_to_cancel: tuple[Request, ...] = tuple( + request + for request in state.active_requests + if self._should_cancel(request) + ) + if not requests_to_cancel: + logger.info( + "Bounded warm shutdown completed without active task cancellation" + ) + return + + pool: object = self._worker.pool + if not isinstance(pool, GeventTaskPool): + logger.error( + "Cannot cancel active tasks during bounded shutdown: " + f"expected gevent pool, got {type(pool).__name__}" + ) + return + + cancellation_pool = _ShutdownCancellationPool( + pool, + self._original_terminate_job, + ) + logger.warning( + "Bounded warm shutdown timeout expired; cancelling " + f"{len(requests_to_cancel)} unacknowledged active task(s)" + ) + for request in requests_to_cancel: + # Request.cancel performs Celery's normal task-ready bookkeeping, + # but the adapter invokes gevent's saved original cancellation. + # The broker connection then closes normally, allowing Kombu to + # restore each unacknowledged reservation for another worker. + try: + request.cancel(cancellation_pool) + except Exception as exc: + # The greenlet is terminated before Celery records its retry + # event. A result-backend outage must not abort this loop and + # leave later active tasks running until ECS force-kills us. + logger.warning( + "Cancelled task but could not record shutdown retry " + f"(task_id={request.id}): {exc}" + ) diff --git a/apps/worker/app/core/worker_bootstrap.py b/apps/worker/app/core/worker_bootstrap.py index 23d48ded8..0947784f6 100644 --- a/apps/worker/app/core/worker_bootstrap.py +++ b/apps/worker/app/core/worker_bootstrap.py @@ -5,13 +5,19 @@ import subprocess import sys -from celery.signals import worker_init, worker_shutdown +from celery.signals import worker_init, worker_shutdown, worker_shutting_down +from celery.worker import WorkController from loguru import logger +from app.core.gevent_worker_shutdown import GeventWorkerShutdownController from shared.core.celery_app import celery_app from shared.core.logging import setup_logging from shared.services.worker_health import start_worker_heartbeat, stop_worker_heartbeat +_worker_shutdown_controller: GeventWorkerShutdownController | None = None +_CHILD_PROCESS_TERM_TIMEOUT_SECONDS: float = 5 +_CHILD_PROCESS_KILL_TIMEOUT_SECONDS: float = 5 + def _register_task_modules() -> None: """Import task modules for Celery side-effect registration.""" @@ -30,44 +36,31 @@ def _stop_child_process( process.terminate() try: - process.wait(timeout=5) + process.wait(timeout=_CHILD_PROCESS_TERM_TIMEOUT_SECONDS) except subprocess.TimeoutExpired: logger.warning(f"{process_name} did not stop after SIGTERM; killing it") process.kill() - process.wait(timeout=5) + process.wait(timeout=_CHILD_PROCESS_KILL_TIMEOUT_SECONDS) @worker_init.connect -def init_worker(**kwargs) -> None: +def init_worker( + sender: WorkController | None = None, + **kwargs: object, +) -> None: """Initialize structured logging and sync Redis when worker process starts.""" + global _worker_shutdown_controller + setup_logging(service_name="knowhere-worker") start_worker_heartbeat() - # Celery gevent cannot cancel greenlets on transport reconnect, so use a no-op. - try: - from celery.concurrency.gevent import TaskPool as GeventTaskPool - - if not hasattr(GeventTaskPool, "_original_terminate_job"): - - def _graceful_terminate_job(self, pid, signal=None): - logger.warning( - f"gevent pool cannot kill greenlet (pid={pid}), " - f"relying on RedisJobLock for dedup on redelivery" - ) - - setattr( - GeventTaskPool, - "_original_terminate_job", - getattr( - GeventTaskPool, - "terminate_job", - None, - ), - ) - GeventTaskPool.terminate_job = _graceful_terminate_job - logger.info("Patched gevent TaskPool.terminate_job for graceful recovery") - except Exception as exc: - logger.warning(f"Could not patch gevent TaskPool: {exc}") + if sender is None: + logger.error("Cannot configure bounded worker shutdown without worker sender") + else: + _worker_shutdown_controller = GeventWorkerShutdownController( + worker=sender, + timeout_seconds=float(celery_app.conf.worker_soft_shutdown_timeout), + ) try: from shared.services.redis.redis_sync_service import SyncRedisServiceFactory @@ -81,9 +74,39 @@ def _graceful_terminate_job(self, pid, signal=None): logger.warning(f"Worker sync Redis init deferred: {exc}") +@worker_shutting_down.connect +def begin_bounded_worker_shutdown( + sender: str | None = None, + sig: str | None = None, + how: str | None = None, + **kwargs: object, +) -> None: + """Bound ECS SIGTERM without entering Celery's gevent-unsafe cold path.""" + if sig != "SIGTERM" or how != "Warm": + return + + shutdown_controller: GeventWorkerShutdownController | None = ( + _worker_shutdown_controller + ) + if shutdown_controller is None: + logger.error("Cannot schedule bounded worker shutdown before worker init") + return + + shutdown_controller.schedule() + + @worker_shutdown.connect -def shutdown_worker(**kwargs) -> None: +def shutdown_worker(**kwargs: object) -> None: """Clean up shared resources on worker shutdown.""" + global _worker_shutdown_controller + + shutdown_controller: GeventWorkerShutdownController | None = ( + _worker_shutdown_controller + ) + _worker_shutdown_controller = None + if shutdown_controller is not None: + shutdown_controller.close() + try: stop_worker_heartbeat() logger.info("Worker heartbeat stopped") diff --git a/apps/worker/tests/contract/test_worker_shutdown_contract.py b/apps/worker/tests/contract/test_worker_shutdown_contract.py index 34e39827e..7c818991d 100644 --- a/apps/worker/tests/contract/test_worker_shutdown_contract.py +++ b/apps/worker/tests/contract/test_worker_shutdown_contract.py @@ -1,16 +1,391 @@ from __future__ import annotations import json +import os +import signal +import subprocess +import sys +import time +from collections.abc import Callable from pathlib import Path +from types import SimpleNamespace +from typing import cast import pytest +from celery import Celery from pydantic import ValidationError +_PROCESS_CONTRACT_QUEUE_NAME: str = "worker_shutdown_contract" +_PROCESS_CONTRACT_TASK_NAME: str = ( + "worker_shutdown_contract.run_blocking_task" +) +_PROCESS_CONTRACT_SOFT_TIMEOUT_SECONDS: float = 0.25 +_PROCESS_CONTRACT_EXIT_DEADLINE_SECONDS: float = 5.0 + + +def _wait_for_started_tasks( + marker_path: Path, + worker_process: subprocess.Popen[str], + task_count: int, + timeout_seconds: float, +) -> bool: + deadline: float = time.monotonic() + timeout_seconds + while time.monotonic() < deadline: + if marker_path.exists() and len( + marker_path.read_text(encoding="utf-8").splitlines() + ) >= task_count: + return True + if worker_process.poll() is not None: + return False + time.sleep(0.05) + return marker_path.exists() and len( + marker_path.read_text(encoding="utf-8").splitlines() + ) >= task_count + + +def _stop_process(worker_process: subprocess.Popen[str]) -> str: + if worker_process.poll() is None: + worker_process.kill() + worker_process.wait(timeout=5) + + if worker_process.stdout is None: + return "" + return worker_process.stdout.read() + + +def _start_shutdown_contract_worker( + repository_root: Path, + broker_directory: Path, + started_marker_path: Path, + heartbeat_path: Path, + concurrency: int, + shutdown_timeout_seconds: float, + fail_result_backend: bool, +) -> subprocess.Popen[str]: + process_environment: dict[str, str] = os.environ.copy() + python_paths: tuple[str, ...] = ( + str(repository_root / "apps/worker"), + str(repository_root / "apps/worker/tests"), + str(repository_root / "packages/shared-python"), + ) + existing_python_path: str | None = process_environment.get("PYTHONPATH") + process_environment["PYTHONPATH"] = os.pathsep.join( + (*python_paths, existing_python_path) + if existing_python_path + else python_paths + ) + process_environment["WORKER_HEARTBEAT_FILE"] = str(heartbeat_path) + process_environment["WORKER_SHUTDOWN_CONTRACT_BROKER_DIRECTORY"] = str( + broker_directory + ) + process_environment["WORKER_SHUTDOWN_CONTRACT_STARTED_MARKER"] = str( + started_marker_path + ) + process_environment["WORKER_SHUTDOWN_CONTRACT_TIMEOUT_SECONDS"] = str( + shutdown_timeout_seconds + ) + if fail_result_backend: + process_environment[ + "WORKER_SHUTDOWN_CONTRACT_FAIL_RESULT_BACKEND" + ] = "1" + else: + process_environment.pop( + "WORKER_SHUTDOWN_CONTRACT_FAIL_RESULT_BACKEND", + None, + ) + + worker_command: list[str] = [ + sys.executable, + "-m", + "celery", + "-A", + "support.worker_shutdown_process_app:celery_app", + "worker", + "--pool=gevent", + f"--concurrency={concurrency}", + "--loglevel=INFO", + "--hostname=shutdown-contract@%h", + "-Q", + _PROCESS_CONTRACT_QUEUE_NAME, + "--without-gossip", + "--without-mingle", + "--without-heartbeat", + ] + return subprocess.Popen( + worker_command, + env=process_environment, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + +def _create_shutdown_contract_producer(broker_directory: Path) -> Celery: + producer_app: Celery = Celery( + "worker-shutdown-contract-producer", + broker="filesystem://", + ) + producer_app.conf.broker_transport_options = { + "data_folder_in": str(broker_directory), + "data_folder_out": str(broker_directory), + "control_folder": str(broker_directory / "control"), + "store_processed": False, + } + return producer_app + + +def _wait_for_worker_exit( + worker_process: subprocess.Popen[str], + repeat_signal: bool = False, +) -> tuple[float, str]: + shutdown_started_at: float = time.monotonic() + os.kill(worker_process.pid, signal.SIGTERM) + if repeat_signal: + time.sleep(0.05) + os.kill(worker_process.pid, signal.SIGTERM) + + try: + worker_process.wait(timeout=_PROCESS_CONTRACT_EXIT_DEADLINE_SECONDS) + except subprocess.TimeoutExpired as exc: + worker_output: str = _stop_process(worker_process) + raise AssertionError( + "warm SIGTERM did not cancel the active gevent task before " + "the scaled ECS stop deadline\n" + f"{worker_output}" + ) from exc + + shutdown_elapsed_seconds: float = time.monotonic() - shutdown_started_at + return shutdown_elapsed_seconds, _stop_process(worker_process) + + +def test_should_cancel_an_active_gevent_task_before_the_ecs_stop_deadline( + tmp_path: Path, + worker_contract_environment: None, +) -> None: + broker_directory: Path = tmp_path / "broker" + broker_directory.mkdir() + started_marker_path: Path = tmp_path / "task-started" + worker_process: subprocess.Popen[str] = _start_shutdown_contract_worker( + repository_root=Path(__file__).resolve().parents[4], + broker_directory=broker_directory, + started_marker_path=started_marker_path, + heartbeat_path=tmp_path / "worker-heartbeat", + concurrency=1, + shutdown_timeout_seconds=_PROCESS_CONTRACT_SOFT_TIMEOUT_SECONDS, + fail_result_backend=False, + ) + producer_app: Celery = _create_shutdown_contract_producer(broker_directory) + + try: + producer_app.send_task( + _PROCESS_CONTRACT_TASK_NAME, + args=[1, 30.0], + queue=_PROCESS_CONTRACT_QUEUE_NAME, + ) + assert _wait_for_started_tasks( + started_marker_path, + worker_process, + task_count=1, + timeout_seconds=15, + ), _stop_process(worker_process) + + shutdown_elapsed_seconds, _ = _wait_for_worker_exit(worker_process) + assert ( + shutdown_elapsed_seconds < _PROCESS_CONTRACT_EXIT_DEADLINE_SECONDS + ) + assert worker_process.returncode == 0 + finally: + producer_app.close() + _stop_process(worker_process) + + +def test_should_continue_cancelling_tasks_when_result_backend_recording_fails( + tmp_path: Path, + worker_contract_environment: None, +) -> None: + broker_directory: Path = tmp_path / "broker" + broker_directory.mkdir() + started_marker_path: Path = tmp_path / "task-started" + worker_process: subprocess.Popen[str] = _start_shutdown_contract_worker( + repository_root=Path(__file__).resolve().parents[4], + broker_directory=broker_directory, + started_marker_path=started_marker_path, + heartbeat_path=tmp_path / "worker-heartbeat", + concurrency=2, + shutdown_timeout_seconds=_PROCESS_CONTRACT_SOFT_TIMEOUT_SECONDS, + fail_result_backend=True, + ) + producer_app: Celery = _create_shutdown_contract_producer(broker_directory) + + try: + for task_number in (1, 2): + producer_app.send_task( + _PROCESS_CONTRACT_TASK_NAME, + args=[task_number, 30.0], + queue=_PROCESS_CONTRACT_QUEUE_NAME, + ) + assert _wait_for_started_tasks( + started_marker_path, + worker_process, + task_count=2, + timeout_seconds=15, + ), _stop_process(worker_process) + + shutdown_elapsed_seconds, worker_output = _wait_for_worker_exit( + worker_process + ) + + assert worker_process.returncode == 0 + assert ( + shutdown_elapsed_seconds < _PROCESS_CONTRACT_EXIT_DEADLINE_SECONDS + ) + assert worker_output.count("could not record shutdown retry") == 2, worker_output + finally: + producer_app.close() + _stop_process(worker_process) + + +def test_should_not_cancel_a_task_that_finishes_before_shutdown_timeout( + tmp_path: Path, + worker_contract_environment: None, +) -> None: + broker_directory: Path = tmp_path / "broker" + broker_directory.mkdir() + started_marker_path: Path = tmp_path / "task-started" + worker_process: subprocess.Popen[str] = _start_shutdown_contract_worker( + repository_root=Path(__file__).resolve().parents[4], + broker_directory=broker_directory, + started_marker_path=started_marker_path, + heartbeat_path=tmp_path / "worker-heartbeat", + concurrency=1, + shutdown_timeout_seconds=0.75, + fail_result_backend=False, + ) + producer_app: Celery = _create_shutdown_contract_producer(broker_directory) + + try: + producer_app.send_task( + _PROCESS_CONTRACT_TASK_NAME, + args=[1, 0.05], + queue=_PROCESS_CONTRACT_QUEUE_NAME, + ) + assert _wait_for_started_tasks( + started_marker_path, + worker_process, + task_count=1, + timeout_seconds=15, + ), _stop_process(worker_process) + time.sleep(0.2) + + _, worker_output = _wait_for_worker_exit(worker_process) + + assert worker_process.returncode == 0 + assert "cancelling 1 unacknowledged active task" not in worker_output + finally: + producer_app.close() + _stop_process(worker_process) + + +def test_should_schedule_only_one_shutdown_timer_for_repeated_sigterm( + tmp_path: Path, + worker_contract_environment: None, +) -> None: + broker_directory: Path = tmp_path / "broker" + broker_directory.mkdir() + started_marker_path: Path = tmp_path / "task-started" + worker_process: subprocess.Popen[str] = _start_shutdown_contract_worker( + repository_root=Path(__file__).resolve().parents[4], + broker_directory=broker_directory, + started_marker_path=started_marker_path, + heartbeat_path=tmp_path / "worker-heartbeat", + concurrency=1, + shutdown_timeout_seconds=_PROCESS_CONTRACT_SOFT_TIMEOUT_SECONDS, + fail_result_backend=False, + ) + producer_app: Celery = _create_shutdown_contract_producer(broker_directory) + + try: + producer_app.send_task( + _PROCESS_CONTRACT_TASK_NAME, + args=[1, 30.0], + queue=_PROCESS_CONTRACT_QUEUE_NAME, + ) + assert _wait_for_started_tasks( + started_marker_path, + worker_process, + task_count=1, + timeout_seconds=15, + ), _stop_process(worker_process) + + _, worker_output = _wait_for_worker_exit( + worker_process, + repeat_signal=True, + ) + + assert worker_process.returncode == 0 + assert worker_output.count( + "Scheduled bounded warm shutdown cancellation" + ) == 1 + finally: + producer_app.close() + _stop_process(worker_process) + + +def test_should_keep_gevent_pool_cancellation_disabled_for_broker_reconnects( + monkeypatch: pytest.MonkeyPatch, + worker_contract_environment: None, +) -> None: + import gevent + from app.core.gevent_worker_shutdown import GeventWorkerShutdownController + from celery.concurrency.gevent import TaskPool as GeventTaskPool + from celery.worker import WorkController + + original_terminate_job: Callable[ + [GeventTaskPool, int, int | None], None + ] = cast( + Callable[[GeventTaskPool, int, int | None], None], + getattr(GeventTaskPool, "_original_terminate_job", None), + ) + if not callable(original_terminate_job): + original_terminate_job = GeventTaskPool.terminate_job + + pool: GeventTaskPool = GeventTaskPool(1) + pool.start() + controller = GeventWorkerShutdownController( + worker=cast( + WorkController, + SimpleNamespace(pool=pool), + ), + timeout_seconds=1, + ) + + def wait_for_reconnect_contract() -> None: + gevent.sleep(10) + + running_greenlet = gevent.spawn(wait_for_reconnect_contract) + pool._pool_map[id(running_greenlet)] = running_greenlet + + try: + pool.terminate_job(id(running_greenlet)) + gevent.sleep(0) + assert running_greenlet.dead is False + finally: + controller.close() + original_terminate_job(pool, id(running_greenlet), None) + gevent.sleep(0) + pool.stop() + monkeypatch.setattr( + GeventTaskPool, + "terminate_job", + original_terminate_job, + ) + def test_should_preserve_fargate_worker_sigterm_shutdown_contract( worker_contract_environment: None, ) -> None: from shared.core.celery_app import celery_app + from app.core import worker_bootstrap repository_root: Path = Path(__file__).resolve().parents[4] task_definition_path: Path = ( @@ -36,6 +411,15 @@ def test_should_preserve_fargate_worker_sigterm_shutdown_contract( assert celery_app.conf.worker_enable_soft_shutdown_on_idle is True assert "REMAP_SIGTERM" not in environment_values assert worker_container["stopTimeout"] == 120 + maximum_child_cleanup_seconds: float = 2 * ( + worker_bootstrap._CHILD_PROCESS_TERM_TIMEOUT_SECONDS + + worker_bootstrap._CHILD_PROCESS_KILL_TIMEOUT_SECONDS + ) + assert ( + celery_app.conf.worker_soft_shutdown_timeout + + maximum_child_cleanup_seconds + < worker_container["stopTimeout"] + ) def test_should_redeliver_interrupted_tasks_before_processing_jobs_expire( diff --git a/apps/worker/tests/support/worker_shutdown_process_app.py b/apps/worker/tests/support/worker_shutdown_process_app.py new file mode 100644 index 000000000..104c0cebd --- /dev/null +++ b/apps/worker/tests/support/worker_shutdown_process_app.py @@ -0,0 +1,92 @@ +"""Minimal Celery application for the worker shutdown process contract.""" + +from __future__ import annotations + +# This process must reproduce the production worker's cooperative runtime. +# Patching after Celery or Redis imports would leave blocking sockets in place. +import gevent.monkey + +gevent.monkey.patch_all() + +import importlib +import os +from pathlib import Path + +import gevent +from celery.worker.request import Request +from kombu import Queue + +# Importing the production bootstrap registers its Celery lifecycle receivers. +# The contract deliberately exercises those receivers instead of a test-only hook. +importlib.import_module("app.core.worker_bootstrap") +from shared.core.celery_app import celery_app + +_BROKER_DIRECTORY_ENVIRONMENT_VARIABLE: str = ( + "WORKER_SHUTDOWN_CONTRACT_BROKER_DIRECTORY" +) +_STARTED_MARKER_ENVIRONMENT_VARIABLE: str = ( + "WORKER_SHUTDOWN_CONTRACT_STARTED_MARKER" +) +_SHUTDOWN_TIMEOUT_ENVIRONMENT_VARIABLE: str = ( + "WORKER_SHUTDOWN_CONTRACT_TIMEOUT_SECONDS" +) +_FAIL_RESULT_BACKEND_ENVIRONMENT_VARIABLE: str = ( + "WORKER_SHUTDOWN_CONTRACT_FAIL_RESULT_BACKEND" +) +_QUEUE_NAME: str = "worker_shutdown_contract" +_TASK_NAME: str = "worker_shutdown_contract.run_blocking_task" + +broker_directory: Path = Path( + os.environ[_BROKER_DIRECTORY_ENVIRONMENT_VARIABLE] +) +started_marker_path: Path = Path( + os.environ[_STARTED_MARKER_ENVIRONMENT_VARIABLE] +) +shutdown_timeout_seconds: float = float( + os.environ[_SHUTDOWN_TIMEOUT_ENVIRONMENT_VARIABLE] +) + +celery_app.conf.broker_url = "filesystem://" +celery_app.conf.result_backend = "cache+memory://" +celery_app.conf.broker_transport_options = { + "data_folder_in": str(broker_directory), + "data_folder_out": str(broker_directory), + "control_folder": str(broker_directory / "control"), + "store_processed": False, +} +celery_app.conf.task_default_queue = _QUEUE_NAME +celery_app.conf.task_queues = (Queue(_QUEUE_NAME),) +celery_app.conf.task_routes = {_TASK_NAME: {"queue": _QUEUE_NAME}} +celery_app.conf.worker_soft_shutdown_timeout = shutdown_timeout_seconds +celery_app.conf.worker_enable_soft_shutdown_on_idle = True + + +def _raise_result_backend_failure(*args: object, **kwargs: object) -> None: + raise RuntimeError("shutdown contract result backend unavailable") + + +if os.getenv(_FAIL_RESULT_BACKEND_ENVIRONMENT_VARIABLE) == "1": + # Request.cancel kills the greenlet before it records retry state. Force a + # bookkeeping failure after that call so the contract proves later tasks + # are still canceled when the result backend is unavailable. + celery_app.backend.mark_as_retry = _raise_result_backend_failure + _original_request_cancel = Request.cancel + + def _cancel_then_fail( + request: Request, + pool: object, + signal: int | None = None, + ) -> None: + _original_request_cancel(request, pool, signal) + raise RuntimeError("shutdown contract retry bookkeeping failed") + + Request.cancel = _cancel_then_fail + + +@celery_app.task(name=_TASK_NAME, acks_late=True) +def run_blocking_task(task_number: int, duration_seconds: float) -> None: + """Remain active until the shutdown contract cancels this greenlet.""" + with started_marker_path.open("a", encoding="utf-8") as marker_file: + marker_file.write(f"{task_number}\n") + marker_file.flush() + gevent.sleep(duration_seconds)