Fix SSHRemoteJobOperator still orphaning the remote job on cancellation

providers/ssh/src/airflow/providers/ssh/utils/remote_job.py

@@ -184,11 +184,22 @@ def build_posix_wrapper_command(
     escaped_command = command.replace("'", "'\"'\"'")
 
     # Launch detached under ``setsid`` so the job is its own session/process-group
-    # leader. ``$!`` is then the leader PID *and* the PGID (verified: setsid does not
-    # fork when started as a background job), recorded synchronously just like before,
-    # so cancellation can signal the whole job tree instead of orphaning the user
-    # command. Without ``setsid`` (some macOS/BSD hosts) ``$!`` is just the wrapper PID
-    # and cancellation degrades to the previous single-process behaviour.
+    # leader, letting cancellation signal the whole job tree instead of orphaning the
+    # user command. We do NOT rely on the launcher's ``$!`` to identify that group:
+    # ``setsid(1)`` forks internally when the process about to exec into it is already
+    # a process-group leader (``setsid(2)`` cannot create a new session from a group
+    # leader), which happens whenever job control is on in the launching shell. When
+    # it forks, ``$!`` is the short-lived setsid parent, not the job's real PGID, and
+    # cancellation signals a dead group. Instead the job script reports its OWN pid:
+    # immediately after ``setsid(2)`` POSIX guarantees ``pid == pgid == sid`` for the
+    # caller, and that identity survives the following exec, so ``$$`` inside the job
+    # script is always the true PGID regardless of whether setsid forked to get there.
+    # The pid file is read only when cancelling (:func:`build_posix_kill_command`),
+    # which happens long after submission, so the job's asynchronous write lands well
+    # before any reader and the launcher does not wait for it (recording the launcher's
+    # ``$!`` would just reintroduce the wrong-pid bug on the fork path). Without
+    # ``setsid`` (some macOS/BSD hosts) ``$$`` is just the job's own PID and
+    # cancellation degrades to the previous single-process behaviour.
     wrapper = f"""set -euo pipefail
 job_dir='{paths.job_dir}'
 log_file='{paths.log_file}'
@@ -202,6 +213,7 @@ def build_posix_wrapper_command(
 
 job_script='
 set +e
+echo -n "$$" > "'"$pid_file"'"
 export LOG_FILE="'"$log_file"'"
 export STATUS_FILE="'"$status_file"'"
 {env_exports}{escaped_command} >>"'"$log_file"'" 2>&1
@@ -216,7 +228,6 @@ def build_posix_wrapper_command(
 else
   nohup bash -c "$job_script" >/dev/null 2>&1 &
 fi
-echo -n $! > "$pid_file"
 echo "{paths.job_id}"
 """
     return wrapper

providers/ssh/tests/unit/ssh/utils/test_remote_job.py

@@ -18,7 +18,10 @@
 from __future__ import annotations
 
 import base64
+import contextlib
 import os
+import pty
+import select
 import shutil
 import subprocess
 import time
@@ -132,16 +135,19 @@ def test_escapes_quotes(self):
         assert wrapper is not None
 
     def test_runs_in_own_process_group(self):
-        """The job launches under setsid (when available); $! is the leader PID/PGID."""
+        """The job launches under setsid (when available) and self-reports its PGID."""
         paths = RemoteJobPaths(job_id="test_job", remote_os="posix")
         wrapper = build_posix_wrapper_command("/path/to/script.sh", paths)
 
         # New session/process group when setsid exists, plain detached run otherwise
         assert "command -v setsid" in wrapper
         assert "setsid bash -c" in wrapper
         assert "nohup bash -c" in wrapper
-        # Leader PID recorded synchronously by the launcher ($! == PGID under setsid)
-        assert 'echo -n $! > "$pid_file"' in wrapper
+        # The job self-reports its own pid ($$ == PGID after setsid), which is correct
+        # even when setsid(1) forks under job control -- unlike the launcher's $!.
+        assert 'echo -n "$$" > "' in wrapper
+        # Launcher must NOT record $! (would be the short-lived setsid parent on a fork).
+        assert 'echo -n $! > "$pid_file"' not in wrapper
 
 
 class TestBuildWindowsWrapperCommand:
@@ -262,46 +268,123 @@ class TestPosixKillBehaviour:
 
     Regression test for the orphaned-process bug: killing only the recorded PID left the
     user command (and its children) running, so the exit_code file was never written and
-    the trigger timed out. The job now runs in its own process group and the kill signals
-    the group.
+    the trigger timed out. The job runs in its own process group and self-reports that
+    group's PGID, so the kill signals the whole group even when setsid(1) forks.
     """
 
+    def _marker(self, tag: str) -> str:
+        # Unique per (test, xdist worker): CI runs these with ``-n auto`` (default
+        # ``load`` distribution), so sibling tests can execute concurrently in separate
+        # workers against the same OS process table. A shared literal would let one
+        # test's ``pgrep -f`` / ``pkill -f`` match or kill another's job. os.getpid()
+        # differs per worker; the tag differs per test.
+        return f"sleep 9{tag}{os.getpid()}"
+
     @staticmethod
     def _group_alive(pgid: int) -> bool:
         # pgrep -g matches by process-group id; rc 0 => at least one member alive.
         return subprocess.run(["pgrep", "-g", str(pgid)], capture_output=True, check=False).returncode == 0
 
-    # setsid only avoids forking when the launching shell is not a process-group leader; on some
-    # CI runners it forks, so the recorded $! is the short-lived setsid parent rather than the job
-    # PGID and the pre-kill pgrep -g finds an empty group. Re-launch on a fresh draw.
-    @pytest.mark.flaky(reruns=5)
-    def test_kill_terminates_whole_job_tree(self, tmp_path):
-        paths = RemoteJobPaths(job_id="killtree", remote_os="posix", base_dir=str(tmp_path / "jobs"))
-        # `sleep 300` runs as a child of the wrapper subshell -> the tree the old kill orphaned.
-        # Run under bash, which is the remote login shell this operator requires (the wrapper
-        # uses `set -o pipefail`); the kill is run the same way below.
-        wrapper = build_posix_wrapper_command("sleep 300", paths)
-        subprocess.run(["bash", "-c", wrapper], check=True, capture_output=True, text=True)
+    @staticmethod
+    def _job_running(marker: str) -> bool:
+        return subprocess.run(["pgrep", "-f", marker], capture_output=True, check=False).returncode == 0
+
+    @staticmethod
+    def _pgid_of(pid: str) -> str:
+        return subprocess.run(
+            ["ps", "-o", "pgid=", "-p", pid], capture_output=True, text=True, check=False
+        ).stdout.strip()
 
-        # The launcher records $! synchronously, so the pid file is present on return.
+    def _await_recorded_pid(self, paths) -> int:
+        # The job writes its pid asynchronously (the launcher does not wait), so poll.
         pid_path = Path(paths.pid_file)
-        assert pid_path.exists(), "job never wrote its pid file"
-        pid_text = pid_path.read_text().strip()
-        assert pid_text, "pid file is empty"
-        pgid = int(pid_text)
+        deadline = time.monotonic() + 5
+        pid_text = ""
+        while time.monotonic() < deadline:
+            pid_text = pid_path.read_text().strip() if pid_path.exists() else ""
+            if pid_text:
+                break
+            time.sleep(0.02)
+        assert pid_text, "job never wrote its pid file"
+        return int(pid_text)
 
+    @staticmethod
+    def _run_bash_mc_under_pty(script: str, marker: bytes, timeout: float = 8.0) -> None:
+        """Run ``bash -mc script`` under a pty we own so job control genuinely activates
+        (bash silently disables ``-m`` without a controlling terminal). Read until the
+        marker, NOT to EOF: the detached job inherits the pty slave as its stdin, so EOF
+        would not arrive until the job itself exits (the full sleep runtime)."""
+        pid, fd = pty.fork()
+        if pid == 0:
+            try:
+                os.execvp("bash", ["bash", "-mc", script])
+            except OSError:
+                os._exit(127)  # never fall through as a duplicate pytest process
         try:
-            assert self._group_alive(pgid), "job tree should be running before kill"
+            deadline = time.monotonic() + timeout
+            buf = b""
+            while time.monotonic() < deadline:
+                r, _, _ = select.select([fd], [], [], 0.2)
+                if fd in r:
+                    try:
+                        chunk = os.read(fd, 4096)
+                    except OSError:
+                        break
+                    if not chunk:
+                        break
+                    buf += chunk
+                    if marker in buf:
+                        break
+        finally:
+            os.close(fd)  # hangs up the pty; the launcher (not the detached job) exits
+            with contextlib.suppress(ChildProcessError):
+                os.waitpid(pid, 0)  # reap the launcher so it does not linger as a zombie
 
+    def _assert_kill_tears_down(self, paths, pgid: int, marker: str) -> None:
+        try:
+            deadline = time.monotonic() + 5
+            while time.monotonic() < deadline and not self._group_alive(pgid):
+                time.sleep(0.02)
+            assert self._group_alive(pgid), "job group should be running before kill"
             subprocess.run(["bash", "-c", build_posix_kill_command(paths.pid_file)], check=True)
-
             deadline = time.monotonic() + 5
-            while time.monotonic() < deadline and self._group_alive(pgid):
+            while time.monotonic() < deadline and self._job_running(marker):
                 time.sleep(0.05)
-            assert not self._group_alive(pgid), "kill left part of the job tree running"
+            assert not self._job_running(marker), "kill left the job running (orphaned)"
         finally:
-            # Belt-and-suspenders: never leave a stray `sleep 300` behind if an assert fails.
             subprocess.run(["bash", "-c", f"kill -9 -{pgid} 2>/dev/null || true"], check=False)
+            subprocess.run(["bash", "-c", f"pkill -9 -f '{marker}' 2>/dev/null || true"], check=False)
+
+    def test_kill_terminates_whole_job_tree(self, tmp_path):
+        """Default path (no job control): the job self-reports its PGID and the kill
+        signals the whole group."""
+        marker = self._marker("1")
+        paths = RemoteJobPaths(job_id="killtree", remote_os="posix", base_dir=str(tmp_path / "jobs"))
+        wrapper = build_posix_wrapper_command(marker, paths)
+        subprocess.run(["bash", "-c", wrapper], check=True, capture_output=True, text=True)
+        pgid = self._await_recorded_pid(paths)
+        self._assert_kill_tears_down(paths, pgid, marker)
+
+    def test_kill_terminates_whole_job_tree_under_job_control(self, tmp_path):
+        """With job control on, setsid(1) forks and the launcher's ``$!`` would name the
+        short-lived setsid parent, not the job -- the condition the old wrapper orphaned
+        the job under. Force it deterministically via a real controlling terminal and
+        assert the recorded pid IS the job's true PGID and the kill reaches the job."""
+        marker = self._marker("2")
+        paths = RemoteJobPaths(job_id="killtree_jc", remote_os="posix", base_dir=str(tmp_path / "jobs"))
+        wrapper = build_posix_wrapper_command(marker, paths)
+        self._run_bash_mc_under_pty(wrapper + "\necho SUBMIT_DONE\n", b"SUBMIT_DONE")
+        pgid = self._await_recorded_pid(paths)
+
+        job_pids = subprocess.run(
+            ["pgrep", "-f", marker], capture_output=True, text=True, check=False
+        ).stdout.split()
+        assert job_pids, "job never started"
+        true_pgid = self._pgid_of(job_pids[0])
+        # Core regression assertion: recorded pid == the job's real PGID. Under the old
+        # $!-based wrapper this differs (setsid forked) and on_kill orphans the job.
+        assert str(pgid) == true_pgid, f"recorded pid {pgid} is not the job PGID {true_pgid}"
+        self._assert_kill_tears_down(paths, pgid, marker)
 
 
 class TestCleanupCommands: