OpenLineage: reuse per-process adapter in dag-state-change pool workers

providers/openlineage/docs/troubleshooting.rst

@@ -43,6 +43,22 @@ as well as the `task_success_overtime <https://airflow.apache.org/docs/apache-ai
 configuration in Airflow config.
 
 
+Scheduler CPU or memory growing steadily while OpenLineage is enabled
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+In provider versions that submit bound adapter methods to the DAG-run event process pool, each pool worker
+built a new OpenLineage client — including a new transport set — for every DAG-run state change, and the
+clients were never closed. With transports that start background worker threads (e.g. the ``datadog``
+transport), each DAG-run event leaked one thread inside the scheduler, so scheduler CPU and memory climbed
+steadily over hours and recovered only on a scheduler restart.
+
+**Possible Solution**
+
+Upgrade to a provider version that reuses a single per-process adapter in the pool workers. If you cannot
+upgrade yet, prefer a transport that does not start background worker threads (e.g. plain ``http``) for the
+affected destination, or restart the scheduler to reclaim the leaked threads as a stopgap.
+
+
 Missing lineage from EmptyOperators
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

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

@@ -113,15 +113,42 @@ def _executor_initializer():
         log.debug("Exception details:", exc_info=True)
 
 
-def _emit_manual_state_change_event(adapter_method, stats_key, **kwargs):
+@cache
+def _get_process_adapter() -> OpenLineageAdapter:
+    """
+    Return the per-process ``OpenLineageAdapter`` used inside pool worker processes.
+
+    Each ``ProcessPoolExecutor`` worker keeps exactly one adapter — and therefore one
+    ``OpenLineageClient`` with one set of transports — for its whole lifetime.
+    """
+    return OpenLineageAdapter()
+
+
+def _run_adapter_method(method_name: str, /, *args, **kwargs):
+    """
+    Run the named ``OpenLineageAdapter`` method on the per-process adapter.
+
+    Module-level so it is picklable across the ProcessPoolExecutor boundary. Bound adapter
+    methods must not be submitted to the pool directly: pickling them serializes the whole
+    adapter, so the worker unpickles a fresh adapter per event and builds a new
+    ``OpenLineageClient`` (with new transports) on every emit. Transports that start
+    background worker threads (e.g. the ``datadog`` transport, which always starts an async
+    HTTP worker thread) are never closed, so this leaks one thread per event and steadily
+    consumes scheduler CPU and memory until restart.
+    """
+    return getattr(_get_process_adapter(), method_name)(*args, **kwargs)
+
+
+def _emit_manual_state_change_event(adapter_method_name: str, stats_key: str, **kwargs):
     """
-    Emit an OL event via the given adapter method and record its serialized size.
+    Emit an OL event via the named 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.
+    externally" emissions. The method is resolved on the per-process adapter so the
+    pool worker reuses one client across events (see ``_run_adapter_method``).
     """
-    event = adapter_method(**kwargs)
+    event = getattr(_get_process_adapter(), adapter_method_name)(**kwargs)
     Stats.gauge(stats_key, len(Serde.to_json(event).encode("utf-8")))
     return event
 
@@ -782,10 +809,10 @@ def _on_task_instance_manual_state_change(
                 return
 
             if ti_state == TaskInstanceState.FAILED:
-                adapter_method = self.adapter.fail_task
+                adapter_method_name = "fail_task"
                 event_type = RunState.FAIL.value.lower()
             elif ti_state in (TaskInstanceState.SUCCESS, TaskInstanceState.SKIPPED):
-                adapter_method = self.adapter.complete_task
+                adapter_method_name = "complete_task"
                 event_type = RunState.COMPLETE.value.lower()
             else:
                 raise ValueError(f"Unsupported ti_state: `{ti_state}`.")
@@ -867,7 +894,7 @@ def _on_task_instance_manual_state_change(
             operator_name = (ti.operator or "unknown").lower()
             self.submit_callable(
                 _emit_manual_state_change_event,
-                adapter_method,
+                adapter_method_name,
                 f"ol.event.size.{event_type}.{operator_name}",
                 **adapter_kwargs,
             )
@@ -979,7 +1006,8 @@ def on_dag_run_running(self, dag_run: DagRun, msg: str) -> None:
             doc, doc_type = get_dag_documentation(dag_run.dag)
 
             self.submit_callable(
-                self.adapter.dag_started,
+                _run_adapter_method,
+                "dag_started",
                 dag_id=dag_run.dag_id,
                 run_id=dag_run.run_id,
                 logical_date=date,
@@ -1031,7 +1059,8 @@ def on_dag_run_success(self, dag_run: DagRun, msg: str) -> None:
             doc, doc_type = get_dag_documentation(dag_run.dag)
 
             self.submit_callable(
-                self.adapter.dag_success,
+                _run_adapter_method,
+                "dag_success",
                 dag_id=dag_run.dag_id,
                 run_id=dag_run.run_id,
                 end_date=dag_run.end_date,
@@ -1082,7 +1111,8 @@ def on_dag_run_failed(self, dag_run: DagRun, msg: str) -> None:
             doc, doc_type = get_dag_documentation(dag_run.dag)
 
             self.submit_callable(
-                self.adapter.dag_failed,
+                _run_adapter_method,
+                "dag_failed",
                 dag_id=dag_run.dag_id,
                 run_id=dag_run.run_id,
                 end_date=dag_run.end_date,

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

@@ -92,15 +92,23 @@ def direct_submit_call(self, callable, *args, **kwargs):
 
     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.
+    The module-level pool wrappers pass adapter method *names* and resolve them
+    on the per-process adapter; here we resolve them on this listener's adapter
+    instead, so assertions against mocked adapter methods keep working. For
+    ``_emit_manual_state_change_event`` this also skips its ``Stats.gauge``
+    side effect (which would try to ``Serde.to_json`` a ``MagicMock`` return).
     """
-    from airflow.providers.openlineage.plugins.listener import _emit_manual_state_change_event
+    from airflow.providers.openlineage.plugins.listener import (
+        _emit_manual_state_change_event,
+        _run_adapter_method,
+    )
 
     if callable is _emit_manual_state_change_event:
-        adapter_method, _stats_key, *_ = args
-        return adapter_method(**kwargs)
+        adapter_method_name, _stats_key, *_ = args
+        return getattr(self.adapter, adapter_method_name)(**kwargs)
+    if callable is _run_adapter_method:
+        adapter_method_name, *rest = args
+        return getattr(self.adapter, adapter_method_name)(*rest, **kwargs)
     return callable(*args, **kwargs)
 
 
@@ -123,6 +131,34 @@ def shutdown(self, *args, **kwargs):
         print("Shutting down")
 
 
+def _probe_process_adapter():
+    """Return (pid, adapter id) from inside a pool worker; module-level so it is picklable."""
+    import os
+
+    from airflow.providers.openlineage.plugins.listener import _get_process_adapter
+
+    return os.getpid(), id(_get_process_adapter())
+
+
+class TestProcessAdapterReuse:
+    def test_process_adapter_reused_across_pool_submissions(self):
+        """
+        A pool worker must reuse one adapter (hence one client/transport set) across events.
+
+        Regression test: submitting bound adapter methods pickled a fresh adapter per event,
+        making the worker build a new OpenLineageClient (and transport worker threads that are
+        never closed) for every DAG-run state change, leaking threads in the scheduler.
+        """
+        from concurrent.futures import ProcessPoolExecutor
+
+        with ProcessPoolExecutor(max_workers=1) as pool:
+            pid_first, adapter_id_first = pool.submit(_probe_process_adapter).result(timeout=60)
+            pid_second, adapter_id_second = pool.submit(_probe_process_adapter).result(timeout=60)
+
+        assert pid_first == pid_second
+        assert adapter_id_first == adapter_id_second
+
+
 class TestExecutorInitializer:
     """Tests for _executor_initializer function."""