vatsrahul1001 commented on code in PR #66412: URL: https://github.com/apache/airflow/pull/66412#discussion_r3188562867
########## airflow-core/tests/unit/jobs/test_trigger_comms_decoder.py: ########## @@ -0,0 +1,279 @@ +# 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. +from __future__ import annotations + +import asyncio +import contextlib +import random +import threading +from socket import socketpair +from typing import Any + +import greenback +import msgspec +import pytest +import pytest_asyncio +from asgiref.sync import async_to_sync + +from airflow.jobs.triggerer_job_runner import TriggerCommsDecoder, messages +from airflow.sdk.execution_time.comms import _RequestFrame, _ResponseFrame + + +def _read_frame_sync(sock) -> _RequestFrame | None: + """Read a length-prefixed msgpack frame from a blocking socket.""" + lb = b"" + while len(lb) < 4: + chunk = sock.recv(4 - len(lb)) + if not chunk: + return None + lb += chunk + n = int.from_bytes(lb, "big") + data = b"" + while len(data) < n: + chunk = sock.recv(n - len(data)) + if not chunk: + return None + data += chunk + return msgspec.msgpack.decode(data, type=_RequestFrame) + + +@pytest_asyncio.fixture +async def decoder_pair(): + """Yield (decoder, server_sock). Caller owns closing.""" + server_sock, client_sock = socketpair() + reader, writer = await asyncio.open_connection(sock=client_sock) + decoder = TriggerCommsDecoder(async_writer=writer, async_reader=reader, socket=client_sock) + await decoder.start_reader() + yield decoder, server_sock + if decoder._reader_task: + if not decoder._reader_task.done(): + decoder._reader_task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await decoder._reader_task + writer.close() + server_sock.close() + + [email protected] [email protected]_timeout(15) +async def test_concurrent_sends_no_deadlock(): + """ + N_ASYNC async asend() calls + N_SYNC send() calls from threads, all concurrent. + Supervisor responds out of order (shuffled by frame ID). + + Without the multiplexing fix this raises RuntimeError('Response read out of order!') + or hangs. With the fix, every caller gets the right response. + """ + server_sock, client_sock = socketpair() + reader, writer = await asyncio.open_connection(sock=client_sock) + decoder = TriggerCommsDecoder(async_writer=writer, async_reader=reader, socket=client_sock) + await decoder.start_reader() + + N_ASYNC, N_SYNC = 10, 5 + N_TOTAL = N_ASYNC + N_SYNC + results: dict[str, Any] = {} + errors: list[Exception] = [] + + def supervisor(): + frames = [] + for _ in range(N_TOTAL): + f = _read_frame_sync(server_sock) + if f is None: + break + frames.append(f) + random.shuffle(frames) + for f in frames: + server_sock.sendall( + _ResponseFrame( + id=f.id, + body={"type": "TriggerStateSync", "to_create": [], "to_cancel": []}, + ).as_bytes() + ) + + sup = threading.Thread(target=supervisor, daemon=True) + sup.start() + + async def async_caller(idx): + try: + results[f"a{idx}"] = await decoder.asend( + messages.TriggerStateChanges(events=None, finished=[idx], failures=None) + ) + except Exception as e: + errors.append(e) + + def sync_caller(idx): + try: + results[f"s{idx}"] = decoder.send( + messages.TriggerStateChanges(events=None, finished=[100 + idx], failures=None) + ) + except Exception as e: + errors.append(e) + + tasks = [asyncio.create_task(async_caller(i)) for i in range(N_ASYNC)] + threads = [threading.Thread(target=sync_caller, args=(i,)) for i in range(N_SYNC)] + for t in threads: + t.start() + + await asyncio.gather(*tasks) + for t in threads: + t.join(timeout=10) + assert not t.is_alive(), f"{t.name} still alive — deadlock" + + sup.join(timeout=5) + + decoder._reader_task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await decoder._reader_task + writer.close() + server_sock.close() + + assert not errors, f"errors: {errors}" + assert len(results) == N_TOTAL + for key, resp in results.items(): + assert isinstance(resp, messages.TriggerStateSync), f"{key}: unexpected {resp!r}" + + [email protected] [email protected]_timeout(15) +async def test_foreign_loop_path(decoder_pair): + """ + asend() called from a foreign event loop (async_to_sync inside a sync_to_async-wrapped + trigger) must bridge back to the main loop via run_coroutine_threadsafe + wrap_future. + This is the code path that caused the original "Response read out of order!" crash. + """ + decoder, server_sock = decoder_pair + + def supervisor(): + f = _read_frame_sync(server_sock) + if f: + server_sock.sendall( + _ResponseFrame( + id=f.id, + body={"type": "TriggerStateSync", "to_create": [], "to_cancel": []}, + ).as_bytes() + ) + + sup = threading.Thread(target=supervisor, daemon=True) + sup.start() + + result = await asyncio.to_thread( + async_to_sync(decoder.asend), + messages.TriggerStateChanges(events=None, finished=[1], failures=None), + ) + + sup.join(timeout=5) + assert isinstance(result, messages.TriggerStateSync) + + [email protected] +async def test_connection_close_cancels_pending(decoder_pair): Review Comment: test_connection_close_cancels_pending, test_unknown_frame_id_logs_warning, and test_send_before_start_reader_raises exist in test_triggerer_job.py as well, should we remove one -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
