Introduce separate channel for trigger workloads

airflow-core/src/airflow/jobs/triggerer_job_runner.py

@@ -18,6 +18,7 @@
 
 import asyncio
 import functools
+import io
 import logging
 import os
 import selectors
@@ -182,6 +183,7 @@ class StartTriggerer(BaseModel):
         """Tell the async trigger runner process to start, and where to send status update messages."""
 
         requests_fd: int
+        trigger_requests_fd: int
         type: Literal["StartTriggerer"] = "StartTriggerer"
 
     class CancelTriggers(BaseModel):
@@ -322,8 +324,11 @@ def start(  # type: ignore[override]
     ):
         proc = super().start(id=job.id, job=job, target=cls.run_in_process, logger=logger, **kwargs)
 
-        msg = messages.StartTriggerer(requests_fd=proc._requests_fd)
-        proc._send(msg)
+        msg = messages.StartTriggerer(
+            requests_fd=proc._requests_fd,
+            trigger_requests_fd=proc._trigger_requests_fd,  # type: ignore[arg-type]
+        )
+        proc._send_first_message(msg)
         return proc
 
     @functools.cached_property
@@ -463,6 +468,9 @@ def emit_metrics(self):
         )
 
     def _send(self, msg: BaseModel):
+        self.trigger_stdin.write(msg.model_dump_json().encode("utf-8") + b"\n")  # type: ignore[union-attr]
+
+    def _send_first_message(self, msg: BaseModel):
         self.stdin.write(msg.model_dump_json().encode("utf-8") + b"\n")
 
     def update_triggers(self, requested_trigger_ids: set[int]):
@@ -678,8 +686,12 @@ def init_comms(self):
         msg = comms_decoder.get_message()
         if not isinstance(msg, messages.StartTriggerer):
             raise RuntimeError(f"Required first message to be a messages.StartTriggerer, it was {msg}")
+
         comms_decoder.request_socket = os.fdopen(msg.requests_fd, "wb", buffering=0)
 
+        binary = os.fdopen(msg.trigger_requests_fd, "rb")
+        handle = io.TextIOWrapper(binary, line_buffering=True)
+        comms_decoder.trigger_input = handle
         task_runner.SUPERVISOR_COMMS = comms_decoder
 
     def run(self):
@@ -750,7 +762,7 @@ async def read_workloads(self, ready_event: asyncio.Event):
         async def connect_stdin() -> asyncio.StreamReader:
             reader = asyncio.StreamReader()
             protocol = asyncio.StreamReaderProtocol(reader)
-            await loop.connect_read_pipe(lambda: protocol, sys.stdin)
+            await loop.connect_read_pipe(lambda: protocol, task_runner.SUPERVISOR_COMMS.trigger_input)
             return reader
 
         stdin = await connect_stdin()

airflow-core/tests/unit/jobs/test_triggerer_job.py

@@ -172,8 +172,10 @@ def builder(job=None):
             job=job,
             pid=process.pid,
             stdin=mocker.Mock(),
+            trigger_stdin=mocker.Mock(),
             process=process,
             requests_fd=-1,
+            trigger_requests_fd=-1,
             capacity=10,
         )
         # Mock the selector
@@ -447,7 +449,7 @@ def is_alive(self):
         return True
 
     supervisor.run()
-    assert supervisor.stdin.write.call_count == 1
+    assert supervisor.trigger_stdin.write.call_count == 1
 
 
 @pytest.mark.execution_timeout(5)
@@ -624,10 +626,10 @@ async def run(self, **args) -> AsyncIterator[TriggerEvent]:
         from airflow.sdk import Variable
         from airflow.sdk.execution_time.xcom import XCom
 
-        conn = await sync_to_async(BaseHook.get_connection)("test_connection")
+        conn = await sync_to_async(BaseHook.get_connection)(f"test_connection_{self.task_id}")
         self.log.info("Loaded conn %s", conn.conn_id)
 
