manual state change should not use fork-execute model on scheduler

Signed-off-by: Maciej Obuchowski <maciej.obuchowski@datadoghq.com>

Author: mobuchowski

providers/openlineage/src/airflow/providers/openlineage/plugins/listener.py

@@ -81,6 +81,19 @@ def _executor_initializer():
     settings.configure_orm()
 
 
+def _emit_manual_state_change_event(adapter_method, stats_key, **kwargs):
+    """
+    Emit an OL event via the given adapter method and record its serialized size.
+
+    Module-level so it is picklable across the ProcessPoolExecutor boundary used by
+    `_on_task_instance_manual_state_change` for scheduler-side "task state changed
+    externally" emissions.
+    """
+    event = adapter_method(**kwargs)
+    Stats.gauge(stats_key, len(Serde.to_json(event).encode("utf-8")))
+    return event
+
+
 class OpenLineageListener:
     """OpenLineage listener sends events on task instance and dag run starts, completes and failures."""
 
@@ -653,6 +666,17 @@ def _on_task_instance_manual_state_change(
         ti_state: TaskInstanceState,
         error: None | str | BaseException = None,
     ) -> None:
+        """
+        Emit an OL event from the scheduler when a TI transitions externally.
+
+        This path is only reached on the scheduler (``process_executor_events ->
+        handle_failure``, or manual UI/API state changes). Emission is routed through
+        the same ``ProcessPoolExecutor`` the DAG-run listeners use rather than through
+        ``_fork_execute``: the pool's ``_executor_initializer`` rebuilds the ORM once
+        per worker, so the child never shares a pooled Postgres SSL connection with
+        the scheduler, and bursts of external-state-change events no longer produce a
+        fork-per-event.
+        """
         self.log.debug("`_on_task_instance_manual_state_change` was called with state: `%s`.", ti_state)
         end_date = timezone.utcnow()
 
@@ -674,22 +698,37 @@ def _on_task_instance_manual_state_change(
             )
             return
 
-        @print_warning(self.log)
-        def on_state_change():
-            date = dagrun.logical_date or dagrun.run_after
-            parent_run_id = self.adapter.build_dag_run_id(
-                dag_id=ti.dag_id,
-                logical_date=date,
-                clear_number=dagrun.clear_number,
-            )
+        try:
+            if not self.executor:
+                self.log.debug("Executor has not started before `_on_task_instance_manual_state_change`")
+                return
+
+            if ti_state == TaskInstanceState.FAILED:
+                adapter_method = self.adapter.fail_task
+                event_type = RunState.FAIL.value.lower()
+            elif ti_state in (TaskInstanceState.SUCCESS, TaskInstanceState.SKIPPED):
+                adapter_method = self.adapter.complete_task
+                event_type = RunState.COMPLETE.value.lower()
+            else:
+                raise ValueError(f"Unsupported ti_state: `{ti_state}`.")
 
+            # Extract primitives from live ORM objects in the parent (scheduler)
+            # before crossing the pool boundary. Passing ORM objects through the pool
+            # pickler loses TaskGroup attributes and crashes event emission -- see
+            # the equivalent note in `on_dag_run_running` (listener.py ~868).
+            date = dagrun.logical_date or dagrun.run_after
             task_uuid = self.adapter.build_task_instance_run_id(
                 dag_id=ti.dag_id,
                 task_id=ti.task_id,
                 try_number=ti.try_number,
                 logical_date=date,
                 map_index=ti.map_index,
             )
+            parent_run_id = self.adapter.build_dag_run_id(
+                dag_id=ti.dag_id,
+                logical_date=date,
+                clear_number=dagrun.clear_number,
+            )
 
             data_interval_start = dagrun.data_interval_start
             if isinstance(data_interval_start, datetime):
@@ -698,21 +737,22 @@ def on_state_change():
             if isinstance(data_interval_end, datetime):
                 data_interval_end = data_interval_end.isoformat()
 
-            dag_tags, owners, doc, doc_type = None, None, None, None
-            airflow_run_facet = {}
+            dag_tags: list | None = None
+            owners: list[str] | None = None
+            doc: str | None = None
+            doc_type: str | None = None
+            airflow_run_facet: dict = {}
             if task:  # on scheduler, we should have access to task
                 doc, doc_type = get_task_documentation(task)
                 dag = getattr(task, "dag")
                 if dag:
                     if not doc:
                         doc, doc_type = get_dag_documentation(dag)
-
                     dag_tags = dag.tags
                     owners = [x.strip() for x in (task if task.owner != "airflow" else dag).owner.split(",")]
-
                     airflow_run_facet = get_airflow_run_facet(dagrun, dag, ti, task, task_uuid)
 
-            adapter_kwargs = {
+            adapter_kwargs: dict = {
                 "run_id": task_uuid,
                 "job_name": get_job_name(ti),
                 "end_time": end_date.isoformat(),
@@ -733,23 +773,20 @@ def on_state_change():
                     **get_airflow_debug_facet(),
                 },
             }
