Support suspending stream consumers

adapters/python/coflux/checkpoint.py

@@ -2,12 +2,20 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from typing import Any, Generic, TypeVar, overload
 
 from .state import get_context
 
 T = TypeVar("T")
 
+# Checkpoint names starting with this are the adapter's own. Reserved as a
+# namespace rather than name by name, so later internal state doesn't need
+# another round of this. Enforced here in ``Checkpoint`` rather than in the
+# context's checkpoint_get/set/reset, which are the path the adapter's own
+# cursors go through.
+RESERVED_PREFIX = "_"
+
 
 class Checkpoint(Generic[T]):
     """A named value that survives across executions of a step.
@@ -33,7 +41,13 @@ def poll_orders():
     a checkpoint as at-least-once and make the code that follows a read safe
     to re-run from it. ``cf.flush()`` gives an explicit boundary where that
     isn't good enough. Whatever is written before an execution suspends,
-    returns or fails is always delivered.
+    returns or fails is delivered.
+
+    Writes are cut into deltas only where the execution could resume from
+    them, so a checkpoint is always read back as part of a state the step was
+    actually in: one written while a stream item is in a loop body's hands is
+    published with the cursor advance that consumes it, and dropped if that
+    iteration never finishes.
 
     A checkpoint is not part of any cache, memo or defer key, and a step that
     resolves from the cache never runs and never sees one.
@@ -48,7 +62,8 @@ def poll_orders():
     nothing is enforced at runtime.
 
     Args:
-        name: Checkpoint name, unique within the step.
+        name: Checkpoint name, unique within the step. Can't start with
+            ``_`` — that prefix is reserved for adapter-managed state.
         default: Value returned when the checkpoint has never been set, or has
             been reset. Client-side only — the server never sees it.
     """
@@ -63,6 +78,12 @@ def __init__(self, name: str) -> None: ...
     # ``-> T`` on ``default`` and ``get()``, which it can't when ``T`` is
     # non-optional and no default was given.
     def __init__(self, name: str, *, default: Any = None) -> None:
+        if name.startswith(RESERVED_PREFIX):
+            raise ValueError(
+                f"checkpoint name {name!r} is reserved: names starting with"
+                f" {RESERVED_PREFIX!r} are used for adapter-managed state,"
+                " such as the cursors behind stream suspension"
+            )
         self._name = name
         self._default = default
 
@@ -93,6 +114,27 @@ def set(self, value: T) -> None:
         """Set the value, replacing anything already there."""
         get_context().checkpoint_set(self._name, value)
 
+    def update(self, fn: Callable[[T], T]) -> T:
+        """Set the value to ``fn(current)``, and return what was stored.
+
+        The read-modify-write that most checkpoints do — advancing a
+        cursor, accumulating a total — without naming the old value::
+
+            n = count.update(lambda x: x + 1)
+
+        ``fn`` receives the declared default when the checkpoint isn't set,
+        exactly as ``get()`` would return it.
+
+        This is a read followed by a write, not an atomic swap: two threads
+        of one execution updating the same checkpoint can still lose one of
+        the updates. That only arises if you share a checkpoint across
+        threads — a task body and a ``cf.stream`` generator, say — in which
+        case guard it yourself.
+        """
+        value = fn(self.get())
+        self.set(value)
+        return value
+
     def reset(self) -> None:
         """Clear the checkpoint, so ``get()`` returns the declared default.
 
@@ -126,7 +168,9 @@ def flush() -> None:
         send_notification()
 
     Not needed before suspending, returning or raising — those are flushed
-    automatically.
+    automatically. Inside a stream loop body it also publishes what is being
+    held for the current item, ahead of the cursor advance that would
+    normally carry it.
     """
     get_context().flush()
 

adapters/python/coflux/context.py

@@ -11,7 +11,7 @@
 from collections.abc import Callable, Iterator
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Any
+from typing import Any, NoReturn
 
 from . import protocol
 from .dispatcher import get_dispatcher
@@ -21,6 +21,7 @@
     ExecutionCrashed,
     ExecutionTimeout,
     InputDismissed,
+    Suspending,
     create_execution_error,
 )
 from .models import Asset, AssetEntry, AssetMetadata, Execution, Input
@@ -83,7 +84,10 @@ def _timeout_to_ms(timeout: float | dt.timedelta | None) -> int | None:
 _group_id: contextvars.ContextVar[int | None] = contextvars.ContextVar(
     "_group_id", default=None
 )
-# Context variable for timeout tracking (not yet enforced)
+# Enclosing `cf.suspense` timeout. Read by `select` when deciding how long
+# to wait before suspending, and by stream subscriptions, where it also
+# switches on cursor tracking so a resumed execution carries on rather than
+# re-reading from the start.
 _timeout: contextvars.ContextVar[float | None] = contextvars.ContextVar(
     "_timeout", default=None
 )
@@ -126,6 +130,23 @@ def __init__(self, execution_id: str, working_dir: Path | None = None):
         # guarded by ``self._lock``.
         self._checkpoint_wire: dict[str, Any] = {}
         self._checkpoint_values: dict[str, Any] = {}