-        variable = await sync_to_async(Variable.get)("test_variable")
+        variable = await sync_to_async(Variable.get)(f"test_variable_{self.task_id}")
         self.log.info("Loaded variable %s", variable)
 
         xcom = await sync_to_async(XCom.get_one)(
@@ -693,8 +695,8 @@ async def test_trigger_can_access_variables_connections_and_xcoms(session, dag_m
     task_instance.trigger_id = trigger_orm.id
 
     # Create the appropriate Connection, Variable and XCom
-    connection = Connection(conn_id="test_connection", conn_type="http")
-    variable = Variable(key="test_variable", val="some_variable_value")
+    connection = Connection(conn_id=f"test_connection_{task_instance.task_id}", conn_type="http")
+    variable = Variable(key=f"test_variable_{task_instance.task_id}", val="some_variable_value")
     XComModel.set(
         key="test_xcom",
         value="some_xcom_value",
@@ -720,7 +722,7 @@ async def test_trigger_can_access_variables_connections_and_xcoms(session, dag_m
     assert task_instance.next_kwargs == {
         "event": {
             "connection": {
-                "conn_id": "test_connection",
+                "conn_id": f"test_connection_{task_instance.task_id}",
                 "conn_type": "http",
                 "description": None,
                 "host": None,
@@ -734,3 +736,80 @@ async def test_trigger_can_access_variables_connections_and_xcoms(session, dag_m
             "xcom": '"some_xcom_value"',
         }
     }
+
+
+@pytest.mark.flaky(reruns=2, reruns_delay=10)
+@pytest.mark.execution_timeout(30)
+@pytest.mark.asyncio
+def test_trigger_multiple_task_instances_correctly_gets_messages_running_in_trigger(session, dag_maker):
+    """
+    Check that the trigger will successfully access Variables, Connections and XComs.
+
+    Running multiple tasks in triggerer to get the correct messages.
+    Create 3 connections, variables, and XComs for 3 tasks. the key will have task_id to
+    assert the correct message received
+    """
+
+    with dag_maker(dag_id="trigger_accessing_variable_connection_and_xcom", session=session):
+        EmptyOperator(task_id="task_one")
+        EmptyOperator(task_id="task_two")
+        EmptyOperator(task_id="task_three")
+    dr = dag_maker.create_dagrun()
+
+    for trigger_id, ti in enumerate(dr.task_instances, start=1):
+        ti.state = TaskInstanceState.DEFERRED
+
+        # Create a Trigger
+        trigger = CustomTrigger(dag_id=dr.dag_id, run_id=dr.run_id, task_id=ti.task_id, map_index=-1)
+        trigger_orm = Trigger(
+            classpath=trigger.serialize()[0],
+            kwargs={"dag_id": dr.dag_id, "run_id": dr.run_id, "task_id": ti.task_id, "map_index": -1},
+        )
+        trigger_orm.id = trigger_id
+        session.add(trigger_orm)
+        session.commit()
+        ti.trigger_id = trigger_orm.id
+
+        # Create the appropriate Connection, Variable and XCom
+        connection = Connection(conn_id=f"test_connection_{ti.task_id}", conn_type="http")
+        variable = Variable(key=f"test_variable_{ti.task_id}", val=f"some_variable_value {trigger_id}")
+        XComModel.set(
+            key="test_xcom",
+            value=f"some_xcom_value {trigger_id}",
+            task_id=ti.task_id,
+            dag_id=dr.dag_id,
+            run_id=dr.run_id,
+            map_index=-1,
+            session=session,
+        )
+        session.add(connection)
+        session.add(variable)
+        session.commit()
+
+    job = Job()
+    session.add(job)
+    session.commit()
+
+    supervisor = DummyTriggerRunnerSupervisor.start(job=job, capacity=4, logger=None)
+    supervisor.run()
+    for trigger_id, ti in enumerate(dr.task_instances, start=1):
+        ti.refresh_from_db()
+        assert ti.state == TaskInstanceState.SCHEDULED
+        assert ti.next_method != "__fail__"
+        assert ti.next_kwargs == {
+            "event": {
+                "connection": {
+                    "conn_id": f"test_connection_{ti.task_id}",
+                    "conn_type": "http",
+                    "description": None,
+                    "host": None,
+                    "schema": None,
+                    "login": None,
+                    "password": None,
+                    "port": None,
+                    "extra": None,
+                },
+                "variable": f"some_variable_value {trigger_id}",
+                "xcom": f'"some_xcom_value {trigger_id}"',
+            }
+        }

task-sdk/src/airflow/sdk/execution_time/supervisor.py

@@ -385,6 +385,16 @@ class WatchedSubprocess:
     _requests_fd: int
     """File descriptor for request handling."""
 
+    """"
+    It require special case to handle the workloads and api calls to api-server in Triggerer, due to mixing up messages
+    a separate channel is used to send the workloads from parent process to child process the child process.
+    connect_stdin will use this channel to read the workloads in read_workload method.
+    trigger_stdin: trigger parent process uses this to send the workloads to child process.
+    _trigger_requests_fd: trigger child process uses this to read the workloads.
+    """
+    _trigger_requests_fd: int | None = None
+    trigger_stdin: BinaryIO | None = None
+
     _num_open_sockets: int = 4
     _exit_code: int | None = attrs.field(default=None, init=False)
 
@@ -413,11 +423,19 @@ def start(
         child_comms, read_msgs = mkpipe()
         child_logs, read_logs = mkpipe()
 
+        if "TriggerRunnerSupervisor" in cls.__name__:
+            trigger_child_stdin, trigger_feed_stdin = mkpipe(remote_read=True)
+        else:
+            trigger_child_stdin, trigger_feed_stdin = None, None
+
         pid = os.fork()
         if pid == 0:
             # Close and delete of the parent end of the sockets.
             cls._close_unused_sockets(feed_stdin, read_stdout, read_stderr, read_msgs, read_logs)
 
+            if trigger_feed_stdin:
+                cls._close_unused_sockets(trigger_feed_stdin)
+
             # Python GC should delete these for us, but lets make double sure that we don't keep anything
             # around in the forked processes, especially things that might involve open files or sockets!
             del constructor_kwargs
@@ -439,6 +457,15 @@ def start(
 
         requests_fd = child_comms.fileno()
 
+        if "TriggerRunnerSupervisor" in cls.__name__:
+            additional_trigger_fds = {
+                "trigger_stdin": trigger_feed_stdin,
+                "trigger_requests_fd": trigger_child_stdin.fileno(),  # type: ignore[union-attr]
+            }
+            cls._close_unused_sockets(trigger_child_stdin)
+        else:
+            additional_trigger_fds = {}
+
         # Close the remaining parent-end of the sockets we've passed to the child via fork. We still have the
         # other end of the pair open
         cls._close_unused_sockets(child_stdin, child_stdout, child_stderr, child_comms, child_logs)
@@ -450,6 +477,7 @@ def start(
             process=psutil.Process(pid),
             requests_fd=requests_fd,
             process_log=logger,
+            **additional_trigger_fds,  # type: ignore[arg-type]
             **constructor_kwargs,
         )
 
@@ -668,6 +696,8 @@ def _check_subprocess_exit(self, raise_on_timeout: bool = False) -> int | None:
                 self._exit_code = self._process.wait(timeout=0)
                 log.debug("%s process exited", type(self).__name__, exit_code=self._exit_code)
                 self._close_unused_sockets(self.stdin)
+                if self.trigger_stdin:
+                    self._close_unused_sockets(self.trigger_stdin)
             except psutil.TimeoutExpired:
                 if raise_on_timeout:
                     raise

task-sdk/src/airflow/sdk/execution_time/task_runner.py

@@ -540,6 +540,13 @@ class CommsDecoder(Generic[ReceiveMsgType, SendMsgType]):
 
     input: TextIO
 
+    """"
+    It require special case to handle the workloads and api calls to api-server, due to mixing up messages
+    a separate channel is used to send the workloads from parent process to child process the child process.
+    connect_stdin will use this channel to read the workloads in read_workload method.
+    """
+    trigger_input: TextIO = attrs.field(init=False, default=None)
+
     request_socket: FileIO = attrs.field(init=False, default=None)
 
     # We could be "clever" here and set the default to this based type parameters and a custom