Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
166 changes: 166 additions & 0 deletions cuprum/_pipeline_stream_fds.py
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
Comment on lines +53 to +74

Copy link
Copy Markdown

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() returns None, 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 the StreamReader buffer 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
 `@contextlib.contextmanager`
-def _paused_reader(reader: asyncio.StreamReader) -> cabc.Iterator[None]:
+def _paused_reader(reader: asyncio.StreamReader) -> cabc.Iterator[bool]:
     """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.
+    transport cannot be paused the yielded value is ``False`` and the resume is
+    never attempted, so the caller can decline the raw-FD hand-off.
     """
     resume = _pause_reader_transport(reader)
     try:
-        yield
+        yield resume is not None
     finally:
         if resume is not None:
             resume()

Then in cuprum/_pipeline_streams.py:

with _paused_reader(reader) as is_paused:
    if not is_paused:
        return False
    ...
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@cuprum/_pipeline_stream_fds.py` around lines 53 - 74, Update
_pause_reader_transport to return an explicit success indicator alongside the
resume callback, distinguishing a completed pause from unsupported transport or
pause errors. Update _paused_reader and _pump_over_raw_fds to consume that
indicator and return False before handing the raw FD to Rust when pausing fails,
while preserving the existing resume cleanup for successful pauses.



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.
Comment thread
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()
187 changes: 69 additions & 118 deletions cuprum/_pipeline_streams.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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
Comment thread
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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The 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 run_in_executor future does not stop the worker thread. On cancellation the finally runs immediately: guard.restore() flips both descriptors back to non-blocking while the Rust pump is mid-read/write, and _paused_reader resumes the event loop on the reader FD that the Rust thread still owns. That yields spurious EAGAIN inside the extension plus two concurrent readers on the same pipe, and leaves an unowned detached thread that will stall loop shutdown.

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

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
loop = asyncio.get_running_loop()
try:
await loop.run_in_executor(None, rust_pump_stream, reader_fd, writer_fd)
finally:
guard.restore()
loop = asyncio.get_running_loop()
pump_future = loop.run_in_executor(
None,
rust_pump_stream,
reader_fd,
writer_fd,
)
try:
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()
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@cuprum/_pipeline_streams.py` around lines 164 - 168, Update the
executor-based pump flow around rust_pump_stream so cancellation does not
immediately run guard.restore(): retain the executor future, await it to
completion when the awaiting task is cancelled, then restore the descriptors and
resume _paused_reader only after the worker thread returns. Add a regression
test that cancels the pumping task mid-transfer and verifies descriptor
restoration occurs only after the worker completes.

Source: Coding guidelines

return True


Expand Down
Loading
Loading