Add micro-benchmarks for the task SDK's supervisor comms

.github/workflows/benchmarks.yml

@@ -0,0 +1,72 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+---
+name: "Benchmarks"
+on:  # yamllint disable-line rule:truthy
+  push:
+    branches: ["main"]
+  pull_request:
+    paths:
+      - "task-sdk/src/airflow/sdk/execution_time/**"
+      - "task-sdk/benchmarks/**"
+      # A dependency bump (msgspec frames every request) changes what is measured.
+      - "task-sdk/pyproject.toml"
+      - "uv.lock"
+      - ".github/workflows/benchmarks.yml"
+  # Lets CodSpeed run the benchmarks on older commits to build its baseline.
+  workflow_dispatch:
+permissions:
+  # All other permissions are set to none by default
+  contents: read
+concurrency:
+  # A group per commit on main, so none of its runs is cancelled or left pending behind a newer one:
+  # CodSpeed compares each PR with the run of its base commit.
+  group: benchmarks-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: true
+
+jobs:
+
+  task-sdk:
+    name: "Task SDK micro-benchmarks"
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    permissions:
+      contents: read
+      # OpenID Connect authentication with CodSpeed, no token secret needed.
+      id-token: write
+    steps:
+      - name: "Checkout ${{ github.ref }} ( ${{ github.sha }} )"
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1  # v7.0.1
+        with:
+          persist-credentials: false
+
+      - name: "Install uv"
+        uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7  # v10.2.0
+        with:
+          version-file: uv.lock
+
+      - name: "Install the benchmark dependencies"
+        run: uv sync --frozen --project task-sdk --group benchmarks
+
+      # Simulation mode counts CPU instructions instead of timing, so results do not depend on
+      # how busy the shared runner is; CodSpeed compares them with the base branch on each PR.
+      - name: "Run the benchmarks"
+        uses: CodSpeedHQ/action@373d6868929f444bc08d901fd0eb0ad52a8875ea  # v5.2.1
+        with:
+          mode: simulation
+          run: uv run --frozen --project task-sdk --group benchmarks pytest task-sdk/benchmarks --codspeed

task-sdk/benchmarks/execution_time/test_comms_benchmark.py

@@ -0,0 +1,105 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Micro-benchmarks of the task's channel to its supervisor.
+
+Every SDK call a task makes (a Variable, a Connection, an XCom, a state store key) is one round
+trip through ``CommsDecoder``. Async tasks make them concurrently on one event loop, and an
+iterated task makes them for every item, so the cost of one round trip multiplies with the
+number of calls. The supervisor side is a thread on the other end of a socket pair that answers
+each request at once, so what is measured is the decoder: framing, locking and the thread hops.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import socket
+import threading
+from collections.abc import Iterator
+
+import msgspec
+import pytest
+import structlog
+
+from airflow.sdk.execution_time.comms import CommsDecoder, GetVariable, _RequestFrame, _ResponseFrame
+
+
+def _recv_exactly(sock: socket.socket, size: int) -> bytes:
+    buffer = bytearray()
+    while len(buffer) < size:
+        chunk = sock.recv(size - len(buffer))
+        if not chunk:
+            raise EOFError
+        buffer.extend(chunk)
+    return bytes(buffer)
+
+
+def _answer_every_request(sock: socket.socket) -> None:
+    """Act as the supervisor: answer each ``GetVariable`` with its key, until the socket closes."""
+    try:
+        while True:
+            length = int.from_bytes(_recv_exactly(sock, 4), "big")
+            request = msgspec.msgpack.decode(_recv_exactly(sock, length), type=_RequestFrame)
+            assert request.body is not None
+            response = {"type": "VariableResult", "key": request.body["key"], "value": "v"}
+            encoded = msgspec.msgpack.encode(_ResponseFrame(request.id, response, None))
+            sock.sendall(len(encoded).to_bytes(4, "big") + encoded)
+    except (EOFError, OSError):
+        pass
+    finally:
+        # Whatever ends this loop, the decoder gets EOF and the benchmark fails instead of waiting forever.
+        sock.close()
+
+
+@pytest.fixture
+def decoder() -> Iterator[CommsDecoder]:
+    task_side, supervisor_side = socket.socketpair()
+    supervisor = threading.Thread(target=_answer_every_request, args=(supervisor_side,), daemon=True)
+    supervisor.start()
+    yield CommsDecoder(socket=task_side, log=structlog.get_logger())
+    task_side.close()
+    supervisor_side.close()
+    supervisor.join(timeout=5)
+
+
+@pytest.fixture
+def loop() -> Iterator[asyncio.AbstractEventLoop]:
+    event_loop = asyncio.new_event_loop()
+    yield event_loop
+    event_loop.close()
+
+
+def test_send_round_trip(benchmark, decoder):
+    """One synchronous request and its response, as a sync task's SDK call makes it."""
+    result = benchmark(lambda: decoder.send(GetVariable(key="key")))
+    assert result.key == "key"
+
+
+def test_asend_round_trip(benchmark, decoder, loop):
+    """One asynchronous request and its response, as an async task's SDK call makes it."""
+    result = benchmark(lambda: loop.run_until_complete(decoder.asend(GetVariable(key="key"))))
+    assert result.key == "key"
+
+
+def test_asend_round_trips_from_concurrent_coroutines(benchmark, decoder, loop):
+    """Fifty coroutines each making one request on the same loop, as iterated async items do."""
+
+    async def fifty_concurrent_requests():
+        return await asyncio.gather(*(decoder.asend(GetVariable(key=f"key{i}")) for i in range(50)))
+
+    results = benchmark(lambda: loop.run_until_complete(fifty_concurrent_requests()))
+    assert [result.key for result in results] == [f"key{i}" for i in range(50)]

task-sdk/pyproject.toml

@@ -229,6 +229,10 @@ dev = [
     "pandas>=2.2.3; python_version >=\"3.13\" and python_version <\"3.14\"",
     "pandas>=2.3.3; python_version >=\"3.14\""
 ]
+benchmarks = [
+    "apache-airflow-devel-common",
+    "pytest-codspeed>=5.0.3",
+]
 docs = [
     "apache-airflow-devel-common[docs]",
 ]