Clean up a Lang-SDK runtime parse whose start fails after the fork

Only the fork itself was guarded. A failure while registering the listeners or sending the start request left the child waiting for that request and both listeners open, and the standalone run does not guard start() either.

Author: jason810496

airflow-core/src/airflow/dag_processing/task_handler_processor.py

@@ -253,23 +253,28 @@ def start(  # type: ignore[override]
             for listener in listeners.values():
                 listener.close()
             raise
-        for channel, listener in listeners.items():
-            proc._open_sockets[listener] = f"{channel}-listener"
-            proc.selector.register(
-                listener,
-                selectors.EVENT_READ,
-                (functools.partial(proc._accept_connection, channel=channel), proc._on_socket_closed),
+        try:
+            for channel, listener in listeners.items():
+                proc._open_sockets[listener] = f"{channel}-listener"
+                proc.selector.register(
+                    listener,
+                    selectors.EVENT_READ,
+                    (functools.partial(proc._accept_connection, channel=channel), proc._on_socket_closed),
+                )
+            proc.send_msg(
+                StartTaskHandlerRuntime(
+                    file=parse_request.file,
+                    bundle_path=bundle_path,
+                    coordinator=coordinator,
+                    comm_address=listeners["comm"].getsockname()[:2],
+                    logs_address=listeners["logs"].getsockname()[:2],
+                ),
+                request_id=0,
             )
-        proc.send_msg(
-            StartTaskHandlerRuntime(
-                file=parse_request.file,
-                bundle_path=bundle_path,
-                coordinator=coordinator,
-                comm_address=listeners["comm"].getsockname()[:2],
-                logs_address=listeners["logs"].getsockname()[:2],
-            ),
-            request_id=0,
-        )
+        except BaseException:
+            proc._kill_runtime()
+            proc.close()
+            raise
         return proc
 
     @classmethod

airflow-core/tests/unit/dag_processing/test_task_handler_processor.py

@@ -250,6 +250,24 @@ def test_a_connection_is_used_once_it_is_verified(self, mock_parse_task_handler,
 
         assert _get_task_ids(proc.parsing_result) == ["extract"]
 
+    @patch.object(
+        LangSDKTaskHandlerProcessorProcess,
+        "send_msg",
+        autospec=True,
+        side_effect=OSError("cannot send the start request"),
+    )
+    def test_a_start_that_fails_after_the_fork_leaves_nothing_behind(self, mock_send_msg, tmp_path):
+        fds_before = _get_open_fds()
+        children_before = {child.pid for child in psutil.Process().children()}
+
+        with selectors.DefaultSelector() as selector:
+            with pytest.raises(OSError, match="cannot send the start request"):
+                _start(tmp_path, selector)
+            assert selector.get_map() == {}
+
+        assert {child.pid for child in psutil.Process().children()} == children_before
+        assert _get_open_fds() <= fds_before
+
     @patch.object(BaseDagFileProcessorProcess, "start", autospec=True, side_effect=OSError("fork failed"))
     def test_a_start_that_fails_before_the_fork_closes_its_listeners(self, mock_start, tmp_path):
         fds_before = _get_open_fds()