Switch the Supervisor/task process from line-based to length-prefixed

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

@@ -77,6 +77,8 @@
 from airflow.utils.sqlalchemy import prohibit_commit, with_row_locks
 
 if TYPE_CHECKING:
+ from socket import socket
+
  from sqlalchemy.orm import Session
 
  from airflow.callbacks.callback_requests import CallbackRequest
@@ -388,7 +390,7 @@ def _service_processor_sockets(self, timeout: float | None = 1.0):
  """
  events = self.selector.select(timeout=timeout)
  for key, _ in events:
- socket_handler = key.data
+ socket_handler, on_close = key.data
 
  # BrokenPipeError should be caught and treated as if the handler returned false, similar
  # to EOF case
@@ -397,8 +399,9 @@ def _service_processor_sockets(self, timeout: float | None = 1.0):
  except BrokenPipeError:
  need_more = False
  if not need_more:
- self.selector.unregister(key.fileobj)
- key.fileobj.close() # type: ignore[union-attr]
+ sock: socket = key.fileobj # type: ignore[assignment]
+ on_close(sock)
+ sock.close()
 
  def _queue_requested_files_for_parsing(self) -> None:
  """Queue any files requested for parsing as requested by users via UI/API."""

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

@@ -68,7 +68,6 @@ class DagFileParseRequest(BaseModel):
  bundle_path: Path
  """Passing bundle path around lets us figure out relative file path."""
 
- requests_fd: int
  callback_requests: list[CallbackRequest] = Field(default_factory=list)
  type: Literal["DagFileParseRequest"] = "DagFileParseRequest"
 
@@ -102,18 +101,16 @@ class DagFileParsingResult(BaseModel):
 def _parse_file_entrypoint():
  import structlog
 
- from airflow.sdk.execution_time import task_runner
+ from airflow.sdk.execution_time import comms, task_runner
 
  # Parse DAG file, send JSON back up!
- comms_decoder = task_runner.CommsDecoder[ToDagProcessor, ToManager](
- input=sys.stdin,
- decoder=TypeAdapter[ToDagProcessor](ToDagProcessor),
+ comms_decoder = comms.CommsDecoder[ToDagProcessor, ToManager](
+ body_decoder=TypeAdapter[ToDagProcessor](ToDagProcessor),
  )
 
- msg = comms_decoder.get_message()
+ msg = comms_decoder._get_response()
  if not isinstance(msg, DagFileParseRequest):
  raise RuntimeError(f"Required first message to be a DagFileParseRequest, it was {msg}")
- comms_decoder.request_socket = os.fdopen(msg.requests_fd, "wb", buffering=0)
 
  task_runner.SUPERVISOR_COMMS = comms_decoder
  log = structlog.get_logger(logger_name="task")
@@ -125,7 +122,7 @@ def _parse_file_entrypoint():
 
  result = _parse_file(msg, log)
  if result is not None:
- comms_decoder.send_request(log, result)
+ comms_decoder.send(result)
 
 
 def _parse_file(msg: DagFileParseRequest, log: FilteringBoundLogger) -> DagFileParsingResult | None:
@@ -266,20 +263,18 @@ def _on_child_started(
  msg = DagFileParseRequest(
  file=os.fspath(path),
  bundle_path=bundle_path,
- requests_fd=self._requests_fd,
  callback_requests=callbacks,
  )
- self.send_msg(msg)
+ self.send_msg(msg, request_id=0)
 
- def _handle_request(self, msg: ToManager, log: FilteringBoundLogger) -> None: # type: ignore[override]
+ def _handle_request(self, msg: ToManager, log: FilteringBoundLogger, req_id: int) -> None: # type: ignore[override]
  from airflow.sdk.api.datamodels._generated import ConnectionResponse, VariableResponse
 
  resp: BaseModel | None = None
  dump_opts = {}
  if isinstance(msg, DagFileParsingResult):
  self.parsing_result = msg
- return
- if isinstance(msg, GetConnection):
+ elif isinstance(msg, GetConnection):
  conn = self.client.connections.get(msg.conn_id)
  if isinstance(conn, ConnectionResponse):
  conn_result = ConnectionResult.from_conn_response(conn)
@@ -301,18 +296,24 @@ def _handle_request(self, msg: ToManager, log: FilteringBoundLogger) -> None: #
  resp = self.client.variables.delete(msg.key)
  else:
  log.error("Unhandled request", msg=msg)
+ self.send_msg(
+ None,
+ request_id=req_id,
+ error=ErrorResponse(
+ detail={"status_code": 400, "message": "Unhandled request"},
+ ),
+ )
  return
 
- if resp:
- self.send_msg(resp, **dump_opts)
+ self.send_msg(resp, request_id=req_id, error=None, **dump_opts)
 
  @property
  def is_ready(self) -> bool:
  if self._check_subprocess_exit() is None:
  # Process still alive, def can't be finished yet
  return False
 
- return self._num_open_sockets == 0
+ return not self._open_sockets
 
  def wait(self) -> int:
  raise NotImplementedError(f"Don't call wait on {type(self).__name__} objects")

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

@@ -28,6 +28,7 @@
 from collections.abc import Generator, Iterable
 from contextlib import suppress
 from datetime import datetime
+from socket import socket
 from traceback import format_exception
 from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal, TypedDict, Union
 
@@ -43,6 +44,7 @@
 from airflow.jobs.job import perform_heartbeat
 from airflow.models.trigger import Trigger
 from airflow.sdk.execution_time.comms import (
+ CommsDecoder,
  ConnectionResult,
  DagRunStateResult,
  DRCount,
@@ -58,6 +60,7 @@
  TICount,
  VariableResult,
  XComResult,
+ _RequestFrame,
 )
 from airflow.sdk.execution_time.supervisor import WatchedSubprocess, make_buffered_socket_reader
 from airflow.stats import Stats
@@ -70,8 +73,6 @@
 from airflow.utils.session import provide_session
 
 if TYPE_CHECKING:
- from socket import socket
-
  from sqlalchemy.orm import Session
  from structlog.typing import FilteringBoundLogger, WrappedLogger
 
@@ -181,7 +182,6 @@ class messages:
  class StartTriggerer(BaseModel):
  """Tell the async trigger runner process to start, and where to send status update messages."""
 
- requests_fd: int
  type: Literal["StartTriggerer"] = "StartTriggerer"
 
  class TriggerStateChanges(BaseModel):
@@ -295,7 +295,7 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
  """
  TriggerRunnerSupervisor is responsible for monitoring the subprocess and marshalling DB access.
 
