diff --git a/adapters/python/coflux/errors.py b/adapters/python/coflux/errors.py index c4a5ebbe..804a3a73 100644 --- a/adapters/python/coflux/errors.py +++ b/adapters/python/coflux/errors.py @@ -78,17 +78,16 @@ def __init__(self, message: str = "worker terminated without reporting a result" class StreamSuperseded(ExecutionTerminated): - """Raised when a stream ended because its producer was superseded. + """Raised when a stream ended because its producer recurred. - The producer suspended, or finished an iteration of a recurrent - target. Neither is a failure — but a stream is owned by exactly one - execution, so the successor registers a *new* stream rather than - continuing this one, and a handle to this stream will never yield - anything further. + The producer finished an iteration of a recurrent target. That isn't + a failure — but each iteration produces its own stream, so a handle + to this one will never yield anything further. - To follow the successor's stream, obtain a fresh handle from the - successor (e.g. re-resolve the producer's result) rather than - re-iterating this one. + To follow the next iteration's stream, obtain a fresh handle from it + rather than re-iterating this one. (A producer that *suspends* does + not end its stream: the resumed execution continues it, and consumers + simply wait.) """ def __init__(self, message: str = "stream producer was superseded"): @@ -215,7 +214,7 @@ def raise_for_close(reason: str, error: dict | None) -> None: raise ExecutionCrashed() if reason == "timeout": raise ExecutionTimeout() - if reason in ("suspended", "recurred"): + if reason == "recurred": raise StreamSuperseded() # Anything else (e.g. "not_found", "already_subscribed", or an diff --git a/adapters/python/coflux/models.py b/adapters/python/coflux/models.py index 80041dbb..3991db4f 100644 --- a/adapters/python/coflux/models.py +++ b/adapters/python/coflux/models.py @@ -298,6 +298,12 @@ class Stream(t.Iterable[T], t.AsyncIterable[T]): starts a fresh subscription from sequence 0, so a stream can be iterated multiple times and each iteration sees the whole sequence. + A stream belongs to the step that produces it. If the producer + suspends, the stream pauses rather than ending, and the execution that + resumes the step appends to it — so iteration simply waits through + the suspension and carries on. Any other way the producer ends closes + the stream. + ``async for`` works too, and does the same thing — the difference is only in how the consumer waits. Prefer it in ``async def`` bodies: the sync iterator blocks its thread between items, which in an async @@ -315,8 +321,8 @@ def __init__( id: str, stride: Stride = (0, None, 1), ): - # Opaque identifier of the form ``_``. - # Users may see this in the CLI/Studio but shouldn't need to parse it. + # Opaque identifier of the form ``:_``. Users may + # see this in the CLI/Studio but shouldn't need to parse it. self._id = id self._stride = stride diff --git a/adapters/python/coflux/protocol.py b/adapters/python/coflux/protocol.py index e7dcf802..7072fafa 100644 --- a/adapters/python/coflux/protocol.py +++ b/adapters/python/coflux/protocol.py @@ -505,16 +505,23 @@ def request_flush(execution_id: str) -> int: return get_protocol().send_request("flush", {"execution_id": execution_id}) -def send_stream_register( +def request_stream_register( execution_id: str, - index: int, + position: int, buffer: int | None = None, timeout_ms: int | None = None, -) -> None: - """Register a stream owned by this execution. - - ``index`` is worker-assigned and monotonic per execution (0, 1, 2, ...); - it identifies the stream within its producer execution. +) -> int: + """Request registration of this execution's ``position``-th stream. + + ``position`` is the order in which this execution registers its + streams (0, 1, 2, ...). The server decides what the registration + means: a new stream of the step, or a resumption of one left paused + by a suspended predecessor at the same position. The response is + ``{"id", "index", "head"}`` — the stream's id + (``:_``), its index within the step (what appends, + closes and demand grants carry from then on), and the last sequence + already in the stream (``-1`` for a new stream; the producer + continues from ``head + 1``). ``buffer`` is the producer-side backpressure budget. ``None`` opts out of backpressure entirely; the server won't issue demand grants and @@ -524,15 +531,15 @@ def send_stream_register( ``timeout_ms`` is the idle-timeout budget (milliseconds). If set, the worker (CLI) force-closes the stream with reason "timeout" when no - item has been appended for that long. Purely informational for the - server; enforcement happens in the worker. + item has been appended for that long. Recorded by the server for + display; enforcement happens in the worker. """ - params: dict[str, Any] = {"execution_id": execution_id, "index": index} + params: dict[str, Any] = {"execution_id": execution_id, "position": position} if buffer is not None: params["buffer"] = buffer if timeout_ms is not None: params["timeout_ms"] = timeout_ms - get_protocol().send_message("stream_register", params) + return get_protocol().send_request("stream_register", params) def send_stream_append( @@ -587,18 +594,18 @@ def send_stream_close( def send_stream_subscribe( execution_id: str, subscription_id: int, - producer_execution_id: str, - index: int, + stream_id: str, from_sequence: int, prefetch: int, stride: dict[str, Any] | None = None, ) -> None: - """Open a consumer subscription to a stream owned by another execution. + """Open a consumer subscription to a stream. ``execution_id`` is the consumer's own execution — the server uses it - to track who's subscribed and where to push items. ``stride`` is an - optional ``{"start": int, "stop": int|None, "step": int}`` dict - restricting which sequence positions are delivered; any chain of + to track who's subscribed and where to push items. ``stream_id`` is + the stream's opaque id from its handle. ``stride`` is an optional + ``{"start": int, "stop": int|None, "step": int}`` dict restricting + which sequence positions are delivered; any chain of slice/partition/stride calls on the handle composes into a single stride before reaching here. @@ -609,8 +616,7 @@ def send_stream_subscribe( params: dict[str, Any] = { "execution_id": execution_id, "subscription_id": subscription_id, - "producer_execution_id": producer_execution_id, - "index": index, + "stream_id": stream_id, "from_sequence": from_sequence, "prefetch": prefetch, } diff --git a/adapters/python/coflux/streams.py b/adapters/python/coflux/streams.py index 4c9c2385..e7f121fb 100644 --- a/adapters/python/coflux/streams.py +++ b/adapters/python/coflux/streams.py @@ -6,6 +6,14 @@ async (``async def`` + ``yield``) generators are supported; async generators get a fresh event loop confined to their worker thread. +A stream belongs to the step, not the execution. Registering is a +round-trip: the server allocates the stream's index within the step and +tells the driver whether this registration resumes a stream a suspended +predecessor left paused — in which case items continue the existing +sequence, and consumers see one unbroken stream across the suspension. +That is what makes ``cf.suspend()`` inside a generator body the way to +write a producer that pauses and carries on. + The consumer side owns a module-level ``StreamRegistry``: open consumer subscriptions are keyed by subscription id. The registry's dispatcher handlers (``stream_items``/``stream_closed``) route incoming pushes from @@ -61,11 +69,18 @@ def stream( handles the registration automatically — you don't need to call ``cf.stream`` explicitly. - Registration happens at call time: the driver thread starts, the - server is told about the stream, and any later serialisation sees a - regular ``Stream`` handle. That means ``cf.stream`` must be called - inside a task or workflow body (where an execution context is - active); calling it from module scope or outside a task raises. + Registration happens at call time: the server is asked to register + the stream (a round-trip, like ``submit``), the driver thread starts, + and any later serialisation sees a regular ``Stream`` handle. That + means ``cf.stream`` must be called inside a task or workflow body + (where an execution context is active); calling it from module scope + or outside a task raises. + + Streams are matched across a suspension by the order they're + registered in, so code that registers several streams before + suspending has to register them in the same order when it resumes — + the same determinism suspend already requires of the code before the + suspend point. Unspecified options inherit from the enclosing task's ``streams=cf.Streams(...)``. Explicit options override per-call. @@ -88,9 +103,9 @@ def stream( ``timedelta``, or ``None`` to disable. Returns: - A ``Stream`` handle referencing the newly registered stream. - It serialises as ``{"type": "stream", "id": ...}`` and is - iterable by downstream tasks. + A ``Stream`` handle referencing the registered stream. It + serialises as ``{"type": "stream", "id": ...}`` and is iterable + by downstream tasks. """ if not (inspect.isgenerator(generator) or inspect.isasyncgen(generator)): raise TypeError( @@ -116,7 +131,10 @@ class StreamDriver: def __init__(self, execution_id: str) -> None: self._execution_id = execution_id - self._next_index = 0 + # Registration order within this execution. The server matches a + # resuming execution's k-th registration onto the paused stream a + # suspended predecessor registered k-th. + self._next_position = 0 self._threads: list[threading.Thread] = [] self._generators: list[Any] = [] self._lock = threading.Lock() @@ -127,6 +145,12 @@ def __init__(self, execution_id: str) -> None: # registration time); the driver never waits. self._demand_cv = threading.Condition() self._demand: dict[int, int | None] = {} + # Credits granted for an index we haven't been told about yet. The + # server grants demand while it handles a registration — before + # the reply carrying the stream's index reaches us — so the grant + # can overtake the reply. Held here until ``register`` learns the + # index, then applied. + self._pending_demand: dict[int, int] = {} self._closing = False self._demand_handler_registered = False self._force_close_handler_registered = False @@ -152,6 +176,11 @@ def register( async generators run inside a fresh event loop confined to that thread. + Registration is a request to the server, which replies with the + stream's id, its index within the step, and the head to sequence + from. A new stream starts at 0; one resumed after a suspend + continues from where the suspended execution left it. + ``buffer`` is the producer-side backpressure budget. ``None`` means unbounded (no flow control); ``0`` means strict lockstep (producer waits for a consumer to ack each item before emitting @@ -162,25 +191,39 @@ def register( worker (CLI) closes the stream with reason "timeout" if no item is appended within that window. ``None`` disables the timeout. - Returns the stream's opaque ``id`` (``_``) - for embedding in the serialized value as a stream reference. + Returns the stream's opaque ``id`` for embedding in the serialized + value as a stream reference. """ self._ensure_demand_handler_registered() self._ensure_force_close_handler_registered() with self._lock: - index = self._next_index - self._next_index += 1 - - with self._demand_cv: - # Unbounded ⇒ driver never waits. Bounded ⇒ starts at 0; the - # server issues a credit grant once demand calculation warrants - # it (or on first consumer subscribing). - self._demand[index] = None if buffer is None else 0 + position = self._next_position + self._next_position += 1 - protocol.send_stream_register( - self._execution_id, index, buffer=buffer, timeout_ms=timeout_ms + request_id = protocol.request_stream_register( + self._execution_id, position, buffer=buffer, timeout_ms=timeout_ms ) + response = get_dispatcher().wait_for_response(request_id) + if response is None: + raise RuntimeError("timed out registering stream") + if response.get("error"): + error = response["error"] + raise RuntimeError(f"{error['code']}: {error['message']}") + registration = response.get("result") or {} + stream_id: str = registration["id"] + index: int = registration["index"] + head: int = registration.get("head", -1) + + with self._demand_cv: + # Unbounded ⇒ driver never waits. Bounded ⇒ starts from whatever + # the server has already granted: it issues a credit grant while + # handling the registration when demand warrants it (a larger + # buffer to pre-warm, or — for a resumed stream — subscribers + # already waiting), and that grant may have arrived ahead of + # the reply. + pending = self._pending_demand.pop(index, 0) + self._demand[index] = None if buffer is None else pending is_async = inspect.isasyncgen(generator) target = self._run_async if is_async else self._run @@ -192,7 +235,7 @@ def register( # context and would lose those settings. parent_context = contextvars.copy_context() thread = threading.Thread( - target=lambda: parent_context.run(target, index, generator), + target=lambda: parent_context.run(target, index, generator, head + 1), name=f"stream-{self._execution_id}-{index}", daemon=False, ) @@ -203,7 +246,7 @@ def register( self._by_index[index] = entry thread.start() - return compose_stream_id(self._execution_id, index) + return stream_id def _ensure_demand_handler_registered(self) -> None: if self._demand_handler_registered: @@ -230,7 +273,11 @@ def _on_stream_demand(self, params: dict[str, Any]) -> None: if index is None or n <= 0: return with self._demand_cv: - current = self._demand.get(index) + if index not in self._demand: + # Overtook the register reply — see ``_pending_demand``. + self._pending_demand[index] = self._pending_demand.get(index, 0) + n + return + current = self._demand[index] if current is None: # Unbounded — nothing to account for. return @@ -301,9 +348,13 @@ def _is_force_closed(self, index: int) -> bool: with self._demand_cv: return index in self._force_closed - def _run(self, index: int, generator: Any) -> None: - """Run one sync generator to exhaustion (or error).""" - sequence = 0 + def _run(self, index: int, generator: Any, start_sequence: int) -> None: + """Run one sync generator to exhaustion (or error). + + ``start_sequence`` is where numbering begins: 0 for a new stream, + or one past the head of a stream this execution is resuming. + """ + sequence = start_sequence try: iterator = iter(generator) while True: @@ -331,6 +382,12 @@ def _run(self, index: int, generator: Any) -> None: # lifecycle closure on execution-end, or has already recorded # the force-close reason (e.g. "timeout"). return + except SystemExit: + # The generator body suspended (``cf.suspend()`` / implicit + # suspense). The stream stays open — paused — for the + # execution that resumes the step to continue, so nothing is + # sent: a close here would end it for every consumer. + return except BaseException as e: # noqa: BLE001 - we propagate all if self._is_force_closed(index): # Worker already recorded the close; don't overwrite. @@ -349,7 +406,7 @@ def _run(self, index: int, generator: Any) -> None: return protocol.send_stream_close(self._execution_id, index) - def _run_async(self, index: int, generator: Any) -> None: + def _run_async(self, index: int, generator: Any, start_sequence: int) -> None: """Run one async generator in a fresh event loop on this thread. The loop handle is recorded so ``close_all`` can schedule aclose() @@ -360,7 +417,7 @@ def _run_async(self, index: int, generator: Any) -> None: asyncio.set_event_loop(loop) async def iterate() -> None: - sequence = 0 + sequence = start_sequence iterator = generator.__aiter__() while True: # The demand wait uses a threading.Condition, which would @@ -387,6 +444,9 @@ async def iterate() -> None: loop.run_until_complete(iterate()) except (GeneratorExit, asyncio.CancelledError): return + except SystemExit: + # Suspended from inside the generator — see ``_run``. + return except BaseException as e: # noqa: BLE001 - we propagate all if self._is_force_closed(index): return @@ -881,25 +941,6 @@ def _stream_registry() -> StreamRegistry: return _registry_instance -def compose_stream_id(execution_id: str, index: int) -> str: - """Build the opaque stream id from its two components. - - Joined with ``_`` because the alternatives are overloaded: ``:`` is - used inside the execution id, ``#`` is used for attempt numbers, ``/`` - separates module/target. Execution ids use only alphanumerics, so - ``rpartition('_')`` is unambiguous on the parse side. - """ - return f"{execution_id}_{index}" - - -def parse_stream_id(id: str) -> tuple[str, int]: - """Reverse of ``compose_stream_id``. Raises ValueError on bad input.""" - exec_id, sep, index = id.rpartition("_") - if not sep or not exec_id: - raise ValueError(f"invalid stream id: {id!r}") - return exec_id, int(index) - - def _open_subscription( stream_id: str, stride: tuple[int, int | None, int], @@ -910,24 +951,20 @@ def _open_subscription( ``stride`` is a ``(start, stop, step)`` tuple — any chain of slice/partition/stride calls on the handle collapses to a single stride before this point. The wire message is the same whichever - reader the caller asked for; only local delivery differs. + reader the caller asked for; only local delivery differs. The stream + id is opaque here — the server resolves it. """ ctx = get_context() execution_id = ctx.execution_id subscription_id, iterator = _stream_registry().allocate(execution_id, factory) - # Split the opaque id for the wire message, which still takes - # producer_execution_id + index positionally. - producer_execution_id, index = parse_stream_id(stream_id) - start, stop, step = stride wire_stride = {"start": start, "stop": stop, "step": step} protocol.send_stream_subscribe( execution_id, subscription_id, - producer_execution_id, - index, + stream_id, 0, _PREFETCH, stride=wire_stride, diff --git a/adapters/python/coflux/target.py b/adapters/python/coflux/target.py index babb062f..7a59876d 100644 --- a/adapters/python/coflux/target.py +++ b/adapters/python/coflux/target.py @@ -324,6 +324,21 @@ def _build_definition( if p.kind != inspect.Parameter.POSITIONAL_OR_KEYWORD: raise TypeError(f"Unsupported parameter type ({p.kind})") parameters_ = [_build_parameter(p) for p in parameters] + # A generator-bodied target's result is its stream handle, never + # ``None``, so it could never recur — and each recurrence would open a + # fresh stream anyway. The continuous form is a suspend inside the + # generator: the resumed execution carries on the same stream. + if ( + recurrent + and not is_stub + and (inspect.isgeneratorfunction(fn) or inspect.isasyncgenfunction(fn)) + ): + raise TypeError( + f"{fn.__name__}: recurrent=True can't be combined with a generator " + "body. To keep producing across pauses, call cf.suspend() inside " + "the generator instead — the resumed execution continues the same " + "stream." + ) return TargetDefinition( type, parameters_, diff --git a/adapters/python/tests/test_stream_driver.py b/adapters/python/tests/test_stream_driver.py new file mode 100644 index 00000000..ffbbfbbc --- /dev/null +++ b/adapters/python/tests/test_stream_driver.py @@ -0,0 +1,232 @@ +"""Producer-side driver tests: registration is a request the server answers +with the stream's identity and resume point, and a suspend from inside a +generator leaves the stream open. + +The driver talks to the server through ``protocol`` and waits on the +``dispatcher``; both are faked here, so no CLI or server is involved. +""" + +from __future__ import annotations + +import threading +from types import SimpleNamespace + +import pytest + +from coflux import protocol, streams +from coflux.streams import StreamDriver + + +class _FakeDispatcher: + """Answers stream_register requests with a canned reply and grants + unlimited demand, so the driver never waits.""" + + def __init__(self, reply): + self.reply = reply + self.handlers = {} + self.closed = False + + def register_notification(self, method, handler): + self.handlers[method] = handler + + def wait_for_response(self, request_id, timeout=None): + return {"id": request_id, "result": dict(self.reply)} + + def is_closed(self): + return self.closed + + +class _Harness: + def __init__(self, reply): + self.dispatcher = _FakeDispatcher(reply) + self.registers = [] + self.appends = [] + self.closes = [] + self.done = threading.Event() + + def request_stream_register( + self, execution_id, position, buffer=None, timeout_ms=None + ): + self.registers.append((execution_id, position, buffer, timeout_ms)) + return len(self.registers) + + def send_stream_append(self, execution_id, index, sequence, value): + self.appends.append((index, sequence, value)) + + def send_stream_close(self, execution_id, index, **error): + self.closes.append((index, error)) + + +@pytest.fixture +def harness(monkeypatch): + def make(reply): + h = _Harness(reply) + monkeypatch.setattr(streams, "get_dispatcher", lambda: h.dispatcher) + monkeypatch.setattr( + protocol, "request_stream_register", h.request_stream_register + ) + monkeypatch.setattr(protocol, "send_stream_append", h.send_stream_append) + monkeypatch.setattr(protocol, "send_stream_close", h.send_stream_close) + monkeypatch.setattr(streams, "serialize_value", lambda value: value) + return h + + return make + + +def test_new_stream_sequences_from_zero(harness): + h = harness({"id": "run1:2_0", "index": 0, "head": -1}) + driver = StreamDriver("run1:2:1") + + def gen(): + yield "a" + yield "b" + + # Unbounded, so the driver doesn't wait for demand grants. + stream_id = driver.register(gen(), buffer=None) + driver.wait_all() + + assert stream_id == "run1:2_0" + assert h.registers == [("run1:2:1", 0, None, None)] + assert h.appends == [(0, 0, "a"), (0, 1, "b")] + assert h.closes == [(0, {})] + + +def test_resumed_stream_continues_from_the_head(harness): + # The server says this registration resumes a paused stream whose + # last item was sequence 4, under the step's index 3. + h = harness({"id": "run1:2_3", "index": 3, "head": 4}) + driver = StreamDriver("run1:2:2") + + def gen(): + yield "e" + yield "f" + + stream_id = driver.register(gen(), buffer=None) + driver.wait_all() + + assert stream_id == "run1:2_3" + assert h.appends == [(3, 5, "e"), (3, 6, "f")] + assert h.closes == [(3, {})] + + +def test_positions_count_registrations_within_the_execution(harness): + h = harness({"id": "run1:2_7", "index": 7, "head": -1}) + driver = StreamDriver("run1:2:1") + + def gen(): + yield from () + + driver.register(gen(), buffer=None) + driver.register(gen(), buffer=None) + driver.wait_all() + + assert [position for _, position, _, _ in h.registers] == [0, 1] + + +def test_suspend_inside_the_generator_sends_no_close(harness): + h = harness({"id": "run1:2_0", "index": 0, "head": -1}) + driver = StreamDriver("run1:2:1") + + def gen(): + yield "a" + # ``cf.suspend()`` ends the calling thread with SystemExit once the + # server has confirmed the suspension. The stream must be left + # paused for the resumed execution, not closed. + raise SystemExit(0) + + driver.register(gen(), buffer=None) + driver.wait_all() + + assert h.appends == [(0, 0, "a")] + assert h.closes == [] + + +def test_generator_error_closes_with_the_error(harness): + h = harness({"id": "run1:2_0", "index": 0, "head": -1}) + driver = StreamDriver("run1:2:1") + + def gen(): + yield "a" + raise ValueError("boom") + + driver.register(gen(), buffer=None) + driver.wait_all() + + assert h.appends == [(0, 0, "a")] + assert len(h.closes) == 1 + index, error = h.closes[0] + assert index == 0 + assert error["error_type"] == "builtins.ValueError" + assert error["error_message"] == "boom" + + +def test_registration_error_is_raised_to_the_caller(harness, monkeypatch): + h = harness({}) + monkeypatch.setattr( + h.dispatcher, + "wait_for_response", + lambda request_id, timeout=None: { + "id": request_id, + "error": {"code": "execution_completed", "message": "finalised"}, + }, + ) + driver = StreamDriver("run1:2:1") + + def gen(): + yield "a" + + with pytest.raises(RuntimeError, match="execution_completed"): + driver.register(gen(), buffer=None) + assert h.appends == [] + + +def test_recurrent_generator_target_is_rejected(): + import coflux as cf + + with pytest.raises(TypeError, match="recurrent=True"): + + @cf.task(recurrent=True) + def ticker(): + yield 1 + + # The non-recurrent form, and recurrence on a plain body, are fine. + @cf.task() + def stream_task(): + yield 1 + + @cf.task(recurrent=True) + def plain_task(): + return None + + _ = SimpleNamespace(stream_task=stream_task, plain_task=plain_task) + + +def test_demand_granted_before_the_register_reply_is_not_lost(harness): + """The server grants demand while handling a registration, so the + grant can reach the adapter before the reply that names the stream's + index. That's the normal case for a resumed stream (its subscribers + are already waiting) and for a pre-warming buffer. The credits must + be held and applied, not dropped as belonging to an unknown stream. + """ + h = harness({"id": "run1:2_0", "index": 0, "head": 3}) + dispatcher = h.dispatcher + original = dispatcher.wait_for_response + + def grant_then_reply(request_id, timeout=None): + # One credit for the item, one for the ``next()`` that ends the + # generator — both delivered ahead of the reply. + dispatcher.handlers["stream_demand"]({"index": 0, "n": 2}) + return original(request_id, timeout) + + dispatcher.wait_for_response = grant_then_reply + driver = StreamDriver("run1:2:2") + + def gen(): + yield "e" + + # Bounded (lockstep), so the driver only proceeds on granted credit. + driver.register(gen(), buffer=0) + driver.wait_all() + + assert h.appends == [(0, 4, "e")] + assert h.closes == [(0, {})] diff --git a/adapters/python/tests/test_stream_iteration.py b/adapters/python/tests/test_stream_iteration.py index 019b5ac6..3213cf4e 100644 --- a/adapters/python/tests/test_stream_iteration.py +++ b/adapters/python/tests/test_stream_iteration.py @@ -75,8 +75,7 @@ def on_subscribe( self, execution_id, subscription_id, - producer_execution_id, - index, + stream_id, from_sequence, prefetch, stride=None, diff --git a/cli/internal/adapter/protocol.go b/cli/internal/adapter/protocol.go index 88bc1789..7e60fe75 100644 --- a/cli/internal/adapter/protocol.go +++ b/cli/internal/adapter/protocol.go @@ -290,18 +290,31 @@ type RegisterGroupParams struct { Name *string `json:"name,omitempty"` } -// StreamRegisterParams for stream_register notification. -// Index is worker-assigned, monotonic per execution — it identifies the -// stream within its producer execution. Buffer is the optional -// backpressure budget; nil means unbounded (no flow control). TimeoutMs -// is the optional idle-timeout budget (milliseconds) — nil disables it. +// StreamRegisterParams for the stream_register request. Position is the +// order in which the execution registered the stream (its k-th, from 0); +// the server allocates the stream's index within its step and replies +// with a StreamRegisterResult. Buffer is the optional backpressure +// budget; nil means unbounded (no flow control). TimeoutMs is the +// optional idle-timeout budget (milliseconds) — nil disables it. type StreamRegisterParams struct { ExecutionID string `json:"execution_id"` - Index int `json:"index"` + Position int `json:"position"` Buffer *int `json:"buffer,omitempty"` TimeoutMs *int `json:"timeout_ms,omitempty"` } +// StreamRegisterResult is the reply to stream_register. ID is the +// stream's id (`:_`); Index is its index within the +// step, which appends, closes and demand grants carry from then on; Head +// is the last sequence already in the stream — -1 for a new stream, or +// the resume point when the registration continues a stream paused by a +// suspend (the producer sequences from Head+1). +type StreamRegisterResult struct { + ID string `json:"id"` + Index int `json:"index"` + Head int `json:"head"` +} + // StreamDemandParams for stream_demand notification pushed CLI → adapter. // Grants the producer “n“ more credits for the given stream. type StreamDemandParams struct { @@ -335,21 +348,21 @@ type StreamCloseError struct { Traceback string `json:"traceback"` } -// StreamSubscribeParams for stream_subscribe notification. `Stride` -// (when present) restricts which sequence positions are delivered: the -// positions `start, start+step, start+2·step, …` up to (but not -// including) `stop`. Any chain of slice/partition/stride calls on the -// consumer side composes into a single stride before the wire. -// `Prefetch` is the adapter's delivery window: the server won't send -// more than this many items beyond what the adapter has acknowledged. +// StreamSubscribeParams for stream_subscribe notification. `StreamID` +// is the stream's id (`:_`). `Stride` (when present) +// restricts which sequence positions are delivered: the positions +// `start, start+step, start+2·step, …` up to (but not including) `stop`. +// Any chain of slice/partition/stride calls on the consumer side +// composes into a single stride before the wire. `Prefetch` is the +// adapter's delivery window: the server won't send more than this many +// items beyond what the adapter has acknowledged. type StreamSubscribeParams struct { - ExecutionID string `json:"execution_id"` // consumer - SubscriptionID int `json:"subscription_id"` - ProducerExecutionID string `json:"producer_execution_id"` - Index int `json:"index"` - FromSequence int `json:"from_sequence"` - Stride map[string]any `json:"stride,omitempty"` - Prefetch int `json:"prefetch"` + ExecutionID string `json:"execution_id"` // consumer + SubscriptionID int `json:"subscription_id"` + StreamID string `json:"stream_id"` + FromSequence int `json:"from_sequence"` + Stride map[string]any `json:"stride,omitempty"` + Prefetch int `json:"prefetch"` } // StreamAckParams for stream_ack notification. `Count` and `Sequence` diff --git a/cli/internal/pool/pool.go b/cli/internal/pool/pool.go index 7ef102c0..c3ffdec1 100644 --- a/cli/internal/pool/pool.go +++ b/cli/internal/pool/pool.go @@ -59,24 +59,26 @@ type ExecutionHandler interface { SubmitInput(ctx context.Context, params *adapter.SubmitInputParams) (string, error) // NotifyTerminated notifies the server that an execution's process has exited NotifyTerminated(ctx context.Context, executionID string) error - // StreamRegister declares a new stream owned by an execution. - // Index is worker-assigned, monotonic per execution — it identifies - // the stream within its producer execution. Buffer is the optional - // backpressure budget; nil means unbounded (no flow control). - // TimeoutMs is the optional idle-timeout budget (milliseconds); - // purely informational for the server (enforced at the worker/CLI). - StreamRegister(ctx context.Context, executionID string, index int, buffer *int, timeoutMs *int) error + // StreamRegister declares an execution's k-th stream (`position`). + // The server allocates the stream's index within its step and decides + // whether the registration resumes a stream paused by a suspend; the + // result carries the stream's id, its index and the head to sequence + // from. Buffer is the optional backpressure budget; nil means + // unbounded (no flow control). TimeoutMs is the optional idle-timeout + // budget (milliseconds); enforced at the worker/CLI, recorded by the + // server for display. + StreamRegister(ctx context.Context, executionID string, position int, buffer *int, timeoutMs *int) (*adapter.StreamRegisterResult, error) // StreamAppend appends an item to a stream. Sequence is worker-assigned, // monotonic per stream — it identifies the item within its stream. StreamAppend(ctx context.Context, executionID string, index int, sequence int, value *adapter.Value) error // StreamClose closes a stream. ``reason`` is "complete" | "errored" | "timeout". // When nil, inferred from ``err`` (nil→complete, non-nil→errored). StreamClose(ctx context.Context, executionID string, index int, err *adapter.StreamCloseError, reason *string) error - // StreamSubscribe opens a consumer subscription to a stream owned - // by another execution. `stride` is an optional + // StreamSubscribe opens a consumer subscription to a stream by its + // id (`:_`). `stride` is an optional // {"start", "stop", "step"} map restricting which positions are // delivered; nil means no filtering. - StreamSubscribe(ctx context.Context, executionID string, subscriptionID int, producerExecutionID string, index int, fromSequence int, stride map[string]any, prefetch int) error + StreamSubscribe(ctx context.Context, executionID string, subscriptionID int, streamID string, fromSequence int, stride map[string]any, prefetch int) error // StreamAck reports consumer progress. `count` and `sequence` are // cumulative — how many items have been processed, and the highest // sequence among them. Frees delivery credit server-side. @@ -355,7 +357,7 @@ loop: // throttle means at most one per window per execution. p.handleCheckpointUpdate(execCtx, executionID, params, logger) - case "submit_execution", "select", "persist_asset", "get_asset", "suspend", "cancel", "download_blob", "upload_blob", "submit_input", "flush": + case "submit_execution", "select", "persist_asset", "get_asset", "suspend", "cancel", "download_blob", "upload_blob", "submit_input", "flush", "stream_register": // Dispatch async: these can block on the server (e.g. a // `select` that waits for a child execution). Blocking the // message loop here would stop us reading the adapter's @@ -368,9 +370,6 @@ loop: case "register_group": p.handleRegisterGroup(execCtx, executionID, params, logger) - case "stream_register": - p.handleStreamRegister(execCtx, executionID, params, logger) - case "stream_append": p.handleStreamAppend(execCtx, executionID, params, logger) @@ -609,22 +608,6 @@ func (p *Pool) handleRegisterGroup(ctx context.Context, executionID string, para } } -func (p *Pool) handleStreamRegister(ctx context.Context, executionID string, params json.RawMessage, logger *slog.Logger) { - var req adapter.StreamRegisterParams - if err := json.Unmarshal(params, &req); err != nil { - logger.Error("failed to parse stream_register message", "error", err) - return - } - - if err := p.handler.StreamRegister(ctx, req.ExecutionID, req.Index, req.Buffer, req.TimeoutMs); err != nil { - logger.Error("failed to register stream", "error", err) - return - } - if req.TimeoutMs != nil { - p.streamTimers.Register(streamKey{req.ExecutionID, req.Index}, *req.TimeoutMs) - } -} - func (p *Pool) handleStreamAppend(ctx context.Context, executionID string, params json.RawMessage, logger *slog.Logger) { var req adapter.StreamAppendParams if err := json.Unmarshal(params, &req); err != nil { @@ -671,8 +654,7 @@ func (p *Pool) handleStreamSubscribe(ctx context.Context, executionID string, pa ctx, req.ExecutionID, req.SubscriptionID, - req.ProducerExecutionID, - req.Index, + req.StreamID, req.FromSequence, req.Stride, req.Prefetch, @@ -746,6 +728,31 @@ func (p *Pool) handleRequest(ctx context.Context, exec *adapter.Executor, method result = out } + case "stream_register": + var req adapter.StreamRegisterParams + if err := json.Unmarshal(params, &req); err != nil { + errInfo = &adapter.ErrorInfo{Code: "parse_error", Message: err.Error()} + break + } + registered, err := p.handler.StreamRegister(ctx, req.ExecutionID, req.Position, req.Buffer, req.TimeoutMs) + if err != nil { + errInfo = &adapter.ErrorInfo{Code: "stream_register_error", Message: err.Error()} + } else { + // The idle timer is keyed by the server-allocated index, which + // is what the adapter's appends and closes carry from here on. + // It's per execution: a resumed stream's timer starts afresh + // with the resuming execution, and a suspension doesn't count + // against it. + if req.TimeoutMs != nil { + p.streamTimers.Register(streamKey{req.ExecutionID, registered.Index}, *req.TimeoutMs) + } + result = map[string]any{ + "id": registered.ID, + "index": registered.Index, + "head": registered.Head, + } + } + case "persist_asset": var req adapter.PersistAssetParams if err := json.Unmarshal(params, &req); err != nil { diff --git a/cli/internal/worker/worker.go b/cli/internal/worker/worker.go index 99393f6c..23cab7d2 100644 --- a/cli/internal/worker/worker.go +++ b/cli/internal/worker/worker.go @@ -97,14 +97,13 @@ type streamSubKey struct { // `delivered` would under-grant it, and could deadlock a consumer that // has already drained its queue and so has nothing left to ack. type streamSubscription struct { - producerExecutionID string - index int - nextSequence int - stride map[string]any - prefetch int - delivered int - ackCount int - ackSequence int + streamID string + nextSequence int + stride map[string]any + prefetch int + delivered int + ackCount int + ackSequence int } type executionState struct { @@ -1430,23 +1429,44 @@ func (w *Worker) RegisterGroup(ctx context.Context, executionID string, groupID return conn.Notify("register_group", executionID, groupID, name) } -func (w *Worker) StreamRegister(ctx context.Context, executionID string, index int, buffer *int, timeoutMs *int) error { - conn, err := w.requireConn() +// StreamRegister declares an execution's k-th stream to the server, which +// allocates the stream's index within its step and decides whether the +// registration resumes a stream paused by a suspend. A request rather +// than a notification: the producer needs the reply (id, index, head) +// before it can embed the handle in its result or append anything. +func (w *Worker) StreamRegister(ctx context.Context, executionID string, position int, buffer *int, timeoutMs *int) (*adapter.StreamRegisterResult, error) { + conn, err := w.waitForConn(ctx) if err != nil { - return err + return nil, err } // The wire protocol takes buffer and timeout_ms positionally; nil - // encodes to JSON null. Server reads [execution_id, index, buffer, - // timeout_ms?]; omitting the trailing timeout_ms keeps compat with - // older server builds that don't read it. - var bufferArg any + // encodes to JSON null. Server reads [execution_id, position, buffer, + // timeout_ms]. + var bufferArg, timeoutArg any if buffer != nil { bufferArg = *buffer } - if timeoutMs == nil { - return conn.Notify("stream_register", executionID, index, bufferArg) + if timeoutMs != nil { + timeoutArg = *timeoutMs + } + result, err := conn.Request(ctx, "stream_register", executionID, position, bufferArg, timeoutArg) + if err != nil { + return nil, err + } + m, ok := result.(map[string]any) + if !ok { + return nil, fmt.Errorf("stream_register: unexpected result type %T", result) + } + id, _ := m["id"].(string) + index, indexOk := m["index"].(float64) + if id == "" || !indexOk { + return nil, fmt.Errorf("stream_register: malformed result %v", result) + } + head := -1 + if h, ok := m["head"].(float64); ok { + head = int(h) } - return conn.Notify("stream_register", executionID, index, bufferArg, *timeoutMs) + return &adapter.StreamRegisterResult{ID: id, Index: int(index), Head: head}, nil } func (w *Worker) StreamAppend(ctx context.Context, executionID string, index int, sequence int, value *adapter.Value) error { @@ -1485,7 +1505,7 @@ func (w *Worker) StreamClose(ctx context.Context, executionID string, index int, return conn.Notify("stream_close", executionID, index, errTuple, *reason) } -func (w *Worker) StreamSubscribe(ctx context.Context, executionID string, subscriptionID int, producerExecutionID string, index int, fromSequence int, stride map[string]any, prefetch int) error { +func (w *Worker) StreamSubscribe(ctx context.Context, executionID string, subscriptionID int, streamID string, fromSequence int, stride map[string]any, prefetch int) error { conn, err := w.requireConn() if err != nil { return err @@ -1496,17 +1516,16 @@ func (w *Worker) StreamSubscribe(ctx context.Context, executionID string, subscr // immediately after this send. w.streamSubsMu.Lock() w.streamSubs[streamSubKey{executionID, subscriptionID}] = &streamSubscription{ - producerExecutionID: producerExecutionID, - index: index, - nextSequence: fromSequence, - stride: stride, - prefetch: prefetch, - ackSequence: fromSequence - 1, + streamID: streamID, + nextSequence: fromSequence, + stride: stride, + prefetch: prefetch, + ackSequence: fromSequence - 1, } w.streamSubsMu.Unlock() - // Params: [subscription_id, consumer_execution_id, producer_execution_id, index, from_sequence, stride, prefetch] - return conn.Notify("stream_subscribe", subscriptionID, executionID, producerExecutionID, index, fromSequence, stride, prefetch) + // Params: [subscription_id, consumer_execution_id, stream_id, from_sequence, stride, prefetch] + return conn.Notify("stream_subscribe", subscriptionID, executionID, streamID, fromSequence, stride, prefetch) } // StreamAck forwards a consumer's cumulative progress to the server, @@ -2153,8 +2172,7 @@ func (w *Worker) resubscribeStreams(known map[string]struct{}) { "acked_seq": r.sub.ackSequence, } err := conn.Notify("stream_subscribe", r.key.subscriptionID, r.key.executionID, - r.sub.producerExecutionID, r.sub.index, r.sub.nextSequence, r.sub.stride, - r.sub.prefetch, progress) + r.sub.streamID, r.sub.nextSequence, r.sub.stride, r.sub.prefetch, progress) if err != nil { // Connection dropped again — the next reconnect retries. w.logger.Debug("stream re-subscribe failed", "execution_id", r.key.executionID, diff --git a/docs/docs/python_reference.md b/docs/docs/python_reference.md index 4b0471f2..559e06dc 100644 --- a/docs/docs/python_reference.md +++ b/docs/docs/python_reference.md @@ -284,6 +284,44 @@ cf.Retries( | `backoff` | `tuple` | `(1, 60)` | Backoff range (min, max) in seconds | | `when` | type, tuple, callable, or `None` | `None` | Exception filter (`None` = retry on any error) | +### `Streams` + +Default stream configuration for a target, passed as `streams=` to `@task` / `@workflow`. Only applies to targets that produce [streams](./streams.md). + +```python +cf.Streams( + buffer: int | None = 0, + timeout: float | timedelta | None = None, +) +``` + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `buffer` | `int \| None` | `0` | How many items the producer may run ahead of its slowest consumer (`0` = lockstep, `None` = no backpressure) | +| `timeout` | `float \| timedelta \| None` | `None` | Idle timeout: the stream is closed if no item is appended within this window | + +## Streams + +Ordered sequences of values produced by one execution and consumed by others as they grow. See [streams](./streams.md). + +### `stream(generator, *, buffer=..., timeout=...)` + +Registers a generator as a stream and returns a `Stream` handle to embed in a return value or pass to another task. A task whose body is itself a generator is registered automatically, with its result being the handle. Unspecified options inherit from the target's `Streams` configuration. Must be called inside a task or workflow body. + +### `Stream` + +A handle to a stream, typed by its items (`Stream[T]`). + +| Method | Description | +|--------|-------------| +| `for item in stream` / `async for` | Iterate from the first item, blocking until each arrives, ending when the stream closes | +| `stream.slice(start, stop=None)` | A view of positions `[start, stop)` | +| `stream.partition(n, i)` | A view of every `n`-th item starting at `i`, for parallel consumers | +| `stream.stride(start=0, stop=None, step=1)` | The general form of the above | +| `stream.id` | An opaque identifier for the stream, as shown in Studio | + +Views compose, and can be passed to other tasks. If the producer raised, iterating raises the same error. + ## Checkpoints State that survives across executions of a step — retries, suspensions, recurrences and re-runs. See [checkpoints](./checkpoints.md). diff --git a/docs/docs/recurring.md b/docs/docs/recurring.md index b80a68c0..9a732117 100644 --- a/docs/docs/recurring.md +++ b/docs/docs/recurring.md @@ -14,6 +14,8 @@ The task recurs as long as it returns `None`. Returning any other value complete Each iteration is a fresh execution starting from the top, so anything that needs to carry forward between them — a cursor, say — belongs in a [checkpoint](./checkpoints.md). +A recurrent task can't have a generator body: each iteration would produce a separate [stream](./streams.md), and the result would never be `None`. For a stream that continues across pauses, suspend from inside the generator instead. + ## Delay By default, recurring tasks restart immediately. Use `delay` to wait between executions: diff --git a/docs/docs/streams.md b/docs/docs/streams.md new file mode 100644 index 00000000..1108d3d6 --- /dev/null +++ b/docs/docs/streams.md @@ -0,0 +1,122 @@ +# Streams + +A stream is an ordered sequence of values that one task produces and others consume as it grows. Where a task result is a single value handed over at the end, a stream lets a consumer start on the first item while the producer is still working on the last. + +## Producing + +The simplest producer is a task whose body is a generator. Its result *is* the stream: + +```python +import coflux as cf + + +@cf.task() +def fetch_pages(url: str): + for page in paginate(url): + yield page +``` + +Calling `fetch_pages(url)` returns a `cf.Stream` handle, and `fetch_pages.submit(url)` returns an execution whose result resolves to one. `async def` generators work the same way. + +To return several streams, or a stream alongside other values, register a generator explicitly with `cf.stream()` and put the handle wherever you like in the return value: + +```python +@cf.task() +def split(url: str): + pages = paginate(url) + return { + "headers": cf.stream(h for h in headers_of(pages)), + "bodies": cf.stream(b for b in bodies_of(pages)), + } +``` + +`cf.stream()` accepts only generators. A list is already a value; return it as one. + +Each item is delivered to consumers as soon as it's produced. The producing execution stays running until every stream it produced has finished, even after it has returned its result. + +## Consuming + +A `cf.Stream` is iterable. Iterating blocks until the next item arrives, and ends when the stream closes: + +```python +@cf.task() +def index(pages: cf.Stream[dict]): + for page in pages: + add_to_index(page) +``` + +Use `async for` inside `async def` bodies. Each iteration of a handle starts from the beginning, so a stream can be read more than once, by more than one consumer, and by a consumer that starts long after the producer finished. + +`stream.slice(start, stop)` restricts iteration to a range of positions, and `stream.partition(n, i)` delivers every `n`-th item starting at `i`, which spreads a stream across parallel consumers. Views compose, and a view can be passed to another task like any handle. + +If the producer raises, iterating raises the same error. If the producer is cancelled or lost, iterating raises the corresponding `ExecutionTerminated` subclass. + +## Backpressure + +By default a producer runs in lockstep with its slowest consumer: it emits one item, waits for a consumer to finish with it, then emits the next. Configure how far ahead it may run with `streams=`: + +```python +@cf.task(streams=cf.Streams(buffer=100)) +def fetch_pages(url: str): ... +``` + +`buffer=N` lets the producer run up to `N` items ahead of the slowest consumer. `buffer=None` disables backpressure. Items are never lost to a slow consumer: they're stored as they're produced, and a consumer that falls behind catches up at its own pace. + +Per-call overrides are available on `cf.stream(generator, buffer=...)`, and `target.with_streams(...)` overrides the configuration for one submission. + +Note that with the default lockstep buffer, a producer whose stream never gains a consumer waits indefinitely, holding its worker slot. Set a `timeout` (below) to bound that. + +## Suspending + +A stream belongs to the *step* that produces it, not to any one execution. The difference shows when a producer [suspends](./suspense.md): its streams pause rather than close, and the execution that resumes the step continues them. Consumers see one unbroken stream. + +This is how to write a producer that keeps going indefinitely without holding a worker slot while it waits: + +```python +cursor = cf.Checkpoint("cursor", default=0) + + +@cf.task() +def tail_events(): + since = cursor.get() + for event in fetch_events(since): + yield event + since = event.id + cursor.set(since) + cf.suspend(60) +``` + +Each resume runs the body from the top, reads the [checkpoint](./checkpoints.md), and carries on the same stream. Consumers just wait through the pause. + +Streams are matched up by the order they're registered in, so code that registers several streams before suspending must register them in the same order when it resumes. That's the same determinism suspend already requires of the code before the suspend point, and it's automatic for a task whose body is a generator. + +Studio shows a stream under its step, listing every attempt that produced into it, and labels each item with the attempt that produced it. + +## When a stream ends + +Only a suspend keeps a stream open across executions. Every other way an execution can end closes the step's open streams: + +- The generator finishing closes its stream normally, and a task completing normally closes anything it left open. +- An exception in the generator closes the stream with that error. A [retry](./retries.md) opens a fresh stream. +- Cancellation, a crash, or a lost worker closes it with that reason. Cancelling a suspended step, which cancels its pending resumption, closes its paused streams too. +- A [recurrent](./recurring.md) task finishing an iteration closes its streams, and the next iteration opens its own. A recurrent task can't have a generator body for this reason: its result would be a stream, so it could never return `None` to recur. Use the suspend form above for a continuous stream. +- Re-running a step from Studio cancels the running attempt, closing its streams, and the new attempt starts fresh. Re-running a *suspended* step instead continues its paused streams, since nothing was producing into them. + +Once closed, a stream is never reopened. A later attempt of the step produces a new stream. + +## Timeouts + +A stream can be given an idle timeout: + +```python +@cf.task(streams=cf.Streams(timeout=30)) +def fetch_pages(url: str): ... +``` + +If the producer doesn't append an item within that window, the stream is closed as timed out and the generator is stopped. The window counts time spent waiting for consumer demand too, so with a lockstep buffer a slow or absent consumer can time out the producer. The producing execution still completes with its value, but is excluded from [caching](./caching.md) and [memoizing](./memoizing.md). + +The timeout is measured from the last item appended, per execution. A suspended step is not idle, so the pause between a suspend and its resumption doesn't count; the timer starts again when the resumed execution registers. There is no consumer-side timeout: iterating a stream waits as long as it takes. + +## Workspaces + +A stream is only written from the [workspace](./concepts.md) it was produced in. Re-running a producer in a derived workspace produces a new stream there rather than continuing the base's, while consumers in a derived workspace can still read a base workspace's stream, the same way they read its results. diff --git a/docs/docs/suspense.md b/docs/docs/suspense.md index 6f198126..3b923f5a 100644 --- a/docs/docs/suspense.md +++ b/docs/docs/suspense.md @@ -6,6 +6,8 @@ The suspense can be either _explicit_ or _implicit_. In either case, it's import State that needs to survive the suspension, but isn't the result of a task, can be kept in a [checkpoint](/checkpoints). +A task that produces a [stream](/streams) can suspend from inside its generator. The stream pauses rather than closing, and the resumed execution continues it, so consumers see one unbroken stream. + Suspense is useful as a way of freeing up resources used by a waiting execution. ## Explicit suspense diff --git a/docs/sidebars.ts b/docs/sidebars.ts index 6c65a353..11377fa6 100644 --- a/docs/sidebars.ts +++ b/docs/sidebars.ts @@ -39,6 +39,7 @@ const sidebars: SidebarsConfig = { "deferring", "suspense", "checkpoints", + "streams", "select", ], }, diff --git a/server/lib/coflux/handlers/worker.ex b/server/lib/coflux/handlers/worker.ex index 8550f395..c1b24817 100644 --- a/server/lib/coflux/handlers/worker.ex +++ b/server/lib/coflux/handlers/worker.ex @@ -293,8 +293,14 @@ defmodule Coflux.Handlers.Worker do {[{:close, 4000, "execution_invalid"}], nil} end + # A request rather than a notification: the server allocates the + # stream's index per step (so its id is unique across attempts and + # workspaces) and decides whether the registration resumes a stream + # paused by a suspend. The producer needs the reply — id, index and + # the head to sequence from — before it can embed the handle in its + # result or append anything. "stream_register" -> - [execution_id, index | rest] = message["params"] + [execution_id, position | rest] = message["params"] buffer = Enum.at(rest, 0) timeout_ms = Enum.at(rest, 1) @@ -302,15 +308,22 @@ defmodule Coflux.Handlers.Worker do case Orchestration.register_stream( state.project_id, execution_id, - index, + position, buffer, timeout_ms, state.session_id ) do - :ok -> {[], state} - # Idempotent — a duplicate register is harmless. - {:error, :already_registered} -> {[], state} - {:error, :not_found} -> {[{:close, 4000, "execution_invalid"}], nil} + {:ok, %{id: id, index: index, head: head}} -> + {[success_message(message["id"], %{"id" => id, "index" => index, "head" => head})], + state} + + # The execution has already been finalised (e.g. it suspended + # and a still-alive driver thread is registering late). + {:error, :completed} -> + {[error_message(message["id"], "execution_completed")], state} + + {:error, :not_found} -> + {[{:close, 4000, "execution_invalid"}], nil} end else {[{:close, 4000, "execution_invalid"}], nil} @@ -344,6 +357,15 @@ defmodule Coflux.Handlers.Worker do {:error, :already_appended} -> {[], state} + # A finalised execution (typically one that suspended, whose + # driver thread is still running down) must not interleave with + # the execution that now owns the stream. + {:error, :completed} -> + {[], state} + + {:error, :not_producer} -> + {[], state} + {:error, :not_found} -> {[{:close, 4000, "execution_invalid"}], nil} end @@ -377,6 +399,7 @@ defmodule Coflux.Handlers.Worker do :ok -> {[], state} {:error, :already_closed} -> {[], state} {:error, :not_registered} -> {[], state} + {:error, :completed} -> {[], state} {:error, :not_found} -> {[{:close, 4000, "execution_invalid"}], nil} end else @@ -387,8 +410,7 @@ defmodule Coflux.Handlers.Worker do [ subscription_id, consumer_execution_id, - producer_execution_id, - index, + stream_id, from_sequence, stride | rest ] = message["params"] @@ -429,8 +451,7 @@ defmodule Coflux.Handlers.Worker do state.session_id, subscription_id, consumer_execution_id, - producer_execution_id, - index, + stream_id, from_sequence, stride, prefetch, @@ -439,17 +460,16 @@ defmodule Coflux.Handlers.Worker do :ok -> {[], state} - # If the stream doesn't exist yet (or producer vanished), push an - # immediate close so the consumer doesn't wait forever. Carry - # the reason atom verbatim — consumers decide how to surface - # "not_found" in their own idiom. - {:error, reason} - when reason in [:stream_not_found, :producer_not_found] -> + # If the stream doesn't exist, push an immediate close so the + # consumer doesn't wait forever. Carry the reason atom verbatim + # — consumers decide how to surface "not_found" in their own + # idiom. + {:error, :stream_not_found} -> {[ command_message("stream_closed", [ consumer_execution_id, subscription_id, - Atom.to_string(reason), + "stream_not_found", nil ]) ], state} @@ -789,14 +809,15 @@ defmodule Coflux.Handlers.Worker do state ) do # `reason` is a string ("complete" / "errored" / "cancelled" / - # "abandoned" / "crashed" / "timeout" / "suspended" / "recurred"); - # `error` is non-nil only for "errored" — the producer's actual - # `{type, message, frames}`. Other reasons travel as the string - # alone; the consumer adapter decides how to represent them. + # "abandoned" / "crashed" / "timeout" / "recurred"); `error` is + # non-nil only for "errored" — the producer's actual `{type, message, + # frames}`. Other reasons travel as the string alone; the consumer + # adapter decides how to represent them. # - # "suspended" / "recurred" mean the producer was superseded rather - # than failing: the successor owns a new stream, so this one is - # terminal even though nothing went wrong. + # "recurred" means the producer finished an iteration rather than + # failing: the next iteration opens its own streams, so this one is + # terminal even though nothing went wrong. A suspend never closes a + # stream — the resuming execution continues it. {[ command_message("stream_closed", [ execution_external_id, diff --git a/server/lib/coflux/orchestration.ex b/server/lib/coflux/orchestration.ex index 6ad74d22..5a98b910 100644 --- a/server/lib/coflux/orchestration.ex +++ b/server/lib/coflux/orchestration.ex @@ -189,14 +189,20 @@ defmodule Coflux.Orchestration do end # Stream producer messages — worker registers a stream, appends items, - # and closes the stream. `index` identifies the stream within its - # producer execution; `sequence` identifies an item within the stream. - # Both are worker-assigned and monotonic from 0. - - def register_stream(project_id, execution_id, index, buffer, timeout_ms, session_id) do + # and closes the stream. + # + # `register_stream` is a request: `position` is the order in which the + # execution registered the stream (its k-th), and the reply carries the + # stream's id, its server-allocated step `index`, and the `head` to + # sequence from (`-1` for a new stream; the last sequence of a paused + # stream the execution is resuming after a suspend). Appends and closes + # then address the stream by that `index`; `sequence` identifies an item + # within the stream, worker-assigned and monotonic. + + def register_stream(project_id, execution_id, position, buffer, timeout_ms, session_id) do call_server( project_id, - {:register_stream, execution_id, index, buffer, timeout_ms, session_id} + {:register_stream, execution_id, position, buffer, timeout_ms, session_id} ) end @@ -209,16 +215,15 @@ defmodule Coflux.Orchestration do end # Stream consumer messages — consumer opens a subscription to receive - # items from a producer's stream; server pushes stream_items / - # stream_closed commands to the consumer's session. + # items from a stream (by its id, `:_`); server pushes + # stream_items / stream_closed commands to the consumer's session. def subscribe_stream( project_id, session_id, subscription_id, consumer_execution_id, - producer_execution_id, - index, + stream_id, from_sequence, filter, prefetch, @@ -226,8 +231,8 @@ defmodule Coflux.Orchestration do ) do call_server( project_id, - {:subscribe_stream, session_id, subscription_id, consumer_execution_id, - producer_execution_id, index, from_sequence, filter, prefetch, progress} + {:subscribe_stream, session_id, subscription_id, consumer_execution_id, stream_id, + from_sequence, filter, prefetch, progress} ) end @@ -304,11 +309,8 @@ defmodule Coflux.Orchestration do call_server(project_id, {:subscribe_run, run_id, pid}) end - def subscribe_stream_topic(project_id, execution_external_id, index, pid) do - call_server( - project_id, - {:subscribe_stream_topic, execution_external_id, index, pid} - ) + def subscribe_stream_topic(project_id, stream_id, pid) do + call_server(project_id, {:subscribe_stream_topic, stream_id, pid}) end def subscribe_targets(project_id, workspace_id, pid) do diff --git a/server/lib/coflux/orchestration/epoch.ex b/server/lib/coflux/orchestration/epoch.ex index a9092186..3703bfe1 100644 --- a/server/lib/coflux/orchestration/epoch.ex +++ b/server/lib/coflux/orchestration/epoch.ex @@ -349,66 +349,88 @@ defmodule Coflux.Orchestration.Epoch do end end) - # Copy streams, their items, and any closure rows. An execution's - # streams may be mid-production (items appended, no closure) — - # carry them forward so consumers can keep reading after rotation. - Enum.each(execution_ids, fn {old_exec_id, new_exec_id} -> + # Copy streams (per step) with their registrations, items and any + # closure rows. A stream may be mid-production, or paused by a + # suspend (items appended, no closure) — carry it forward so + # consumers can keep reading, and a resuming execution can keep + # appending, after rotation. + Enum.each(step_ids, fn {old_step_id, new_step_id} -> {:ok, streams} = query( source_db, - "SELECT `index`, buffer, timeout_ms, created_at FROM streams WHERE execution_id = ?1", - {old_exec_id} + "SELECT id, workspace_id, `index`, position, created_at FROM streams WHERE step_id = ?1", + {old_step_id} ) - Enum.each(streams, fn {index, buffer, timeout_ms, stream_created_at} -> - {:ok, _} = + Enum.each(streams, fn {old_stream_id, workspace_id, index, position, + stream_created_at} -> + {:ok, new_stream_id} = insert_one(target_db, :streams, %{ - execution_id: new_exec_id, + step_id: new_step_id, + workspace_id: remap_workspace_id(source_db, target_db, workspace_id), index: index, - buffer: buffer, - timeout_ms: timeout_ms, + position: position, created_at: stream_created_at }) + {:ok, registrations} = + query( + source_db, + """ + SELECT execution_id, buffer, timeout_ms, created_at + FROM stream_registrations + WHERE stream_id = ?1 + """, + {old_stream_id} + ) + + Enum.each(registrations, fn {old_exec_id, buffer, timeout_ms, + registration_created_at} -> + {:ok, _} = + insert_one(target_db, :stream_registrations, %{ + stream_id: new_stream_id, + execution_id: Map.fetch!(execution_ids, old_exec_id), + buffer: buffer, + timeout_ms: timeout_ms, + created_at: registration_created_at + }) + end) + {:ok, items} = query( source_db, """ - SELECT sequence, value_id, created_at + SELECT sequence, value_id, execution_id, created_at FROM stream_items - WHERE execution_id = ?1 AND `index` = ?2 + WHERE stream_id = ?1 """, - {old_exec_id, index} + {old_stream_id} ) - Enum.each(items, fn {sequence, value_id, item_created_at} -> - new_value_id = ensure_value(source_db, target_db, value_id) - + Enum.each(items, fn {sequence, value_id, old_exec_id, item_created_at} -> {:ok, _} = insert_one(target_db, :stream_items, %{ - execution_id: new_exec_id, - index: index, + stream_id: new_stream_id, sequence: sequence, - value_id: new_value_id, + value_id: ensure_value(source_db, target_db, value_id), + execution_id: Map.fetch!(execution_ids, old_exec_id), created_at: item_created_at }) end) case query_one( source_db, - "SELECT reason, error_id, created_at FROM stream_closures WHERE execution_id = ?1 AND `index` = ?2", - {old_exec_id, index} + "SELECT reason, error_id, execution_id, created_at FROM stream_closures WHERE stream_id = ?1", + {old_stream_id} ) do - {:ok, {reason, error_id, closure_created_at}} -> - new_error_id = - if error_id, do: ensure_error(source_db, target_db, error_id) - + {:ok, {reason, error_id, old_exec_id, closure_created_at}} -> {:ok, _} = insert_one(target_db, :stream_closures, %{ - execution_id: new_exec_id, - index: index, + stream_id: new_stream_id, reason: reason, - error_id: new_error_id, + error_id: + if(error_id, do: ensure_error(source_db, target_db, error_id)), + execution_id: Map.fetch!(execution_ids, old_exec_id), created_at: closure_created_at }) @@ -605,28 +627,44 @@ defmodule Coflux.Orchestration.Epoch do {:ok, stream_deps} = query( source_db, - "SELECT stream_ref_id, stream_index, created_at FROM stream_dependencies WHERE execution_id = ?1", + "SELECT stream_ref_id, created_at FROM stream_dependencies WHERE execution_id = ?1", {old_exec_id} ) - Enum.each(stream_deps, fn {old_ref_id, stream_index, created_at} -> - new_ref_id = ensure_execution_ref(source_db, target_db, old_ref_id) + Enum.each(stream_deps, fn {old_ref_id, created_at} -> + new_ref_id = ensure_stream_ref(source_db, target_db, old_ref_id) {:ok, _} = insert_one( target_db, :stream_dependencies, - %{ - execution_id: new_exec_id, - stream_ref_id: new_ref_id, - stream_index: stream_index, - created_at: created_at - }, + %{execution_id: new_exec_id, stream_ref_id: new_ref_id, created_at: created_at}, on_conflict: "DO NOTHING" ) end) end + defp ensure_stream_ref(source_db, target_db, old_ref_id) do + {:ok, {run_ext_id, step_number, index, module, target}} = + query_one!( + source_db, + "SELECT run_external_id, step_number, `index`, module, target FROM stream_refs WHERE id = ?1", + {old_ref_id} + ) + + {:ok, ref_id} = + Coflux.Orchestration.Streams.get_or_create_stream_ref( + target_db, + run_ext_id, + step_number, + index, + module, + target + ) + + ref_id + end + defp copy_execution_inputs(source_db, target_db, old_exec_id, new_exec_id, new_run_id) do # Copy execution_inputs records (which executions submitted which inputs) {:ok, exec_inputs} = diff --git a/server/lib/coflux/orchestration/runs.ex b/server/lib/coflux/orchestration/runs.ex index 102315c5..393590f2 100644 --- a/server/lib/coflux/orchestration/runs.ex +++ b/server/lib/coflux/orchestration/runs.ex @@ -640,22 +640,6 @@ defmodule Coflux.Orchestration.Runs do end) end - def record_stream_dependency(db, execution_id, stream_ref_id, stream_index) do - with_transaction(db, fn -> - insert_one( - db, - :stream_dependencies, - %{ - execution_id: execution_id, - stream_ref_id: stream_ref_id, - stream_index: stream_index, - created_at: current_timestamp() - }, - on_conflict: "DO NOTHING" - ) - end) - end - def get_unassigned_executions(db) do query( db, @@ -1030,28 +1014,6 @@ defmodule Coflux.Orchestration.Runs do end end - def get_run_stream_dependencies(db, run_id) do - case query( - db, - """ - SELECT d.execution_id, d.stream_ref_id, d.stream_index - FROM stream_dependencies AS d - INNER JOIN executions AS e ON e.id = d.execution_id - INNER JOIN steps AS s ON s.id = e.step_id - WHERE s.run_id = ?1 - """, - {run_id} - ) do - {:ok, rows} -> - {:ok, - Enum.group_by( - rows, - fn {execution_id, _ref_id, _index} -> execution_id end, - fn {_execution_id, ref_id, index} -> {ref_id, index} end - )} - end - end - def get_step_assignments(db, step_id) do case query( db, diff --git a/server/lib/coflux/orchestration/server.ex b/server/lib/coflux/orchestration/server.ex index fdec9545..3ab20e7f 100644 --- a/server/lib/coflux/orchestration/server.ex +++ b/server/lib/coflux/orchestration/server.ex @@ -140,24 +140,32 @@ defmodule Coflux.Orchestration.Server do # them from the DB. # # stream_subscriptions: {consumer_execution_id, subscription_id} -> - # %{consumer_execution_external_id, producer_execution_id, - # index, cursor, stride, prefetch, delivered, acked_count, - # acked_seq, pending_close} - # stream_subscribers: {producer_execution_id, index} -> MapSet of + # %{consumer_execution_external_id, stream_id, cursor, + # stride, prefetch, delivered, acked_count, acked_seq, + # pending_close} + # stream_subscribers: stream_id -> MapSet of # {consumer_execution_id, subscription_id} stream_subscriptions: %{}, stream_subscribers: %{}, # Per-stream producer state for backpressure. Only present # when the producer opted in by registering with a non-nil - # buffer. Keyed by {producer_execution_id, index}. + # buffer. Keyed by stream_id (the `streams` row). # - # %{buffer, demand_granted, session_id, execution_external_id} + # %{buffer, demand_granted, session_id, execution_external_id, + # index} # # * buffer — configured backpressure budget - # * demand_granted — cumulative credits sent so far + # * demand_granted — cumulative credits sent so far, in + # sequence space (so a producer + # resuming a paused stream starts at + # head + 1) # * session_id — where to route stream_demand - # * execution_external_id — external id for the command wire + # * execution_external_id — the current producer, for the wire + # * index — the stream's step index, for the wire + # + # A stream that spans a suspension changes producer: the + # resuming execution's registration replaces this entry. # # The watermark the budget is measured against is the # *slowest* subscriber's acknowledged position, recomputed @@ -1588,7 +1596,11 @@ defmodule Coflux.Orchestration.Server do if base_workspace_id == workspace_id || is_workspace_ancestor?(state, base_workspace_id, workspace_id) do - state = cancel_active_step_executions(state, step.id, workspace_id) + # A live attempt is cancelled and its streams closed, so the new + # attempt starts fresh. A pending successor of a suspended attempt + # never produced anything, so its cancellation leaves the paused + # streams open and the new attempt continues them. + state = cancel_active_step_executions(state, step.id, workspace_id, streams: :registered) {:ok, _execution_id, attempt, state} = rerun_step(state, step, workspace_id, created_by: access[:principal_id]) @@ -1933,42 +1945,69 @@ defmodule Coflux.Orchestration.Server do end end + # A producer declares its k-th stream (`position`). The server decides + # whether that continues a paused stream of the step — one left open by a + # suspended execution — or opens a new one, and replies with the stream's + # id, its step index and the head it should sequence from. def handle_call( - {:register_stream, execution_external_id, index, buffer, timeout_ms, session_external_id}, + {:register_stream, execution_external_id, position, buffer, timeout_ms, + session_external_id}, _from, state ) do - case Map.fetch(state.execution_ids, execution_external_id) do - {:ok, execution_id} -> - case Streams.register_stream(state.db, execution_id, index, buffer, timeout_ms) do - {:ok, created_at} -> - # Resolve the session's external id to the internal one — - # send_session (which delivers stream_demand) indexes by the - # internal id. - internal_session_id = - Map.get(state.session_ids, session_external_id) + with {:ok, execution_id} <- + Map.fetch(state.execution_ids, execution_external_id) |> ok_or(:not_found), + {:ok, false} <- Results.has_completion?(state.db, execution_id) do + {:ok, {step_id, workspace_id, _attempt}} = + Runs.get_execution_location(state.db, execution_id) + + {:ok, registration} = + Streams.register( + state.db, + step_id, + workspace_id, + execution_id, + position, + buffer, + timeout_ms + ) - state = - state - |> maybe_init_stream_producer( - execution_id, - execution_external_id, - index, - buffer, - internal_session_id - ) - |> notify_stream_opened(execution_id, index, buffer, timeout_ms, created_at) - |> maybe_send_initial_demand(execution_id, index) - |> flush_notifications() + {:ok, stream} = Streams.get_stream(state.db, registration.id) - {:reply, :ok, state} + external_id = + stream_external_id(stream.run_external_id, stream.step_number, stream.index) - {:error, :already_registered} -> - {:reply, {:error, :already_registered}, state} + state = + if registration.created_at do + # Resolve the session's external id to the internal one — + # send_session (which delivers stream_demand) indexes by the + # internal id. + internal_session_id = Map.get(state.session_ids, session_external_id) + + state + |> init_stream_producer( + stream, + execution_external_id, + buffer, + registration.head, + internal_session_id + ) + |> notify_stream_registered(stream, execution_id, registration, buffer, timeout_ms) + # Lockstep (buffer=0) stays paused until a consumer attaches; a + # larger buffer lets the producer pre-warm. A resuming producer + # picks up whatever demand its subscribers have already built. + |> refresh_stream_demand(stream.id) + |> flush_notifications() + else + # The same execution registering the same position again — + # idempotent, nothing new to announce. + state end - :error -> - {:reply, {:error, :not_found}, state} + {:reply, {:ok, %{id: external_id, index: stream.index, head: registration.head}}, state} + else + {:ok, true} -> {:reply, {:error, :completed}, state} + {:error, reason} -> {:reply, {:error, reason}, state} end end @@ -1977,112 +2016,97 @@ defmodule Coflux.Orchestration.Server do _from, state ) do - case Map.fetch(state.execution_ids, execution_external_id) do - {:ok, execution_id} -> - case Streams.append_item( - state.db, - execution_id, - index, - sequence, - normalize_value(value) - ) do - {:ok, created_at} -> - # If we came out of a server restart with no in-memory - # producer state for this stream, rebuild it now from the - # persisted buffer so subsequent consumer advances can - # refresh demand. The appending session is the producer. - producer_session_id = - case find_session_for_execution(state, execution_external_id) do - {:ok, sid} -> sid - :error -> nil - end - - state = - state - |> ensure_stream_producer( - execution_id, - execution_external_id, - index, - producer_session_id - ) - |> push_stream_item(execution_id, index, sequence, value) - |> notify_stream_item_appended( - execution_id, - index, - sequence, - value, - created_at - ) - |> flush_notifications() - - {:reply, :ok, state} - - {:error, reason} -> - {:reply, {:error, reason}, state} + # Appends are refused once the execution has a completion: after a + # suspend the resuming execution owns the stream, and a still-alive + # predecessor must not interleave with it. + with {:ok, execution_id} <- + Map.fetch(state.execution_ids, execution_external_id) |> ok_or(:not_found), + {:ok, stream_id} <- resolve_step_stream(state.db, execution_id, index), + {:ok, false} <- Results.has_completion?(state.db, execution_id), + {:ok, created_at} <- + Streams.append_item( + state.db, + stream_id, + execution_id, + sequence, + normalize_value(value) + ) do + # If we came out of a server restart with no in-memory producer + # state for this stream, rebuild it now from the persisted config so + # subsequent consumer advances can refresh demand. The appending + # session is the producer. + producer_session_id = + case find_session_for_execution(state, execution_external_id) do + {:ok, sid} -> sid + :error -> nil end - :error -> - {:reply, {:error, :not_found}, state} + state = + state + |> ensure_stream_producer(stream_id, producer_session_id) + |> push_stream_item(stream_id, sequence, value) + |> notify_stream_item_appended(stream_id, execution_id, sequence, value, created_at) + |> flush_notifications() + + {:reply, :ok, state} + else + {:ok, true} -> {:reply, {:error, :completed}, state} + {:error, reason} -> {:reply, {:error, reason}, state} end end def handle_call({:close_stream, execution_external_id, index, close_spec}, _from, state) do - case Map.fetch(state.execution_ids, execution_external_id) do - {:ok, execution_id} -> - {spec, reason, error} = - case close_spec do - nil -> - {:complete, :complete, nil} - - :timeout -> - {:timeout, :timeout, nil} + with {:ok, execution_id} <- + Map.fetch(state.execution_ids, execution_external_id) |> ok_or(:not_found), + {:ok, stream_id} <- resolve_step_stream(state.db, execution_id, index), + {:ok, false} <- Results.has_completion?(state.db, execution_id) do + {spec, reason, error} = + case close_spec do + nil -> + {:complete, :complete, nil} - {type, message, frames} -> - {{:errored, type, message, frames}, :errored, {type, message, frames}} - end + :timeout -> + {:timeout, :timeout, nil} - case Streams.close_stream(state.db, execution_id, index, spec) do - {:ok, closed_at} -> - state = - state - |> push_stream_closed(execution_id, index, reason, error) - |> notify_stream_closed(execution_id, index, reason, error, closed_at) - |> drop_stream_producer({execution_id, index}) - |> flush_notifications() + {type, message, frames} -> + {{:errored, type, message, frames}, :errored, {type, message, frames}} + end - {:reply, :ok, state} + case Streams.close_stream(state.db, stream_id, execution_id, spec) do + {:ok, closed_at} -> + state = + state + |> push_stream_closed(stream_id, reason, error) + |> notify_stream_closed(stream_id, execution_id, reason, error, closed_at) + |> drop_stream_producer(stream_id) + |> flush_notifications() - {:error, reason} -> - {:reply, {:error, reason}, state} - end + {:reply, :ok, state} - :error -> - {:reply, {:error, :not_found}, state} + {:error, reason} -> + {:reply, {:error, reason}, state} + end + else + {:ok, true} -> {:reply, {:error, :completed}, state} + {:error, reason} -> {:reply, {:error, reason}, state} end end def handle_call( {:subscribe_stream, session_external_id, subscription_id, consumer_execution_external_id, - producer_execution_external_id, index, from_sequence, stride, prefetch, progress}, + stream_external_id, from_sequence, stride, prefetch, progress}, _from, state ) do - # Producer may already have terminated — resolve from DB (active epoch) - # rather than the in-memory active-execution cache. - producer_result = - case resolve_internal_execution_id(state, producer_execution_external_id) do - {:ok, id} -> {:ok, id} - {:error, :not_found} -> {:error, :producer_not_found} - end - with {:ok, _session_id} <- Map.fetch(state.session_ids, session_external_id) |> ok_or(:session_not_found), {:ok, consumer_execution_id} <- Map.fetch(state.execution_ids, consumer_execution_external_id) |> ok_or(:consumer_not_found), - {:ok, producer_execution_id} <- producer_result, - {:ok, true} <- Streams.exists?(state.db, producer_execution_id, index), + # The stream's run may have been rotated into an older epoch — + # resolving by id copies it forward if so. + {:ok, stream_id} <- resolve_stream_id(state, stream_external_id), key = {consumer_execution_id, subscription_id}, false <- Map.has_key?(state.stream_subscriptions, key) do # `progress` is nil for a fresh subscribe and carries the CLI's @@ -2110,8 +2134,7 @@ defmodule Coflux.Orchestration.Server do subscription = %{ consumer_execution_external_id: consumer_execution_external_id, - producer_execution_id: producer_execution_id, - index: index, + stream_id: stream_id, cursor: from_sequence, stride: stride, prefetch: prefetch, @@ -2125,38 +2148,18 @@ defmodule Coflux.Orchestration.Server do state |> Map.update!(:stream_subscriptions, &Map.put(&1, key, subscription)) |> Map.update!(:stream_subscribers, fn m -> - Map.update( - m, - {producer_execution_id, index}, - MapSet.new([key]), - &MapSet.put(&1, key) - ) + Map.update(m, stream_id, MapSet.new([key]), &MapSet.put(&1, key)) end) # Post-restart recovery: producer state may be missing. The - # producer's session isn't necessarily the one the subscribe - # came from — look it up across sessions by external execution - # id. - producer_session_id = - case find_session_for_execution(state, producer_execution_external_id) do - {:ok, sid} -> sid - :error -> nil - end - - state = - ensure_stream_producer( - state, - producer_execution_id, - producer_execution_external_id, - index, - producer_session_id - ) + # producer's session isn't necessarily the one the subscribe came + # from — look it up from the stream's current producer. + state = ensure_stream_producer(state, stream_id, producer_session_id(state, stream_id)) # First subscriber (or a later one whose cursor exceeds the prior # max) may unblock the producer — recompute demand before pushing # backlog so any delivered items keep the credit maths honest. - state = - refresh_stream_demand(state, {producer_execution_id, index}) + state = refresh_stream_demand(state, stream_id) # If the stream has already closed, record that as pending first so # the pump can emit it — but only once the backlog it's allowed to @@ -2165,29 +2168,27 @@ defmodule Coflux.Orchestration.Server do state = mark_closed_if_closed(state, key) state = pump_subscription(state, key) - # Record the subscribe as a lineage edge (consumer -> producer stream). - # Done unconditionally on subscribe, independent of whether items end - # up being read. Uses execution_refs so the edge survives epoch - # rotation. - {:ok, stream_ref_id} = - Runs.create_execution_ref_for(state.db, producer_execution_id) + # Record the subscribe as a lineage edge (consumer -> stream). Done + # unconditionally on subscribe, independent of whether items end up + # being read. Uses stream_refs so the edge survives epoch rotation. + {:ok, stream_ref_id} = Streams.create_stream_ref_for(state.db, stream_id) {:ok, inserted_id} = - Runs.record_stream_dependency(state.db, consumer_execution_id, stream_ref_id, index) + Streams.record_dependency(state.db, consumer_execution_id, stream_ref_id) state = if inserted_id do {:ok, {run_external_id}} = Runs.get_external_run_id_for_execution(state.db, consumer_execution_id) - {producer_ext_id, _module, _target} = - producer_metadata = resolve_execution_ref(state.db, stream_ref_id) + {:ok, {stream_run_ext_id, step_number, index, module, target}} = + Streams.get_stream_ref(state.db, stream_ref_id) notify_listeners( state, {:run, run_external_id}, - {:stream_dependency, consumer_execution_external_id, producer_ext_id, index, - producer_metadata} + {:stream_dependency, consumer_execution_external_id, + stream_external_id(stream_run_ext_id, step_number, index), module, target} ) else state @@ -2197,8 +2198,8 @@ defmodule Coflux.Orchestration.Server do {:reply, :ok, state} else - {:ok, false} -> {:reply, {:error, :stream_not_found}, state} true -> {:reply, {:error, :already_subscribed}, state} + {:error, :not_found} -> {:reply, {:error, :stream_not_found}, state} {:error, reason} -> {:reply, {:error, reason}, state} end end @@ -2234,7 +2235,7 @@ defmodule Coflux.Orchestration.Server do ) state = pump_subscription(state, key) - state = refresh_stream_demand(state, {sub.producer_execution_id, sub.index}) + state = refresh_stream_demand(state, sub.stream_id) {:reply, :ok, flush_notifications(state)} else @@ -2956,15 +2957,10 @@ defmodule Coflux.Orchestration.Server do end end - def handle_call( - {:subscribe_stream_topic, execution_external_id, index, pid}, - _from, - state - ) do - case build_stream_topic_initial(state, execution_external_id, index) do + def handle_call({:subscribe_stream_topic, stream_external_id, pid}, _from, state) do + case build_stream_topic_initial(state, stream_external_id) do {:ok, initial} -> - {:ok, ref, state} = - add_listener(state, {:stream, execution_external_id, index}, pid) + {:ok, ref, state} = add_listener(state, {:stream, stream_external_id}, pid) {:reply, {:ok, initial, ref}, state} @@ -3019,44 +3015,29 @@ defmodule Coflux.Orchestration.Server do {:reply, :ok, state} end - defp maybe_init_stream_producer( - state, - _execution_id, - _execution_external_id, - _index, - nil, - _session_id - ) do + # Producer-side backpressure state for a stream, (re)initialised on + # every registration. A resuming execution replaces the entry left by + # the suspended one: the session changes, and its credit starts at the + # head it was told to sequence from. + defp init_stream_producer(state, stream, _execution_external_id, nil, _head, _session_id) do # buffer=nil means the producer has opted out of backpressure — no # tracking required on the server side. It'll emit freely and the - # adapter's driver never waits. - state + # adapter's driver never waits. Drop anything a previous producer of + # the stream left behind. + drop_stream_producer(state, stream.id) end - defp maybe_init_stream_producer( - state, - execution_id, - execution_external_id, - index, - buffer, - session_id - ) + defp init_stream_producer(state, stream, execution_external_id, buffer, head, session_id) when is_integer(buffer) and buffer >= 0 do - put_in(state.stream_producers[{execution_id, index}], %{ + put_in(state.stream_producers[stream.id], %{ buffer: buffer, - demand_granted: 0, + demand_granted: head + 1, session_id: session_id, - execution_external_id: execution_external_id + execution_external_id: execution_external_id, + index: stream.index }) end - defp maybe_send_initial_demand(state, execution_id, index) do - # At registration time there are no subscribers yet. Allow the - # producer to pre-warm up to `buffer` items; lockstep (buffer=0) - # stays paused until a consumer attaches. - refresh_stream_demand(state, {execution_id, index}) - end - # Recompute the target demand for one stream and, if it's grown, # send a delta grant to the producer's session. # @@ -3083,8 +3064,8 @@ defmodule Coflux.Orchestration.Server do # mark. The target *rising* does have to be noticed, though, which is # why drop_subscription refreshes: losing the slowest subscriber # raises the minimum, and nothing else would recompute it. - defp refresh_stream_demand(state, key) do - case Map.fetch(state.stream_producers, key) do + defp refresh_stream_demand(state, stream_id) do + case Map.fetch(state.stream_producers, stream_id) do :error -> state @@ -3104,41 +3085,41 @@ defmodule Coflux.Orchestration.Server do if not Map.has_key?(state.sessions, producer.session_id) do state else - refresh_stream_demand_for(state, key, producer) + refresh_stream_demand_for(state, stream_id, producer) end end end - defp refresh_stream_demand_for(state, {_execution_id, index} = key, producer) do - has_subscribers = has_stream_subscribers?(state, key) - watermark = slowest_ack_watermark(state, key) + defp refresh_stream_demand_for(state, stream_id, producer) do + has_subscribers = has_stream_subscribers?(state, stream_id) + watermark = slowest_ack_watermark(state, stream_id) bump = if has_subscribers, do: 1, else: 0 target = watermark + producer.buffer + bump delta = target - producer.demand_granted if delta > 0 do state - |> put_in([Access.key(:stream_producers), key, :demand_granted], target) + |> put_in([Access.key(:stream_producers), stream_id, :demand_granted], target) |> send_session( producer.session_id, - {:stream_demand, producer.execution_external_id, index, delta} + {:stream_demand, producer.execution_external_id, producer.index, delta} ) else state end end - defp has_stream_subscribers?(state, key) do - case Map.get(state.stream_subscribers, key) do + defp has_stream_subscribers?(state, stream_id) do + case Map.get(state.stream_subscribers, stream_id) do nil -> false set -> MapSet.size(set) > 0 end end - defp slowest_ack_watermark(state, key) do + defp slowest_ack_watermark(state, stream_id) do watermarks = state.stream_subscribers - |> Map.get(key, MapSet.new()) + |> Map.get(stream_id, MapSet.new()) |> Enum.flat_map(fn sub_key -> case Map.get(state.stream_subscriptions, sub_key) do nil -> [] @@ -3168,58 +3149,85 @@ defmodule Coflux.Orchestration.Server do defp ack_watermark(%{acked_seq: acked_seq}), do: acked_seq + 1 - defp drop_stream_producer(state, key) do - Map.update!(state, :stream_producers, &Map.delete(&1, key)) + defp drop_stream_producer(state, stream_id) do + Map.update!(state, :stream_producers, &Map.delete(&1, stream_id)) end # Lazily rebuild stream_producer state from the DB if it's missing. # Used after server restart — in-memory producer state is gone but - # the ``streams`` table still has the buffer. We rebuild on first - # append or subscribe for a given stream, recovering flow control. + # the registration still has the config. We rebuild on first append or + # subscribe for a given stream, recovering flow control. # # ``session_id`` is the internal id of the producer's current session; # supply ``nil`` if not known, in which case demand grants will be # deferred until the session is resolvable. - defp ensure_stream_producer( - state, - execution_id, - execution_external_id, - index, - session_id - ) do - key = {execution_id, index} + defp ensure_stream_producer(state, stream_id, session_id) do + if Map.has_key?(state.stream_producers, stream_id) do + state + else + case Streams.get_config(state.db, stream_id) do + {:ok, {nil, _timeout_ms}} -> + # Stream opted out of backpressure; nothing to track. + state - cond do - Map.has_key?(state.stream_producers, key) -> - state + {:ok, {buffer, _timeout_ms}} when is_integer(buffer) -> + # Reconstruct state. demand_granted starts at items already + # produced — we assume earlier-us granted enough for those, + # and rely on the producer having kept its local credit + # counter consistent. + {:ok, head} = Streams.get_stream_head(state.db, stream_id) + {:ok, stream} = Streams.get_stream(state.db, stream_id) + + put_in(state.stream_producers[stream_id], %{ + buffer: buffer, + demand_granted: head + 1, + session_id: session_id, + execution_external_id: producer_external_id(state.db, stream_id), + index: stream.index + }) - true -> - case Streams.get_buffer(state.db, execution_id, index) do - {:ok, nil} -> - # Stream opted out of backpressure; nothing to track. - state + {:error, :not_found} -> + state + end + end + end - {:ok, buffer} when is_integer(buffer) -> - # Reconstruct state. demand_granted starts at items already - # produced — we assume earlier-us granted enough for those, - # and rely on the producer having kept its local credit - # counter consistent. - {:ok, head} = Streams.get_stream_head(state.db, execution_id, index) - items_produced = if head < 0, do: 0, else: head + 1 - - put_in(state.stream_producers[key], %{ - buffer: buffer, - demand_granted: items_produced, - session_id: session_id, - execution_external_id: execution_external_id - }) + # External id of the stream's current producer (its latest registrant), + # or nil if it has none. + defp producer_external_id(db, stream_id) do + with {:ok, execution_id} <- Streams.get_producer(db, stream_id), + {:ok, {r, s, a}} <- Runs.get_execution_key(db, execution_id) do + execution_external_id(r, s, a) + else + _ -> nil + end + end - {:error, :not_found} -> - state + # The session the stream's current producer is running on, if it's live. + defp producer_session_id(state, stream_id) do + case producer_external_id(state.db, stream_id) do + nil -> + nil + + ext_id -> + case find_session_for_execution(state, ext_id) do + {:ok, sid} -> sid + :error -> nil end end end + # The stream a producer means by `index`: the one with that step index + # on the producer's own step. + defp resolve_step_stream(db, execution_id, index) do + {:ok, {step_id, _workspace_id, _attempt}} = Runs.get_execution_location(db, execution_id) + + case Streams.get_stream_by_step_index(db, step_id, index) do + {:ok, stream_id} -> {:ok, stream_id} + {:error, :not_found} -> {:error, :not_registered} + end + end + def handle_cast({:unsubscribe, ref}, state) do Process.demonitor(ref, [:flush]) @@ -4162,14 +4170,21 @@ defmodule Coflux.Orchestration.Server do end # Cancel a single execution: record :cancelled, abort if assigned, cancel descendants. - defp do_cancel_execution(state, execution_id, workspace_id) do + # + # `streams: :step` (the default) closes every open stream of the step in + # the workspace, so a consumer waiting on a paused stream is released when + # the pending successor is cancelled. `streams: :registered` closes only + # the streams the cancelled execution itself produced into — used by + # re-run, where cancelling a never-started successor must leave the + # paused stream for the new attempt to continue. + defp do_cancel_execution(state, execution_id, workspace_id, opts \\ []) do # Write the completion row (kind = cancelled) and fire notifications. # The result row is left untouched: if the worker already produced a # value, it stays; otherwise nothing is recorded. UI shows "cancelled" # via the completion kind, with any prior result visible in the # sidebar. state = - case record_and_notify_result(state, execution_id, :cancelled, nil) do + case record_and_notify_result(state, execution_id, :cancelled, nil, nil, opts) do {:ok, state} -> state {:error, :already_recorded} -> state {:error, :already_completed} -> state @@ -4179,7 +4194,8 @@ defmodule Coflux.Orchestration.Server do # subsequent `append_item` from the producer will fail with `:closed`, # signalling the worker to stop. Recorded as :lifecycle — consumers # derive the ExecutionCancelled error from the recorded result. - state = close_open_streams(state, execution_id) + state = + close_open_streams(state, execution_id, :lifecycle, Keyword.get(opts, :streams, :step)) state = case Runs.get_execution_key(state.db, execution_id) do @@ -4266,12 +4282,12 @@ defmodule Coflux.Orchestration.Server do end # Cancel all active (unresolved) executions for a step in a workspace. - defp cancel_active_step_executions(state, step_id, workspace_id) do + defp cancel_active_step_executions(state, step_id, workspace_id, opts) do {:ok, active_execution_ids} = Runs.get_active_execution_ids_for_step(state.db, step_id, workspace_id) Enum.reduce(active_execution_ids, state, fn exec_id, state -> - do_cancel_execution(state, exec_id, workspace_id) + do_cancel_execution(state, exec_id, workspace_id, opts) end) end @@ -5403,7 +5419,9 @@ defmodule Coflux.Orchestration.Server do {:ok, steps} = Runs.get_run_steps(db, run.id) {:ok, run_executions} = Runs.get_run_executions(db, run.id) {:ok, run_dependencies} = Runs.get_run_dependencies(db, run.id) - {:ok, run_stream_dependencies} = Runs.get_run_stream_dependencies(db, run.id) + {:ok, run_stream_dependencies} = Streams.get_run_dependencies(db, run.id) + {:ok, run_streams} = Streams.get_streams_for_run(db, run.id) + streams_by_step = build_run_streams(db, run_streams) {:ok, run_children} = Runs.get_run_children(db, run.id) {:ok, groups} = Runs.get_groups_for_run(db, run.id) {:ok, run_metric_defs} = Runs.get_run_metric_definitions(db, run.id) @@ -5569,6 +5587,7 @@ defmodule Coflux.Orchestration.Server do created_at: step.created_at, arguments: arguments, requires: requires, + streams: Map.get(streams_by_step, step.id, %{}), executions: run_executions |> Enum.filter(&(elem(&1, 1) == step.id)) @@ -5621,11 +5640,12 @@ defmodule Coflux.Orchestration.Server do stream_deps = run_stream_dependencies |> Map.get(execution_id, []) - |> Map.new(fn {stream_ref_id, stream_index} -> - {producer_ext_id, _module, _target} = - execution = resolve_execution_ref(db, stream_ref_id) + |> Map.new(fn stream_ref_id -> + {:ok, {stream_run_ext_id, step_number, index, module, target}} = + Streams.get_stream_ref(db, stream_ref_id) - {"#{producer_ext_id}:#{stream_index}", {:stream, stream_index, execution}} + id = stream_external_id(stream_run_ext_id, step_number, index) + {id, {:stream, id, module, target}} end) dependencies = @@ -5640,8 +5660,6 @@ defmodule Coflux.Orchestration.Server do ) ) - streams = streams_with_resolved_reasons(db, execution_id) - {:ok, {checkpoints_before, checkpoints_after}} = Checkpoints.get_execution_snapshots( db, @@ -5670,7 +5688,6 @@ defmodule Coflux.Orchestration.Server do result_created_by: result_created_by, children: Map.get(run_children, execution_id, []), metric_definitions: Map.get(metric_definitions_by_execution, execution_id, %{}), - streams: streams, checkpoints: %{ before: enrich_checkpoints(checkpoints_before, db), after: enrich_checkpoints(checkpoints_after, db) @@ -6176,6 +6193,53 @@ defmodule Coflux.Orchestration.Server do end end + defp stream_external_id(run_external_id, step_number, index) do + "#{run_external_id}:#{step_number}_#{index}" + end + + # `:_`. Run ids are alphanumeric and step numbers are + # integers, so the last `_` unambiguously separates the index. + defp parse_stream_external_id(id) when is_binary(id) do + case String.split(id, "_") do + parts when length(parts) >= 2 -> + {index_s, prefix_parts} = List.pop_at(parts, -1) + + with {index, ""} when index >= 0 <- Integer.parse(index_s), + {:ok, run_external_id, step_number} <- parse_step_id(Enum.join(prefix_parts, "_")) do + {:ok, run_external_id, step_number, index} + else + _ -> {:error, :invalid_format} + end + + _ -> + {:error, :invalid_format} + end + end + + defp parse_stream_external_id(_), do: {:error, :invalid_format} + + # Resolve a stream's external id to its row in the active epoch, copying + # its run forward from an archived epoch if that's where it lives. + defp resolve_stream_id(state, external_id) do + with {:ok, run_ext_id, step_number, index} <- parse_stream_external_id(external_id) do + case Streams.get_stream_id_by_key(state.db, run_ext_id, step_number, index) do + {:ok, id} -> + {:ok, id} + + {:error, :not_found} -> + case find_and_copy_run_from_archives(state, run_ext_id) do + {:ok, _remap} -> + Streams.get_stream_id_by_key(state.db, run_ext_id, step_number, index) + + :not_found -> + {:error, :not_found} + end + end + else + {:error, :invalid_format} -> {:error, :not_found} + end + end + defp parse_step_id(step_id) do case String.split(step_id, ":", parts: 2) do [run_external_id, step_number_s] -> @@ -6358,7 +6422,18 @@ defmodule Coflux.Orchestration.Server do # The completion row is written later via complete_execution (triggered by # notify_terminated for worker-involved cases, or by the server-initiated # paths directly when no worker is involved). - defp record_and_notify_result(state, execution_id, result, _module, created_by \\ nil) do + # `opts[:streams]` selects which open streams an immediately-completing + # result closes (see `close_open_streams/4`): `:step` by default, or + # `:registered` when a re-run cancels a never-started successor and the + # paused streams must survive for the new attempt. + defp record_and_notify_result( + state, + execution_id, + result, + _module, + created_by \\ nil, + opts \\ [] + ) do result = case result do {:value, value} -> {:value, normalize_value(value)} @@ -6384,7 +6459,7 @@ defmodule Coflux.Orchestration.Server do state = if writes_completion_immediately?(result) do state - |> maybe_close_open_streams(result, execution_id) + |> maybe_close_open_streams(result, execution_id, Keyword.get(opts, :streams, :step)) |> fire_completion_notification(execution_id, timestamp) else state @@ -6397,27 +6472,24 @@ defmodule Coflux.Orchestration.Server do end end - # Which of the immediately-completing shapes should close the - # execution's open streams. Deliberately the same set that - # derive_lifecycle_info can name — closing a stream whose reason we - # can't derive would push a nil reason, which consumers coerce to a - # clean "complete" and silently accept as a truncated stream. + # Which of the immediately-completing shapes should close the step's + # open streams. Deliberately the same set that derive_lifecycle_info + # can name — closing a stream whose reason we can't derive would push a + # nil reason, which consumers coerce to a clean "complete" and silently + # accept as a truncated stream. # # :deferred / :cached / :spawned never reach here with open streams # (the execution was superseded before its body ran, so it appended # nothing). # - # :suspended / :recurred *are* included. A stream is owned by exactly - # one execution (see the `streams` table invariants), so a successor - # does not continue its predecessor's stream — it registers its own - # under a new execution id, and consumer references are concrete to - # the original. That makes the predecessor's stream terminal at the - # moment it suspends or recurs. Leaving it open would strand every - # attached consumer: no further item can ever be appended to it, and - # no later event would close it. - defp maybe_close_open_streams(state, result, execution_id) do + # :suspended is deliberately *not* included: a suspend pauses the + # step's streams, and the execution that resumes the step continues + # them, so consumers see one unbroken sequence. :recurred *is* + # included — a recurrent iteration finishing is a completion, and the + # next iteration opens its own streams. + defp maybe_close_open_streams(state, result, execution_id, scope) do if closes_streams_on_completion?(result) do - close_open_streams(state, execution_id) + close_open_streams(state, execution_id, :lifecycle, scope) else state end @@ -6427,7 +6499,6 @@ defmodule Coflux.Orchestration.Server do defp closes_streams_on_completion?({:abandoned, _}), do: true defp closes_streams_on_completion?({:crashed, _}), do: true defp closes_streams_on_completion?({:timeout, _}), do: true - defp closes_streams_on_completion?({:suspended, _}), do: true defp closes_streams_on_completion?({:recurred, _}), do: true defp closes_streams_on_completion?(_), do: false @@ -6587,15 +6658,17 @@ defmodule Coflux.Orchestration.Server do end # Value result + drain: dispatch on stream closure outcomes. - # * any owned stream closed `:errored` → `:stream_errored` (retried) - # * else any owned stream closed `:timeout` → `:partial` (not retried, - # not cacheable) + # * any stream this execution closed `:errored` → `:stream_errored` + # (retried) + # * else any stream it closed `:timeout` → `:stream_timeout` (not + # retried, not cacheable) # * else `:succeeded` - # `close_open_streams` runs first so any still-open streams get a - # `:lifecycle` row (which doesn't influence the dispatch — only the - # explicit `:errored`/`:timeout` reasons do). + # `close_open_streams` runs first so any of the step's streams still + # open get a `:complete` row — the step finished, so they're done. That + # doesn't influence the dispatch; only the explicit `:errored` / + # `:timeout` reasons do. defp finalize_success_completion(state, execution_id) do - state = close_open_streams(state, execution_id) + state = close_open_streams(state, execution_id, :complete) {:ok, summary} = Streams.get_closure_summary_for_execution(state.db, execution_id) @@ -6752,38 +6825,47 @@ defmodule Coflux.Orchestration.Server do end end - # Closes every stream owned by `execution_id` that doesn't yet have a - # closure row, and pushes a `stream_closed` notification to every active - # subscriber. Streams already closed by the producer (clean or errored) - # are left untouched. + # Closes the step's open streams on behalf of `execution_id`, and pushes + # a `stream_closed` notification to every active subscriber. Streams + # already closed by the producer (clean or errored) are left untouched. + # + # `scope` picks which streams: `:step` (default) is every open stream + # of the execution's step in its workspace — including paused streams + # left by an earlier suspended attempt, which nothing else would close. + # `:registered` is only the streams this execution produced into; used + # when a re-run cancels a pending successor so the paused streams + # survive for the new attempt. # - # The closure is recorded on disk with reason :lifecycle (no error on - # the closure row). On the wire, we resolve that to a specific reason - # (:cancelled / :abandoned / :crashed / :timeout / :errored) by looking - # at the execution's completion — consumers then decide how to handle - # each case. - defp close_open_streams(state, execution_id, spec \\ :lifecycle) do - {:ok, indexes} = Streams.get_open_streams_for_execution(state.db, execution_id) + # `spec` is how the closure is recorded: `:lifecycle` (the reason is + # derived on read from the closing execution's completion), `:timeout` + # or `:complete`. + defp close_open_streams(state, execution_id, spec \\ :lifecycle, scope \\ :step) do + {:ok, stream_ids} = + case scope do + :step -> + {:ok, {step_id, workspace_id, _attempt}} = + Runs.get_execution_location(state.db, execution_id) + + Streams.get_open_stream_ids_for_step(state.db, step_id, workspace_id) + + :registered -> + Streams.get_open_stream_ids_for_execution(state.db, execution_id) + end {push_reason, push_error} = case spec do :lifecycle -> derive_lifecycle_info(state.db, execution_id) :timeout -> {:timeout, nil} + :complete -> {:complete, nil} end - Enum.reduce(indexes, state, fn index, state -> - case Streams.close_stream(state.db, execution_id, index, spec) do + Enum.reduce(stream_ids, state, fn stream_id, state -> + case Streams.close_stream(state.db, stream_id, execution_id, spec) do {:ok, closed_at} -> state - |> push_stream_closed(execution_id, index, push_reason, push_error) - |> notify_stream_closed( - execution_id, - index, - push_reason, - push_error, - closed_at - ) - |> drop_stream_producer({execution_id, index}) + |> push_stream_closed(stream_id, push_reason, push_error) + |> notify_stream_closed(stream_id, execution_id, push_reason, push_error, closed_at) + |> drop_stream_producer(stream_id) {:error, :already_closed} -> state @@ -6791,32 +6873,66 @@ defmodule Coflux.Orchestration.Server do end) end - # Returns the streams list for `execution_id` for the run topic's - # initial state. Shape: - # `{index, buffer, opened_at, closed_at | nil, reason | nil, error | nil}` - # - # DB closures are recorded with reason :complete / :errored / :lifecycle. - # For :lifecycle we resolve the actual cause (from the execution's - # completion kind) and surface that directly as the reason — - # :cancelled / :abandoned / :crashed / :timeout / :errored — so Studio - # doesn't have to deal with a generic "lifecycle" bucket. `error` is - # non-nil only when the reason is :errored. - defp streams_with_resolved_reasons(db, execution_id) do - {:ok, rows} = Streams.get_streams_with_closures_for_execution(db, execution_id) - - Enum.map(rows, fn - {index, buffer, timeout_ms, opened_at, nil, nil, nil} -> - {index, buffer, timeout_ms, opened_at, nil, nil, nil} - - {index, buffer, timeout_ms, opened_at, closed_at, :lifecycle, _} -> - {resolved_reason, resolved_error} = derive_lifecycle_info(db, execution_id) - {index, buffer, timeout_ms, opened_at, closed_at, resolved_reason, resolved_error} - - {index, buffer, timeout_ms, opened_at, closed_at, reason, error} -> - {index, buffer, timeout_ms, opened_at, closed_at, reason, error} + # Streams for the run topic's initial state, grouped by step id: + # `%{step_id => %{index => entry}}`. Each entry is the same shape the + # `:stream_registered` / `:stream_closed` notifications build up, with + # `:lifecycle` closures resolved to their specific cause against the + # closing execution. + defp build_run_streams(db, run_streams) do + run_streams + |> Enum.group_by(& &1.step_id) + |> Map.new(fn {step_id, streams} -> + {step_id, + Map.new(streams, fn stream -> + {:ok, registrations} = Streams.get_registrations(db, stream.id) + {buffer, timeout_ms} = latest_registration_config(registrations) + + {:ok, workspace_external_id} = + Workspaces.get_workspace_external_id(db, stream.workspace_id) + + {reason, error, closed_by_attempt} = + if stream.closed_at do + {reason, error} = + resolve_closure_reason(db, stream.reason, stream.error, stream.closed_by) + + {:ok, {_r, _s, attempt}} = Runs.get_execution_key(db, stream.closed_by) + {reason, error, attempt} + else + {nil, nil, nil} + end + + {stream.index, + %{ + id: stream_external_id(stream.run_external_id, stream.step_number, stream.index), + index: stream.index, + position: stream.position, + workspace_id: workspace_external_id, + buffer: buffer, + timeout_ms: timeout_ms, + opened_at: stream.created_at, + attempts: Enum.map(registrations, fn {_id, attempt, _b, _t, _c} -> attempt end), + closed_at: stream.closed_at, + closed_by: closed_by_attempt, + reason: reason, + error: error + }} + end)} end) end + defp latest_registration_config([]), do: {nil, nil} + + defp latest_registration_config(registrations) do + {_id, _attempt, buffer, timeout_ms, _created_at} = List.last(registrations) + {buffer, timeout_ms} + end + + # A `:lifecycle` closure's meaning comes from the closing execution's + # completion; any other reason is stored directly. + defp resolve_closure_reason(db, :lifecycle, _stored_error, closed_by), + do: derive_lifecycle_info(db, closed_by) + + defp resolve_closure_reason(_db, reason, stored_error, _closed_by), do: {reason, stored_error} # Derive a semantic reason + optional error for a lifecycle stream # closure, from the execution's completion kind. Used when pushing # closures to live consumers and when late subscribers attach to @@ -6849,15 +6965,11 @@ defmodule Coflux.Orchestration.Server do {:ok, {:timeout, _, _, _, _}} -> {:timeout, nil} - # The producer didn't fail — it suspended, or finished a recurrent - # iteration. Either way *this* stream is finished: the successor - # registers its own under a new execution id and consumer - # references are concrete to the original, so nothing more can - # arrive here. Reported distinctly rather than as :abandoned so a - # truncated-by-recurrence stream doesn't read as a worker failure. - {:ok, {:suspended, _, _, _, _}} -> - {:suspended, nil} - + # The producer didn't fail — it finished a recurrent iteration, and + # the next iteration opens its own streams. Reported distinctly + # rather than as :abandoned so a truncated-by-recurrence stream + # doesn't read as a worker failure. (A suspend never closes a + # stream, so it never appears here.) {:ok, {:recurred, _, _, _, _}} -> {:recurred, nil} @@ -6986,19 +7098,43 @@ defmodule Coflux.Orchestration.Server do {:ok, state} has_result? -> - # Mid-drain (value result recorded, completion pending): a - # wall-clock timeout here means the drain was cut short. Close the - # remaining open streams as :timeout so complete_execution promotes - # the completion to :stream_timeout — otherwise the kill would land - # as a clean :succeeded with silently truncated streams. - state = - if result == :timeout do - close_open_streams(state, execution_id, :timeout) - else - state - end + # Mid-drain: the value result is recorded and the completion is + # pending while the execution's streams drain. + cond do + # A wall-clock timeout here means the drain was cut short. Close + # the remaining open streams as :timeout so complete_execution + # promotes the completion to :stream_timeout — otherwise the + # kill would land as a clean :succeeded with silently truncated + # streams. + result == :timeout -> + {:ok, close_open_streams(state, execution_id, :timeout)} + + # A generator-bodied producer suspends from inside its body, + # after its value (the stream handle) was recorded. Write the + # completion so the successor is scheduled and the streams stay + # paused for it. The run topic's `:result` isn't re-fired: the + # value stands, and the completion carries the suspension. + match?({:suspended, _, _}, result) -> + {:ok, step} = Runs.get_step_for_execution(state.db, execution_id) + {:ok, workspace_id} = Runs.get_workspace_id_for_execution(state.db, execution_id) + + {successor_id, _recurred?, state} = + decide_and_create_successor(state, execution_id, step, workspace_id, result) + + case Results.record_completion(state.db, execution_id, :suspended, + successor_id: successor_id, + created_by: created_by + ) do + {:ok, completion_at} -> + {:ok, fire_completion_notification(state, execution_id, completion_at)} - {:ok, state} + {:error, :already_completed} -> + {:ok, state} + end + + true -> + {:ok, state} + end true -> {:ok, step} = Runs.get_step_for_execution(state.db, execution_id) @@ -8285,46 +8421,64 @@ defmodule Coflux.Orchestration.Server do # --- Stream topic notifications (for Studio subscribers) --- # These flow through `notify_listeners` → the run topic and the - # per-stream `{:stream, execution_ext_id, index}` inspection topic, - # distinct from the session-directed `push_stream_*` helpers which - # target subscribed consumer sessions' WebSockets. + # per-stream `{:stream, stream_external_id}` inspection topic, distinct + # from the session-directed `push_stream_*` helpers which target + # subscribed consumer sessions' WebSockets. # Bounded tail of items held by the stream inspection topic. Long # streams don't need to materialise every item — the UI loads older # items on demand. @stream_topic_tail_size 200 - defp notify_stream_opened(state, execution_id, index, buffer, timeout_ms, created_at) do - {:ok, {r, _s, _a}} = Runs.get_execution_key(state.db, execution_id) - {:ok, execution_ext_id} = execution_external_id_for(state.db, execution_id) + # An execution registered on a stream: either opening it or resuming it + # after a suspend. The run topic keeps streams under their step, with + # the attempts that have produced into each. + defp notify_stream_registered(state, stream, execution_id, registration, buffer, timeout_ms) do + {:ok, {_r, _s, attempt}} = Runs.get_execution_key(state.db, execution_id) + + {:ok, workspace_external_id} = + Workspaces.get_workspace_external_id(state.db, stream.workspace_id) + + ext_id = stream_external_id(stream.run_external_id, stream.step_number, stream.index) + + state = + notify_listeners( + state, + {:run, stream.run_external_id}, + {:stream_registered, stream.step_number, stream.index, + %{ + id: ext_id, + position: stream.position, + workspace_id: workspace_external_id, + attempt: attempt, + buffer: buffer, + timeout_ms: timeout_ms, + continued: registration.continued, + opened_at: stream.created_at + }} + ) notify_listeners( state, - {:run, r}, - {:stream_opened, execution_ext_id, index, buffer, timeout_ms, created_at} + {:stream, ext_id}, + {:registered, attempt, buffer, timeout_ms, registration.created_at} ) end - defp notify_stream_item_appended( - state, - execution_id, - index, - sequence, - value, - created_at - ) do - {:ok, execution_ext_id} = execution_external_id_for(state.db, execution_id) - topic = {:stream, execution_ext_id, index} + defp notify_stream_item_appended(state, stream_id, execution_id, sequence, value, created_at) do + {:ok, ext_id} = stream_external_id_for(state.db, stream_id) + topic = {:stream, ext_id} # Skip the build_value (which hits the DB to resolve refs) when the # inspection topic has no active subscribers. if Map.has_key?(state.topics, topic) do + {:ok, {_r, _s, attempt}} = Runs.get_execution_key(state.db, execution_id) resolved = build_value(normalize_value(value), state.db) notify_listeners( state, topic, - {:item_appended, sequence, resolved, created_at} + {:item_appended, sequence, resolved, attempt, created_at} ) else state @@ -8333,90 +8487,101 @@ defmodule Coflux.Orchestration.Server do # Fire the stream-closed notification on run + stream topics. `reason` # is a semantic atom from the full set (:complete / :errored / - # :cancelled / :abandoned / :crashed / :timeout) — Studio renders each - # directly in UI-appropriate language, rather than displaying a - # fabricated exception type. `error` is non-nil only for :errored. - defp notify_stream_closed(state, execution_id, index, reason, error, closed_at) do - {:ok, {r, _s, _a}} = Runs.get_execution_key(state.db, execution_id) - {:ok, execution_ext_id} = execution_external_id_for(state.db, execution_id) + # :cancelled / :abandoned / :crashed / :timeout / :recurred) — Studio + # renders each directly in UI-appropriate language, rather than + # displaying a fabricated exception type. `error` is non-nil only for + # :errored. `execution_id` is the execution that closed the stream. + defp notify_stream_closed(state, stream_id, execution_id, reason, error, closed_at) do + {:ok, stream} = Streams.get_stream(state.db, stream_id) + ext_id = stream_external_id(stream.run_external_id, stream.step_number, stream.index) + {:ok, {_r, _s, attempt}} = Runs.get_execution_key(state.db, execution_id) encoded_error = encode_stream_error_summary(error) - reason_str = if reason, do: Atom.to_string(reason) state = notify_listeners( state, - {:run, r}, - {:stream_closed, execution_ext_id, index, reason_str, encoded_error, closed_at} + {:run, stream.run_external_id}, + {:stream_closed, stream.step_number, stream.index, reason_str, encoded_error, attempt, + closed_at} ) notify_listeners( state, - {:stream, execution_ext_id, index}, - {:closed, reason_str, encoded_error, closed_at} + {:stream, ext_id}, + {:closed, reason_str, encoded_error, attempt, closed_at} ) end - # Build the initial state for a newly-opened stream inspection topic. - # Returns {:ok, state} with producer metadata, opened/closed timestamps, - # closure info (with lifecycle errors already derived), bounded tail of - # items, and the total item count. - defp build_stream_topic_initial(state, execution_ext_id, index) do - with {:ok, execution_id} <- resolve_internal_execution_id(state, execution_ext_id), - {:ok, true} <- Streams.exists?(state.db, execution_id, index), - {:ok, opened_at} <- Streams.get_opened_at(state.db, execution_id, index), - {:ok, buffer} <- Streams.get_buffer(state.db, execution_id, index), - {:ok, timeout_ms} <- Streams.get_timeout_ms(state.db, execution_id, index), + # Build the initial state for a newly-opened stream inspection topic: + # the step it belongs to, its config and producers, closure info (with + # lifecycle reasons already derived), a bounded tail of items, and the + # total item count. + defp build_stream_topic_initial(state, stream_external_id) do + with {:ok, stream_id} <- resolve_stream_id(state, stream_external_id), + {:ok, stream} <- Streams.get_stream(state.db, stream_id), + {:ok, registrations} <- Streams.get_registrations(state.db, stream_id), {:ok, {items, total_count}} <- - Streams.get_stream_tail(state.db, execution_id, index, @stream_topic_tail_size) do + Streams.get_stream_tail(state.db, stream_id, @stream_topic_tail_size) do # Keep the tuple shape here — the topic module runs TopicUtils.build_value # on each item's value to produce the JSON-encodable form, matching # how live :item_appended notifications are handled. resolved_items = - Enum.map(items, fn {sequence, value, created_at} -> - {sequence, build_value(value, state.db), created_at} + Enum.map(items, fn {sequence, value, attempt, created_at} -> + {sequence, build_value(value, state.db), attempt, created_at} end) - closure = build_stream_topic_closure(state, execution_id, index) + {buffer, timeout_ms} = latest_registration_config(registrations) + + {:ok, workspace_external_id} = + Workspaces.get_workspace_external_id(state.db, stream.workspace_id) {:ok, %{ - producer: build_stream_producer(state.db, execution_ext_id, execution_id), + id: stream_external_id, + step: %{ + stepId: "#{stream.run_external_id}:#{stream.step_number}", + module: stream.module, + target: stream.target + }, + workspaceId: workspace_external_id, + index: stream.index, + position: stream.position, buffer: buffer, timeoutMs: timeout_ms, - openedAt: opened_at, - closure: closure, + openedAt: stream.created_at, + attempts: Enum.map(registrations, fn {_id, attempt, _b, _t, _c} -> attempt end), + closure: build_stream_topic_closure(state, stream_id), items: resolved_items, totalCount: total_count, tailSize: @stream_topic_tail_size }} else - {:ok, false} -> {:error, :not_found} - {:error, reason} -> {:error, reason} + {:error, _reason} -> {:error, :not_found} end end - defp build_stream_topic_closure(state, execution_id, index) do - case Streams.get_stream_closure(state.db, execution_id, index) do + defp build_stream_topic_closure(state, stream_id) do + case Streams.get_stream_closure(state.db, stream_id) do {:ok, nil} -> nil - {:ok, {reason, stored_error, closed_at}} -> - # DB stores `:lifecycle` for closures driven by the producer - # execution ending; on read we resolve that to the specific - # cause (:cancelled / :abandoned / :crashed / :timeout / - # :errored) so clients don't need to know about the internal + {:ok, {reason, stored_error, closed_by, closed_at}} -> + # DB stores `:lifecycle` for closures driven by an execution + # ending; on read we resolve that to the specific cause + # (:cancelled / :abandoned / :crashed / :timeout / :errored / + # :recurred) so clients don't need to know about the internal # bucket. `error` only accompanies a genuine :errored close. {effective_reason, effective_error} = - case reason do - :lifecycle -> derive_lifecycle_info(state.db, execution_id) - _ -> {reason, stored_error} - end + resolve_closure_reason(state.db, reason, stored_error, closed_by) + + {:ok, {_r, _s, attempt}} = Runs.get_execution_key(state.db, closed_by) %{ reason: if(effective_reason, do: Atom.to_string(effective_reason)), error: encode_stream_error_summary(effective_error), + attempt: attempt, closedAt: closed_at } end @@ -8435,18 +8600,13 @@ defmodule Coflux.Orchestration.Server do } end - defp build_stream_producer(db, execution_ext_id, execution_id) do - # The external execution_id already encodes run + step + attempt, so - # the producer reference only carries identifier + module/target — - # matching ExecutionReference on the wire. - {:ok, step} = Runs.get_step_for_execution(db, execution_id) - Coflux.TopicUtils.build_execution({execution_ext_id, step.module, step.target}) - end + defp stream_external_id_for(db, stream_id) do + case Streams.get_stream(db, stream_id) do + {:ok, stream} -> + {:ok, stream_external_id(stream.run_external_id, stream.step_number, stream.index)} - defp execution_external_id_for(db, execution_id) do - case Runs.get_execution_key(db, execution_id) do - {:ok, {r, s, a}} -> {:ok, execution_external_id(r, s, a)} - err -> err + err -> + err end end @@ -8513,13 +8673,7 @@ defmodule Coflux.Orchestration.Server do state else {:ok, items} = - Streams.get_stream_items( - state.db, - sub.producer_execution_id, - sub.index, - sub.cursor, - @backlog_page_size - ) + Streams.get_stream_items(state.db, sub.stream_id, sub.cursor, @backlog_page_size) if items == [] do state @@ -8613,7 +8767,7 @@ defmodule Coflux.Orchestration.Server do # The push moved this consumer's cursor forward, which may have moved # the slowest-subscriber watermark and unblocked the producer. - refresh_stream_demand(state, {sub.producer_execution_id, sub.index}) + refresh_stream_demand(state, sub.stream_id) end # Resolve the consumer's current session and send, skipping if the @@ -8637,10 +8791,8 @@ defmodule Coflux.Orchestration.Server do # stride or credit semantics must be made there *and* mirrored here. # Anything this path declines to send is left untouched and durable, # so the pump re-reads it in order — declining is always safe. - defp push_stream_item(state, producer_execution_id, index, sequence, value) do - stream_key = {producer_execution_id, index} - - subscribers = Map.get(state.stream_subscribers, stream_key, MapSet.new()) + defp push_stream_item(state, stream_id, sequence, value) do + subscribers = Map.get(state.stream_subscribers, stream_id, MapSet.new()) state = Enum.reduce(subscribers, state, fn key, state -> @@ -8714,7 +8866,7 @@ defmodule Coflux.Orchestration.Server do # Subscriber cursors may have advanced — recompute demand once per # stream (cheaper than once per subscriber, same result). - refresh_stream_demand(state, stream_key) + refresh_stream_demand(state, stream_id) end # On close, tell every subscriber. `reason` is a semantic atom @@ -8728,9 +8880,8 @@ defmodule Coflux.Orchestration.Server do # stream head, and emitting the close now would land it ahead of the # items the consumer hasn't been given room for yet. Each subscription # emits its close once it has drained. - defp push_stream_closed(state, producer_execution_id, index, reason, error) do - subscribers = - Map.get(state.stream_subscribers, {producer_execution_id, index}, MapSet.new()) + defp push_stream_closed(state, stream_id, reason, error) do + subscribers = Map.get(state.stream_subscribers, stream_id, MapSet.new()) reason_str = if reason, do: Atom.to_string(reason) encoded_error = encode_stream_error(error) @@ -8738,7 +8889,7 @@ defmodule Coflux.Orchestration.Server do # The stream is closed in the DB by the time we get here, so the head # is final. Read it once and record it on each pending close rather # than re-querying per subscriber per ack for the rest of the drain. - {:ok, head} = Streams.get_stream_head(state.db, producer_execution_id, index) + {:ok, head} = Streams.get_stream_head(state.db, stream_id) Enum.reduce(subscribers, state, fn key, state -> state @@ -8827,24 +8978,21 @@ defmodule Coflux.Orchestration.Server do end defp do_mark_closed(state, sub, key) do - case Streams.get_stream_closure(state.db, sub.producer_execution_id, sub.index) do + case Streams.get_stream_closure(state.db, sub.stream_id) do {:ok, nil} -> state - {:ok, {reason, stored_error, _closed_at}} -> + {:ok, {reason, stored_error, closed_by, _closed_at}} -> # Resolve :lifecycle to the specific cause for the wire — same # treatment as live closures so late subscribers don't get a # less-informative signal than those attached at close time. {effective_reason, effective_error} = - case reason do - :lifecycle -> derive_lifecycle_info(state.db, sub.producer_execution_id) - _ -> {reason, stored_error} - end + resolve_closure_reason(state.db, reason, stored_error, closed_by) reason_str = if effective_reason, do: Atom.to_string(effective_reason) # Already closed, so the head is final — see mark_pending_close. - {:ok, head} = Streams.get_stream_head(state.db, sub.producer_execution_id, sub.index) + {:ok, head} = Streams.get_stream_head(state.db, sub.stream_id) mark_pending_close( state, @@ -8862,7 +9010,7 @@ defmodule Coflux.Orchestration.Server do state {:ok, sub} -> - stream_key = {sub.producer_execution_id, sub.index} + stream_key = sub.stream_id state |> Map.update!(:stream_subscriptions, &Map.delete(&1, key)) diff --git a/server/lib/coflux/orchestration/streams.ex b/server/lib/coflux/orchestration/streams.ex index ccbd69cd..d8f91179 100644 --- a/server/lib/coflux/orchestration/streams.ex +++ b/server/lib/coflux/orchestration/streams.ex @@ -1,138 +1,469 @@ defmodule Coflux.Orchestration.Streams do @moduledoc """ - Storage for execution-produced streams. + Storage for step-produced streams. - A stream is an ordered, append-only sequence of values produced by an - execution. Each stream is identified by `(execution_id, index)` where - `index` is assigned monotonically by the worker during return-value - serialisation — the worker mints ids locally, no server round-trip. + A stream is an ordered, append-only sequence of values produced by a + step within a workspace. It is identified externally by + `:_`, where `index` is allocated here, per step, so + the id is unique across attempts and across workspaces. - Items within a stream are identified by `sequence` — a 0-based, - monotonically increasing per-item counter. + Streams outlive the execution that opens them in exactly one case: a + suspend. An execution that suspends leaves its open streams *paused*, + and the execution that resumes the step continues them. Registration + matches a resuming execution's k-th stream onto the paused stream whose + opener registered it k-th (`position`). Every other way an execution can + end closes the step's open streams — that lifecycle logic lives in + `Server`, not here. - The SQL column is quoted with backticks (``` `index` ```) throughout - queries because `INDEX` is a SQLite keyword; at the Elixir level we - just pass `:index` as a map key — the Store helper handles quoting - for inserts. + Items within a stream are identified by `sequence` — a 0-based, + monotonically increasing per-item counter that continues across a + suspension. Each item records the execution that appended it. Invariants enforced here (and by schema FKs): - * A stream is owned by exactly one execution (its producer). * Items are append-only with monotonic `sequence` starting at 0. - * A closure is terminal — no items may be appended after one is recorded. - * On execution completion / cancel / crash, every owned stream that lacks - a closure receives one (clean, cancelled, or crashed). Enforced by the - lifecycle code in `Server`, not by this module. - * Re-running a producer execution creates fresh streams (new attempt ⇒ - new execution_id ⇒ new rows). Consumer refs pin to the original streams. + * A closure is terminal — no items may be appended after one is + recorded. The closure records the execution that closed the stream. + * Only an execution registered on a stream may append to it. + * The latest registrant is the stream's producer, and its registration + carries the config (buffer, timeout) in force. * Consumer cursors are kept in-memory only; re-run consumers subscribe fresh from sequence 0. + + The SQL column is quoted with backticks (``` `index` ```) throughout + because `INDEX` is a SQLite keyword; at the Elixir level we pass `:index` + as a map key — the Store helper handles quoting for inserts. """ import Coflux.Store - alias Coflux.Orchestration.{Errors, Values} - - # Registers a new stream owned by `execution_id` at `index` (monotonic - # per-execution, worker-assigned). ``buffer`` is the persisted flow- - # control budget — ``nil`` means no backpressure, integer N means the - # producer may be up to N items ahead of the slowest consumer. - # ``timeout_ms`` is the idle-timeout budget (milliseconds) — ``nil`` - # disables the timeout. The server only stores it (for display in - # Studio); enforcement happens at the worker (CLI). Returns - # ``{:error, :already_registered}`` if the index was already used. - def register_stream(db, execution_id, index, buffer \\ nil, timeout_ms \\ nil) do - now = current_timestamp() - - case insert_one(db, :streams, %{ - execution_id: execution_id, + alias Coflux.Orchestration.{Errors, Results, Values} + + # --- Registration --- + + # Registers `execution_id` as a producer of the stream it opened at + # `position` (its k-th registration). Either continues the step's paused + # stream at that position, or opens a new one with the next index for + # the step. + # + # A stream is paused when it is open and its latest registrant has + # completed as suspended. An open stream whose producer is still live is + # never matched (a recurrent step's successor can be dispatched while + # its predecessor is still draining an explicit stream), and a closed + # stream is never continued. + # + # Returns `{:ok, %{id, index, head, continued, created_at}}`. `head` is + # the highest sequence already in the stream (`-1` when empty), so the + # producer knows where to resume numbering. Registering the same + # position twice for one execution is idempotent and returns the + # existing stream. + def register(db, step_id, workspace_id, execution_id, position, buffer, timeout_ms) do + with_transaction(db, fn -> + now = current_timestamp() + + case get_registered_stream(db, execution_id, position) do + {:ok, {stream_id, index}} -> + {:ok, head} = get_stream_head(db, stream_id) + {:ok, %{id: stream_id, index: index, head: head, continued: false, created_at: nil}} + + {:ok, nil} -> + case find_paused_stream(db, step_id, workspace_id, position) do + {:ok, {stream_id, index}} -> + insert_registration(db, stream_id, execution_id, buffer, timeout_ms, now) + {:ok, head} = get_stream_head(db, stream_id) + {:ok, %{id: stream_id, index: index, head: head, continued: true, created_at: now}} + + {:ok, nil} -> + {:ok, index} = next_index(db, step_id) + + {:ok, stream_id} = + insert_one(db, :streams, %{ + step_id: step_id, + workspace_id: workspace_id, + index: index, + position: position, + created_at: now + }) + + insert_registration(db, stream_id, execution_id, buffer, timeout_ms, now) + {:ok, %{id: stream_id, index: index, head: -1, continued: false, created_at: now}} + end + end + end) + end + + defp insert_registration(db, stream_id, execution_id, buffer, timeout_ms, now) do + {:ok, _} = + insert_one(db, :stream_registrations, %{ + stream_id: stream_id, + execution_id: execution_id, + buffer: buffer, + timeout_ms: timeout_ms, + created_at: now + }) + end + + defp get_registered_stream(db, execution_id, position) do + query_one( + db, + """ + SELECT s.id, s.`index` + FROM stream_registrations AS r + INNER JOIN streams AS s ON s.id = r.stream_id + WHERE r.execution_id = ?1 AND s.position = ?2 + """, + {execution_id, position} + ) + end + + defp find_paused_stream(db, step_id, workspace_id, position) do + case query_one( + db, + """ + SELECT s.id, s.`index` + FROM streams AS s + LEFT JOIN stream_closures AS c ON c.stream_id = s.id + WHERE s.step_id = ?1 AND s.workspace_id = ?2 AND s.position = ?3 + AND c.stream_id IS NULL + ORDER BY s.`index` DESC + LIMIT 1 + """, + {step_id, workspace_id, position} + ) do + {:ok, nil} -> + {:ok, nil} + + {:ok, {stream_id, index}} -> + if paused?(db, stream_id) do + {:ok, {stream_id, index}} + else + {:ok, nil} + end + end + end + + # An open stream is paused when its latest registrant completed as + # suspended. No completion means the producer is live. + defp paused?(db, stream_id) do + suspended = Results.atom_kind(:suspended) + + case query_one( + db, + """ + SELECT c.kind + FROM stream_registrations AS r + INNER JOIN executions AS e ON e.id = r.execution_id + LEFT JOIN completions AS c ON c.execution_id = r.execution_id + WHERE r.stream_id = ?1 + ORDER BY e.attempt DESC + LIMIT 1 + """, + {stream_id} + ) do + {:ok, {^suspended}} -> true + {:ok, _} -> false + end + end + + defp next_index(db, step_id) do + case query_one( + db, + "SELECT COALESCE(MAX(`index`) + 1, 0) FROM streams WHERE step_id = ?1", + {step_id} + ) do + {:ok, {index}} -> {:ok, index} + end + end + + # --- Lookup --- + + def get_stream_by_step_index(db, step_id, index) do + case query_one( + db, + "SELECT id FROM streams WHERE step_id = ?1 AND `index` = ?2", + {step_id, index} + ) do + {:ok, {id}} -> {:ok, id} + {:ok, nil} -> {:error, :not_found} + end + end + + def get_stream_id_by_key(db, run_external_id, step_number, index) do + case query_one( + db, + """ + SELECT st.id + FROM streams AS st + INNER JOIN steps AS s ON s.id = st.step_id + INNER JOIN runs AS r ON r.id = s.run_id + WHERE r.external_id = ?1 AND s.number = ?2 AND st.`index` = ?3 + """, + {run_external_id, step_number, index} + ) do + {:ok, {id}} -> {:ok, id} + {:ok, nil} -> {:error, :not_found} + end + end + + # Everything needed to describe a stream: its identity, where it lives, + # and the step it belongs to. + def get_stream(db, stream_id) do + case query_one( + db, + """ + SELECT st.step_id, st.workspace_id, st.`index`, st.position, st.created_at, + r.external_id, s.number, s.module, s.target + FROM streams AS st + INNER JOIN steps AS s ON s.id = st.step_id + INNER JOIN runs AS r ON r.id = s.run_id + WHERE st.id = ?1 + """, + {stream_id} + ) do + {:ok, nil} -> + {:error, :not_found} + + {:ok, + {step_id, workspace_id, index, position, created_at, run_external_id, step_number, module, + target}} -> + {:ok, + %{ + id: stream_id, + step_id: step_id, + workspace_id: workspace_id, index: index, - buffer: buffer, - timeout_ms: timeout_ms, - created_at: now - }) do - {:ok, _} -> {:ok, now} - {:error, "UNIQUE constraint failed: " <> _} -> {:error, :already_registered} + position: position, + created_at: created_at, + run_external_id: run_external_id, + step_number: step_number, + module: module, + target: target + }} + end + end + + def exists?(db, stream_id) do + case query_one(db, "SELECT 1 FROM streams WHERE id = ?1", {stream_id}) do + {:ok, nil} -> {:ok, false} + {:ok, {1}} -> {:ok, true} end end - # Returns the persisted buffer for a stream. Result is ``{:ok, buffer}`` - # where ``buffer`` is either an integer or ``nil`` (no backpressure). - # ``{:error, :not_found}`` if the stream doesn't exist. - def get_buffer(db, execution_id, index) do + # The config in force: the latest registration's. `{:ok, {buffer, + # timeout_ms}}`, or `{:error, :not_found}` for an unknown stream. + def get_config(db, stream_id) do case query_one( db, - "SELECT buffer FROM streams WHERE execution_id = ?1 AND `index` = ?2", - {execution_id, index} + """ + SELECT r.buffer, r.timeout_ms + FROM stream_registrations AS r + INNER JOIN executions AS e ON e.id = r.execution_id + WHERE r.stream_id = ?1 + ORDER BY e.attempt DESC + LIMIT 1 + """, + {stream_id} ) do {:ok, nil} -> {:error, :not_found} - {:ok, {buffer}} -> {:ok, buffer} + {:ok, {buffer, timeout_ms}} -> {:ok, {buffer, timeout_ms}} end end - # Returns the persisted timeout (milliseconds) for a stream. - # ``{:ok, nil}`` means no timeout. ``{:error, :not_found}`` if the - # stream doesn't exist. - def get_timeout_ms(db, execution_id, index) do + # The stream's producer: the execution behind the latest registration. + def get_producer(db, stream_id) do case query_one( db, - "SELECT timeout_ms FROM streams WHERE execution_id = ?1 AND `index` = ?2", - {execution_id, index} + """ + SELECT r.execution_id + FROM stream_registrations AS r + INNER JOIN executions AS e ON e.id = r.execution_id + WHERE r.stream_id = ?1 + ORDER BY e.attempt DESC + LIMIT 1 + """, + {stream_id} ) do {:ok, nil} -> {:error, :not_found} - {:ok, {timeout_ms}} -> {:ok, timeout_ms} + {:ok, {execution_id}} -> {:ok, execution_id} + end + end + + def registered?(db, stream_id, execution_id) do + case query_one( + db, + "SELECT 1 FROM stream_registrations WHERE stream_id = ?1 AND execution_id = ?2", + {stream_id, execution_id} + ) do + {:ok, nil} -> {:ok, false} + {:ok, {1}} -> {:ok, true} end end - # Appends an item at `sequence` to the stream. Caller supplies the sequence - # (worker-assigned, monotonic). Returns: + # Every execution that has produced into the stream, in attempt order: + # `[{execution_id, attempt, buffer, timeout_ms, created_at}, ...]`. + def get_registrations(db, stream_id) do + query( + db, + """ + SELECT r.execution_id, e.attempt, r.buffer, r.timeout_ms, r.created_at + FROM stream_registrations AS r + INNER JOIN executions AS e ON e.id = r.execution_id + WHERE r.stream_id = ?1 + ORDER BY e.attempt + """, + {stream_id} + ) + end + + # --- Items --- + + # Appends an item at `sequence`. The caller supplies the sequence + # (worker-assigned, monotonic, continuing from the head a resuming + # producer was told about). Returns: # * `{:error, :not_registered}` if the stream doesn't exist # * `{:error, :closed}` if the stream has a closure row - # * `{:error, :already_appended}` if sequence collides with an existing item - def append_item(db, execution_id, index, sequence, value) do + # * `{:error, :not_producer}` if `execution_id` isn't registered on it + # * `{:error, :already_appended}` if the sequence collides + def append_item(db, stream_id, execution_id, sequence, value) do with_transaction(db, fn -> - case has_closure?(db, execution_id, index) do - {:ok, true} -> - {:error, :closed} - - {:ok, false} -> - case exists?(db, execution_id, index) do - {:ok, false} -> - {:error, :not_registered} - - {:ok, true} -> - {:ok, value_id} = Values.get_or_create_value(db, value) - now = current_timestamp() - - case insert_one(db, :stream_items, %{ - execution_id: execution_id, - index: index, - sequence: sequence, - value_id: value_id, - created_at: now - }) do - {:ok, _} -> {:ok, now} - {:error, "UNIQUE constraint failed: " <> _} -> {:error, :already_appended} - end - end + with {:ok, true} <- exists_or(db, stream_id), + {:ok, false} <- closed_or(db, stream_id), + {:ok, true} <- registered_or(db, stream_id, execution_id) do + {:ok, value_id} = Values.get_or_create_value(db, value) + now = current_timestamp() + + case insert_one(db, :stream_items, %{ + stream_id: stream_id, + sequence: sequence, + value_id: value_id, + execution_id: execution_id, + created_at: now + }) do + {:ok, _} -> {:ok, now} + {:error, "UNIQUE constraint failed: " <> _} -> {:error, :already_appended} + end end end) end - # Closes the stream. `spec` describes *why* it closed: + defp exists_or(db, stream_id) do + case exists?(db, stream_id) do + {:ok, true} -> {:ok, true} + {:ok, false} -> {:error, :not_registered} + end + end + + defp closed_or(db, stream_id) do + case has_closure?(db, stream_id) do + {:ok, false} -> {:ok, false} + {:ok, true} -> {:error, :closed} + end + end + + defp registered_or(db, stream_id, execution_id) do + case registered?(db, stream_id, execution_id) do + {:ok, true} -> {:ok, true} + {:ok, false} -> {:error, :not_producer} + end + end + + # Fetches up to `max_items` items starting at `from_sequence`, as + # `[{sequence, value, created_at}, ...]` in sequence order. The caller + # (Server) layers stride logic on top. + def get_stream_items(db, stream_id, from_sequence, max_items) do + case query( + db, + """ + SELECT sequence, value_id, created_at + FROM stream_items + WHERE stream_id = ?1 AND sequence >= ?2 + ORDER BY sequence + LIMIT ?3 + """, + {stream_id, from_sequence, max_items} + ) do + {:ok, rows} -> + items = + Enum.map(rows, fn {sequence, value_id, created_at} -> + {:ok, value} = Values.get_value_by_id(db, value_id) + {sequence, value, created_at} + end) + + {:ok, items} + end + end + + # Highest sequence recorded for the stream, or `-1` if empty. + def get_stream_head(db, stream_id) do + case query_one( + db, + "SELECT MAX(sequence) FROM stream_items WHERE stream_id = ?1", + {stream_id} + ) do + {:ok, {nil}} -> {:ok, -1} + {:ok, {sequence}} -> {:ok, sequence} + end + end + + # The last `max_items` items in sequence order, each with the attempt + # that appended it, alongside the total item count. Used by the + # inspection topic to bootstrap its bounded tail without materialising + # the full log. + def get_stream_tail(db, stream_id, max_items) do + {:ok, {total_count}} = + query_one( + db, + "SELECT COUNT(*) FROM stream_items WHERE stream_id = ?1", + {stream_id} + ) + + case query( + db, + """ + SELECT i.sequence, i.value_id, e.attempt, i.created_at + FROM stream_items AS i + INNER JOIN executions AS e ON e.id = i.execution_id + WHERE i.stream_id = ?1 + ORDER BY i.sequence DESC + LIMIT ?2 + """, + {stream_id, max_items} + ) do + {:ok, rows} -> + items = + rows + |> Enum.reverse() + |> Enum.map(fn {sequence, value_id, attempt, created_at} -> + {:ok, value} = Values.get_value_by_id(db, value_id) + {sequence, value, attempt, created_at} + end) + + {:ok, {items, total_count}} + end + end + + # --- Closure --- + + # Closure reason codes — kept in sync with the CHECK constraint in 4.sql. + @reason_complete 0 + @reason_errored 1 + @reason_lifecycle 2 + @reason_timeout 3 + + # Closes the stream on behalf of `execution_id`. `spec` describes *why*: # - # * `:complete` — producer finished normally - # * `{:errored, type, message, frames}` — producer raised an error; the + # * `:complete` — the producer finished normally, or the step + # completed successfully with the stream still open + # * `{:errored, type, message, frames}` — the producer raised; the # error is stored via the errors table, same as Results - # * `:lifecycle` — closed implicitly because the producer execution - # ended (cancel/crash/abandon/error). No error is recorded here — - # callers that need to surface an error derive it from the - # execution's recorded result at read time. + # * `:lifecycle` — closed implicitly because an execution of the step + # ended (cancel/crash/abandon/error/recur). No error is recorded + # here — callers derive one from the closing execution's completion. # * `:timeout` — the worker closed the stream because its idle # timeout elapsed without a new item being appended. - def close_stream(db, execution_id, index, spec \\ :complete) do + def close_stream(db, stream_id, execution_id, spec) do with_transaction(db, fn -> - case exists?(db, execution_id, index) do + case exists?(db, stream_id) do {:ok, false} -> {:error, :not_registered} @@ -141,10 +472,10 @@ defmodule Coflux.Orchestration.Streams do {reason, error_id} = resolve_close_spec(db, spec) case insert_one(db, :stream_closures, %{ - execution_id: execution_id, - index: index, + stream_id: stream_id, reason: reason, error_id: error_id, + execution_id: execution_id, created_at: now }) do {:ok, _} -> {:ok, now} @@ -154,12 +485,6 @@ defmodule Coflux.Orchestration.Streams do end) end - # Closure reason codes — kept in sync with the CHECK constraint in 4.sql. - @reason_complete 0 - @reason_errored 1 - @reason_lifecycle 2 - @reason_timeout 3 - defp resolve_close_spec(_db, :complete), do: {@reason_complete, nil} defp resolve_close_spec(_db, :lifecycle), do: {@reason_lifecycle, nil} defp resolve_close_spec(_db, :timeout), do: {@reason_timeout, nil} @@ -169,68 +494,88 @@ defmodule Coflux.Orchestration.Streams do {@reason_errored, error_id} end - # Atom form of the reason integer — used by callers that want to decide - # whether to derive an error from the execution's result (:lifecycle) - # or use the stored one (:errored / :complete / :timeout). def reason_from_int(@reason_complete), do: :complete def reason_from_int(@reason_errored), do: :errored def reason_from_int(@reason_lifecycle), do: :lifecycle def reason_from_int(@reason_timeout), do: :timeout - def exists?(db, execution_id, index) do - case query_one( - db, - "SELECT 1 FROM streams WHERE execution_id = ?1 AND `index` = ?2", - {execution_id, index} - ) do + def has_closure?(db, stream_id) do + case query_one(db, "SELECT 1 FROM stream_closures WHERE stream_id = ?1", {stream_id}) do {:ok, nil} -> {:ok, false} {:ok, {1}} -> {:ok, true} end end - # Returns the stream's registration timestamp, or `{:error, :not_found}`. - def get_opened_at(db, execution_id, index) do + # Closure info, or `{:ok, nil}` if the stream is still open. Closure + # info is `{reason, error | nil, closing_execution_id, created_at}`, + # where `reason` is :complete | :errored | :lifecycle | :timeout and + # `error` is the `{type, message, frames}` triple for :errored only — + # for :lifecycle, callers derive it from the closing execution's + # completion. + def get_stream_closure(db, stream_id) do case query_one( db, - "SELECT created_at FROM streams WHERE execution_id = ?1 AND `index` = ?2", - {execution_id, index} + "SELECT reason, error_id, execution_id, created_at FROM stream_closures WHERE stream_id = ?1", + {stream_id} ) do - {:ok, nil} -> {:error, :not_found} - {:ok, {created_at}} -> {:ok, created_at} + {:ok, nil} -> + {:ok, nil} + + {:ok, {reason_int, nil, execution_id, created_at}} -> + {:ok, {reason_from_int(reason_int), nil, execution_id, created_at}} + + {:ok, {reason_int, error_id, execution_id, created_at}} -> + {:ok, error} = Errors.get_by_id(db, error_id) + {:ok, {reason_from_int(reason_int), error, execution_id, created_at}} end end - def has_closure?(db, execution_id, index) do - case query_one( + # Ids of the step's open streams in `workspace_id`, in index order. The + # lifecycle code closes these when an execution of the step ends any + # way other than suspending. + def get_open_stream_ids_for_step(db, step_id, workspace_id) do + case query( db, - "SELECT 1 FROM stream_closures WHERE execution_id = ?1 AND `index` = ?2", - {execution_id, index} + """ + SELECT s.id + FROM streams AS s + LEFT JOIN stream_closures AS c ON c.stream_id = s.id + WHERE s.step_id = ?1 AND s.workspace_id = ?2 AND c.stream_id IS NULL + ORDER BY s.`index` + """, + {step_id, workspace_id} ) do - {:ok, nil} -> {:ok, false} - {:ok, {1}} -> {:ok, true} + {:ok, rows} -> {:ok, Enum.map(rows, fn {id} -> id end)} end end - # Returns `{:ok, [index, ...]}` for every stream owned by `execution_id`, - # in index order. - def get_streams_for_execution(db, execution_id) do + # Ids of the open streams `execution_id` is registered on, in index + # order. Narrower than the step's open streams: a pending execution that + # never registered anything has none. + def get_open_stream_ids_for_execution(db, execution_id) do case query( db, - "SELECT `index` FROM streams WHERE execution_id = ?1 ORDER BY `index`", + """ + SELECT s.id + FROM stream_registrations AS r + INNER JOIN streams AS s ON s.id = r.stream_id + LEFT JOIN stream_closures AS c ON c.stream_id = s.id + WHERE r.execution_id = ?1 AND c.stream_id IS NULL + ORDER BY s.`index` + """, {execution_id} ) do - {:ok, rows} -> - {:ok, Enum.map(rows, fn {index} -> index end)} + {:ok, rows} -> {:ok, Enum.map(rows, fn {id} -> id end)} end end - # Returns a summary of how the streams owned by `execution_id` closed. - # Used by `complete_execution` to decide whether to promote a value-result - # to `:stream_errored` / `:partial`. + # Summary of the closures `execution_id` wrote. Used by + # `complete_execution` to decide whether to promote a value result to + # `:stream_errored` / `:stream_timeout`. # # Shape: `{:ok, %{errored: integer | nil, timed_out: boolean}}` - # * `errored` — the `errors.id` for the *first* errored stream closure - # (in stream-index order), or `nil` if none errored + # * `errored` — the `errors.id` for the *first* errored closure (in + # stream-index order), or `nil` if none errored # * `timed_out` — true if any stream closed via idle timeout # # Lifecycle / complete closures are ignored: the former inherit the @@ -239,10 +584,11 @@ defmodule Coflux.Orchestration.Streams do case query( db, """ - SELECT reason, error_id - FROM stream_closures - WHERE execution_id = ?1 - ORDER BY `index` + SELECT c.reason, c.error_id + FROM stream_closures AS c + INNER JOIN streams AS s ON s.id = c.stream_id + WHERE c.execution_id = ?1 + ORDER BY s.`index` """, {execution_id} ) do @@ -260,165 +606,151 @@ defmodule Coflux.Orchestration.Streams do end end - # Returns indexes of streams owned by `execution_id` that don't yet have - # a closure row. Used by the lifecycle code to discover which streams to - # close on completion / cancel / crash. - def get_open_streams_for_execution(db, execution_id) do + # --- Run/topic views --- + + # One map per stream belonging to any step of the run, for the run + # topic's initial state. Closure reasons are returned raw (`:lifecycle` + # unresolved) — the server resolves them against the closing execution. + def get_streams_for_run(db, run_id) do case query( db, """ - SELECT s.`index` - FROM streams AS s - LEFT JOIN stream_closures AS c - ON c.execution_id = s.execution_id AND c.`index` = s.`index` - WHERE s.execution_id = ?1 AND c.execution_id IS NULL - ORDER BY s.`index` + SELECT st.id, st.step_id, s.number, r.external_id, st.workspace_id, st.`index`, + st.position, st.created_at, c.created_at, c.reason, c.error_id, c.execution_id + FROM streams AS st + INNER JOIN steps AS s ON s.id = st.step_id + INNER JOIN runs AS r ON r.id = s.run_id + LEFT JOIN stream_closures AS c ON c.stream_id = st.id + WHERE s.run_id = ?1 + ORDER BY s.number, st.`index` """, - {execution_id} + {run_id} ) do {:ok, rows} -> - {:ok, Enum.map(rows, fn {index} -> index end)} + streams = + Enum.map(rows, fn {id, step_id, step_number, run_external_id, workspace_id, index, + position, created_at, closed_at, reason_int, error_id, closed_by} -> + error = if error_id, do: get_error(db, error_id) + + %{ + id: id, + step_id: step_id, + step_number: step_number, + run_external_id: run_external_id, + workspace_id: workspace_id, + index: index, + position: position, + created_at: created_at, + closed_at: closed_at, + reason: if(reason_int, do: reason_from_int(reason_int)), + error: error, + closed_by: closed_by + } + end) + + {:ok, streams} end end - # Returns closure info or `{:ok, nil}` if the stream is still open. - # Closure info: `{reason, error | nil, created_at}` where - # * reason is :complete | :errored | :lifecycle - # * error is the `{type, message, frames}` triple for :errored, nil - # otherwise (callers derive it from the execution's result on - # :lifecycle) - def get_stream_closure(db, execution_id, index) do - case query_one( - db, - "SELECT reason, error_id, created_at FROM stream_closures WHERE execution_id = ?1 AND `index` = ?2", - {execution_id, index} - ) do - {:ok, nil} -> - {:ok, nil} + defp get_error(db, error_id) do + {:ok, error} = Errors.get_by_id(db, error_id) + error + end - {:ok, {reason_int, nil, created_at}} -> - {:ok, {reason_from_int(reason_int), nil, created_at}} + # --- Refs and lineage --- - {:ok, {reason_int, error_id, created_at}} -> - {:ok, error} = Errors.get_by_id(db, error_id) - {:ok, {reason_from_int(reason_int), error, created_at}} - end - end + def get_or_create_stream_ref(db, run_external_id, step_number, index, module, target) do + {:ok, _} = + insert_one( + db, + :stream_refs, + %{ + run_external_id: run_external_id, + step_number: step_number, + index: index, + module: module, + target: target + }, + on_conflict: "DO NOTHING" + ) - # Fetches up to `max_items` items from the stream starting at `from_sequence`. - # Returns `{:ok, [{sequence, value, created_at}, ...]}` in sequence order. - # The caller (Server) layers filter logic (slice / partition) on top of this. - def get_stream_items(db, execution_id, index, from_sequence, max_items) do - case query( + case query_one( db, """ - SELECT sequence, value_id, created_at - FROM stream_items - WHERE execution_id = ?1 AND `index` = ?2 AND sequence >= ?3 - ORDER BY sequence - LIMIT ?4 + SELECT id + FROM stream_refs + WHERE run_external_id = ?1 AND step_number = ?2 AND `index` = ?3 """, - {execution_id, index, from_sequence, max_items} + {run_external_id, step_number, index} ) do - {:ok, rows} -> - items = - Enum.map(rows, fn {sequence, value_id, created_at} -> - {:ok, value} = Values.get_value_by_id(db, value_id) - {sequence, value, created_at} - end) - - {:ok, items} + {:ok, {id}} -> {:ok, id} end end - # Returns one row per stream owned by `execution_id`: - # `{index, buffer, timeout_ms, created_at, closed_at | nil, reason | nil, error | nil}`. - # * buffer is the persisted backpressure budget (integer or nil) - # * timeout_ms is the persisted idle-timeout budget (integer or nil) - # * reason is :complete | :errored | :lifecycle | :timeout when closed, nil when open - # * error is the stored `{type, message, frames}` triple for :errored - # closures only — callers that need to surface an error for a - # :lifecycle closure derive it from the execution's result. - # Used when populating the topic state for a run. - def get_streams_with_closures_for_execution(db, execution_id) do - case query( + def get_stream_ref(db, ref_id) do + case query_one( db, - """ - SELECT s.`index`, s.buffer, s.timeout_ms, s.created_at, c.created_at, c.reason, c.error_id - FROM streams AS s - LEFT JOIN stream_closures AS c - ON c.execution_id = s.execution_id AND c.`index` = s.`index` - WHERE s.execution_id = ?1 - ORDER BY s.`index` - """, - {execution_id} + "SELECT run_external_id, step_number, `index`, module, target FROM stream_refs WHERE id = ?1", + {ref_id} ) do - {:ok, rows} -> - streams = - Enum.map(rows, fn - {index, buffer, timeout_ms, created_at, nil, nil, nil} -> - {index, buffer, timeout_ms, created_at, nil, nil, nil} - - {index, buffer, timeout_ms, created_at, closed_at, reason_int, nil} -> - {index, buffer, timeout_ms, created_at, closed_at, reason_from_int(reason_int), nil} - - {index, buffer, timeout_ms, created_at, closed_at, reason_int, error_id} -> - {:ok, error} = Errors.get_by_id(db, error_id) - - {index, buffer, timeout_ms, created_at, closed_at, reason_from_int(reason_int), - error} - end) + {:ok, {run_external_id, step_number, index, module, target}} -> + {:ok, {run_external_id, step_number, index, module, target}} - {:ok, streams} + {:ok, nil} -> + {:error, :not_found} end end - # Returns the highest sequence recorded for the stream, or `-1` if empty. - # Used by the worker protocol to report "head" for flow control without - # requiring the caller to scan all items. - def get_stream_head(db, execution_id, index) do - case query_one( - db, - "SELECT MAX(sequence) FROM stream_items WHERE execution_id = ?1 AND `index` = ?2", - {execution_id, index} - ) do - {:ok, {nil}} -> {:ok, -1} - {:ok, {sequence}} -> {:ok, sequence} + def create_stream_ref_for(db, stream_id) do + case get_stream(db, stream_id) do + {:ok, stream} -> + get_or_create_stream_ref( + db, + stream.run_external_id, + stream.step_number, + stream.index, + stream.module, + stream.target + ) + + {:error, :not_found} -> + {:error, :not_found} end end - # Returns the last `max_items` items of the stream, in sequence order, - # alongside the total item count. Used by the inspection topic to - # bootstrap its bounded tail buffer without materialising the full log. - def get_stream_tail(db, execution_id, index, max_items) do - {:ok, {total_count}} = - query_one( + # Records that `execution_id` subscribed to the stream. Returns + # `{:ok, id}` for a new edge, `{:ok, nil}` if it already existed. + def record_dependency(db, execution_id, stream_ref_id) do + with_transaction(db, fn -> + insert_one( db, - "SELECT COUNT(*) FROM stream_items WHERE execution_id = ?1 AND `index` = ?2", - {execution_id, index} + :stream_dependencies, + %{ + execution_id: execution_id, + stream_ref_id: stream_ref_id, + created_at: current_timestamp() + }, + on_conflict: "DO NOTHING" ) + end) + end + # `%{execution_id => [stream_ref_id, ...]}` for every consumer execution + # in the run. + def get_run_dependencies(db, run_id) do case query( db, """ - SELECT sequence, value_id, created_at - FROM stream_items - WHERE execution_id = ?1 AND `index` = ?2 - ORDER BY sequence DESC - LIMIT ?3 + SELECT d.execution_id, d.stream_ref_id + FROM stream_dependencies AS d + INNER JOIN executions AS e ON e.id = d.execution_id + INNER JOIN steps AS s ON s.id = e.step_id + WHERE s.run_id = ?1 """, - {execution_id, index, max_items} + {run_id} ) do {:ok, rows} -> - items = - rows - |> Enum.reverse() - |> Enum.map(fn {sequence, value_id, created_at} -> - {:ok, value} = Values.get_value_by_id(db, value_id) - {sequence, value, created_at} - end) - - {:ok, {items, total_count}} + {:ok, Enum.group_by(rows, &elem(&1, 0), &elem(&1, 1))} end end diff --git a/server/lib/coflux/topics/run.ex b/server/lib/coflux/topics/run.ex index 0d64ae0f..5b7418d8 100644 --- a/server/lib/coflux/topics/run.ex +++ b/server/lib/coflux/topics/run.ex @@ -65,6 +65,7 @@ defmodule Coflux.Topics.Run do createdAt: step.created_at, arguments: Enum.map(step.arguments, &build_value/1), requires: step.requires, + streams: %{}, executions: %{} }) else @@ -105,7 +106,6 @@ defmodule Coflux.Topics.Run do inputs: %{}, result: nil, metrics: %{}, - streams: %{}, # Nothing has run yet, so what the execution will start from is also # what it currently holds. checkpoints: build_checkpoints(checkpoints, checkpoints) @@ -183,26 +183,13 @@ defmodule Coflux.Topics.Run do defp process_notification( topic, - {:stream_dependency, execution_external_id, producer_execution_id, index, - producer_metadata} + {:stream_dependency, execution_external_id, stream_id, module, target} ) do - dependency = %{ - type: "stream", - execution: build_execution(producer_metadata), - index: index - } + dependency = %{type: "stream", streamId: stream_id, module: module, target: target} - update_execution( - topic, - execution_external_id, - fn topic, base_path -> - Topic.merge( - topic, - base_path ++ [:dependencies, "#{producer_execution_id}:#{index}"], - dependency - ) - end - ) + update_execution(topic, execution_external_id, fn topic, base_path -> + Topic.merge(topic, base_path ++ [:dependencies, stream_id], dependency) + end) end defp process_notification(topic, {:child, parent_execution_external_id, child}) do @@ -249,34 +236,69 @@ defmodule Coflux.Topics.Run do end) end - defp process_notification( - topic, - {:stream_opened, execution_external_id, index, buffer, timeout_ms, created_at} - ) do - update_execution(topic, execution_external_id, fn topic, base_path -> - Topic.set(topic, base_path ++ [:streams, Integer.to_string(index)], %{ - buffer: buffer, - timeoutMs: timeout_ms, - openedAt: created_at, - closedAt: nil, - reason: nil, - error: nil - }) - end) + # An execution registered on one of the step's streams — opening it, or + # resuming it after a suspend (`continued`). Streams live under the step + # in the topic, with the attempts that have produced into each. + defp process_notification(topic, {:stream_registered, step_number, index, info}) do + step_key = "#{topic.state.external_run_id}:#{step_number}" + index_key = Integer.to_string(index) + path = [:steps, step_key, :streams, index_key] + + cond do + info.workspace_id not in topic.state.workspace_ids -> + topic + + not Map.has_key?(topic.value.steps, step_key) -> + topic + + true -> + case get_in(topic.value, path) do + nil -> + Topic.set(topic, path, %{ + id: info.id, + index: index, + position: info.position, + workspaceId: info.workspace_id, + buffer: info.buffer, + timeoutMs: info.timeout_ms, + openedAt: info.opened_at, + attempts: [info.attempt], + closedAt: nil, + closedBy: nil, + reason: nil, + error: nil + }) + + existing -> + attempts = + if info.attempt in existing.attempts, + do: existing.attempts, + else: existing.attempts ++ [info.attempt] + + topic + |> Topic.set(path ++ [:attempts], attempts) + |> Topic.set(path ++ [:buffer], info.buffer) + |> Topic.set(path ++ [:timeoutMs], info.timeout_ms) + end + end end defp process_notification( topic, - {:stream_closed, execution_external_id, index, reason, error, closed_at} + {:stream_closed, step_number, index, reason, error, attempt, closed_at} ) do - index_key = Integer.to_string(index) + step_key = "#{topic.state.external_run_id}:#{step_number}" + path = [:steps, step_key, :streams, Integer.to_string(index)] - update_execution(topic, execution_external_id, fn topic, base_path -> + if get_in(topic.value, path) do topic - |> Topic.set(base_path ++ [:streams, index_key, :closedAt], closed_at) - |> Topic.set(base_path ++ [:streams, index_key, :reason], reason) - |> Topic.set(base_path ++ [:streams, index_key, :error], error) - end) + |> Topic.set(path ++ [:closedAt], closed_at) + |> Topic.set(path ++ [:closedBy], attempt) + |> Topic.set(path ++ [:reason], reason) + |> Topic.set(path ++ [:error], error) + else + topic + end end defp process_notification( @@ -414,6 +436,7 @@ defmodule Coflux.Topics.Run do createdAt: step.created_at, arguments: Enum.map(step.arguments, &build_value/1), requires: step.requires, + streams: build_streams(step.streams, workspace_ids), executions: step.executions |> Enum.filter(fn {_, execution} -> @@ -455,7 +478,6 @@ defmodule Coflux.Topics.Run do upper: def_data.upper }} end), - streams: build_streams(execution.streams), checkpoints: build_checkpoints( execution.checkpoints.before, @@ -479,13 +501,8 @@ defmodule Coflux.Topics.Run do {id, {:asset, asset}} -> {id, %{type: "asset", assetId: id, asset: build_asset(asset)}} - {id, {:stream, index, execution}} -> - {id, - %{ - type: "stream", - execution: build_execution(execution), - index: index - }} + {id, {:stream, stream_id, module, target}} -> + {id, %{type: "stream", streamId: stream_id, module: module, target: target}} end) end @@ -618,43 +635,36 @@ defmodule Coflux.Topics.Run do Map.new(values, fn {name, value} -> {name, build_value(value)} end) end - defp build_streams(streams) do - Map.new(streams, fn - {index, buffer, timeout_ms, opened_at, nil, nil, nil} -> - {Integer.to_string(index), - %{ - buffer: buffer, - timeoutMs: timeout_ms, - openedAt: opened_at, - closedAt: nil, - reason: nil, - error: nil - }} - - {index, buffer, timeout_ms, opened_at, closed_at, reason, nil} -> - {Integer.to_string(index), - %{ - buffer: buffer, - timeoutMs: timeout_ms, - openedAt: opened_at, - closedAt: closed_at, - reason: Atom.to_string(reason), - error: nil - }} - - {index, buffer, timeout_ms, opened_at, closed_at, reason, {type, message, frames}} -> - {Integer.to_string(index), - %{ - buffer: buffer, - timeoutMs: timeout_ms, - openedAt: opened_at, - closedAt: closed_at, - reason: Atom.to_string(reason), - error: %{type: type, message: message, frames: build_frames(frames)} - }} + # Streams belong to the step. Only those in a workspace the topic is + # showing are included — a re-run in another workspace opens its own. + defp build_streams(streams, workspace_ids) do + streams + |> Enum.filter(fn {_index, stream} -> stream.workspace_id in workspace_ids end) + |> Map.new(fn {index, stream} -> + {Integer.to_string(index), + %{ + id: stream.id, + index: stream.index, + position: stream.position, + workspaceId: stream.workspace_id, + buffer: stream.buffer, + timeoutMs: stream.timeout_ms, + openedAt: stream.opened_at, + attempts: stream.attempts, + closedAt: stream.closed_at, + closedBy: stream.closed_by, + reason: if(stream.reason, do: Atom.to_string(stream.reason)), + error: build_stream_error(stream.error) + }} end) end + defp build_stream_error(nil), do: nil + + defp build_stream_error({type, message, frames}) do + %{type: type, message: message, frames: build_frames(frames)} + end + defp execution_attempt({ext_id, _module, _target}) do ext_id |> String.split(":") |> List.last() |> String.to_integer() end diff --git a/server/lib/coflux/topics/stream.ex b/server/lib/coflux/topics/stream.ex index 331fc14e..9698e24a 100644 --- a/server/lib/coflux/topics/stream.ex +++ b/server/lib/coflux/topics/stream.ex @@ -1,10 +1,11 @@ defmodule Coflux.Topics.Stream do @moduledoc """ - Inspection topic for a single stream, keyed by the stream's opaque id - (``_``). Used by the Studio UI when a - user opens a stream dialog — the topic keeps a bounded tail of items - (with resolved values) plus closure state, and receives live updates - as items are appended or the stream is closed. + Inspection topic for a single stream, keyed by the stream's id + (``:_``). Used by the Studio UI when a user opens a + stream dialog — the topic keeps a bounded tail of items (with resolved + values) plus closure state, and receives live updates as items are + appended, as executions register on the stream (opening it, or resuming + it after a suspend), or as the stream is closed. """ use Topical.Topic, route: ["streams", :id] @@ -17,32 +18,23 @@ defmodule Coflux.Topics.Stream do def init(params) do project_id = Map.fetch!(params, :project) + id = Map.fetch!(params, :id) - case parse_id(Map.fetch!(params, :id)) do - {:ok, execution_id, index} -> - do_init(project_id, execution_id, index) - - :error -> - {:error, :not_found} - end - end - - defp do_init(project_id, execution_id, index) do - case Orchestration.subscribe_stream_topic( - project_id, - execution_id, - index, - self() - ) do + case Orchestration.subscribe_stream_topic(project_id, id, self()) do {:ok, initial, ref} -> {:ok, Topic.new( %{ - producer: initial.producer, + id: initial.id, + step: initial.step, + workspaceId: initial.workspaceId, + index: initial.index, + position: initial.position, buffer: initial.buffer, timeoutMs: initial.timeoutMs, openedAt: initial.openedAt, - closure: build_closure(initial.closure), + attempts: initial.attempts, + closure: initial.closure, items: Enum.map(initial.items, &build_item/1), totalCount: initial.totalCount, tailSize: initial.tailSize @@ -60,10 +52,10 @@ defmodule Coflux.Topics.Stream do {:ok, topic} end - defp process_notification({:item_appended, sequence, value, created_at}, topic) do + defp process_notification({:item_appended, sequence, value, attempt, created_at}, topic) do tail_size = topic.state.tail_size || 200 - item = build_item({sequence, value, created_at}) + item = build_item({sequence, value, attempt, created_at}) existing = topic.value.items # Keep items bounded: drop the head once we're at capacity. @@ -80,49 +72,32 @@ defmodule Coflux.Topics.Stream do |> Topic.set([:totalCount], topic.value.totalCount + 1) end - defp process_notification({:closed, reason, error, closed_at}, topic) do - closure = %{reason: reason, error: error, closedAt: closed_at} - Topic.set(topic, [:closure], closure) - end + # An execution registered on the stream — the latest registration's + # config is the one in force. + defp process_notification({:registered, attempt, buffer, timeout_ms, _created_at}, topic) do + attempts = topic.value.attempts - defp build_item({sequence, value, created_at}) do - %{ - sequence: sequence, - value: TopicUtils.build_value(value), - createdAt: created_at - } + topic + |> Topic.set([:attempts], if(attempt in attempts, do: attempts, else: attempts ++ [attempt])) + |> Topic.set([:buffer], buffer) + |> Topic.set([:timeoutMs], timeout_ms) end - defp build_closure(nil), do: nil - - defp build_closure(%{reason: reason, error: error, closedAt: closed_at}) do - %{ + defp process_notification({:closed, reason, error, attempt, closed_at}, topic) do + Topic.set(topic, [:closure], %{ reason: reason, error: error, + attempt: attempt, closedAt: closed_at - } + }) end - # Split an opaque stream id back into (execution_id, index). The - # separator is `_` — execution ids use alphanumerics + `:`, so the last - # `_` unambiguously marks the index suffix. - defp parse_id(id) when is_binary(id) do - case String.split(id, "_") do - parts when length(parts) >= 2 -> - {index_str, execution_parts} = List.pop_at(parts, -1) - - with {index, ""} when index >= 0 <- Integer.parse(index_str), - execution_id when execution_id != "" <- - Enum.join(execution_parts, "_") do - {:ok, execution_id, index} - else - _ -> :error - end - - _ -> - :error - end + defp build_item({sequence, value, attempt, created_at}) do + %{ + sequence: sequence, + value: TopicUtils.build_value(value), + attempt: attempt, + createdAt: created_at + } end - - defp parse_id(_), do: :error end diff --git a/server/priv/migrations/orchestration/4.sql b/server/priv/migrations/orchestration/4.sql index bb8c87ed..3dc8458f 100644 --- a/server/priv/migrations/orchestration/4.sql +++ b/server/priv/migrations/orchestration/4.sql @@ -127,91 +127,138 @@ ALTER TABLE workflows ADD COLUMN streams_timeout_ms INTEGER; ALTER TABLE steps ADD COLUMN streams_buffer INTEGER; ALTER TABLE steps ADD COLUMN streams_timeout_ms INTEGER; --- Streams — ordered, append-only sequences of values produced by an --- execution. Each stream is identified by (execution_id, index), where --- `index` is assigned monotonically by the worker when serialising the --- execution's return value. The worker manages allocation locally, so --- no server round-trip is needed to mint an id. The column is quoted --- with backticks throughout because INDEX is a SQLite keyword. +-- Streams — ordered, append-only sequences of values produced by a step. +-- +-- A stream belongs to a step within a workspace, not to an execution. An +-- execution that suspends leaves its open streams *paused*, and the +-- execution that resumes the step appends to them, so consumers see one +-- unbroken sequence across the suspension. Every other way an execution +-- can end closes the step's open streams (see the lifecycle code in +-- `Server`), so a retry, a recurrence or a re-run after a normal completion +-- opens fresh streams. +-- +-- `index` is allocated by the server per step, monotonically from 0, which +-- makes a stream's id `:_` unique across attempts and +-- workspaces. `position` is the order in which the opening execution +-- registered the stream (0 for a generator-bodied task's own stream). It is +-- how a resuming execution's registrations are matched back onto the paused +-- streams: its k-th registration continues the paused stream whose opener +-- registered it k-th. That relies on the same determinism suspend already +-- requires of the code before the suspend point. -- -- Invariants: --- • A stream is owned by exactly one execution (its producer). -- • stream_items are append-only with monotonic sequence starting at 0. --- • stream_closures are terminal — no items may be appended after closure. --- • On execution completion / cancellation / crash, every owned stream --- that lacks a closure receives one (clean, cancelled, or crashed). --- • Re-running a producer execution creates fresh streams (new attempt ⇒ --- new execution_id ⇒ new rows). Consumer references are concrete to --- the original streams. +-- Each records the execution that appended it; a stream that spans a +-- suspension has items from several attempts. +-- • stream_closures are terminal — no items may be appended after +-- closure. Each records the execution that closed the stream, or whose +-- completion did. +-- • A stream is appended to by one live execution at a time, always in +-- the stream's own workspace. Consumers may read from any workspace. +-- • stream_registrations records every execution that produced into a +-- stream, with the config it registered. The latest registration's +-- config is the one in force; the latest registrant is the producer. -- • Consumer cursors are kept in-memory only; re-run consumers subscribe -- fresh from sequence 0. +-- +-- The column is quoted with backticks throughout because INDEX is a SQLite +-- keyword. CREATE TABLE streams ( - execution_id INTEGER NOT NULL, + id INTEGER PRIMARY KEY, + step_id INTEGER NOT NULL, + workspace_id INTEGER NOT NULL, `index` INTEGER NOT NULL, + position INTEGER NOT NULL, + created_at INTEGER NOT NULL, + UNIQUE (step_id, `index`), + FOREIGN KEY (step_id) REFERENCES steps ON DELETE CASCADE, + FOREIGN KEY (workspace_id) REFERENCES workspaces ON DELETE RESTRICT +) STRICT; + +-- Finding the paused stream a resuming execution should continue is a +-- lookup by step, workspace and position. +CREATE INDEX idx_streams_step_workspace_position ON streams(step_id, workspace_id, position); + +CREATE TABLE stream_registrations ( + stream_id INTEGER NOT NULL, + execution_id INTEGER NOT NULL, -- Producer-side backpressure budget. NULL opts out of flow control - -- (producer emits freely). Integer N means the producer may run up - -- to N items ahead of the slowest consumer's acknowledged position; - -- N=0 is strict lockstep. - -- Persisted so the server can reconstruct per-stream flow-control - -- state on restart and so Studio can display the configuration. + -- (producer emits freely). Integer N means the producer may run up to N + -- items ahead of the slowest consumer's acknowledged position; N=0 is + -- strict lockstep. buffer INTEGER, -- Idle-timeout budget in milliseconds. NULL disables the timeout. - -- Enforced at the worker (CLI) level; persisted here only so Studio - -- can display the configured value. + -- Enforced at the worker (CLI) level, per execution; persisted here so + -- Studio can display the configured value. timeout_ms INTEGER, created_at INTEGER NOT NULL, - PRIMARY KEY (execution_id, `index`), + PRIMARY KEY (stream_id, execution_id), + FOREIGN KEY (stream_id) REFERENCES streams ON DELETE CASCADE, FOREIGN KEY (execution_id) REFERENCES executions ON DELETE CASCADE ) STRICT; +CREATE INDEX idx_stream_registrations_execution_id ON stream_registrations(execution_id); + CREATE TABLE stream_items ( - execution_id INTEGER NOT NULL, - `index` INTEGER NOT NULL, + stream_id INTEGER NOT NULL, sequence INTEGER NOT NULL, value_id INTEGER NOT NULL, + execution_id INTEGER NOT NULL, created_at INTEGER NOT NULL, - PRIMARY KEY (execution_id, `index`, sequence), - FOREIGN KEY (execution_id, `index`) REFERENCES streams (execution_id, `index`) ON DELETE CASCADE, - FOREIGN KEY (value_id) REFERENCES values_ ON DELETE RESTRICT + PRIMARY KEY (stream_id, sequence), + FOREIGN KEY (stream_id) REFERENCES streams ON DELETE CASCADE, + FOREIGN KEY (value_id) REFERENCES values_ ON DELETE RESTRICT, + FOREIGN KEY (execution_id) REFERENCES executions ON DELETE CASCADE ) STRICT; -- Closure of a stream. `reason` records *why* it closed: --- 0 = complete — producer finished normally (no error) --- 1 = errored — producer raised an error (stored in errors via error_id) --- 2 = lifecycle — closed implicitly because the producer execution ended --- (cancel/crash/abandon/error). The specific error is --- derived on read by looking up the execution's completion, --- so we don't duplicate that state here. +-- 0 = complete — the producer finished normally, or the step completed +-- with the stream still open (no error) +-- 1 = errored — the producer raised an error (stored in errors via +-- error_id) +-- 2 = lifecycle — closed implicitly because an execution of the step +-- ended (cancel/crash/abandon/error/recur). The specific +-- reason is derived on read from the closing execution's +-- completion, so we don't duplicate that state here. -- 3 = timeout — closed by the worker because the configured idle -- timeout elapsed without a new item being appended. +-- `execution_id` is the execution that closed the stream, or whose +-- completion did. CREATE TABLE stream_closures ( - execution_id INTEGER NOT NULL, - `index` INTEGER NOT NULL, + stream_id INTEGER PRIMARY KEY, reason INTEGER NOT NULL, error_id INTEGER, + execution_id INTEGER NOT NULL, created_at INTEGER NOT NULL, - PRIMARY KEY (execution_id, `index`), - FOREIGN KEY (execution_id, `index`) REFERENCES streams (execution_id, `index`) ON DELETE CASCADE, + FOREIGN KEY (stream_id) REFERENCES streams ON DELETE CASCADE, FOREIGN KEY (error_id) REFERENCES errors ON DELETE RESTRICT, + FOREIGN KEY (execution_id) REFERENCES executions ON DELETE CASCADE, CHECK ((reason = 1) = (error_id IS NOT NULL)) ) STRICT; --- Track stream subscriptions as a lineage edge between executions, mirroring --- result_dependencies / asset_dependencies. A row is written when a consumer --- subscribes to a producer's stream (regardless of whether items are read), --- so data lineage is preserved even for subscriptions that yield no values. --- --- The producer side is referenced via `execution_refs` (not the live --- `executions` row) so the edge survives epoch rotation, and by `stream_index` --- so we can distinguish between multiple streams produced by the same --- execution. +-- Stable reference to a stream that survives epoch rotation (the stream's +-- run may live in an older epoch's database), mirroring execution_refs. +CREATE TABLE stream_refs ( + id INTEGER PRIMARY KEY, + run_external_id TEXT NOT NULL, + step_number INTEGER NOT NULL, + `index` INTEGER NOT NULL, + module TEXT, + target TEXT, + UNIQUE (run_external_id, step_number, `index`) +) STRICT; + +-- Track stream subscriptions as a lineage edge between a consumer execution +-- and a stream, mirroring result_dependencies / asset_dependencies. A row +-- is written when a consumer subscribes (regardless of whether items are +-- read), so data lineage is preserved even for subscriptions that yield no +-- values. CREATE TABLE stream_dependencies ( execution_id INTEGER NOT NULL, stream_ref_id INTEGER NOT NULL, - stream_index INTEGER NOT NULL, created_at INTEGER NOT NULL, - PRIMARY KEY (execution_id, stream_ref_id, stream_index), + PRIMARY KEY (execution_id, stream_ref_id), FOREIGN KEY (execution_id) REFERENCES executions ON DELETE CASCADE, - FOREIGN KEY (stream_ref_id) REFERENCES execution_refs ON DELETE RESTRICT + FOREIGN KEY (stream_ref_id) REFERENCES stream_refs ON DELETE RESTRICT ) STRICT; diff --git a/server/test/coflux/streams_test.exs b/server/test/coflux/streams_test.exs new file mode 100644 index 00000000..9378776c --- /dev/null +++ b/server/test/coflux/streams_test.exs @@ -0,0 +1,254 @@ +defmodule Coflux.StreamsTest do + use ExUnit.Case, async: true + + alias Coflux.Orchestration.{Results, Streams} + alias Coflux.Store.Migrations + alias Exqlite.Sqlite3 + + @base_ws 1 + @child_ws 2 + @step 1 + + setup do + {:ok, db} = Sqlite3.open(":memory:") + :ok = Migrations.run(db, "orchestration") + + create_workspace(db, @base_ws, "base") + create_workspace(db, @child_ws, "child") + create_run(db, 1, "r1") + create_step(db, @step, 1, 0) + + {:ok, db: db} + end + + defp val(data), do: {:raw, data, []} + + describe "register/7" do + test "opens a new stream per position with step-allocated indexes", %{db: db} do + create_execution(db, 1, @step, 1, @base_ws) + + assert {:ok, %{index: 0, head: -1, continued: false} = first} = register(db, 1, 0) + assert {:ok, %{index: 1, head: -1, continued: false} = second} = register(db, 1, 1) + assert first.id != second.id + + assert {:ok, ^first} = + Streams.register(db, @step, @base_ws, 1, 0, 0, nil) |> strip_created() + end + + test "resumes a stream paused by a suspend, continuing the sequence", %{db: db} do + create_execution(db, 1, @step, 1, @base_ws) + {:ok, %{id: stream_id}} = register(db, 1, 0) + {:ok, _} = Streams.append_item(db, stream_id, 1, 0, val("a")) + {:ok, _} = Streams.append_item(db, stream_id, 1, 1, val("b")) + suspend(db, 1, 2) + + create_execution(db, 2, @step, 2, @base_ws) + assert {:ok, %{id: ^stream_id, index: 0, head: 1, continued: true}} = register(db, 2, 0) + + assert {:ok, _} = Streams.append_item(db, stream_id, 2, 2, val("c")) + + assert {:ok, [{0, _, _}, {1, _, _}, {2, _, _}]} = + Streams.get_stream_items(db, stream_id, 0, 10) + + assert {:ok, {0, nil}} = Streams.get_config(db, stream_id) + assert {:ok, 2} = Streams.get_producer(db, stream_id) + end + + test "the latest registration's config is the one in force", %{db: db} do + create_execution(db, 1, @step, 1, @base_ws) + {:ok, %{id: stream_id}} = Streams.register(db, @step, @base_ws, 1, 0, 0, nil) + suspend(db, 1, 2) + + create_execution(db, 2, @step, 2, @base_ws) + + {:ok, %{id: ^stream_id, continued: true}} = + Streams.register(db, @step, @base_ws, 2, 0, 5, 1000) + + assert {:ok, {5, 1000}} = Streams.get_config(db, stream_id) + end + + test "does not continue a stream whose producer is still live", %{db: db} do + create_execution(db, 1, @step, 1, @base_ws) + {:ok, %{id: first}} = register(db, 1, 0) + + create_execution(db, 2, @step, 2, @base_ws) + assert {:ok, %{index: 1, continued: false} = second} = register(db, 2, 0) + assert second.id != first + end + + test "does not continue a closed stream", %{db: db} do + create_execution(db, 1, @step, 1, @base_ws) + {:ok, %{id: first}} = register(db, 1, 0) + {:ok, _} = Streams.close_stream(db, first, 1, :complete) + suspend(db, 1, 2) + + create_execution(db, 2, @step, 2, @base_ws) + assert {:ok, %{index: 1, continued: false} = second} = register(db, 2, 0) + assert second.id != first + end + + test "does not continue a stream from another workspace", %{db: db} do + create_execution(db, 1, @step, 1, @base_ws) + {:ok, %{id: first}} = register(db, 1, 0) + suspend(db, 1, 2) + + create_execution(db, 2, @step, 2, @child_ws) + + assert {:ok, %{index: 1, continued: false} = second} = + Streams.register(db, @step, @child_ws, 2, 0, 0, nil) + + assert second.id != first + end + + test "matches paused streams by the opener's registration position", %{db: db} do + create_execution(db, 1, @step, 1, @base_ws) + {:ok, %{id: s0}} = register(db, 1, 0) + {:ok, %{id: s1}} = register(db, 1, 1) + {:ok, %{id: s2}} = register(db, 1, 2) + {:ok, _} = Streams.close_stream(db, s1, 1, :complete) + suspend(db, 1, 2) + + create_execution(db, 2, @step, 2, @base_ws) + assert {:ok, %{id: ^s0, continued: true}} = register(db, 2, 0) + assert {:ok, %{index: 3, continued: false} = fresh} = register(db, 2, 1) + assert fresh.id != s1 + assert {:ok, %{id: ^s2, continued: true}} = register(db, 2, 2) + end + + test "survives a resuming execution that suspends before reaching later positions", %{db: db} do + create_execution(db, 1, @step, 1, @base_ws) + {:ok, %{id: s0}} = register(db, 1, 0) + {:ok, %{id: s1}} = register(db, 1, 1) + suspend(db, 1, 2) + + create_execution(db, 2, @step, 2, @base_ws) + {:ok, %{id: ^s0, continued: true}} = register(db, 2, 0) + suspend(db, 2, 3) + + create_execution(db, 3, @step, 3, @base_ws) + assert {:ok, %{id: ^s0, continued: true}} = register(db, 3, 0) + assert {:ok, %{id: ^s1, continued: true}} = register(db, 3, 1) + end + end + + describe "append_item/5" do + test "rejects appends from an execution that isn't registered", %{db: db} do + create_execution(db, 1, @step, 1, @base_ws) + create_execution(db, 2, @step, 2, @base_ws) + {:ok, %{id: stream_id}} = register(db, 1, 0) + + assert {:error, :not_producer} = Streams.append_item(db, stream_id, 2, 0, val("x")) + end + + test "rejects appends after closure and duplicate sequences", %{db: db} do + create_execution(db, 1, @step, 1, @base_ws) + {:ok, %{id: stream_id}} = register(db, 1, 0) + {:ok, _} = Streams.append_item(db, stream_id, 1, 0, val("x")) + + assert {:error, :already_appended} = Streams.append_item(db, stream_id, 1, 0, val("y")) + {:ok, _} = Streams.close_stream(db, stream_id, 1, :complete) + assert {:error, :closed} = Streams.append_item(db, stream_id, 1, 1, val("z")) + assert {:error, :already_closed} = Streams.close_stream(db, stream_id, 1, :complete) + end + end + + describe "open streams" do + test "distinguishes the step's open streams from an execution's own", %{db: db} do + create_execution(db, 1, @step, 1, @base_ws) + {:ok, %{id: s0}} = register(db, 1, 0) + {:ok, %{id: s1}} = register(db, 1, 1) + suspend(db, 1, 2) + + # The pending successor never registered anything. + create_execution(db, 2, @step, 2, @base_ws) + + assert {:ok, [^s0, ^s1]} = Streams.get_open_stream_ids_for_step(db, @step, @base_ws) + assert {:ok, []} = Streams.get_open_stream_ids_for_step(db, @step, @child_ws) + assert {:ok, []} = Streams.get_open_stream_ids_for_execution(db, 2) + assert {:ok, [^s0, ^s1]} = Streams.get_open_stream_ids_for_execution(db, 1) + end + + test "summarises closures by the execution that wrote them", %{db: db} do + create_execution(db, 1, @step, 1, @base_ws) + {:ok, %{id: s0}} = register(db, 1, 0) + {:ok, %{id: s1}} = register(db, 1, 1) + {:ok, _} = Streams.close_stream(db, s0, 1, :timeout) + {:ok, _} = Streams.close_stream(db, s1, 1, {:errored, "E", "boom", []}) + + assert {:ok, %{timed_out: true, errored: error_id}} = + Streams.get_closure_summary_for_execution(db, 1) + + assert is_integer(error_id) + + assert {:ok, %{timed_out: false, errored: nil}} = + Streams.get_closure_summary_for_execution(db, 2) + + assert {:ok, {:errored, {"E", "boom", []}, 1, _}} = Streams.get_stream_closure(db, s1) + end + end + + describe "refs" do + test "resolves a stream by its key and creates a stable ref", %{db: db} do + create_execution(db, 1, @step, 1, @base_ws) + {:ok, %{id: stream_id}} = register(db, 1, 0) + + assert {:ok, ^stream_id} = Streams.get_stream_id_by_key(db, "r1", 0, 0) + assert {:error, :not_found} = Streams.get_stream_id_by_key(db, "r1", 0, 1) + assert {:ok, ref_id} = Streams.create_stream_ref_for(db, stream_id) + assert {:ok, ^ref_id} = Streams.create_stream_ref_for(db, stream_id) + assert {:ok, {"r1", 0, 0, "module", "target"}} = Streams.get_stream_ref(db, ref_id) + assert {:ok, id} = Streams.record_dependency(db, 1, ref_id) + assert is_integer(id) + assert {:ok, nil} = Streams.record_dependency(db, 1, ref_id) + assert {:ok, %{1 => [^ref_id]}} = Streams.get_run_dependencies(db, 1) + end + end + + # --- helpers --- + + defp register(db, execution_id, position) do + Streams.register(db, @step, @base_ws, execution_id, position, 0, nil) |> strip_created() + end + + defp strip_created({:ok, registration}), do: {:ok, Map.delete(registration, :created_at)} + + defp suspend(db, execution_id, successor_id) do + {:ok, _} = Results.record_completion(db, execution_id, :suspended, successor_id: nil) + _ = successor_id + end + + defp create_workspace(db, id, external_id) do + :ok = + Sqlite3.execute( + db, + "INSERT INTO workspaces (id, external_id) VALUES (#{id}, '#{external_id}')" + ) + end + + defp create_run(db, id, external_id) do + :ok = + Sqlite3.execute( + db, + "INSERT INTO runs (id, external_id, created_at) VALUES (#{id}, '#{external_id}', 0)" + ) + end + + defp create_step(db, id, run_id, number) do + :ok = + Sqlite3.execute(db, """ + INSERT INTO steps ( + id, number, run_id, module, target, type, priority, wait_for, + retry_limit, retry_backoff_min, retry_backoff_max, created_at + ) + VALUES (#{id}, #{number}, #{run_id}, 'module', 'target', 0, 0, 0, 0, 0, 0, 0) + """) + end + + defp create_execution(db, id, step_id, attempt, workspace_id) do + :ok = + Sqlite3.execute(db, """ + INSERT INTO executions (id, step_id, attempt, workspace_id, created_at) + VALUES (#{id}, #{step_id}, #{attempt}, #{workspace_id}, 0) + """) + end +end diff --git a/tests/support/executor.py b/tests/support/executor.py index aff1493f..6f3da300 100644 --- a/tests/support/executor.py +++ b/tests/support/executor.py @@ -131,12 +131,24 @@ def _request(self, msg): self._next_request_id += 1 msg["id"] = rid self.send(msg) - while True: - incoming = self.recv() - if incoming.get("id") == rid: - return incoming - # Park non-matching messages (typically notifications). - self._buffer.append(incoming) + # Hold non-matching messages (typically notifications) aside and + # restore them afterwards. They must not go back onto `_buffer` + # mid-loop: `recv` serves the buffer before the socket, so a + # re-buffered message would be popped straight back out and the + # loop would spin without ever reading the reply. A stream_demand + # grant arriving ahead of a stream_register reply is the everyday + # case. + held = [] + try: + while True: + incoming = self.recv() + if incoming.get("id") == rid: + self._buffer[:0] = held + return incoming + held.append(incoming) + except (TimeoutError, ConnectionError): + self._buffer[:0] = held + raise def submit_task(self, execution_id, module, target, arguments, **kwargs): """Submit a child task execution and return the target execution ID.""" @@ -292,15 +304,23 @@ def resolve_input( # --- Stream producer helpers --- - def stream_register(self, execution_id, index, buffer=None, timeout_ms=None): - """Notify that a new stream exists. ``buffer`` enables - backpressure; ``timeout_ms`` enables idle-timeout enforcement - at the worker.""" - self.send( + def stream_register(self, execution_id, position, buffer=None, timeout_ms=None): + """Register the execution's ``position``-th stream and return the + server's reply: ``{"id", "index", "head"}``. ``buffer`` enables + backpressure; ``timeout_ms`` enables idle-timeout enforcement at + the worker. + + For a first attempt the index equals the position, so tests that + don't care about the reply can keep addressing the stream by the + position they registered.""" + resp = self._request( protocol.stream_register( - execution_id, index, buffer=buffer, timeout_ms=timeout_ms + execution_id, position, buffer=buffer, timeout_ms=timeout_ms ) ) + if resp.get("error"): + raise RuntimeError(f"stream_register error: {resp['error']}") + return resp["result"] def stream_append(self, execution_id, index, sequence, value, format="json"): """Append an item (raw JSON value) to a stream.""" @@ -318,29 +338,36 @@ def stream_subscribe( self, execution_id, subscription_id, - producer_execution_id, - index, + stream_id=None, + *, + producer_execution_id=None, + index=None, from_sequence=0, stride=None, prefetch=protocol.DEFAULT_PREFETCH, ): - """Subscribe to a stream. ``stride`` is an optional - ``{"start", "stop", "step"}`` dict restricting which positions - are delivered — built via ``protocol.stride`` / - ``slice_stride`` / ``partition_stride``. ``None`` means no - filtering (identity stride). + """Subscribe to a stream, by ``stream_id`` (the ``id`` from a + ``stream_register`` reply) or by ``producer_execution_id`` + + ``index`` — the stream at that index on the producer's step. + + ``stride`` is an optional ``{"start", "stop", "step"}`` dict + restricting which positions are delivered — built via + ``protocol.stride`` / ``slice_stride`` / ``partition_stride``. + ``None`` means no filtering (identity stride). ``prefetch`` is the delivery window — the server pushes at most this many items beyond what's been acknowledged. The default is large enough to be invisible; lower it to exercise credit-gated delivery.""" + if stream_id is None: + assert producer_execution_id is not None and index is not None + stream_id = protocol.stream_id_for(producer_execution_id, index) self._sub_execution_ids[subscription_id] = execution_id self.send( protocol.stream_subscribe( execution_id, subscription_id, - producer_execution_id, - index, + stream_id, from_sequence=from_sequence, stride=stride, prefetch=prefetch, @@ -401,32 +428,41 @@ def drain_stream(self, subscription_id, timeout=10, ack=True): for asserting that the window actually stops delivery. """ items = [] + # Messages for other subscriptions (or other methods) are held + # aside and restored on exit — see `_request` for why they can't + # be appended back onto `_buffer` while the loop is still reading. + held = [] deadline = time.time() + timeout - while True: - remaining = max(0.01, deadline - time.time()) - msg = self.recv(timeout=remaining) - method = msg.get("method") - params = msg.get("params", {}) - if params.get("subscription_id") != subscription_id or method not in ( - "stream_items", - "stream_closed", - ): - self._buffer.append(msg) - continue - if method == "stream_items": - batch = params.get("items", []) - items.extend(batch) - count, sequence = self._record_items(subscription_id, batch) - if ack and batch: - self.stream_ack( - self._sub_execution_ids[subscription_id], - subscription_id, - count, - sequence, - ) - continue - # stream_closed — terminal - return items, params + try: + while True: + remaining = max(0.01, deadline - time.time()) + msg = self.recv(timeout=remaining) + method = msg.get("method") + params = msg.get("params", {}) + if params.get("subscription_id") != subscription_id or method not in ( + "stream_items", + "stream_closed", + ): + held.append(msg) + continue + if method == "stream_items": + batch = params.get("items", []) + items.extend(batch) + count, sequence = self._record_items(subscription_id, batch) + if ack and batch: + self.stream_ack( + self._sub_execution_ids[subscription_id], + subscription_id, + count, + sequence, + ) + continue + # stream_closed — terminal + self._buffer[:0] = held + return items, params + except (TimeoutError, ConnectionError): + self._buffer[:0] = held + raise def _record_items(self, subscription_id, batch): """Fold a delivered batch into this subscription's cumulative diff --git a/tests/support/protocol.py b/tests/support/protocol.py index 1d95dcaa..c09edd24 100644 --- a/tests/support/protocol.py +++ b/tests/support/protocol.py @@ -260,8 +260,14 @@ def register_group_notification(execution_id, group_id, name=None): # --- Stream messages (producer side: adapter → server) --- -def stream_register(execution_id, index, buffer=None, timeout_ms=None): - params = {"execution_id": execution_id, "index": index} +def stream_register(execution_id, position, buffer=None, timeout_ms=None): + """Register the execution's ``position``-th stream. A request: the + server replies with ``{"id", "index", "head"}`` — the stream's id, + its index within the step (what appends and closes carry), and the + last sequence already in it (``-1`` for a new stream, or the resume + point when the registration continues a stream a suspended + predecessor left paused).""" + params = {"execution_id": execution_id, "position": position} if buffer is not None: params["buffer"] = buffer if timeout_ms is not None: @@ -269,10 +275,18 @@ def stream_register(execution_id, index, buffer=None, timeout_ms=None): return {"method": "stream_register", "params": params} +def stream_id_for(execution_id, index): + """The id of the stream at ``index`` on the step ``execution_id`` + belongs to: ``:_``. Streams belong to steps, so + the attempt in the execution id is dropped.""" + step_id = execution_id.rsplit(":", 1)[0] + return f"{step_id}_{index}" + + def stream_append(execution_id, index, sequence, value, format="json"): """Append an item to a stream. ``value`` is the raw JSON value. - ``index`` identifies the stream within its execution; ``sequence`` + ``index`` identifies the stream within its step; ``sequence`` identifies the item within the stream. Builds a Value wire-form message with an empty references list. Tests that need references should build the Value dict manually. @@ -313,13 +327,12 @@ def stream_close(execution_id, index, error=None): def stream_subscribe( execution_id, subscription_id, - producer_execution_id, - index, + stream_id, from_sequence=0, stride=None, prefetch=DEFAULT_PREFETCH, ): - """Subscribe to a stream. + """Subscribe to a stream by its id (see ``stream_id_for``). ``prefetch`` bounds how many items the server will push beyond what has been acknowledged via ``stream_ack``. @@ -331,8 +344,7 @@ def stream_subscribe( params = { "execution_id": execution_id, "subscription_id": subscription_id, - "producer_execution_id": producer_execution_id, - "index": index, + "stream_id": stream_id, "from_sequence": from_sequence, "prefetch": prefetch, } diff --git a/tests/test_epochs.py b/tests/test_epochs.py index 97447a96..e4d443ce 100644 --- a/tests/test_epochs.py +++ b/tests/test_epochs.py @@ -1,5 +1,7 @@ """Tests for epoch database rotation and cross-epoch behavior.""" +import time + from support import cli from support.helpers import api_post, managed_worker, poll_result from support.manifest import task, workflow @@ -435,3 +437,76 @@ def test_idempotency_across_epoch_boundary(isolated_server, tmp_path): run_id2 = resp2["runId"] assert run_id1 == run_id2 + + +def test_stream_readable_across_epoch_boundary(isolated_server, tmp_path): + """A stream produced before rotation can still be read after it. + + Resolving the stream id copies its run forward — streams, items and + closure included — so a late consumer replays the backlog. + """ + server, host, project_id = isolated_server + targets = [workflow("test", "producer"), workflow("test", "consumer")] + + with managed_worker(targets, host, tmp_path, concurrency=2) as executor: + resp = cli.submit("test/producer", host=host) + prod = executor.next_execute() + stream = prod.conn.stream_register(prod.execution_id, 0) + prod.conn.stream_append(prod.execution_id, 0, 0, "a") + prod.conn.stream_append(prod.execution_id, 0, 1, "b") + prod.conn.stream_close(prod.execution_id, 0) + prod.conn.complete(prod.execution_id, value="done") + assert poll_result(resp["runId"], host)["value"]["data"] == "done" + + _rotate_epoch(server.port, project_id) + + cli.submit("test/consumer", host=host) + cons = executor.next_execute() + cons.conn.stream_subscribe( + cons.execution_id, subscription_id=1, stream_id=stream["id"] + ) + items, closed = cons.conn.drain_stream(subscription_id=1) + assert [item[1]["value"] for item in items] == ["a", "b"] + assert closed["reason"] == "complete" + cons.conn.complete(cons.execution_id) + + +def test_paused_stream_continued_across_epoch_boundary(isolated_server, tmp_path): + """A stream left paused by a suspend survives rotation as paused. + + Its registrations and the suspended completion are carried forward, so + a re-run in the new epoch continues the same stream from its head. + """ + server, host, project_id = isolated_server + targets = [workflow("test", "producer"), workflow("test", "consumer")] + + with managed_worker(targets, host, tmp_path, concurrency=2) as executor: + resp = cli.submit("test/producer", host=host) + prod = executor.next_execute() + stream = prod.conn.stream_register(prod.execution_id, 0) + prod.conn.stream_append(prod.execution_id, 0, 0, "a") + # Keep the successor pending, so the step is paused across rotation. + prod.conn.suspend( + prod.execution_id, execute_after=int(time.time() * 1000) + 3_600_000 + ) + + _rotate_epoch(server.port, project_id) + + cli.runs_rerun(resp["stepId"], host=host) + prod2 = executor.next_execute() + resumed = prod2.conn.stream_register(prod2.execution_id, 0) + assert resumed == {"id": stream["id"], "index": 0, "head": 0} + prod2.conn.stream_append(prod2.execution_id, 0, 1, "b") + prod2.conn.stream_close(prod2.execution_id, 0) + prod2.conn.complete(prod2.execution_id, value="done") + assert poll_result(resp["runId"], host)["value"]["data"] == "done" + + cli.submit("test/consumer", host=host) + cons = executor.next_execute() + cons.conn.stream_subscribe( + cons.execution_id, subscription_id=1, stream_id=stream["id"] + ) + items, closed = cons.conn.drain_stream(subscription_id=1) + assert [item[1]["value"] for item in items] == ["a", "b"] + assert closed["reason"] == "complete" + cons.conn.complete(cons.execution_id) diff --git a/tests/test_streams.py b/tests/test_streams.py index 06f1fb58..c6257925 100644 --- a/tests/test_streams.py +++ b/tests/test_streams.py @@ -17,6 +17,7 @@ import time import pytest +from support import cli from support.manifest import task, workflow from support.protocol import ( execution_handle, @@ -285,12 +286,13 @@ def test_subscribe_to_unknown_producer_closes_immediately(worker): _items, closed = cons_ex.conn.drain_stream(subscription_id=1) cons_ex.conn.complete(cons_ex.execution_id) - assert closed.get("reason") == "producer_not_found" + assert closed.get("reason") == "stream_not_found" assert closed.get("error") is None def test_topic_exposes_stream_state(worker): - """Studio topic gets `streams` per execution: opened, closed, error.""" + """Studio topic gets `streams` per step: opened, closed, error, and the + attempts that produced into each.""" targets = [workflow("test", "producer")] with worker(targets) as ctx: @@ -308,12 +310,16 @@ def test_topic_exposes_stream_state(worker): ctx.result(prod_resp["runId"]) snapshot = ctx.inspect(prod_resp["runId"]) - # The run snapshot has a `steps → {run:step → {executions → {attempt → {...}}}}` shape. - step = next(iter(snapshot["steps"].values())) - execution = next(iter(step["executions"].values())) - streams = execution["streams"] + # The run snapshot has a `steps → {run:step → {streams, executions}}` + # shape: streams belong to the step, not to an attempt. + step_id, step = next(iter(snapshot["steps"].items())) + streams = step["streams"] assert "0" in streams and "1" in streams + assert streams["0"]["id"] == f"{step_id}_0" + assert streams["1"]["id"] == f"{step_id}_1" + assert streams["0"]["attempts"] == [1] + assert streams["0"]["closedBy"] == 1 assert streams["0"]["openedAt"] is not None assert streams["0"]["closedAt"] is not None assert streams["0"]["reason"] == "complete" @@ -1131,45 +1137,64 @@ def test_slowest_consumer_leaving_releases_producer(worker): ctx.result(prod_resp["runId"]) -def test_suspended_producer_closes_stream_as_superseded(worker): - """A producer that suspends leaves its stream terminal. +def test_suspended_producer_stream_is_continued_by_resumed_execution(worker): + """A producer that suspends leaves its stream paused, not closed. - A stream is owned by exactly one execution, so the resumed execution - registers its own under a new id rather than continuing this one. - The attached consumer must therefore be told the stream is over — - and told *why*, since "suspended" is not a failure and shouldn't - read like one. + Streams belong to the step: the execution that resumes it registers at + the same position, is handed the same stream (and where its sequence + left off), and appends to it. A consumer sees one unbroken stream — + including one that subscribed while the step was suspended. """ targets = [workflow("test", "producer"), workflow("test", "consumer")] with worker(targets, concurrency=3) as ctx: ctx.submit("test", "producer") prod_ex = ctx.executor.next_execute() - prod_ex.conn.stream_register(prod_ex.execution_id, 0) + first = prod_ex.conn.stream_register(prod_ex.execution_id, 0) + assert first["index"] == 0 and first["head"] == -1 prod_ex.conn.stream_append(prod_ex.execution_id, 0, 0, "v0") ctx.submit("test", "consumer") cons_ex = ctx.executor.next_execute() cons_ex.conn.stream_subscribe( - cons_ex.execution_id, - subscription_id=1, - producer_execution_id=prod_ex.execution_id, - index=0, + cons_ex.execution_id, subscription_id=1, stream_id=first["id"] ) items = cons_ex.conn.recv_push("stream_items", subscription_id=1, timeout=2) assert [i[0] for i in items["items"]] == [0] prod_ex.conn.suspend(prod_ex.execution_id) - closed = cons_ex.conn.recv_push("stream_closed", subscription_id=1, timeout=2) - assert closed["reason"] == "suspended" - assert closed.get("error") is None - - # The resumed execution is a distinct execution owning a distinct - # stream — which is exactly why the old one had to be closed. + # The resumed execution is a distinct execution — but it continues + # the same stream from where the suspended one left it. prod2 = ctx.executor.next_execute() assert prod2.execution_id != prod_ex.execution_id - prod2.conn.stream_register(prod2.execution_id, 0) + + # Nothing closed for the live consumer in the meantime. + with pytest.raises(TimeoutError): + cons_ex.conn.recv_push("stream_closed", subscription_id=1, timeout=0.5) + + # A consumer subscribing while the step is paused gets the backlog + # and then waits, rather than being told the stream is over. + cons_ex.conn.stream_subscribe( + cons_ex.execution_id, subscription_id=2, stream_id=first["id"] + ) + backlog = cons_ex.conn.recv_push("stream_items", subscription_id=2, timeout=2) + assert [i[0] for i in backlog["items"]] == [0] + with pytest.raises(TimeoutError): + cons_ex.conn.recv_push("stream_closed", subscription_id=2, timeout=0.5) + + resumed = prod2.conn.stream_register(prod2.execution_id, 0) + assert resumed == {"id": first["id"], "index": 0, "head": 0} + + prod2.conn.stream_append(prod2.execution_id, 0, 1, "v1") + items = cons_ex.conn.recv_push("stream_items", subscription_id=1, timeout=2) + assert [(i[0], i[1]["value"]) for i in items["items"]] == [(1, "v1")] + items = cons_ex.conn.recv_push("stream_items", subscription_id=2, timeout=2) + assert [(i[0], i[1]["value"]) for i in items["items"]] == [(1, "v1")] + + prod2.conn.stream_close(prod2.execution_id, 0) + closed = cons_ex.conn.recv_push("stream_closed", subscription_id=1, timeout=2) + assert closed["reason"] == "complete" prod2.conn.complete(prod2.execution_id, value="done") cons_ex.conn.complete(cons_ex.execution_id) @@ -1215,17 +1240,160 @@ def test_recurrent_producer_closes_stream_each_iteration(worker): assert closed["reason"] == "recurred" # The next iteration is a separate execution with a separate - # stream; the old handle does not follow it. + # stream — the step's next index, rather than a continuation of the + # closed one; the old handle does not follow it. tick2 = ctx.executor.next_execute() assert tick2.execution_id != tick1.execution_id - tick2.conn.stream_register(tick2.execution_id, 0) - tick2.conn.stream_append(tick2.execution_id, 0, 0, "tick-1") + fresh = tick2.conn.stream_register(tick2.execution_id, 0) + assert fresh["index"] == 1 and fresh["head"] == -1 + tick2.conn.stream_append(tick2.execution_id, fresh["index"], 0, "tick-1") with pytest.raises(TimeoutError): cons_ex.conn.recv_push("stream_items", subscription_id=1, timeout=0.5) cons_ex.conn.complete(cons_ex.execution_id) +def _far_future_ms(): + """An execute_after that keeps a suspended step's successor pending + for the rest of the test, so the step can be re-run or cancelled + while paused.""" + return int(time.time() * 1000) + 3_600_000 + + +def test_rerun_of_suspended_step_continues_stream(worker): + """Re-running a suspended step cancels its pending successor — which + never produced anything — and the new attempt continues the paused + stream. A manual "resume now" doesn't reset it. + """ + targets = [workflow("test", "producer"), workflow("test", "consumer")] + + with worker(targets, concurrency=3) as ctx: + resp = ctx.submit("test", "producer") + prod_ex = ctx.executor.next_execute() + first = prod_ex.conn.stream_register(prod_ex.execution_id, 0) + prod_ex.conn.stream_append(prod_ex.execution_id, 0, 0, "v0") + + ctx.submit("test", "consumer") + cons_ex = ctx.executor.next_execute() + cons_ex.conn.stream_subscribe( + cons_ex.execution_id, subscription_id=1, stream_id=first["id"] + ) + cons_ex.conn.recv_push("stream_items", subscription_id=1, timeout=2) + + prod_ex.conn.suspend(prod_ex.execution_id, execute_after=_far_future_ms()) + + ctx.rerun(resp["stepId"]) + prod2 = ctx.executor.next_execute() + with pytest.raises(TimeoutError): + cons_ex.conn.recv_push("stream_closed", subscription_id=1, timeout=0.5) + + resumed = prod2.conn.stream_register(prod2.execution_id, 0) + assert resumed == {"id": first["id"], "index": 0, "head": 0} + + prod2.conn.stream_append(prod2.execution_id, 0, 1, "v1") + items = cons_ex.conn.recv_push("stream_items", subscription_id=1, timeout=2) + assert [(i[0], i[1]["value"]) for i in items["items"]] == [(1, "v1")] + + prod2.conn.stream_close(prod2.execution_id, 0) + prod2.conn.complete(prod2.execution_id, value="done") + cons_ex.conn.complete(cons_ex.execution_id) + + +def test_cancelling_suspended_step_closes_paused_stream(worker): + """Cancelling the pending successor of a suspended producer ends the + step, so the paused stream is closed for its consumers — as + "cancelled", even though the cancelled execution itself never + produced into it. + """ + targets = [workflow("test", "producer"), workflow("test", "consumer")] + + with worker(targets, concurrency=3) as ctx: + resp = ctx.submit("test", "producer") + prod_ex = ctx.executor.next_execute() + first = prod_ex.conn.stream_register(prod_ex.execution_id, 0) + prod_ex.conn.stream_append(prod_ex.execution_id, 0, 0, "v0") + + ctx.submit("test", "consumer") + cons_ex = ctx.executor.next_execute() + cons_ex.conn.stream_subscribe( + cons_ex.execution_id, subscription_id=1, stream_id=first["id"] + ) + cons_ex.conn.recv_push("stream_items", subscription_id=1, timeout=2) + + prod_ex.conn.suspend(prod_ex.execution_id, execute_after=_far_future_ms()) + + # The pending successor is the step's next attempt. + ctx.cancel(f"{resp['stepId']}:2") + + closed = cons_ex.conn.recv_push("stream_closed", subscription_id=1, timeout=2) + assert closed["reason"] == "cancelled" + assert closed.get("error") is None + + cons_ex.conn.complete(cons_ex.execution_id) + + +def test_rerun_of_live_producer_starts_fresh_stream(worker): + """Re-running a step whose attempt is mid-stream cancels that attempt, + which closes its stream; the new attempt opens a fresh stream under + the step's next index rather than continuing the cancelled one. + """ + targets = [workflow("test", "producer"), workflow("test", "consumer")] + + with worker(targets, concurrency=3) as ctx: + resp = ctx.submit("test", "producer") + prod_ex = ctx.executor.next_execute() + first = prod_ex.conn.stream_register(prod_ex.execution_id, 0) + prod_ex.conn.stream_append(prod_ex.execution_id, 0, 0, "v0") + + ctx.submit("test", "consumer") + cons_ex = ctx.executor.next_execute() + cons_ex.conn.stream_subscribe( + cons_ex.execution_id, subscription_id=1, stream_id=first["id"] + ) + cons_ex.conn.recv_push("stream_items", subscription_id=1, timeout=2) + + ctx.rerun(resp["stepId"]) + + closed = cons_ex.conn.recv_push("stream_closed", subscription_id=1, timeout=2) + assert closed["reason"] == "cancelled" + + prod2 = ctx.executor.next_execute() + fresh = prod2.conn.stream_register(prod2.execution_id, 0) + assert fresh["index"] == 1 and fresh["head"] == -1 + assert fresh["id"] != first["id"] + prod2.conn.stream_close(prod2.execution_id, fresh["index"]) + prod2.conn.complete(prod2.execution_id, value="done") + + cons_ex.conn.complete(cons_ex.execution_id) + + +def test_stream_is_not_continued_across_workspaces(worker): + """A derived-workspace re-run of a suspended step opens its own stream + rather than appending to the base's paused one. Writes never cross + workspaces (a derived consumer could still read the base's stream). + """ + targets = [workflow("test", "producer")] + + with worker(targets, workspace="base", concurrency=2) as ctx_base: + resp = ctx_base.submit("test", "producer") + prod_ex = ctx_base.executor.next_execute() + first = prod_ex.conn.stream_register(prod_ex.execution_id, 0) + prod_ex.conn.stream_append(prod_ex.execution_id, 0, 0, "v0") + prod_ex.conn.suspend(prod_ex.execution_id, execute_after=_far_future_ms()) + saved_host = ctx_base.host + + cli.workspaces_create("derived", base="base", host=saved_host, workspace="derived") + + with worker(targets, workspace="derived") as ctx_derived: + ctx_derived.rerun(resp["stepId"]) + ex1 = ctx_derived.executor.next_execute() + fresh = ex1.conn.stream_register(ex1.execution_id, 0) + assert fresh["index"] == 1 and fresh["head"] == -1 + assert fresh["id"] != first["id"] + ex1.conn.stream_close(ex1.execution_id, fresh["index"]) + ex1.conn.complete(ex1.execution_id, value="done") + + # --- Idle timeout ------------------------------------------------------- @@ -1345,8 +1513,7 @@ def test_timeout_visible_in_topic(worker): snapshot = ctx.inspect(prod_resp["runId"]) step = next(iter(snapshot["steps"].values())) - execution = next(iter(step["executions"].values())) - stream = execution["streams"]["0"] + stream = step["streams"]["0"] assert stream["timeoutMs"] == 120 assert stream["reason"] == "timeout" assert stream["error"] is None @@ -1371,11 +1538,10 @@ def test_clean_stream_keeps_completion_succeeded(worker): prod_ex.conn.complete(prod_ex.execution_id, value=1) ctx.result(prod_resp["runId"]) - snapshot = ctx.inspect(prod_resp["runId"]) - execution = next( - iter(snapshot["steps"][f"{prod_resp['runId']}:1"]["executions"].values()) - ) - assert execution["completion"]["kind"] == "succeeded" + # The value is recorded before the completion, which lands when the + # worker reports the process exited — so wait for it rather than + # reading the snapshot straight away. + assert _wait_for_completion(ctx, prod_resp["runId"], 1) == "succeeded" def test_stream_error_promotes_completion_to_stream_errored(worker): @@ -1399,11 +1565,10 @@ def test_stream_error_promotes_completion_to_stream_errored(worker): prod_ex.conn.complete(prod_ex.execution_id, value=1) ctx.result(prod_resp["runId"]) - snapshot = ctx.inspect(prod_resp["runId"]) - execution = next( - iter(snapshot["steps"][f"{prod_resp['runId']}:1"]["executions"].values()) - ) - assert execution["completion"]["kind"] == "stream_errored" + # The value is recorded before the completion, which lands when the + # worker reports the process exited — so wait for it rather than + # reading the snapshot straight away. + assert _wait_for_completion(ctx, prod_resp["runId"], 1) == "stream_errored" def test_stream_timeout_promotes_completion_to_stream_timeout(worker): @@ -1427,11 +1592,10 @@ def test_stream_timeout_promotes_completion_to_stream_timeout(worker): prod_ex.conn.complete(prod_ex.execution_id, value=1) ctx.result(prod_resp["runId"]) - snapshot = ctx.inspect(prod_resp["runId"]) - execution = next( - iter(snapshot["steps"][f"{prod_resp['runId']}:1"]["executions"].values()) - ) - assert execution["completion"]["kind"] == "stream_timeout" + # The value is recorded before the completion, which lands when the + # worker reports the process exited — so wait for it rather than + # reading the snapshot straight away. + assert _wait_for_completion(ctx, prod_resp["runId"], 1) == "stream_timeout" def test_stream_error_outranks_timeout(worker): @@ -1462,11 +1626,10 @@ def test_stream_error_outranks_timeout(worker): prod_ex.conn.complete(prod_ex.execution_id, value=1) ctx.result(prod_resp["runId"]) - snapshot = ctx.inspect(prod_resp["runId"]) - execution = next( - iter(snapshot["steps"][f"{prod_resp['runId']}:1"]["executions"].values()) - ) - assert execution["completion"]["kind"] == "stream_errored" + # The value is recorded before the completion, which lands when the + # worker reports the process exited — so wait for it rather than + # reading the snapshot straight away. + assert _wait_for_completion(ctx, prod_resp["runId"], 1) == "stream_errored" def _wait_for_completion(ctx, run_id, step_num, timeout=5):