-
Notifications
You must be signed in to change notification settings - Fork 0
Hypothesis fault-injection for the Rust pump's FD lifecycle (#74) #244
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,166 @@ | ||
| """Raw file-descriptor lifecycle for the Rust inter-stage pump. | ||
|
|
||
| When the Rust backend pumps one pipeline stage's stdout into the next stage's | ||
| stdin it takes over the raw pipe descriptors from asyncio for the duration of | ||
| the transfer. That hand-off has several partial-failure paths: extracting the | ||
| descriptor from an asyncio transport, pausing the reader transport so asyncio | ||
| does not race the Rust pump, and switching the descriptors to blocking mode | ||
| (then restoring their prior mode). This module isolates that lifecycle behind | ||
| two seams — the [`_BlockingModeGuard`][cuprum._pipeline_stream_fds._BlockingModeGuard] | ||
| FD-state object and the [`_paused_reader`][cuprum._pipeline_stream_fds._paused_reader] | ||
| context manager — so it can be fault-injected without a live pump. | ||
| """ | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import contextlib | ||
| import dataclasses as dc | ||
| import os | ||
| import typing as typ | ||
|
|
||
| if typ.TYPE_CHECKING: | ||
| import asyncio | ||
| import collections.abc as cabc | ||
|
|
||
|
|
||
| def _fd_from_transport(transport: object | None) -> int | None: | ||
| """Extract a raw FD via ``transport.get_extra_info('pipe').fileno()``.""" | ||
| get_extra = getattr(transport, "get_extra_info", None) | ||
| if get_extra is None: | ||
| return None | ||
| pipe: object | None = get_extra("pipe") | ||
| fileno = getattr(pipe, "fileno", None) if pipe is not None else None | ||
| if fileno is None: | ||
| return None | ||
| try: | ||
| return int(fileno()) | ||
| except (OSError, ValueError, TypeError, AttributeError): | ||
| return None | ||
|
|
||
|
|
||
| def _extract_stream_fd( | ||
| stream: asyncio.StreamReader | asyncio.StreamWriter | None, | ||
| ) -> int | None: | ||
| """Extract a raw FD from an asyncio stream via its transport.""" | ||
| if stream is None: | ||
| return None | ||
| transport = getattr(stream, "transport", None) | ||
| if transport is None: | ||
| transport = getattr(stream, "_transport", None) | ||
| return _fd_from_transport(transport) | ||
|
|
||
|
|
||
| def _pause_reader_transport( | ||
| reader: asyncio.StreamReader, | ||
| ) -> cabc.Callable[[], None] | None: | ||
| """Pause reader transport callbacks while Rust pump owns the raw FD.""" | ||
| transport = getattr(reader, "transport", None) | ||
| if transport is None: | ||
| transport = getattr(reader, "_transport", None) | ||
| pause_reading = getattr(transport, "pause_reading", None) | ||
| resume_reading = getattr(transport, "resume_reading", None) | ||
| if not callable(pause_reading) or not callable(resume_reading): | ||
| return None | ||
| try: | ||
| pause_reading() | ||
| except (RuntimeError, OSError): | ||
| return None | ||
|
|
||
| def _resume() -> None: | ||
| """Resume the paused reader transport, ignoring teardown errors.""" | ||
| with contextlib.suppress(RuntimeError, OSError): | ||
| resume_reading() | ||
|
|
||
| return _resume | ||
|
|
||
|
|
||
| def _set_stream_fds_blocking(*, reader_fd: int, writer_fd: int) -> tuple[bool, bool]: | ||
| """Switch pipe FDs to blocking mode and return their prior state.""" | ||
| reader_was_blocking = os.get_blocking(reader_fd) | ||
| writer_was_blocking = os.get_blocking(writer_fd) | ||
| reader_changed = False | ||
| try: | ||
| if not reader_was_blocking: | ||
| os.set_blocking(reader_fd, True) | ||
| reader_changed = True | ||
| if not writer_was_blocking: | ||
| os.set_blocking(writer_fd, True) | ||
| except OSError: | ||
| if reader_changed: | ||
| with contextlib.suppress(OSError, ValueError): | ||
| os.set_blocking(reader_fd, reader_was_blocking) | ||
| raise | ||
| return reader_was_blocking, writer_was_blocking | ||
|
|
||
|
|
||
| def _restore_stream_fd_blocking( | ||
| *, | ||
| reader_fd: int, | ||
| writer_fd: int, | ||
| reader_was_blocking: bool, | ||
| writer_was_blocking: bool, | ||
| ) -> None: | ||
| """Restore pipe FD blocking mode captured before Rust pumping.""" | ||
| with contextlib.suppress(OSError, ValueError): | ||
| os.set_blocking(reader_fd, reader_was_blocking) | ||
| with contextlib.suppress(OSError, ValueError): | ||
| os.set_blocking(writer_fd, writer_was_blocking) | ||
|
|
||
|
|
||
| @dc.dataclass(slots=True) | ||
| class _BlockingModeGuard: | ||
| """The prior blocking mode of a reader/writer FD pair, restorable on demand. | ||
|
leynos marked this conversation as resolved.
|
||
|
|
||
| This is the FD-state object behind the Rust pump's blocking lifecycle: it | ||
| captures each descriptor's blocking mode when the pump takes over the raw | ||
| FDs and restores it afterwards, so no descriptor is left in the temporary | ||
| blocking mode the Rust pump requires. | ||
| """ | ||
|
|
||
| reader_fd: int | ||
| writer_fd: int | ||
| reader_was_blocking: bool | ||
| writer_was_blocking: bool | ||
|
|
||
| @classmethod | ||
| def engage(cls, *, reader_fd: int, writer_fd: int) -> _BlockingModeGuard: | ||
| """Switch both FDs to blocking mode, capturing their prior state. | ||
|
|
||
| Raises ``OSError`` — after rolling back any partial change so no | ||
| descriptor is left switched — when either FD cannot be made blocking. | ||
| """ | ||
| reader_was_blocking, writer_was_blocking = _set_stream_fds_blocking( | ||
| reader_fd=reader_fd, | ||
| writer_fd=writer_fd, | ||
| ) | ||
| return cls( | ||
| reader_fd=reader_fd, | ||
| writer_fd=writer_fd, | ||
| reader_was_blocking=reader_was_blocking, | ||
| writer_was_blocking=writer_was_blocking, | ||
| ) | ||
|
|
||
| def restore(self) -> None: | ||
| """Restore both FDs to the blocking mode captured by ``engage``.""" | ||
| _restore_stream_fd_blocking( | ||
| reader_fd=self.reader_fd, | ||
| writer_fd=self.writer_fd, | ||
| reader_was_blocking=self.reader_was_blocking, | ||
| writer_was_blocking=self.writer_was_blocking, | ||
| ) | ||
|
|
||
|
|
||
| @contextlib.contextmanager | ||
| def _paused_reader(reader: asyncio.StreamReader) -> cabc.Iterator[None]: | ||
| """Pause the reader transport for the block, always resuming it on exit. | ||
|
|
||
| Wraps ``_pause_reader_transport`` so the resume call cannot be skipped on | ||
| any exit path — normal return, exception, or cancellation. When the | ||
| transport cannot be paused the resume is simply never attempted. | ||
| """ | ||
| resume = _pause_reader_transport(reader) | ||
| try: | ||
| yield | ||
| finally: | ||
| if resume is not None: | ||
| resume() | ||
| Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -17,7 +17,6 @@ | |||||||||||||||||||||||||||||||||||||||||||||
| import asyncio | ||||||||||||||||||||||||||||||||||||||||||||||
| import contextlib | ||||||||||||||||||||||||||||||||||||||||||||||
| import dataclasses as dc | ||||||||||||||||||||||||||||||||||||||||||||||
| import os | ||||||||||||||||||||||||||||||||||||||||||||||
| import typing as typ | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| from cuprum._backend import StreamBackend, get_stream_backend | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
@@ -30,6 +29,11 @@ | |||||||||||||||||||||||||||||||||||||||||||||
| from cuprum._pipeline_stage_streams import ( | ||||||||||||||||||||||||||||||||||||||||||||||
| _create_stage_capture_tasks as _create_stage_capture_tasks, | ||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||
| from cuprum._pipeline_stream_fds import ( | ||||||||||||||||||||||||||||||||||||||||||||||
| _BlockingModeGuard, | ||||||||||||||||||||||||||||||||||||||||||||||
| _extract_stream_fd, | ||||||||||||||||||||||||||||||||||||||||||||||
| _paused_reader, | ||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||
| from cuprum._streams import _close_stream_writer, _pump_stream | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| if typ.TYPE_CHECKING: | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
@@ -79,142 +83,89 @@ def reset_pump_stream_dispatch_for_testing() -> None: | |||||||||||||||||||||||||||||||||||||||||||||
| configure_pump_stream_dispatch_for_testing() | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| def _fd_from_transport(transport: object | None) -> int | None: | ||||||||||||||||||||||||||||||||||||||||||||||
| """Extract a raw FD via ``transport.get_extra_info('pipe').fileno()``.""" | ||||||||||||||||||||||||||||||||||||||||||||||
| get_extra = getattr(transport, "get_extra_info", None) | ||||||||||||||||||||||||||||||||||||||||||||||
| if get_extra is None: | ||||||||||||||||||||||||||||||||||||||||||||||
| return None | ||||||||||||||||||||||||||||||||||||||||||||||
| pipe: object | None = get_extra("pipe") | ||||||||||||||||||||||||||||||||||||||||||||||
| fileno = getattr(pipe, "fileno", None) if pipe is not None else None | ||||||||||||||||||||||||||||||||||||||||||||||
| if fileno is None: | ||||||||||||||||||||||||||||||||||||||||||||||
| return None | ||||||||||||||||||||||||||||||||||||||||||||||
| try: | ||||||||||||||||||||||||||||||||||||||||||||||
| return int(fileno()) | ||||||||||||||||||||||||||||||||||||||||||||||
| except (OSError, ValueError, TypeError, AttributeError): | ||||||||||||||||||||||||||||||||||||||||||||||
| return None | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| def _extract_stream_fd( | ||||||||||||||||||||||||||||||||||||||||||||||
| stream: asyncio.StreamReader | asyncio.StreamWriter | None, | ||||||||||||||||||||||||||||||||||||||||||||||
| ) -> int | None: | ||||||||||||||||||||||||||||||||||||||||||||||
| """Extract a raw FD from an asyncio stream via its transport.""" | ||||||||||||||||||||||||||||||||||||||||||||||
| if stream is None: | ||||||||||||||||||||||||||||||||||||||||||||||
| return None | ||||||||||||||||||||||||||||||||||||||||||||||
| transport = getattr(stream, "transport", None) | ||||||||||||||||||||||||||||||||||||||||||||||
| if transport is None: | ||||||||||||||||||||||||||||||||||||||||||||||
| transport = getattr(stream, "_transport", None) | ||||||||||||||||||||||||||||||||||||||||||||||
| return _fd_from_transport(transport) | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| def _pause_reader_transport( | ||||||||||||||||||||||||||||||||||||||||||||||
| reader: asyncio.StreamReader, | ||||||||||||||||||||||||||||||||||||||||||||||
| ) -> cabc.Callable[[], None] | None: | ||||||||||||||||||||||||||||||||||||||||||||||
| """Pause reader transport callbacks while Rust pump owns the raw FD.""" | ||||||||||||||||||||||||||||||||||||||||||||||
| transport = getattr(reader, "transport", None) | ||||||||||||||||||||||||||||||||||||||||||||||
| if transport is None: | ||||||||||||||||||||||||||||||||||||||||||||||
| transport = getattr(reader, "_transport", None) | ||||||||||||||||||||||||||||||||||||||||||||||
| pause_reading = getattr(transport, "pause_reading", None) | ||||||||||||||||||||||||||||||||||||||||||||||
| resume_reading = getattr(transport, "resume_reading", None) | ||||||||||||||||||||||||||||||||||||||||||||||
| if not callable(pause_reading) or not callable(resume_reading): | ||||||||||||||||||||||||||||||||||||||||||||||
| return None | ||||||||||||||||||||||||||||||||||||||||||||||
| try: | ||||||||||||||||||||||||||||||||||||||||||||||
| pause_reading() | ||||||||||||||||||||||||||||||||||||||||||||||
| except (RuntimeError, OSError): | ||||||||||||||||||||||||||||||||||||||||||||||
| return None | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| def _resume() -> None: | ||||||||||||||||||||||||||||||||||||||||||||||
| """Resume the paused reader transport, ignoring teardown errors.""" | ||||||||||||||||||||||||||||||||||||||||||||||
| with contextlib.suppress(RuntimeError, OSError): | ||||||||||||||||||||||||||||||||||||||||||||||
| resume_reading() | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| return _resume | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| def _set_stream_fds_blocking(*, reader_fd: int, writer_fd: int) -> tuple[bool, bool]: | ||||||||||||||||||||||||||||||||||||||||||||||
| """Switch pipe FDs to blocking mode and return their prior state.""" | ||||||||||||||||||||||||||||||||||||||||||||||
| reader_was_blocking = os.get_blocking(reader_fd) | ||||||||||||||||||||||||||||||||||||||||||||||
| writer_was_blocking = os.get_blocking(writer_fd) | ||||||||||||||||||||||||||||||||||||||||||||||
| reader_changed = False | ||||||||||||||||||||||||||||||||||||||||||||||
| try: | ||||||||||||||||||||||||||||||||||||||||||||||
| if not reader_was_blocking: | ||||||||||||||||||||||||||||||||||||||||||||||
| os.set_blocking(reader_fd, True) | ||||||||||||||||||||||||||||||||||||||||||||||
| reader_changed = True | ||||||||||||||||||||||||||||||||||||||||||||||
| if not writer_was_blocking: | ||||||||||||||||||||||||||||||||||||||||||||||
| os.set_blocking(writer_fd, True) | ||||||||||||||||||||||||||||||||||||||||||||||
| except OSError: | ||||||||||||||||||||||||||||||||||||||||||||||
| if reader_changed: | ||||||||||||||||||||||||||||||||||||||||||||||
| with contextlib.suppress(OSError, ValueError): | ||||||||||||||||||||||||||||||||||||||||||||||
| os.set_blocking(reader_fd, reader_was_blocking) | ||||||||||||||||||||||||||||||||||||||||||||||
| raise | ||||||||||||||||||||||||||||||||||||||||||||||
| return reader_was_blocking, writer_was_blocking | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| def _restore_stream_fd_blocking( | ||||||||||||||||||||||||||||||||||||||||||||||
| async def _run_rust_pump( | ||||||||||||||||||||||||||||||||||||||||||||||
| *, | ||||||||||||||||||||||||||||||||||||||||||||||
| reader: asyncio.StreamReader, | ||||||||||||||||||||||||||||||||||||||||||||||
| writer: asyncio.StreamWriter | None, | ||||||||||||||||||||||||||||||||||||||||||||||
| reader_fd: int, | ||||||||||||||||||||||||||||||||||||||||||||||
| writer_fd: int, | ||||||||||||||||||||||||||||||||||||||||||||||
| reader_was_blocking: bool, | ||||||||||||||||||||||||||||||||||||||||||||||
| writer_was_blocking: bool, | ||||||||||||||||||||||||||||||||||||||||||||||
| ) -> None: | ||||||||||||||||||||||||||||||||||||||||||||||
| """Restore pipe FD blocking mode captured before Rust pumping.""" | ||||||||||||||||||||||||||||||||||||||||||||||
| with contextlib.suppress(OSError, ValueError): | ||||||||||||||||||||||||||||||||||||||||||||||
| os.set_blocking(reader_fd, reader_was_blocking) | ||||||||||||||||||||||||||||||||||||||||||||||
| with contextlib.suppress(OSError, ValueError): | ||||||||||||||||||||||||||||||||||||||||||||||
| os.set_blocking(writer_fd, writer_was_blocking) | ||||||||||||||||||||||||||||||||||||||||||||||
| ) -> bool: | ||||||||||||||||||||||||||||||||||||||||||||||
| """Run the Rust pump for one pipe hop; return ``True`` when handled.""" | ||||||||||||||||||||||||||||||||||||||||||||||
| handled = await _pump_over_raw_fds( | ||||||||||||||||||||||||||||||||||||||||||||||
| reader=reader, | ||||||||||||||||||||||||||||||||||||||||||||||
| writer=writer, | ||||||||||||||||||||||||||||||||||||||||||||||
| reader_fd=reader_fd, | ||||||||||||||||||||||||||||||||||||||||||||||
| writer_fd=writer_fd, | ||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||
| if not handled: | ||||||||||||||||||||||||||||||||||||||||||||||
| return False | ||||||||||||||||||||||||||||||||||||||||||||||
| # The Rust pump closes the writer FD on return (drop semantics). | ||||||||||||||||||||||||||||||||||||||||||||||
| # Suppress OSError so the asyncio transport close does not raise | ||||||||||||||||||||||||||||||||||||||||||||||
| # EBADF when the descriptor is already gone. | ||||||||||||||||||||||||||||||||||||||||||||||
| with contextlib.suppress(OSError): | ||||||||||||||||||||||||||||||||||||||||||||||
| await _close_stream_writer(writer) | ||||||||||||||||||||||||||||||||||||||||||||||
| return True | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| async def _run_rust_pump( | ||||||||||||||||||||||||||||||||||||||||||||||
| async def _pump_over_raw_fds( | ||||||||||||||||||||||||||||||||||||||||||||||
| *, | ||||||||||||||||||||||||||||||||||||||||||||||
| reader: asyncio.StreamReader, | ||||||||||||||||||||||||||||||||||||||||||||||
| writer: asyncio.StreamWriter | None, | ||||||||||||||||||||||||||||||||||||||||||||||
| reader_fd: int, | ||||||||||||||||||||||||||||||||||||||||||||||
| writer_fd: int, | ||||||||||||||||||||||||||||||||||||||||||||||
| ) -> bool: | ||||||||||||||||||||||||||||||||||||||||||||||
| """Run the Rust pump for one pipe hop; return ``True`` when handled.""" | ||||||||||||||||||||||||||||||||||||||||||||||
| # Flush any bytes asyncio already buffered in the StreamReader | ||||||||||||||||||||||||||||||||||||||||||||||
| # before the Rust pump takes over the raw file descriptor. | ||||||||||||||||||||||||||||||||||||||||||||||
| resume_reader = _pause_reader_transport(reader) | ||||||||||||||||||||||||||||||||||||||||||||||
| try: | ||||||||||||||||||||||||||||||||||||||||||||||
| await _drain_reader_buffer(reader, writer) | ||||||||||||||||||||||||||||||||||||||||||||||
| except Exception: | ||||||||||||||||||||||||||||||||||||||||||||||
| if resume_reader is not None: | ||||||||||||||||||||||||||||||||||||||||||||||
| resume_reader() | ||||||||||||||||||||||||||||||||||||||||||||||
| raise | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| try: | ||||||||||||||||||||||||||||||||||||||||||||||
| reader_was_blocking, writer_was_blocking = _set_stream_fds_blocking( | ||||||||||||||||||||||||||||||||||||||||||||||
| """Drive the Rust pump over the raw FDs, managing the FD lifecycle. | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| Pauses the reader transport (always resuming on exit via ``_paused_reader``) | ||||||||||||||||||||||||||||||||||||||||||||||
| and switches the FDs to blocking mode under a ``_BlockingModeGuard`` that | ||||||||||||||||||||||||||||||||||||||||||||||
| restores their prior mode. Returns ``False`` — the signal to fall back to | ||||||||||||||||||||||||||||||||||||||||||||||
|
leynos marked this conversation as resolved.
|
||||||||||||||||||||||||||||||||||||||||||||||
| the Python pump — when the FDs cannot be made blocking; the writer close is | ||||||||||||||||||||||||||||||||||||||||||||||
| left to the caller. | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| Examples | ||||||||||||||||||||||||||||||||||||||||||||||
| -------- | ||||||||||||||||||||||||||||||||||||||||||||||
| A successful hop reports that Rust handled it, leaving the caller to close | ||||||||||||||||||||||||||||||||||||||||||||||
| the writer:: | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| handled = await _pump_over_raw_fds( | ||||||||||||||||||||||||||||||||||||||||||||||
| reader=reader, | ||||||||||||||||||||||||||||||||||||||||||||||
| writer=writer, | ||||||||||||||||||||||||||||||||||||||||||||||
| reader_fd=reader_fd, | ||||||||||||||||||||||||||||||||||||||||||||||
| writer_fd=writer_fd, | ||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||
| except OSError: | ||||||||||||||||||||||||||||||||||||||||||||||
| if resume_reader is not None: | ||||||||||||||||||||||||||||||||||||||||||||||
| resume_reader() | ||||||||||||||||||||||||||||||||||||||||||||||
| return False | ||||||||||||||||||||||||||||||||||||||||||||||
| assert handled is True | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| from cuprum._streams_rs import rust_pump_stream | ||||||||||||||||||||||||||||||||||||||||||||||
| When the descriptors cannot be switched to blocking mode the call reports | ||||||||||||||||||||||||||||||||||||||||||||||
| ``False`` instead of raising, and the reader transport has already been | ||||||||||||||||||||||||||||||||||||||||||||||
| resumed, so the caller falls back to the Python pump:: | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| loop = asyncio.get_running_loop() | ||||||||||||||||||||||||||||||||||||||||||||||
| try: | ||||||||||||||||||||||||||||||||||||||||||||||
| await loop.run_in_executor( | ||||||||||||||||||||||||||||||||||||||||||||||
| None, | ||||||||||||||||||||||||||||||||||||||||||||||
| rust_pump_stream, | ||||||||||||||||||||||||||||||||||||||||||||||
| reader_fd, | ||||||||||||||||||||||||||||||||||||||||||||||
| writer_fd, | ||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||
| finally: | ||||||||||||||||||||||||||||||||||||||||||||||
| _restore_stream_fd_blocking( | ||||||||||||||||||||||||||||||||||||||||||||||
| handled = await _pump_over_raw_fds( | ||||||||||||||||||||||||||||||||||||||||||||||
| reader=reader, | ||||||||||||||||||||||||||||||||||||||||||||||
| writer=writer, | ||||||||||||||||||||||||||||||||||||||||||||||
| reader_fd=reader_fd, | ||||||||||||||||||||||||||||||||||||||||||||||
| writer_fd=writer_fd, | ||||||||||||||||||||||||||||||||||||||||||||||
| reader_was_blocking=reader_was_blocking, | ||||||||||||||||||||||||||||||||||||||||||||||
| writer_was_blocking=writer_was_blocking, | ||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||
| if resume_reader is not None: | ||||||||||||||||||||||||||||||||||||||||||||||
| resume_reader() | ||||||||||||||||||||||||||||||||||||||||||||||
| # The Rust pump closes the writer FD on return (drop semantics). | ||||||||||||||||||||||||||||||||||||||||||||||
| # Suppress OSError so the asyncio transport close does not raise | ||||||||||||||||||||||||||||||||||||||||||||||
| # EBADF when the descriptor is already gone. | ||||||||||||||||||||||||||||||||||||||||||||||
| with contextlib.suppress(OSError): | ||||||||||||||||||||||||||||||||||||||||||||||
| await _close_stream_writer(writer) | ||||||||||||||||||||||||||||||||||||||||||||||
| if not handled: | ||||||||||||||||||||||||||||||||||||||||||||||
| await _run_python_pump(reader, writer) | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| """ | ||||||||||||||||||||||||||||||||||||||||||||||
| with _paused_reader(reader): | ||||||||||||||||||||||||||||||||||||||||||||||
| # Flush any bytes asyncio already buffered in the StreamReader before | ||||||||||||||||||||||||||||||||||||||||||||||
| # the Rust pump takes over the raw file descriptor. | ||||||||||||||||||||||||||||||||||||||||||||||
| await _drain_reader_buffer(reader, writer) | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| try: | ||||||||||||||||||||||||||||||||||||||||||||||
| guard = _BlockingModeGuard.engage(reader_fd=reader_fd, writer_fd=writer_fd) | ||||||||||||||||||||||||||||||||||||||||||||||
| except OSError: | ||||||||||||||||||||||||||||||||||||||||||||||
| return False | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| from cuprum._streams_rs import rust_pump_stream | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
| loop = asyncio.get_running_loop() | ||||||||||||||||||||||||||||||||||||||||||||||
| try: | ||||||||||||||||||||||||||||||||||||||||||||||
| await loop.run_in_executor(None, rust_pump_stream, reader_fd, writer_fd) | ||||||||||||||||||||||||||||||||||||||||||||||
| finally: | ||||||||||||||||||||||||||||||||||||||||||||||
| guard.restore() | ||||||||||||||||||||||||||||||||||||||||||||||
|
Comment on lines
+164
to
+168
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift Bind the FD lifecycle to the executor thread, not to the awaiting task. Cancelling the Await the pump future to completion even when cancelled, so restore and resume only run once the thread has released the descriptors. 🧵 Keep the descriptors alive until the worker finishes loop = asyncio.get_running_loop()
+ pump_future = loop.run_in_executor(
+ None,
+ rust_pump_stream,
+ reader_fd,
+ writer_fd,
+ )
try:
- await loop.run_in_executor(None, rust_pump_stream, reader_fd, writer_fd)
+ await asyncio.shield(pump_future)
+ except asyncio.CancelledError:
+ # The worker still owns both descriptors; restoring their blocking
+ # mode or resuming the reader now would corrupt the in-flight
+ # transfer, so wait for the thread to release them first.
+ await asyncio.wait([pump_future])
+ raise
finally:
guard.restore()Add a regression test that cancels the pumping task mid-transfer and asserts the descriptors are restored only after the worker returns. 📝 Committable suggestion
Suggested change
🤖 Prompt for AI AgentsSource: Coding guidelines |
||||||||||||||||||||||||||||||||||||||||||||||
| return True | ||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Surface pause failure instead of silently racing asyncio against the Rust pump.
A failed or unsupported
pause_reading()returnsNone, which is indistinguishable from a successful pause._pump_over_raw_fds(cuprum/_pipeline_streams.py Lines 152-169) then hands the raw reader FD to the Rust thread while the event loop is still reading that same descriptor, so bytes get split between theStreamReaderbuffer and the Rust pump — silent, non-deterministic data loss on the hop. The blocking-mode seam already signals failure and falls back to Python; make this seam equally honest.🛠️ Propagate pause success through the seam
Then in
cuprum/_pipeline_streams.py:🤖 Prompt for AI Agents