- This class (which runs in the main process) is responsible for querying the DB, sending RunTrigger
+ This class (which runs in the main/sync process) is responsible for querying the DB, sending RunTrigger
  workload messages to the subprocess, and collecting results and updating them in the DB.
  """
 
@@ -342,8 +342,8 @@ 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)
+ msg = messages.StartTriggerer()
+ proc.send_msg(msg, request_id=0)
  return proc
 
  @functools.cached_property
@@ -355,7 +355,7 @@ def client(self) -> Client:
  client.base_url = "http://in-process.invalid./" # type: ignore[assignment]
  return client
 
- def _handle_request(self, msg: ToTriggerSupervisor, log: FilteringBoundLogger) -> None: # type: ignore[override]
+ def _handle_request(self, msg: ToTriggerSupervisor, log: FilteringBoundLogger, req_id: int) -> None: # type: ignore[override]
  from airflow.sdk.api.datamodels._generated import (
  ConnectionResponse,
  TaskStatesResponse,
@@ -454,8 +454,7 @@ def _handle_request(self, msg: ToTriggerSupervisor, log: FilteringBoundLogger) -
  else:
  raise ValueError(f"Unknown message type {type(msg)}")
 
- if resp:
- self.send_msg(resp, **dump_opts)
+ self.send_msg(resp, request_id=req_id, error=None, **dump_opts)
 
  def run(self) -> None:
  """Run synchronously and handle all database reads/writes."""
@@ -628,7 +627,7 @@ def _register_pipe_readers(self, stdout: socket, stderr: socket, requests: socke
  ),
  )
 
- def _process_log_messages_from_subprocess(self) -> Generator[None, bytes, None]:
+ def _process_log_messages_from_subprocess(self) -> Generator[None, bytes | bytearray, None]:
  import msgspec
  from structlog.stdlib import NAME_TO_LEVEL
 
@@ -691,14 +690,60 @@ class TriggerDetails(TypedDict):
  events: int
 
 
+@attrs.define(kw_only=True)
+class TriggerCommsDecoder(CommsDecoder[ToTriggerRunner, ToTriggerSupervisor]):
+ _async_writer: asyncio.StreamWriter = attrs.field(alias="async_writer")
+ _async_reader: asyncio.StreamReader = attrs.field(alias="async_reader")
+
+ body_decoder: TypeAdapter[ToTriggerRunner] = attrs.field(
+ factory=lambda: TypeAdapter(ToTriggerRunner), repr=False
+ )
+
+ _lock: asyncio.Lock = attrs.field(factory=asyncio.Lock, repr=False)
+
+ def _read_frame(self):
+ from asgiref.sync import async_to_sync
+
+ return async_to_sync(self._aread_frame)()
+
+ def send(self, msg: ToTriggerSupervisor) -> ToTriggerRunner | None:
+ from asgiref.sync import async_to_sync
+
+ return async_to_sync(self.asend)(msg)
+
+ async def _aread_frame(self):
+ len_bytes = await self._async_reader.readexactly(4)
+ len = int.from_bytes(len_bytes, byteorder="big")
+ if len >= 2**32:
+ raise OverflowError(f"Refusing to receive messages larger than 4GiB {len=}")
+
+ buffer = await self._async_reader.readexactly(len)
+ return self.resp_decoder.decode(buffer)
+
+ async def _aget_response(self, expect_id: int) -> ToTriggerRunner | None:
+ frame = await self._aread_frame()
+ if frame.id != expect_id:
+ # Given the lock we take out in `asend`, this _shouldn't_ be possible, but I'd rather fail with
+ # this explicit error return the wrong type of message back to a Trigger
+ raise RuntimeError(f"Response read out of order! Got {frame.id=}, {expect_id=}")
+ return self._from_frame(frame)
+
+ async def asend(self, msg: ToTriggerSupervisor) -> ToTriggerRunner | None:
+ frame = _RequestFrame(id=next(self.id_counter), body=msg.model_dump())
+ bytes = frame.as_bytes()
+
+ async with self._lock:
+ self._async_writer.write(bytes)
+
+ return await self._aget_response(frame.id)
+
+
 class TriggerRunner:
  """
  Runtime environment for all triggers.
 
