diff --git a/providers/edge3/src/airflow/providers/edge3/cli/worker.py b/providers/edge3/src/airflow/providers/edge3/cli/worker.py index 2e4010e9dc778..f5396127219eb 100644 --- a/providers/edge3/src/airflow/providers/edge3/cli/worker.py +++ b/providers/edge3/src/airflow/providers/edge3/cli/worker.py @@ -23,6 +23,7 @@ import traceback from asyncio import Task, create_task, get_running_loop, sleep from collections.abc import Awaitable, Callable +from contextlib import suppress from datetime import datetime from functools import cached_property from http import HTTPStatus @@ -78,6 +79,25 @@ def _edge_hostname() -> str: return os.environ.get("HOSTNAME", getfqdn()) +def _reset_parent_signal_state() -> None: + """ + Detach a forked child from the parent's asyncio signal plumbing. + + The parent installs asyncio signal handlers for SIGTERM/SIGINT/SIG_STATUS + via ``loop.add_signal_handler``, which internally uses + ``signal.set_wakeup_fd`` on one end of a shared socketpair. On Linux + ``fork()`` duplicates that fd into the child; signals delivered to the + child then write bytes into the socketpair the parent is reading from, + re-firing the parent's handlers. Reset the inherited state before the + child runs any supervised code. + """ + signal.set_wakeup_fd(-1) + signal.signal(signal.SIGTERM, signal.SIG_DFL) + signal.signal(signal.SIGINT, signal.SIG_DFL) + with suppress(ValueError, OSError): + signal.signal(SIG_STATUS, signal.SIG_DFL) + + class EdgeWorker: """Runner instance which executes the Edge Worker.""" @@ -194,13 +214,14 @@ def signal_drain(self): logger.info("Request to shut down Edge Worker received, waiting for jobs to complete.") def shutdown_handler(self): + if self.drain: + return self.drain = True - msg = "SIGTERM received. Sending SIGTERM to all jobs and quit" - logger.info(msg) + logger.info("SIGTERM received. Sending SIGTERM to all jobs and quit") for job in self.jobs: if job.process.pid: - os.setpgid(job.process.pid, 0) - os.kill(job.process.pid, signal.SIGTERM) + with suppress(ProcessLookupError, PermissionError): + os.kill(job.process.pid, signal.SIGTERM) async def _get_sysinfo(self) -> dict[str, str | int | float | datetime]: """Produce the sysinfo from worker to post to central site.""" @@ -240,6 +261,8 @@ def _get_state(self) -> EdgeWorkerState: return EdgeWorkerState.IDLE def _run_job_via_supervisor(self, workload: ExecuteTask, results_queue: Queue) -> int: + _reset_parent_signal_state() + from airflow.sdk.execution_time.supervisor import supervise # Ignore ctrl-c in this process -- we don't want to kill _this_ one. we let tasks run to completion diff --git a/providers/edge3/tests/unit/edge3/cli/test_worker.py b/providers/edge3/tests/unit/edge3/cli/test_worker.py index dcfbe6614ef3d..62565cdf98cac 100644 --- a/providers/edge3/tests/unit/edge3/cli/test_worker.py +++ b/providers/edge3/tests/unit/edge3/cli/test_worker.py @@ -19,6 +19,7 @@ import contextlib import importlib import json +import signal from datetime import datetime from io import StringIO from multiprocessing import Process, Queue @@ -36,6 +37,7 @@ from airflow.providers.common.compat.sdk import timezone from airflow.providers.edge3.cli import edge_command from airflow.providers.edge3.cli.dataclasses import Job +from airflow.providers.edge3.cli.signalling import SIG_STATUS from airflow.providers.edge3.cli.worker import EdgeWorker from airflow.providers.edge3.models.edge_worker import ( EdgeWorkerModel, @@ -572,3 +574,121 @@ def test_list_edge_workers(self, mock_edgeworker: EdgeWorkerModel): ]: assert key in edge_workers[0] assert any("test_edge_worker" in h["worker_name"] for h in edge_workers) + + +class TestSignalHandling: + """Regression tests for the SIGTERM-storm fix (Bug B).""" + + # Override the module-level ``pytestmark = [pytest.mark.asyncio]`` — these + # tests are intentionally synchronous. + pytestmark: list = [] + + @staticmethod + def _make_workload(): + # Pydantic coerces MOCK_COMMAND (dict) into a real ExecuteTask via + # EdgeJobFetched's field validator, giving us a workload with + # attribute access (``.ti.dag_id`` etc.) that the code path needs. + return EdgeJobFetched( + dag_id="test", + task_id="test1", + run_id="test", + map_index=-1, + try_number=1, + concurrency_slots=1, + command=MOCK_COMMAND, # type: ignore[arg-type] + ).command + + @staticmethod + def _make_job(tmp_path: Path) -> Job: + logfile = tmp_path / "job.log" + logfile.touch() + proc = mock.Mock() + proc.pid = 12345 + return Job( + edge_job=EdgeJobFetched( + dag_id="test", + task_id="test1", + run_id="test", + map_index=-1, + try_number=1, + concurrency_slots=1, + command=MOCK_COMMAND, # type: ignore[arg-type] + ), + process=proc, + logfile=logfile, + logsize=0, + ) + + @pytest.fixture + def worker_with_one_job(self, tmp_path: Path): + worker = EdgeWorker(str(tmp_path / "mock.pid"), "mock", None, 8) + EdgeWorker.jobs = [self._make_job(tmp_path)] + EdgeWorker.drain = False + try: + yield worker + finally: + EdgeWorker.jobs = [] + EdgeWorker.drain = False + + def test_reset_parent_signal_state_clears_all_handlers(self): + from airflow.providers.edge3.cli.worker import _reset_parent_signal_state + + original = { + signal.SIGTERM: signal.getsignal(signal.SIGTERM), + signal.SIGINT: signal.getsignal(signal.SIGINT), + SIG_STATUS: signal.getsignal(SIG_STATUS), + } + sentinel = lambda *_: None + signal.signal(signal.SIGTERM, sentinel) + signal.signal(signal.SIGINT, sentinel) + signal.signal(SIG_STATUS, sentinel) + try: + _reset_parent_signal_state() + assert signal.getsignal(signal.SIGTERM) is signal.SIG_DFL + assert signal.getsignal(signal.SIGINT) is signal.SIG_DFL + assert signal.getsignal(SIG_STATUS) is signal.SIG_DFL + assert signal.set_wakeup_fd(-1) == -1 + finally: + for sig, prev in original.items(): + signal.signal(sig, prev) + + def test_run_job_via_supervisor_resets_signals_before_supervise(self, tmp_path): + """Reset must run first: before ``os.setpgrp`` and before ``supervise``.""" + worker = EdgeWorker(str(tmp_path / "mock.pid"), "mock", None, 8) + worker.__dict__["_execution_api_server_url"] = "https://mock-server/execution" + order = mock.MagicMock() + with ( + mock.patch( + "airflow.providers.edge3.cli.worker._reset_parent_signal_state", + side_effect=lambda: order("reset"), + ), + mock.patch("os.setpgrp", side_effect=lambda: order("setpgrp")), + mock.patch( + "airflow.sdk.execution_time.supervisor.supervise", + side_effect=lambda **_: order("supervise"), + ), + ): + rc = worker._run_job_via_supervisor( + workload=self._make_workload(), + results_queue=mock.MagicMock(), + ) + assert rc == 0 + assert [c.args[0] for c in order.call_args_list] == ["reset", "setpgrp", "supervise"] + + def test_shutdown_handler_is_idempotent(self, worker_with_one_job): + with mock.patch("os.kill") as m_kill: + worker_with_one_job.shutdown_handler() + worker_with_one_job.shutdown_handler() + worker_with_one_job.shutdown_handler() + assert m_kill.call_count == 1 + assert worker_with_one_job.drain is True + + def test_shutdown_handler_does_not_call_setpgid(self, worker_with_one_job): + with mock.patch("os.setpgid") as m_setpgid, mock.patch("os.kill"): + worker_with_one_job.shutdown_handler() + m_setpgid.assert_not_called() + + def test_shutdown_handler_tolerates_dead_child(self, worker_with_one_job): + with mock.patch("os.kill", side_effect=ProcessLookupError): + worker_with_one_job.shutdown_handler() + assert worker_with_one_job.drain is True