[verified] feat: complete Ralph task workstreams

packages/gateway/meshnet_gateway/prefill_backpressure.py

@@ -0,0 +1,130 @@
+"""Bounded, ordered prefill transfer primitives.
+
+Prefill chunks mutate the downstream shard's session cache, so they must reach a
+route in order.  This deliberately uses a serial acknowledgement window: it is
+the safe default for both current peers and old peers which do not advertise a
+windowing capability.  The configured in-flight limit is still explicit so a
+future ordered transport can widen the window without changing callers.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from threading import Event
+from typing import Callable, Iterable, TypeVar
+
+
+DEFAULT_PREFILL_CHUNK_TOKENS = 128
+DEFAULT_PREFILL_MAX_IN_FLIGHT = 1
+DEFAULT_PREFILL_MAX_CHUNK_BYTES = 8 * 1024 * 1024
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+@dataclass(frozen=True)
+class PrefillTransferLimits:
+    """Configuration for one ordered prefill seam."""
+
+    chunk_tokens: int = DEFAULT_PREFILL_CHUNK_TOKENS
+    max_in_flight: int = DEFAULT_PREFILL_MAX_IN_FLIGHT
+    max_chunk_bytes: int = DEFAULT_PREFILL_MAX_CHUNK_BYTES
+
+    @property
+    def effective_in_flight(self) -> int:
+        """Current peers require ordered session-cache mutation, hence one ack."""
+        return 1
+
+    @property
+    def max_buffered_bytes(self) -> int:
+        """Hard accounting bound, including any future wider ack window."""
+        return self.max_chunk_bytes * self.max_in_flight
+
+    @classmethod
+    def from_env(cls) -> "PrefillTransferLimits":
+        # MESHNET_CHUNK_TOKENS was the pre-DIP-007 name.  Keep it as a fallback
+        # so existing deployments retain their chunk shape while upgrading.
+        return cls(
+            chunk_tokens=_positive_env(
+                "MESHNET_PREFILL_CHUNK_TOKENS",
+                _positive_env("MESHNET_CHUNK_TOKENS", DEFAULT_PREFILL_CHUNK_TOKENS),
+            ),
+            max_in_flight=_positive_env(
+                "MESHNET_PREFILL_MAX_IN_FLIGHT", DEFAULT_PREFILL_MAX_IN_FLIGHT,
+            ),
+            max_chunk_bytes=_positive_env(
+                "MESHNET_PREFILL_MAX_CHUNK_BYTES", DEFAULT_PREFILL_MAX_CHUNK_BYTES,
+            ),
+        )
+
+
+class BoundedPrefillSender:
+    """Send lazily-produced chunks with bounded ownership and ordered acks."""
+
+    def __init__(self, limits: PrefillTransferLimits) -> None:
+        self.limits = limits
+        self.buffered_bytes = 0
+        self.peak_buffered_bytes = 0
+        self.in_flight = 0
+        self.peak_in_flight = 0
+        self.closed = False
+
+    def send(
+        self,
+        chunks: Iterable[T],
+        *,
+        body_size: Callable[[T], int],
+        forward: Callable[[T], R],
+        cancelled: Event | None = None,
+    ) -> list[R]:
+        """Forward chunks in source order, releasing each body after its ack.
+
+        ``forward`` is synchronous by design: a slow consumer therefore blocks
+        production of the next chunk instead of accumulating an unbounded queue.
+        Every retained body is dropped on cancellation or route failure.
+        """
+        results: list[R] = []
+        try:
+            for chunk in chunks:
+                if self.closed or (cancelled is not None and cancelled.is_set()):
+                    break
+                size = body_size(chunk)
+                if size < 0:
+                    raise ValueError("prefill chunk size cannot be negative")
+                if size > self.limits.max_chunk_bytes:
+                    raise ValueError(
+                        f"prefill chunk exceeds {self.limits.max_chunk_bytes} byte limit"
+                    )
+                self.buffered_bytes += size
+                self.in_flight += 1
+                self.peak_buffered_bytes = max(self.peak_buffered_bytes, self.buffered_bytes)
+                self.peak_in_flight = max(self.peak_in_flight, self.in_flight)
+                try:
+                    results.append(forward(chunk))
+                finally:
+                    # Do not retain a body while waiting for the next chunk.
+                    self.buffered_bytes -= size
+                    self.in_flight -= 1
+        except BaseException:
+            self.close()
+            raise
+        return results
+
+    def close(self) -> None:
+        """Release accounting after cancellation or route failure.
+
+        The sender deliberately owns no queued chunk references; callers must
+        discard their iterator on close rather than trying to drain it.
+        """
+        self.buffered_bytes = 0
+        self.in_flight = 0
+        self.closed = True
+
+
+def _positive_env(name: str, default: int) -> int:
+    try:
+        value = int(os.environ.get(name, default))
+    except (TypeError, ValueError):
+        return default
+    return value if value > 0 else default