-
             if ti_state == TaskInstanceState.FAILED:
-                event_type = RunState.FAIL.value.lower()
-                redacted_event = self.adapter.fail_task(**adapter_kwargs, error=error)
-            elif ti_state in (TaskInstanceState.SUCCESS, TaskInstanceState.SKIPPED):
-                event_type = RunState.COMPLETE.value.lower()
-                redacted_event = self.adapter.complete_task(**adapter_kwargs)
-            else:
-                raise ValueError(f"Unsupported ti_state: `{ti_state}`.")
+                adapter_kwargs["error"] = error
 
-            operator_name = ti.operator.lower()
-            Stats.gauge(
-                f"ol.event.size.{event_type}.{operator_name}",
-                len(Serde.to_json(redacted_event).encode("utf-8")),
+            self.submit_callable(
+                _emit_manual_state_change_event,
+                adapter_method,
+                f"ol.event.size.{event_type}.{ti.operator.lower()}",
+                **adapter_kwargs,
+            )
+        except BaseException as e:
+            self.log.warning(
+                "OpenLineage received exception in method `_on_task_instance_manual_state_change`",
+                exc_info=e,
             )
-
-        self._execute(on_state_change, "on_state_change", use_fork=True)
 
     def _execute(self, callable, callable_name: str, use_fork: bool = False):
         if use_fork:
@@ -787,7 +824,19 @@ def _fork_execute(self, callable, callable_name: str):
             self.log.debug("Process with pid %s finished - parent", pid)
         else:
             setproctitle(getproctitle() + " - OpenLineage - " + callable_name)
-            if not AIRFLOW_V_3_0_PLUS:
+            # Rebuild the ORM in the forked child so it does not share pooled
+            # Postgres connections (and in-flight SSL session state) with the
+            # parent. Without this, the child and parent can both write on an
+            # inherited SSL socket, desynchronising the TLS sequence counter
+            # and producing "SSL error: decryption failed or bad record mac"
+            # on the parent's next query (see #47580 and the DB-access crash
+            # it introduced: on AF3+ workers the Task SDK sets SQL_ALCHEMY_CONN
+            # to "airflow-db-not-allowed:///", which would make configure_orm
+            # fail). settings.engine is only populated in processes where the
+            # ORM was successfully configured -- i.e., the scheduler on AF3+
+            # and both scheduler and workers on AF2 -- so gating on it gives
+            # us fork-safety where we can have it and a no-op where we can't.
+            if settings.engine is not None:
                 configure_orm(disable_connection_pool=True)
             self.log.debug("Executing OpenLineage process - %s - pid %s", callable_name, os.getpid())
             callable()

providers/openlineage/tests/unit/openlineage/plugins/test_listener.py

@@ -85,6 +85,23 @@ def regular_call(self, callable, callable_name, use_fork):
     callable()
 
 
+def direct_submit_call(self, callable, *args, **kwargs):
+    """Synchronous stand-in for ``OpenLineageListener.submit_callable``.
+
+    Bypasses the ``ProcessPoolExecutor`` so tests can assert against mocked
+    adapter methods without hitting pickling of ``unittest.mock.Mock``.
+    When the submitted callable is ``_emit_manual_state_change_event``, skip
+    its ``Stats.gauge`` side effect (which would try to ``Serde.to_json`` a
+    ``MagicMock`` return value) and invoke the adapter method directly.
+    """
+    from airflow.providers.openlineage.plugins.listener import _emit_manual_state_change_event
+
+    if callable is _emit_manual_state_change_event:
+        adapter_method, _stats_key, *_ = args
+        return adapter_method(**kwargs)
+    return callable(*args, **kwargs)
+
+
 class MockExecutor:
     def __init__(self, *args, **kwargs):
         self.submitted = False
@@ -1463,7 +1480,8 @@ def test_adapter_fail_task_is_called_with_dag_description_when_task_doc_is_empty
     @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_run_facet")
     @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet")
     @mock.patch(
-        "airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call
+        "airflow.providers.openlineage.plugins.listener.OpenLineageListener.submit_callable",
+        new=direct_submit_call,
     )
     def test_adapter_fail_task_is_called_with_proper_arguments_for_db_task_instance_model(
         self,
@@ -1482,6 +1500,7 @@ def test_adapter_fail_task_is_called_with_proper_arguments_for_db_task_instance_
         time_machine.move_to(timezone.datetime(2023, 1, 3, 13, 1, 1), tick=False)
 
         listener, task_instance = self._create_listener_and_task_instance(runtime_ti=False)
+        listener._executor = MagicMock()  # satisfy `if not self.executor` guard
         mock_get_airflow_run_facet.return_value = {"airflow": 3}
         mock_get_task_parent_run_facet.return_value = {"parent": 4}
         mock_debug_facet.return_value = {"debug": "packages"}
@@ -1649,7 +1668,8 @@ def test_adapter_complete_task_is_called_with_dag_description_when_task_doc_is_e
     @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet")
     @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet")
     @mock.patch(
-        "airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call
+        "airflow.providers.openlineage.plugins.listener.OpenLineageListener.submit_callable",
+        new=direct_submit_call,
     )
     def test_adapter_complete_task_is_called_with_proper_arguments_for_db_task_instance_model(
         self, mock_get_task_parent_run_facet, mock_debug_facet, mock_debug_mode, mock_emit, time_machine
@@ -1662,6 +1682,7 @@ def test_adapter_complete_task_is_called_with_proper_arguments_for_db_task_insta
         time_machine.move_to(timezone.datetime(2023, 1, 3, 13, 1, 1), tick=False)
 
         listener, task_instance = self._create_listener_and_task_instance(runtime_ti=False)
+        listener._executor = MagicMock()  # satisfy `if not self.executor` guard
         delattr(task_instance, "task")  # Test api server path, where task is not available
         mock_get_task_parent_run_facet.return_value = {"parent": 4}
         mock_debug_facet.return_value = {"debug": "packages"}
@@ -1856,7 +1877,8 @@ def test_listener_on_task_instance_skipped_do_not_call_adapter_when_disabled_ope
     @mock.patch("airflow.providers.openlineage.plugins.listener.get_airflow_debug_facet")
     @mock.patch("airflow.providers.openlineage.plugins.listener.get_task_parent_run_facet")
     @mock.patch(
-        "airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute", new=regular_call
+        "airflow.providers.openlineage.plugins.listener.OpenLineageListener.submit_callable",
+        new=direct_submit_call,
     )
     def test_adapter_complete_task_is_called_with_proper_arguments_for_db_task_instance_model_on_skip(
         self, mock_get_task_parent_run_facet, mock_debug_facet, mock_debug_mode, mock_emit, time_machine
@@ -1869,6 +1891,7 @@ def test_adapter_complete_task_is_called_with_proper_arguments_for_db_task_insta
         time_machine.move_to(timezone.datetime(2023, 1, 3, 13, 1, 1), tick=False)
 
         listener, task_instance = self._create_listener_and_task_instance(runtime_ti=False)
+        listener._executor = MagicMock()  # satisfy `if not self.executor` guard
         delattr(task_instance, "task")  # Test api server path, where task is not available
         mock_get_task_parent_run_facet.return_value = {"parent": 4}
         mock_debug_facet.return_value = {"debug": "packages"}
@@ -1980,6 +2003,55 @@ def set_result(*args, **kwargs):
         listener.log.warning.assert_called_once()
 
 
+class TestOpenLineageListenerForkExecute:
+    """Regression tests for `OpenLineageListener._fork_execute`.
+
+    On processes where the ORM is configured (scheduler; workers on AF2), the
+    forked child must rebuild the engine so it does not share pooled Postgres
+    connections with the parent -- otherwise an inherited SSL socket gets
+    written by both processes and the parent's next query dies with
+    ``SSL error: decryption failed or bad record mac``.
+
+    On AF3+ workers the Task SDK sets SQL_ALCHEMY_CONN to
+    ``airflow-db-not-allowed:///``; ``configure_orm`` would raise there, so
+    the child must skip the rebuild when ``settings.engine`` is ``None``.
+    """
+
+    @staticmethod
+    def _run_child(engine_value):
+        listener = OpenLineageListener()
+        called = MagicMock()
+        with (
+            patch("airflow.providers.openlineage.plugins.listener.os.fork", return_value=0),
+            patch("airflow.providers.openlineage.plugins.listener.os._exit") as mock_exit,
+            patch("airflow.providers.openlineage.plugins.listener.configure_orm") as mock_configure_orm,
+            patch("airflow.providers.openlineage.plugins.listener.setproctitle"),
+            patch(
+                "airflow.providers.openlineage.plugins.listener.getproctitle",
+                return_value="test",
+            ),
+            patch("airflow.providers.openlineage.plugins.listener.settings") as mock_settings,
+        ):
+            mock_settings.engine = engine_value
+            listener._fork_execute(called, "on_failure")
+        return mock_configure_orm, called, mock_exit
+
+    def test_child_rebuilds_orm_when_engine_is_configured(self):
+        mock_configure_orm, called, mock_exit = self._run_child(engine_value=MagicMock())
+        mock_configure_orm.assert_called_once_with(disable_connection_pool=True)
+        called.assert_called_once()
+        mock_exit.assert_called_once_with(0)
+
+    def test_child_skips_orm_rebuild_when_engine_is_none(self):
+        # On AF3+ workers the metadata DB is intentionally unreachable and
+        # configure_orm would raise on the sentinel URL. The callable must
+        # still run and the child must still exit cleanly.
+        mock_configure_orm, called, mock_exit = self._run_child(engine_value=None)
+        mock_configure_orm.assert_not_called()
+        called.assert_called_once()
+        mock_exit.assert_called_once_with(0)
+
+
 @pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Airflow 2 tests")
 class TestOpenLineageSelectiveEnableAirflow2:
     def setup_method(self):