- Mainly runs inside its own thread, where it hands control off to an asyncio
- event loop, but is also sometimes interacted with from the main thread
- (where all the DB queries are done). All communication between threads is
- done via Deques.
+ Mainly runs inside its own process, where it hands control off to an asyncio
+ event loop. All communication between this and it's (sync) supervisor is done via sockets
  """
 
  # Maps trigger IDs to their running tasks and other info
@@ -726,10 +771,7 @@ class TriggerRunner:
  # TODO: connect this to the parent process
  log: FilteringBoundLogger = structlog.get_logger()
 
- requests_sock: asyncio.StreamWriter
- response_sock: asyncio.StreamReader
-
- decoder: TypeAdapter[ToTriggerRunner]
+ comms_decoder: TriggerCommsDecoder
 
  def __init__(self):
  super().__init__()
@@ -740,7 +782,6 @@ def __init__(self):
  self.events = deque()
  self.failed_triggers = deque()
  self.job_id = None
- self.decoder = TypeAdapter(ToTriggerRunner)
 
  def run(self):
  """Sync entrypoint - just run a run in an async loop."""
@@ -796,36 +837,21 @@ async def init_comms(self):
  """
  from airflow.sdk.execution_time import task_runner
 
- loop = asyncio.get_event_loop()
+ # Yes, we read and write to stdin! It's a socket, not a normal stdin.
+ reader, writer = await asyncio.open_connection(sock=socket(fileno=0))
 
- comms_decoder = task_runner.CommsDecoder[ToTriggerRunner, ToTriggerSupervisor](
-input=sys.stdin,
-decoder=self.decoder,
+self.comms_decoder = TriggerCommsDecoder(
+async_writer=writer,
+async_reader=reader,
  )
 
- task_runner.SUPERVISOR_COMMS = comms_decoder
-
- async def connect_stdin() -> asyncio.StreamReader:
- reader = asyncio.StreamReader()
- protocol = asyncio.StreamReaderProtocol(reader)
- await loop.connect_read_pipe(lambda: protocol, sys.stdin)
- return reader
-
- self.response_sock = await connect_stdin()
+ task_runner.SUPERVISOR_COMMS = self.comms_decoder
 
-line = await self.response_sock.readline()
+msg = await self.comms_decoder._aget_response(expect_id=0)
 
- msg = self.decoder.validate_json(line)
  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)
- writer_transport, writer_protocol = await loop.connect_write_pipe(
- lambda: asyncio.streams.FlowControlMixin(loop=loop),
- comms_decoder.request_socket,
- )
- self.requests_sock = asyncio.streams.StreamWriter(writer_transport, writer_protocol, None, loop)
-
  async def create_triggers(self):
  """Drain the to_create queue and create all new triggers that have been requested in the DB."""
  while self.to_create:
@@ -934,8 +960,6 @@ async def cleanup_finished_triggers(self) -> list[int]:
  return finished_ids
 
  async def sync_state_to_supervisor(self, finished_ids: list[int]):
- from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
-
  # Copy out of our deques in threadsafe manner to sync state with parent
  events_to_send = []
  while self.events:
@@ -961,19 +985,17 @@ async def sync_state_to_supervisor(self, finished_ids: list[int]):
  if not finished_ids:
  msg.finished = None
 
- # Block triggers from making any requests for the duration of this
- async with SUPERVISOR_COMMS.lock:
- # Tell the monitor that we've finished triggers so it can update things
- self.requests_sock.write(msg.model_dump_json(exclude_none=True).encode() + b"\n")
- line = await self.response_sock.readline()
-
- if line == b"": # EoF received!
+ # Tell the monitor that we've finished triggers so it can update things
+ try:
+ resp = await self.comms_decoder.asend(msg)
+ except asyncio.IncompleteReadError:
  if task := asyncio.current_task():
  task.cancel("EOF - shutting down")
+ return
+ raise
 
- resp = self.decoder.validate_json(line)
  if not isinstance(resp, messages.TriggerStateSync):
- raise RuntimeError(f"Expected to get a TriggerStateSync message, instead we got f{type(msg)}")
+ raise RuntimeError(f"Expected to get a TriggerStateSync message, instead we got {type(msg)}")
  self.to_create.extend(resp.to_create)
  self.to_cancel.extend(resp.to_cancel)