+        # Writes held back until the execution is somewhere it could resume
+        # from. ``_checkpoint_holds`` counts the non-replayable inputs
+        # currently in the body's hands (see ``hold_checkpoints``); while it
+        # is non-zero, writes accumulate here instead of going on the wire.
+        # A name is either set or reset, never both — the later write
+        # replaces the earlier one, so the delta only describes the net
+        # effect, exactly as the worker-side throttle coalesces them.
+        self._checkpoint_holds = 0
+        self._checkpoint_pending_set: dict[str, Any] = {}
+        self._checkpoint_pending_reset: set[str] = set()
+        # Occurrence counts for auto-named stream cursors, so two loops
+        # over an identical view of the same stream get distinct
+        # checkpoints. Keyed by the content-addressed base name; the
+        # count is deterministic across attempts as long as subscriptions
+        # are opened in the same order, which is the determinism suspend
+        # already requires.
+        self._cursor_occurrences: dict[str, int] = {}
 
     def set_default_streams(self, streams: Streams | None) -> None:
         """Record the decorator's stream config so ``cf.stream(...)`` can
@@ -624,25 +645,136 @@ def checkpoint_set(self, name: str, value: Any) -> None:
         with self._lock:
             self._checkpoint_values[name] = value
             self._checkpoint_wire.pop(name, None)
-        protocol.send_checkpoint_update(
-            self.execution_id, set_={name: serialize_value(value)}
-        )
+        # Serialised here rather than at publication time: the value is the
+        # one the caller passed, and holding a reference to a mutable object
+        # would record whatever it became later instead.
+        self._record_checkpoint_delta(set_={name: serialize_value(value)})
 
     def checkpoint_reset(self, name: str) -> None:
         with self._lock:
             self._checkpoint_values.pop(name, None)
             self._checkpoint_wire.pop(name, None)
-        protocol.send_checkpoint_update(self.execution_id, reset=[name])
+        self._record_checkpoint_delta(reset=[name])
+
+    def hold_checkpoints(self) -> None:
+        """Note that a non-replayable input is in the body's hands.
+
+        Checkpoint state is one snapshot of a step's progress rather than a
+        set of independent cells — the server stores an execution's row-set
+        as a complete snapshot and applies each delta in a transaction. What
+        decides whether that snapshot is *coherent* is where the deltas get
+        cut, and one cut at an arbitrary point describes a state the
+        execution was never in.
+
+        That only matters for state derived from something a replay can't
+        re-read. A result resolves again; a stream item does not — its
+        position lives in a cursor, and if the cursor and whatever the body
+        derived from the item reach the server separately, a crash in
+        between leaves the successor counting on from a position it never
+        actually reached.
+
+        So writes made while an item is in hand are held, and published in
+        the same delta as the cursor advance that retires it. The pair moves
+        together or not at all, and a replay repeats whole items rather than
+        fractions of one.
+
+        Balanced by ``release_checkpoints``. Nested holds — a body iterating
+        two streams — publish at the outermost release, the only point at
+        which every cursor involved is up to date.
+        """
+        with self._lock:
+            self._checkpoint_holds += 1
+
+    def release_checkpoints(self, *, publish: bool) -> None:
+        """Retire a hold taken by ``hold_checkpoints``.
+
+        ``publish`` says whether the item the hold covered was consumed. On
+        the way out of a completed iteration it is true, and the held writes
+        go out with the cursor advance. Where the item is abandoned instead
+        — ``break``, an exception, a dropped iterator — it is false: the
+        cursor was never advanced, so the item will be delivered again, and
+        anything derived from it must not be recorded or the replay counts
+        it twice.
+
+        Discarding drops the whole pending delta, including writes made
+        under an enclosing hold. That is not over-eager: an enclosing hold
+        means that iteration has not advanced its own cursor either, so
+        everything pending derives from an item that is still unconsumed.
+        """
+        with self._lock:
+            if not self._checkpoint_holds:
+                return
+            self._checkpoint_holds -= 1
+            held = self._checkpoint_holds
+            if not publish:
+                self._checkpoint_pending_set.clear()
+                self._checkpoint_pending_reset.clear()
+        if publish and not held:
+            self._publish_checkpoints()
+
+    def _record_checkpoint_delta(
+        self,
+        set_: dict[str, Any] | None = None,
+        reset: list[str] | None = None,
+    ) -> None:
+        """Put a write on the wire, or hold it for the next safe point."""
+        with self._lock:
+            if self._checkpoint_holds:
+                for name, value in (set_ or {}).items():
+                    self._checkpoint_pending_set[name] = value
+                    self._checkpoint_pending_reset.discard(name)
+                for name in reset or []:
+                    self._checkpoint_pending_reset.add(name)
+                    self._checkpoint_pending_set.pop(name, None)
+                return
+        protocol.send_checkpoint_update(self.execution_id, set_=set_, reset=reset)
+
+    def _publish_checkpoints(self) -> None:
+        """Send whatever is being held, as a single delta."""
+        with self._lock:
+            set_ = self._checkpoint_pending_set
+            reset = self._checkpoint_pending_reset
+            self._checkpoint_pending_set = {}
+            self._checkpoint_pending_reset = set()
+        if set_ or reset:
+            protocol.send_checkpoint_update(
+                self.execution_id,
+                set_=set_ or None,
+                # Sorted only so the delta is deterministic; the server
+                # applies the whole thing at once either way.
+                reset=sorted(reset) or None,
+            )
 
     def flush(self) -> None:
-        """Block until buffered state has reached the server."""
+        """Block until buffered state has reached the server.
+
+        Publishes anything currently held first. An explicit flush is the
+        caller declaring this point consistent, which is what makes it the
+        escape hatch for state that has to be durable before a side effect —
+        including inside a loop body, where the runtime would otherwise wait
+        for the iteration to end. The cursor advance is still to come at
+        that point, so a flush there deliberately records derived state
+        ahead of the position it came from.
+        """
+        self._publish_checkpoints()
         request_id = protocol.request_flush(self.execution_id)
         self._wait_response(request_id)
 
     def suspend_execution(
-        self, delay: float | dt.timedelta | dt.datetime | None = None
-    ) -> None:
-        """Suspend the current execution, optionally resuming after a delay."""
+        self,
+        delay: float | dt.timedelta | dt.datetime | None = None,
+        stream_wait: tuple[str, int] | None = None,
+    ) -> NoReturn:
+        """Signal that this execution should suspend.
+
+        Raises rather than performing the handshake here. The server
+        records a suspension as a completion, and a completed execution's
+        checkpoint writes are rejected — so everything that runs while the
+        body unwinds (``finally`` blocks, cancelled tasks, generator
+        cleanup) has to happen *before* the request is sent, or its state
+        is silently dropped. ``finish_suspension`` completes it once the
+        body is done.
+        """
         execute_after = None
         if isinstance(delay, dt.datetime):
             execute_after = int(delay.timestamp() * 1000)
@@ -657,12 +789,84 @@ def suspend_execution(
                 ).timestamp()
                 * 1000
             )
-        request_id = protocol.request_suspend(self.execution_id, execute_after)
+        raise Suspending(execute_after, stream_wait)
+
+    def finish_suspension(
+        self,
+        execute_after: int | None,
+        stream_wait: tuple[str, int] | None = None,
+    ) -> None:
+        """Complete a suspension once the body has unwound. Never returns.
+
+        Stops any in-flight stream producers and joins their driver threads
+        first, so their cleanup runs while the execution is still live.
+
+        Winding the generators down does *not* close their streams: the
+        driver skips ``send_stream_close`` on ``GeneratorExit``, and the
+        server leaves a suspended execution's streams paused rather than
+        closing them, so the execution that resumes the step continues
+        them.
+
+        ``stream_wait`` gates the successor on a stream reaching a
+        sequence, for a consumer that suspended partway through iterating.
+        """
+        try:
+            self.close_streams()
+            self.wait_streams()
+        except Exception:  # noqa: BLE001, S110
+            # Best-effort teardown — the suspension below is what matters.
+            pass
+        request_id = protocol.request_suspend(
+            self.execution_id, execute_after, stream_wait
+        )
         self._wait_response(request_id)
         # Suspension confirmed. Block until the server aborts this execution.
         get_dispatcher().wait_closed()
         raise SystemExit(0)
 
+    def take_stream_suspension(
+        self,
+    ) -> tuple[int | None, tuple[str, int] | None] | None:
+        """Claim a suspension requested from inside a generator body.
+
+        Returns ``(execute_after, stream_wait)`` — either of which may
+        itself be ``None`` — or ``None`` when no generator asked to
+        suspend. The executor checks this after its streams have drained.
+        """
+        return self._stream_driver.take_suspension()
+
+    def stream_available(self, stream_id: str, sequence: int) -> bool:
+        """Whether ``stream_id`` has reached ``sequence``, or has closed.
+
+        The consumer can't answer this itself. Its queue is fed
+        asynchronously, so an empty one means "nothing has arrived yet",
+        not "the stream has nothing" — checking locally right after
+        subscribing always finds it empty, whatever the stream holds.
+
+        A poll, never a suspension: the server reports what it knows and
+        the decision of what to do about it stays here, so a suspension
+        still unwinds the body before the handshake.
+        """
+        request_id = protocol.request_select(
+            self.execution_id,
+            [{"type": "stream", "id": stream_id, "sequence": sequence}],
+            timeout_ms=0,
+            suspend=False,
+        )
+        return self._wait_response(request_id) is not None
+
+    @property
+    def suspense_timeout(self) -> float | None:
+        """The enclosing ``cf.suspense`` timeout, or ``None`` outside one."""
+        return _timeout.get()
+
+    def next_cursor_occurrence(self, name: str) -> int:
+        """Count of prior subscriptions in this execution sharing ``name``."""
+        with self._lock:
+            occurrence = self._cursor_occurrences.get(name, 0)
+            self._cursor_occurrences[name] = occurrence + 1
+            return occurrence
+
     def _parse_response(self, msg: dict) -> Any:
         """Extract the result from a response message, raising on error."""
         if msg.get("error"):