Fix Edge worker SIGTERM storm during API server outage

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

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