diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a575ce6c0..d8031de09d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -78,6 +78,18 @@ to include examples, links to docs, or any other relevant information. which the handler maps onto the store's accessor for that owner. Configure the front with `data_converter=` to run a payload codec on the caller side, so records are encoded before they leave the process. +- Added experimental External Workflow Streams in + `temporalio.contrib.external_workflow_streams`. Workflow stream payloads are + stored in a configured external backend instead of Temporal History, with a + Redis Streams provider included. Workflows subscribe with `external_stream`, + external processes publish with `ExternalStreamProducer`, and Workers are + configured with `external_stream_backend`. +- Added the output direction for External Workflow Streams. + Workflows publish with `external_output_stream`, Activities and external + processes use `ExternalOutputStreamProducer`, and external consumers resume + through `ExternalOutputStreamClient`. Workflow output is staged outside + History and becomes readable only after its compact Workflow Task marker is + committed. ### Changed diff --git a/pyproject.toml b/pyproject.toml index fcf82b5266..2c08b9c06d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -95,6 +95,7 @@ dev = [ "pytest-rerunfailures>=16.1", "pytest-xdist>=3.6,<4", "moto[s3,server]>=5", + "redis>=5,<9", "langgraph>=1.1.0", "langsmith>=0.7.34,<0.9", "deepagents>=0.6.12,<0.7; python_version >= '3.11'", diff --git a/temporalio/bridge/src/worker.rs b/temporalio/bridge/src/worker.rs index e74698e944..c3416d948a 100644 --- a/temporalio/bridge/src/worker.rs +++ b/temporalio/bridge/src/worker.rs @@ -17,9 +17,9 @@ use temporalio_common::protos::temporal::api::history::v1::History; use temporalio_common::protos::temporal::api::worker::v1::{PluginInfo, StorageDriverInfo}; use temporalio_sdk_core::replay::{HistoryForReplay, ReplayWorkerInput}; use temporalio_sdk_core::{ - PollError, SlotInfo, SlotInfoTrait, SlotKind, SlotKindType, SlotMarkUsedContext, - SlotReleaseContext, SlotReservationContext, SlotSupplier as SlotSupplierTrait, - SlotSupplierPermit, WorkflowErrorType, + ExternalStreamReadyResult, ExternalStreamRunStatus, PollError, SlotInfo, SlotInfoTrait, + SlotKind, SlotKindType, SlotMarkUsedContext, SlotReleaseContext, SlotReservationContext, + SlotSupplier as SlotSupplierTrait, SlotSupplierPermit, WorkflowErrorType, }; use tokio::sync::mpsc::{channel, Sender}; use tokio_stream::wrappers::ReceiverStream; @@ -671,6 +671,63 @@ impl WorkerRef { Ok(()) } + /// Tell Core a record is buffered for one external stream wait. + /// + /// **Synchronous and blocking**, like `record_activity_heartbeat`, but unlike it + /// **acknowledged**: the caller is a watcher whose next action -- keep waiting, re-probe, send + /// a wake Signal, or tear itself down -- depends entirely on which of the five results comes + /// back, so a fire-and-forget send would leave it unable to tell "Core will activate" from + /// "nothing here will". + /// + /// Synchronous rather than a Python awaitable on purpose. There is one watcher per + /// subscription and they do not share an event loop, and a future produced by + /// `future_into_py` is bound to the loop that created it -- so an awaitable here would work + /// from the worker's own loop and hang from anywhere else. Blocking with the GIL released is + /// safe: the answer comes off Core's serialized local-input lane, which does no I/O. + /// + /// Returns one of `Accepted`, `Stale`, `Parked`, `NoOpenWorkflowTask`, `RunNotFound`. + fn notify_external_stream_ready( + &self, + py: Python<'_>, + run_id: String, + wait_id: u32, + wait_generation: u64, + ) -> PyResult<&'static str> { + let worker = self.worker.as_ref().unwrap().clone(); + // Deliberately *not* `enter_sync!`: entering the tokio handle would put this thread in an + // async context, and `Handle::block_on` panics there. + let handle = self.runtime.core.tokio_handle(); + let result = py.detach(move || { + handle.block_on(worker.notify_external_stream_ready(&run_id, wait_id, wait_generation)) + }); + Ok(match result { + ExternalStreamReadyResult::Accepted => "Accepted", + ExternalStreamReadyResult::Stale => "Stale", + ExternalStreamReadyResult::Parked => "Parked", + ExternalStreamReadyResult::NoOpenWorkflowTask => "NoOpenWorkflowTask", + ExternalStreamReadyResult::RunNotFound => "RunNotFound", + }) + } + + /// Ask what state a run's external stream wait set is in, changing nothing. + /// + /// Deliberately not the readiness call: readiness means "a record is buffered", so probing + /// with it would assert something false and manufacture a spurious activation on the way out + /// of a shutting-down worker. + /// + /// Returns one of `WftOpen`, `Parked`, `NoOpenWorkflowTask`, `RunNotFound`. + fn external_stream_run_status(&self, py: Python<'_>, run_id: String) -> PyResult<&'static str> { + let worker = self.worker.as_ref().unwrap().clone(); + let handle = self.runtime.core.tokio_handle(); + let status = py.detach(move || handle.block_on(worker.external_stream_run_status(&run_id))); + Ok(match status { + ExternalStreamRunStatus::WftOpen => "WftOpen", + ExternalStreamRunStatus::Parked => "Parked", + ExternalStreamRunStatus::NoOpenWorkflowTask => "NoOpenWorkflowTask", + ExternalStreamRunStatus::RunNotFound => "RunNotFound", + }) + } + fn request_workflow_eviction(&self, run_id: &str) -> PyResult<()> { enter_sync!(self.runtime); self.worker diff --git a/temporalio/bridge/worker.py b/temporalio/bridge/worker.py index 4b7f55d092..bff886a4f8 100644 --- a/temporalio/bridge/worker.py +++ b/temporalio/bridge/worker.py @@ -5,6 +5,8 @@ from __future__ import annotations +import asyncio +import enum from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass from typing import ( @@ -32,6 +34,62 @@ from temporalio.worker._command_aware_visitor import CommandAwarePayloadVisitor +class ExternalStreamReadyResult(enum.Enum): + """What Core did with an external stream readiness notification. + + Five values, not two, because a watcher does something different with each + and an operator should conclude something different from each. + + ``NO_OPEN_WORKFLOW_TASK`` and ``RUN_NOT_FOUND`` are separate on purpose. A + Run cached between Workflow Tasks is the healthy post-completion state; + reporting it as a missing Run would both corrupt the metric and tell the + watcher to tear itself down while it is still needed. + """ + + ACCEPTED = "Accepted" + """Serialized into an open Workflow Task. Core will activate; do nothing.""" + + STALE = "Stale" + """The wait moved on. Re-probe; do **not** send a Signal.""" + + PARKED = "Parked" + """A confirmed park generation exists. Send the reserved wake Signal.""" + + NO_OPEN_WORKFLOW_TASK = "NoOpenWorkflowTask" + """Cached with waits registered, no task open. Signal, and **keep** the watcher.""" + + RUN_NOT_FOUND = "RunNotFound" + """Absent from this Worker's cache. Signal, then tear the watcher down.""" + + @property + def needs_wake_signal(self) -> bool: + """Whether local readiness could not be delivered and a Signal is owed.""" + return self in ( + ExternalStreamReadyResult.PARKED, + ExternalStreamReadyResult.NO_OPEN_WORKFLOW_TASK, + ExternalStreamReadyResult.RUN_NOT_FOUND, + ) + + +class ExternalStreamRunStatus(enum.Enum): + """What state a Run's external stream wait set is in. + + The answer to a **read-only** probe. Asking changes nothing. + """ + + WFT_OPEN = "WftOpen" + """A Workflow Task is open, retained by the wait set.""" + + PARKED = "Parked" + """The complete wait set is parked.""" + + NO_OPEN_WORKFLOW_TASK = "NoOpenWorkflowTask" + """Cached with active waits, but no Workflow Task open.""" + + RUN_NOT_FOUND = "RunNotFound" + """Absent from this Worker's cache.""" + + @dataclass class WorkerConfig: """Python representation of the Rust struct for configuring a worker.""" @@ -259,6 +317,64 @@ def record_activity_heartbeat( """Record an activity heartbeat.""" self._ref.record_activity_heartbeat(comp.SerializeToString()) # type: ignore[reportOptionalMemberAccess] + def notify_external_stream_ready_sync( + self, run_id: str, wait_id: int, wait_generation: int + ) -> ExternalStreamReadyResult: + """Tell Core a record is **buffered** for one external stream wait. + + Buffered, not merely available: a notification for a record the caller + has not yet read into its buffer would produce an activation whose + drain must block, which is the deadlock hazard the out-of-thread buffer + exists to remove. + + Acknowledged rather than fire-and-forget, because what the caller does + next depends entirely on the answer -- see + :py:class:`ExternalStreamReadyResult`. + + Blocks while Core answers on its serialized local-input lane. Safe from + any thread, and safe to block in -- the lane does no I/O. + + .. warning:: + Do **not** call this from the event loop running the Worker. That + loop is what drives the poll this answer comes from, so blocking it + here waits for something only it can produce. Use + :py:meth:`notify_external_stream_ready` from a loop. + """ + return ExternalStreamReadyResult( + self._ref.notify_external_stream_ready( # type: ignore[reportOptionalMemberAccess] + run_id, wait_id, wait_generation + ) + ) + + async def notify_external_stream_ready( + self, run_id: str, wait_id: int, wait_generation: int + ) -> ExternalStreamReadyResult: + """:py:meth:`notify_external_stream_ready_sync`, off the calling loop. + + There is one watcher per subscription and they do not share a loop, so + this runs the blocking call in a thread rather than binding a future to + whichever loop happened to create it. + """ + return await asyncio.to_thread( + self.notify_external_stream_ready_sync, run_id, wait_id, wait_generation + ) + + def external_stream_run_status_sync(self, run_id: str) -> ExternalStreamRunStatus: + """Ask what state a Run's external stream wait set is in, changing nothing. + + A separate call from :py:meth:`notify_external_stream_ready` on purpose: + readiness means "a record is buffered", so probing with it would assert + something false and manufacture a spurious Workflow Task on the way out + of a shutting-down Worker. + """ + return ExternalStreamRunStatus( + self._ref.external_stream_run_status(run_id) # type: ignore[reportOptionalMemberAccess] + ) + + async def external_stream_run_status(self, run_id: str) -> ExternalStreamRunStatus: + """:py:meth:`external_stream_run_status_sync`, off the calling loop.""" + return await asyncio.to_thread(self.external_stream_run_status_sync, run_id) + def request_workflow_eviction(self, run_id: str) -> None: """Request a workflow be evicted.""" self._ref.request_workflow_eviction(run_id) # type: ignore[reportOptionalMemberAccess] diff --git a/temporalio/contrib/external_workflow_streams/__init__.py b/temporalio/contrib/external_workflow_streams/__init__.py new file mode 100644 index 0000000000..80acb17aea --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/__init__.py @@ -0,0 +1,177 @@ +"""External Workflow Streams for Temporal Workflows. + +.. warning:: + This package is experimental and may change in future versions. + +High-volume stream payloads live in a pluggable external backend such as Redis +Streams and never enter Temporal History. Deterministic replay is preserved +with compact markers recording consumed offset ranges and the +availability/blocking boundaries the original execution observed. + +Workflow code creates topics with :data:`external_stream` and subscribes to +them. A Worker receives a :class:`StreamBackend` through +``Worker(external_stream_backend=...)``. External processes publish with an +:class:`ExternalStreamProducer` bound to the Workflow's +:class:`WorkflowChainKey`. + +The complementary output direction uses :data:`external_output_stream` inside +Workflow code, :class:`ExternalOutputStreamProducer` in Activities and external +processes, and :class:`ExternalOutputStreamClient` for resumable external +consumers. Output payloads are staged behind a pending provider barrier until a +compact marker in Temporal History proves their producing Workflow Task was +accepted. + +This is a **mirror image** of the shipped +:py:mod:`temporalio.contrib.workflow_streams`, not a replacement for it, and +the two coexist (ADR-001). No name here may begin with +``__temporal_workflow_stream``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from temporalio.contrib.external_workflow_streams._api import ( + ExternalStreamOptions, + ExternalStreamSubscription, + ExternalStreamTopic, + external_stream, + merge, +) +from temporalio.contrib.external_workflow_streams._backend import ( + AppendConflictError, + ParkIntent, + ParkIntentRemoval, + StreamBackend, + StreamDirection, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec +from temporalio.contrib.external_workflow_streams._errors import ( + ConcurrentStreamConsumerError, + ExternalStreamCapacityError, + StreamDecodeError, + StreamError, + StreamIntegrityError, + StreamStorageError, +) +from temporalio.contrib.external_workflow_streams._output_api import ( + ExternalOutputStreamOptions, + ExternalOutputStreamTopic, + external_output_stream, +) +from temporalio.contrib.external_workflow_streams._output_backend import ( + OutputReadResult, + OutputStage, + OutputStageConflictError, + OutputStageManifest, + OutputStageNotFoundError, + OutputStageResolutionError, + OutputStageStatus, + OutputStreamBackend, + OutputStreamRecord, + PendingOutputBarrier, + StagedOutputRecord, +) +from temporalio.contrib.external_workflow_streams._output_client import ( + ExternalOutputStreamClient, + ExternalOutputStreamClientTopic, + ExternalOutputStreamItem, +) +from temporalio.contrib.external_workflow_streams._output_producer import ( + ExternalOutputStreamProducer, + ExternalOutputStreamProducerTopic, + OutputAppendNotAcknowledgedError, +) +from temporalio.contrib.external_workflow_streams._producer import ( + AppendNotAcknowledgedError, + ChainKeyMismatchError, + ExternalStreamProducer, + ExternalStreamProducerTopic, + PrecedingWriteFailedError, + WakeNotAcknowledgedError, + WorkflowChainKey, +) +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Cursor, + IdempotencyKey, + Offset, + OffsetComparator, + RecordKind, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._wake import WakeRequest + +if TYPE_CHECKING: + from temporalio.contrib.external_workflow_streams._redis import RedisStreamBackend + +__all__ = [ + "AFTER", + "BEGINNING", + "AppendConflictError", + "AppendNotAcknowledgedError", + "ChainKeyMismatchError", + "ConcurrentStreamConsumerError", + "Cursor", + "ExternalStreamCapacityError", + "ExternalOutputStreamClient", + "ExternalOutputStreamClientTopic", + "ExternalOutputStreamItem", + "ExternalOutputStreamOptions", + "ExternalOutputStreamProducer", + "ExternalOutputStreamProducerTopic", + "ExternalOutputStreamTopic", + "ExternalStreamOptions", + "ExternalStreamProducer", + "ExternalStreamProducerTopic", + "ExternalStreamSubscription", + "ExternalStreamTopic", + "IdempotencyKey", + "Offset", + "OffsetComparator", + "OutputAppendNotAcknowledgedError", + "OutputReadResult", + "OutputStage", + "OutputStageConflictError", + "OutputStageManifest", + "OutputStageNotFoundError", + "OutputStageResolutionError", + "OutputStageStatus", + "OutputStreamBackend", + "OutputStreamRecord", + "ParkIntent", + "ParkIntentRemoval", + "PrecedingWriteFailedError", + "RecordKind", + "RedisStreamBackend", + "StreamBackend", + "StreamDecodeError", + "StreamError", + "StreamIntegrityError", + "StreamDirection", + "StreamKey", + "StreamPayloadCodec", + "StreamRecord", + "StreamStorageError", + "StagedOutputRecord", + "PendingOutputBarrier", + "WakeNotAcknowledgedError", + "WakeRequest", + "WorkflowChainKey", + "external_stream", + "external_output_stream", + "merge", +] + + +def __getattr__(name: str) -> Any: + """Lazily expose the Redis provider without importing it in Workflows.""" + if name == "RedisStreamBackend": + from temporalio.contrib.external_workflow_streams._redis import ( + RedisStreamBackend, + ) + + return RedisStreamBackend + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/temporalio/contrib/external_workflow_streams/_annotation.py b/temporalio/contrib/external_workflow_streams/_annotation.py new file mode 100644 index 0000000000..d5c55bb536 --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_annotation.py @@ -0,0 +1,726 @@ +"""The replay annotation codec (P5). + +The opaque bytes Python encodes, Core stores in a marker, and Python decodes on +replay. **Core never parses any of it** -- it appends observation deltas to a +byte buffer and hands the result back. + +That last fact drives the encoding's shape. Core accumulates by byte +concatenation, so an annotation is a sequence of self-delimiting **frames**: + +.. code-block:: text + + annotation := schema_version, header_frame + , (bindings_frame | segment_frame)*, terminal_frame + + header := streams[] // wait_id -> binding + bindings := streams[] // the same, for waits added later + binding := (stream_key, start_cursor, provider_id + , provider_format_version) + segment := run*, segment_end_reason + run := (wait_id, first_offset, last_offset, count, control_positions) + terminal := blocked_snapshot[] // wait_id -> BEGINNING | AFTER(offset) + +The provider identity is carried with each wait so every recorded cursor and +range can be checked against the backend configured on the replaying Worker +before that backend interprets it. + +A subscription may be created at **any** activation of a retained Workflow +Task, which is later than the header frame that already went to Core. Core +appends bytes and never rewrites them, so a header cannot be extended in place; +the binding rides its own frame instead, emitted with the delta of the +activation that registered the wait and before the segment that first records a +run for it. Decoding merges every bindings frame into ``header.streams``, so +what replay reads is one complete ``wait_id -> binding`` table however late a +wait joined. Without this a wait registered after the first delta reaches the +marker as runs and a terminal entry with no stream key, no backend, and no start +cursor -- and replay of *unchanged* code fails as "the Workflow did not create" +that wait. + +Concatenating the deltas of one Workflow Task therefore *is* the annotation, +with no reassembly step that could disagree with Core's. + +Two properties the encoding is built around: + +- **No field is per-record.** A run costs two offsets, a count, and a sparse + control list whether it covers ten records or a hundred thousand. This is + what makes marker size scale with cross-stream schedule transitions rather + than with item count. +- **Both run endpoints are recorded, not a start plus a count.** Backend + offsets are ordered but not dense, so ``(first, count)`` does not determine + where a run ends -- and a deletion inside the range would be undetectable. +""" + +from __future__ import annotations + +import enum +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Final + +import temporalio.exceptions +from temporalio.contrib.external_workflow_streams._backend import StreamKey +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Cursor, + Offset, +) + +__all__ = [ + "MAX_ANNOTATION_BYTES", + "Annotation", + "AnnotationAccumulator", + "AnnotationDecodeError", + "AnnotationHeader", + "Run", + "Segment", + "SegmentEndReason", + "StreamBinding", + "decode_annotation", + "encode_annotation", + "encode_bindings", + "encoded_run_size", + "encoded_segment_size", +] + +SCHEMA_VERSION: Final = 1 +"""Leads the encoding, so a marker written by an older SDK stays readable. + +It is also the extension point that keeps ADR-003's accepted risk bounded: a +per-record content-hash mode can be added later without a format break. + +There are no legacy decoders. The feature is private and unreleased, so the +format has one current schema and rejects everything else. +""" + +MAX_ANNOTATION_BYTES: Final = 64 * 1024 +"""Hard cap on one marker's annotation, well below the server's event limit. + +A constant rather than a guideline: it is enforced while encoding, so an +annotation can never exceed it. The runtime rolls the Workflow Task over +instead of growing the marker. +""" + +ROLLOVER_HIGH_WATER: Final = 0.75 +"""Fraction of the budget at which the encoder asks for a rollover. + +The margin exists because the request has to travel to Core and come back: +a rollover asked for at 100% would arrive too late to prevent the overflow it +was meant to avoid. +""" + +_FRAME_HEADER: Final = 0x01 +_FRAME_SEGMENT: Final = 0x02 +_FRAME_TERMINAL: Final = 0x03 +_FRAME_BINDINGS: Final = 0x04 + +_CURSOR_BEGINNING: Final = 0x00 +_CURSOR_AFTER: Final = 0x01 + + +@enum.unique +class SegmentEndReason(enum.IntEnum): + """Why one activation's segment ended. + + Recorded because replay must reproduce not only *what* was delivered but + *where the runtime returned control with nothing further available* -- + those boundaries are what make ``wait_condition`` predicates fire the same + number of times. + """ + + NO_DATA_AVAILABLE = 1 + BATCH_LIMIT = 2 + FENCE_REACHED = 3 + BUDGET_ROLLOVER = 4 + """This batch continues in the following marker.""" + + +class AnnotationDecodeError(Exception): + """The annotation bytes are not a well-formed annotation.""" + + +class AnnotationBudgetExceeded(temporalio.exceptions.ApplicationError): + """Encoding would push the annotation past :data:`MAX_ANNOTATION_BYTES`. + + Unreachable if the three preventions hold, and each closes a way it used to + be reachable: + + - the closing frames are **reserved** rather than checked, so the terminal + and any late bindings frame always fit (see + :attr:`AnnotationAccumulator.reserved`); + - the runtime **stops delivering** rather than growing a segment it could not + then record, and asks Core to roll the Workflow Task over; + - ``subscribe()`` **refuses** a subscription set whose own header and + terminal could not fit an empty annotation, at the point the Workflow makes + it. + + An :py:class:`~temporalio.exceptions.ApplicationError` marked + non-retryable, and that is the substance of this class rather than a detail. + A plain exception here fails the *Workflow Task*, and the server retries + Workflow Task failures forever: the encoding that overflowed overflows again + on every retry, so the Workflow is stuck permanently with no marker, no + terminal, and no rollover ever requested. ADR-007 rejects exactly that + check-and-fail behaviour. Failing the Workflow instead is still bad news, but + it is bounded, visible, and reported once. + """ + + def __init__(self, message: str) -> None: + super().__init__(message, type="AnnotationBudgetExceeded", non_retryable=True) + + +# --- the decoded shape ------------------------------------------------------ + + +@dataclass(frozen=True) +class StreamBinding: + """What a ``wait_id`` was subscribed to, where it started, and through what. + + ``start_cursor`` is explicit rather than re-derived: it is what makes an + annotation with no segments at all -- a subscription to an empty stream -- + a complete replay instruction. + + ``stream_key`` is chosen by Workflow code, so a mismatch is nondeterminism. + ``provider_id`` and ``provider_format_version`` describe the backend + configured on the Worker, so a mismatch is a deployment problem: the + Workflow is unchanged and the store need not be damaged. + """ + + stream_key: StreamKey + start_cursor: Cursor + provider_id: str + """The configured backend's provider when this wait was recorded.""" + provider_format_version: int + """That provider's on-the-wire format version when this wait was recorded. + + Checked on replay rather than merely stored: a backend implementation that + keeps its provider id while changing how it lays records out would otherwise + be read as though nothing had happened. + """ + + +@dataclass(frozen=True) +class Run: + """A maximal set of consecutive deliveries from **one** stream. + + Alternating streams produce one run per delivery; a single-stream batch of + 100,000 records produces one run. That difference is the whole reason the + schedule is encoded as runs. + """ + + wait_id: int + first_offset: Offset + last_offset: Offset + count: int + control_positions: tuple[int, ...] = () + """Sparse relative indices within the run at which a control record sat. + + Sparse, not one kind tag per record: fences are rare by construction, so + this keeps a run's control encoding proportional to the number of fences + rather than to the number of records. + """ + + def __post_init__(self) -> None: + """Validate the run length and sparse control positions.""" + if self.count < 1: + raise ValueError(f"a run covers at least one record, got {self.count}") + if any(not 0 <= p < self.count for p in self.control_positions): + raise ValueError( + f"control positions {self.control_positions} fall outside a run " + f"of {self.count} record(s)" + ) + if list(self.control_positions) != sorted(set(self.control_positions)): + raise ValueError( + f"control positions must be strictly increasing, got " + f"{self.control_positions}" + ) + + +@dataclass(frozen=True) +class Segment: + """One original activation's worth of deliveries. + + ``runs`` may be empty and that is meaningful: an activation that drained + and found nothing still ran one event-loop drain, and replay must reproduce + it or conditions fire a different number of times. + """ + + runs: tuple[Run, ...] + end_reason: SegmentEndReason + + +@dataclass(frozen=True) +class AnnotationHeader: + """The bindings, and nothing else. + + Holds **every** wait's binding once decoded, including the ones that arrived + in a later bindings frame because their ``subscribe()`` call ran after the + header had already gone to Core. Where a binding was carried is an encoding + detail; the decoded table is flat, so replay never has to ask when a wait + joined. + + There is deliberately no annotation-wide provider here. One existed through + schema version 1 and was taken from whichever subscription happened to be + registered first, which made it wrong for every other wait in a + multi-backend annotation. + """ + + streams: dict[int, StreamBinding] = field(default_factory=dict) + schema_version: int = SCHEMA_VERSION + + +@dataclass(frozen=True) +class Annotation: + """A complete marker annotation. + + ``terminal`` is ``None`` only for an annotation still being accumulated. + Core refuses to write a marker for one without a terminal, so a decoded + annotation always has one. + """ + + header: AnnotationHeader + segments: tuple[Segment, ...] = () + terminal: dict[int, Cursor] | None = None + + +# --- primitives ------------------------------------------------------------- + + +def _put_uvarint(out: bytearray, value: int) -> None: + if value < 0: + raise ValueError(f"cannot encode a negative value: {value}") + while True: + byte = value & 0x7F + value >>= 7 + out.append(byte | (0x80 if value else 0)) + if not value: + return + + +def _put_str(out: bytearray, value: str) -> None: + encoded = value.encode() + _put_uvarint(out, len(encoded)) + out += encoded + + +def _put_cursor(out: bytearray, cursor: Cursor) -> None: + """Cursors and offsets are distinct types in the encoding, on purpose. + + A boundary and a record identity are not interchangeable, and an encoding + that spelled them the same way would let one be read as the other. + """ + if cursor.is_beginning: + out.append(_CURSOR_BEGINNING) + return + assert cursor.offset is not None + out.append(_CURSOR_AFTER) + _put_str(out, cursor.offset.serialize()) + + +def _put_stream_key(out: bytearray, key: StreamKey) -> None: + _put_str(out, key.namespace) + _put_str(out, key.workflow_id) + _put_str(out, key.first_execution_run_id) + _put_str(out, key.stream_name) + + +class _Reader: + def __init__(self, data: bytes) -> None: + self._data = data + self._at = 0 + + @property + def exhausted(self) -> bool: + return self._at >= len(self._data) + + def _take(self, n: int) -> bytes: + if self._at + n > len(self._data): + raise AnnotationDecodeError( + f"annotation truncated: wanted {n} byte(s) at offset {self._at}, " + f"only {len(self._data) - self._at} remain" + ) + chunk = self._data[self._at : self._at + n] + self._at += n + return chunk + + def byte(self) -> int: + return self._take(1)[0] + + def uvarint(self) -> int: + value = 0 + shift = 0 + while True: + byte = self.byte() + value |= (byte & 0x7F) << shift + if not byte & 0x80: + return value + shift += 7 + if shift > 63: + raise AnnotationDecodeError("varint is too long to be valid") + + def string(self) -> str: + return self._take(self.uvarint()).decode() + + def cursor(self) -> Cursor: + kind = self.byte() + if kind == _CURSOR_BEGINNING: + return BEGINNING + if kind == _CURSOR_AFTER: + return AFTER(Offset.deserialize(self.string())) + raise AnnotationDecodeError(f"unknown cursor kind {kind:#x}") + + def stream_key(self) -> StreamKey: + return StreamKey(self.string(), self.string(), self.string(), self.string()) + + def bindings(self) -> dict[int, StreamBinding]: + streams: dict[int, StreamBinding] = {} + for _ in range(self.uvarint()): + # Read into locals rather than nesting the calls in the constructor: + # argument evaluation order is not the thing that should decide which + # field a byte lands in. + wait_id = self.uvarint() + stream_key = self.stream_key() + start_cursor = self.cursor() + provider_id = self.string() + provider_format_version = self.uvarint() + streams[wait_id] = StreamBinding( + stream_key, + start_cursor, + provider_id, + provider_format_version, + ) + return streams + + +# --- frame encoding --------------------------------------------------------- + + +def _put_bindings(out: bytearray, streams: Mapping[int, StreamBinding]) -> None: + _put_uvarint(out, len(streams)) + for wait_id in sorted(streams): + binding = streams[wait_id] + _put_uvarint(out, wait_id) + _put_stream_key(out, binding.stream_key) + _put_cursor(out, binding.start_cursor) + _put_str(out, binding.provider_id) + _put_uvarint(out, binding.provider_format_version) + + +def encode_header(header: AnnotationHeader) -> bytes: + """The schema version and header frame, which lead every annotation.""" + out = bytearray() + _put_uvarint(out, header.schema_version) + out.append(_FRAME_HEADER) + _put_bindings(out, header.streams) + return bytes(out) + + +def encode_bindings(streams: Mapping[int, StreamBinding]) -> bytes: + """Bindings for waits that were registered after the header was emitted. + + The same body as the header frame under a different tag. A second *header* + frame would have been ambiguous -- it also carries the schema version, and + a decoder would have to decide whether the later one replaced the earlier -- + whereas a bindings frame says exactly one thing: these waits exist too. + """ + out = bytearray() + out.append(_FRAME_BINDINGS) + _put_bindings(out, streams) + return bytes(out) + + +def _put_run(out: bytearray, run: Run) -> None: + _put_uvarint(out, run.wait_id) + _put_str(out, run.first_offset.serialize()) + _put_str(out, run.last_offset.serialize()) + _put_uvarint(out, run.count) + _put_uvarint(out, len(run.control_positions)) + previous = 0 + for position in run.control_positions: + # Delta-encoded: control positions ascend, so the gaps are smaller + # numbers than the absolute indices and cost fewer varint bytes. + _put_uvarint(out, position - previous) + previous = position + + +def encode_segment(segment: Segment) -> bytes: + out = bytearray() + out.append(_FRAME_SEGMENT) + _put_uvarint(out, len(segment.runs)) + for run in segment.runs: + _put_run(out, run) + out.append(int(segment.end_reason)) + return bytes(out) + + +def encoded_run_size(run: Run) -> int: + """What one run costs inside a segment frame, in bytes. + + Exposed because the runtime has to know what an *open* segment costs before + it is closed: a segment frame that no longer fits the annotation's remaining + budget cannot be deferred to the next annotation -- its deliveries happened + in this Workflow Task and the marker for that task is where replay must find + them -- so the only place to act is before the records that would grow it are + handed over. Measured rather than estimated per record, because a run's cost + is dominated by two provider-supplied offset strings whose length this side + does not choose. + """ + out = bytearray() + _put_run(out, run) + return len(out) + + +def encoded_segment_size(run_sizes: Sequence[int]) -> int: + """What a segment frame costs, given what each of its runs costs. + + Takes the per-run sizes rather than the runs so that a caller extending one + run at a time re-measures only the run it changed. + """ + out = bytearray() + _put_uvarint(out, len(run_sizes)) + # Frame tag, the run count, the runs themselves, and the end-reason byte. + return 1 + len(out) + sum(run_sizes) + 1 + + +def encode_terminal(blocked: dict[int, Cursor]) -> bytes: + out = bytearray() + out.append(_FRAME_TERMINAL) + _put_uvarint(out, len(blocked)) + for wait_id in sorted(blocked): + _put_uvarint(out, wait_id) + _put_cursor(out, blocked[wait_id]) + return bytes(out) + + +def encode_annotation(annotation: Annotation) -> bytes: + """A whole annotation, as Core would have accumulated it.""" + parts = [encode_header(annotation.header)] + parts.extend(encode_segment(s) for s in annotation.segments) + if annotation.terminal is not None: + parts.append(encode_terminal(annotation.terminal)) + return b"".join(parts) + + +def decode_annotation(data: bytes) -> Annotation: + """Parses a complete annotation. + + Accepts an annotation with no terminal so a partially accumulated one can + be inspected; Core is what refuses to *write* a marker for one. + """ + reader = _Reader(data) + schema_version = reader.uvarint() + if schema_version != SCHEMA_VERSION: + raise AnnotationDecodeError( + f"annotation schema version {schema_version} is not supported by this " + f"SDK, which writes and reads version {SCHEMA_VERSION}" + ) + + if reader.byte() != _FRAME_HEADER: + raise AnnotationDecodeError("an annotation must begin with its header frame") + streams = reader.bindings() + + segments: list[Segment] = [] + terminal: dict[int, Cursor] | None = None + while not reader.exhausted: + frame = reader.byte() + if frame == _FRAME_BINDINGS: + if terminal is not None: + raise AnnotationDecodeError("a bindings frame follows the terminal") + for wait_id, binding in reader.bindings().items(): + if wait_id in streams: + # A wait is bound once. A second binding for the same id + # would leave replay choosing between two stream keys, and + # whichever it chose could be the one the records were not + # written to. + raise AnnotationDecodeError( + f"external stream wait {wait_id} is bound twice in one " + "annotation" + ) + streams[wait_id] = binding + elif frame == _FRAME_SEGMENT: + if terminal is not None: + raise AnnotationDecodeError("a segment frame follows the terminal") + runs = [] + for _ in range(reader.uvarint()): + wait_id = reader.uvarint() + first = Offset.deserialize(reader.string()) + last = Offset.deserialize(reader.string()) + count = reader.uvarint() + positions: list[int] = [] + previous = 0 + for _ in range(reader.uvarint()): + previous += reader.uvarint() + positions.append(previous) + runs.append(Run(wait_id, first, last, count, tuple(positions))) + segments.append(Segment(tuple(runs), SegmentEndReason(reader.byte()))) + elif frame == _FRAME_TERMINAL: + if terminal is not None: + raise AnnotationDecodeError("an annotation has one terminal, not two") + terminal = {} + for _ in range(reader.uvarint()): + # Read into locals rather than `terminal[reader.uvarint()] = + # reader.cursor()`: Python evaluates the assigned value before + # the subscript, which would read the two fields backwards. + wait_id = reader.uvarint() + terminal[wait_id] = reader.cursor() + else: + raise AnnotationDecodeError(f"unknown frame tag {frame:#x}") + + return Annotation( + AnnotationHeader(streams, schema_version), tuple(segments), terminal + ) + + +# --- accumulation and the byte budget --------------------------------------- + + +class AnnotationAccumulator: + """Builds one Workflow Task's annotation, one observation delta at a time. + + Mirrors what Core holds: every delta this emits is appended verbatim to + ``ExternalWaitSet.replay_annotation``, so the concatenation of everything + emitted here *is* the annotation the marker carries. There is no separate + assembly step that could disagree. + + The budget is enforced here, while encoding, rather than checked afterwards + -- by the time an oversized annotation exists it is too late to do anything + but discard work. + """ + + def __init__( + self, + header: AnnotationHeader, + *, + max_bytes: int = MAX_ANNOTATION_BYTES, + high_water: float = ROLLOVER_HIGH_WATER, + ) -> None: + """Start an accumulator with its immutable header and byte limits.""" + self._max_bytes = max_bytes + self._high_water_bytes = int(max_bytes * high_water) + self._emitted: list[bytes] = [] + self._size = 0 + self._closing = 0 + self._spill = 0 + self._terminated = False + self._emit(encode_header(header), closing=True) + + @property + def size(self) -> int: + """Bytes accumulated so far, which is what the marker will carry.""" + return self._size + + @property + def reserved(self) -> int: + """Bytes held back from the ordinary run of segments. + + Two parts, and they are held back against different things: + + - **closing** -- the terminal, plus a bindings frame for any wait + registered since the header went out. Neither may ever be refused: both + record something that has already happened, and an annotation whose + terminal does not fit is one Core writes *without* a terminal -- + durable, and undecodable past the frame after it. Only they may spend + this. + - **spill** -- a margin a segment frame may overrun into. The runtime + stops delivering before a segment it could not record, but it prices a + record it has not seen yet, and a provider chooses how long its offsets + are. The margin is what turns a misprice into a rollover instead of a + refusal; an annotation that dips into it has already asked Core to end + the Workflow Task. + """ + return self._closing + self._spill + + def reserve(self, closing: int, *, spill: int = 0) -> None: + """Sets what closing costs, and how much a segment may overrun.""" + self._closing = closing + self._spill = spill + + @property + def headroom(self) -> int: + """Bytes available to a segment before it starts spending the margin. + + What the runtime measures affordability against. Zero does not mean the + annotation is full; it means the next segment overruns into the spill and + the Workflow Task has to end. + """ + return max(0, self._max_bytes - self._closing - self._spill - self._size) + + def fits(self, frame_bytes: int) -> bool: + """Whether a segment of this size fits without spending the margin.""" + return frame_bytes <= self.headroom + + @property + def request_rollover(self) -> bool: + """Whether the next progress report should ask Core to roll over. + + Set once the high-water mark is passed. Core then rolls the task over + *without* a finalization round trip, because the progress report + carrying this flag already carried the terminal. + + The high-water mark alone is not the whole condition -- it is a fraction + of the budget, and a frame can be larger than the fraction that is left. + The runtime adds the other half by refusing to grow a segment it could not + then record; see ``WorkflowStreamRuntime.request_rollover``. + """ + return self._size >= self._high_water_bytes + + @property + def terminated(self) -> bool: + """Whether a terminal frame has already been emitted.""" + return self._terminated + + def accumulated(self) -> bytes: + """Everything emitted so far, concatenated -- what Core now holds.""" + return b"".join(self._emitted) + + def add_bindings(self, streams: Mapping[int, StreamBinding]) -> bytes: + """Encodes bindings for waits registered after the header went out. + + Its own frame rather than an amended header: Core appends the deltas it + is given and never rewrites what it already holds, so the only way a + binding decided later can reach the marker is to be appended after the + bytes that preceded it. + """ + if self._terminated: + raise ValueError("cannot bind a wait after the terminal") + if not streams: + raise ValueError("a bindings frame binds at least one wait") + # A closing frame for budget purposes: the wait it binds was admitted by + # `register`, which charged this frame into the reserve at the time, and + # refusing it now would leave the wait in the terminal with no stream key, + # no backend, and no start cursor -- which replay reports as a wait the + # Workflow never created. + return self._emit(encode_bindings(streams), closing=True) + + def add_segment(self, segment: Segment) -> bytes: + """Encodes one activation's segment and returns it as a delta.""" + if self._terminated: + raise ValueError("cannot add a segment after the terminal") + return self._emit(encode_segment(segment)) + + def add_terminal(self, blocked: dict[int, Cursor]) -> bytes: + """Encodes the blocked snapshot that closes this annotation.""" + if self._terminated: + raise ValueError("an annotation has one terminal, not two") + delta = self._emit(encode_terminal(blocked), closing=True) + self._terminated = True + return delta + + def _emit(self, frame: bytes, *, closing: bool = False) -> bytes: + # A closing frame -- the header, a bindings frame, the terminal -- may + # spend everything, including the margin, because it is what makes the + # annotation readable at all. A segment may spend the margin but not the + # closing reserve: overrunning into the margin is a rollover, overrunning + # past it is the failure this is the last line against. + cap = self._max_bytes if closing else self._max_bytes - self._closing + if self._size + len(frame) > cap: + raise AnnotationBudgetExceeded( + f"encoding {len(frame)} more byte(s) would take the annotation to " + f"{self._size + len(frame)}, past the {cap}-byte limit " + f"({self._max_bytes}-byte budget less {self._closing} reserved for " + "the frames that close it); the runtime should have stopped " + "delivering and asked Core to roll the Workflow Task over before " + "this" + ) + self._emitted.append(frame) + self._size += len(frame) + return frame diff --git a/temporalio/contrib/external_workflow_streams/_api.py b/temporalio/contrib/external_workflow_streams/_api.py new file mode 100644 index 0000000000..2ae02d3914 --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_api.py @@ -0,0 +1,898 @@ +"""The Workflow-facing API (P9). + +.. code-block:: python + + streams = external_stream.with_options(idle_timeout=timedelta(seconds=1)) + tokens = streams.topic("tokens", type=str) + + async for token in tokens.subscribe(): + process(token) + +Workflow code never constructs or imports the backend. Its provider instance +holds connections and credentials, lives on the Worker outside the sandbox, +and is reached only through an opaque handle. + +Everything here is a mirror image of the shipped +:py:mod:`temporalio.contrib.workflow_streams`, not a second implementation of +it, so no name may collide -- and in particular no name here may begin with +``__temporal_workflow_stream``. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncGenerator, AsyncIterator, Sequence +from dataclasses import dataclass, field, replace +from datetime import timedelta +from typing import Any, Generic, Protocol + +import temporalio.workflow +from temporalio.contrib.external_workflow_streams._backend import StreamKey +from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec +from temporalio.contrib.external_workflow_streams._errors import ( + ConcurrentStreamConsumerError, + StreamError, + classify_read_failure, +) +from temporalio.contrib.external_workflow_streams._record import ( + StreamRecord, +) +from temporalio.types import AnyType + +__all__ = [ + "merge", + "ExternalStreamOptions", + "ExternalStreamSubscription", + "ExternalStreamTopic", + "external_stream", +] + +DEFAULT_IDLE_TIMEOUT = timedelta(seconds=1) +"""How long the complete blocked set waits before parking. + +A property of the **set**, not of one subscription: one idle stream must not +park a Workflow Task another stream is still driving. +""" + +MAX_RECORDS_PER_ACTIVATION = 256 +"""How many records one activation may hand to Workflow code. + +Without a cap, a producer that keeps the buffer non-empty makes one +``activate()`` call never return: the iterator re-fills before every record and +always finds one. Activations run on a thread-pool executor under a **2-second +deadlock timeout**, so the Workflow Task fails -- and every retry hits the same +producer, so the Workflow is stuck permanently rather than merely slowed. Driving +the real iterator against a never-empty buffer consumed 316,086 records in two +seconds and blocked zero times. + +A **record count, never a duration**. A time-based cap would be nondeterministic, +and because the boundary each activation stopped at is recorded in the +annotation, a replay of the same records would cut the segments somewhere else +and diverge from the live run. +""" + +#: Where per-Run subscription state hangs off the Workflow instance. Reserved, +#: and deliberately not in the `__temporal_workflow_stream*` namespace the +#: shipped contrib feature already owns. +_RUN_STATE_ATTR = "__temporal_external_stream_state" + + +class ExternalStreamRuntime(Protocol): + """What Workflow code needs from the Worker, and nothing more. + + An opaque handle across the sandbox boundary. Workflow code never sees the + provider instance configured on the Worker. + """ + + def stream_key(self, stream_name: str) -> StreamKey: + """The full stream identity for a name, from the running Run's chain.""" + ... + + def register( + self, + *, + wait_id: int, + stream_key: StreamKey, + idle_timeout: timedelta, + ) -> None: + """Registers a wait with the Worker's subscription manager. + + ``idle_timeout`` is the *configured* value for this one subscription. + The runtime reduces the quiescent set's values to one; this side has no + business doing that reduction, because it can only see one member. + """ + ... + + def drain(self, wait_id: int, max_records: int | None = None) -> list[StreamRecord]: + """Pops buffered records. Performs no I/O.""" + ... + + def delivery_budget_remaining(self) -> int: + """How many more records this activation may hand to Workflow code. + + Lives on the runtime rather than on a subscription because the budget is + an *activation* budget: :py:func:`merge` consumes from several + subscriptions in one activation, and a per-subscription counter would let + *n* streams run *n* times as long. + """ + ... + + def codec_for( + self, value_type: type | None, wait_id: int + ) -> StreamPayloadCodec[Any]: + """The Workflow's DataConverter, bound to a topic's declared type. + + ``wait_id`` selects the *context* the converter carries. It is this + Run's own for a live record, and the one the marker recorded for a wait + a replayed record arrives on -- which an offline ``Replayer``, running + under its own namespace, is not otherwise able to reproduce. + """ + ... + + def new_readiness_future(self) -> asyncio.Future[None]: + """A future the readiness activation handler will resolve. + + Created by the runtime rather than here because it must belong to the + Workflow's own deterministic event loop, which this module has no + business reaching into. + """ + ... + + def record_delivery(self, wait_id: int, record: StreamRecord) -> None: + """Notes one record reaching the Workflow, in observed global order. + + Called for control records too. They are never yielded, but they occupy + offsets inside a run, so the annotation has to know about them. + """ + ... + + def record_consumption(self, wait_id: int, record: StreamRecord) -> None: + """Notes one delivered record leaving the subscription's ready list.""" + ... + + def note_blocked(self, wait_id: int, blocked: bool) -> None: + """Records whether Workflow code is currently waiting on this wait.""" + ... + + def unsubscribe(self, wait_id: int) -> None: + """Ends a wait and stops the Worker serving it. + + Keeps the wait's recorded state -- replay and a Continue-As-New + successor both still need it -- and only makes it unblockable. + """ + ... + + def register_pending(self, wait_id: int, future: asyncio.Future[None]) -> None: + """Registers the future the readiness activation will resolve.""" + ... + + def discard_pending(self, wait_id: int) -> None: + """Forgets a future that is no longer being awaited.""" + ... + + +@dataclass +class _RunState: + """Per-Run subscription bookkeeping. + + Lives on the Workflow instance rather than in a module global, so it shares + the instance's lifetime exactly -- a module global would outlive an evicted + Run and hand its wait ids to the next one. + """ + + runtime: ExternalStreamRuntime | None = None + next_wait_id: int = 1 + #: `wait_id -> Future`, resolved by the readiness activation handler. + pending: dict[int, Any] = field(default_factory=dict) + + +def _run_state() -> _RunState: + instance = temporalio.workflow.instance() + state = getattr(instance, _RUN_STATE_ATTR, None) + if state is None: + state = _RunState() + setattr(instance, _RUN_STATE_ATTR, state) + return state + + +def _install_runtime( # pyright: ignore[reportUnusedFunction] + instance: Any, runtime: ExternalStreamRuntime +) -> None: + """Gives a Workflow instance its handle to the Worker's manager. + + Called by the Worker when it creates the instance. + """ + state = getattr(instance, _RUN_STATE_ATTR, None) + if state is None: + state = _RunState() + setattr(instance, _RUN_STATE_ATTR, state) + state.runtime = runtime + + +@dataclass(frozen=True) +class ExternalStreamOptions: + """The entry point, and the options every topic under it inherits.""" + + idle_timeout: timedelta = DEFAULT_IDLE_TIMEOUT + + def with_options( + self, *, idle_timeout: timedelta | None = None + ) -> ExternalStreamOptions: + """A copy with the given options replaced. + + Args: + idle_timeout: How long the complete blocked set waits with no record + on any active subscription before parking. + + Raises: + ValueError: The timeout is not positive. Rejected rather than + coerced -- zero would park instantly and a negative value means + nothing at all, so neither can be what a caller intended. + """ + if idle_timeout is not None and idle_timeout <= timedelta(0): + raise ValueError( + f"idle_timeout must be positive, got {idle_timeout}. A " + "non-positive timeout is a configuration error rather than a " + "request to park immediately." + ) + return replace( + self, + idle_timeout=self.idle_timeout if idle_timeout is None else idle_timeout, + ) + + def topic( + self, name: str, *, type: type[AnyType] | None = None + ) -> ExternalStreamTopic[Any]: + """A handle for one stream. + + Args: + name: The stream name. The only place it appears. + type: The value type, used as the decode hint. + """ + if not name: + raise ValueError("a topic needs a non-empty name") + return ExternalStreamTopic(name=name, value_type=type, options=self) + + +@dataclass(frozen=True) +class ExternalStreamTopic(Generic[AnyType]): + """One stream, from the Workflow's side. + + Has ``subscribe`` and no ``publish``: the consumer and producer handles are + distinct types, not one object passed across a process boundary. + """ + + name: str + value_type: type[AnyType] | None + options: ExternalStreamOptions + + def subscribe(self) -> ExternalStreamSubscription[AnyType]: + """Starts a new subscription and returns its async iterator. + + Each call is an **independent** subscription with its own ``wait_id``, + cursor, and park intent -- even two calls naming the same stream. + Delivery is broadcast, so each sees every record from its own cursor. + + ``wait_id`` comes from a per-Run counter in call order, which puts it in + the same hazard class as timers and activities: inserting, removing, or + reordering a ``subscribe()`` call renumbers every later wait in the Run + and must be gated behind ``workflow.patched()``. + """ + state = _run_state() + if state.runtime is None: + raise RuntimeError( + "external streams are not configured on this Worker; pass " + "external_stream_backend=... to the Worker" + ) + wait_id = state.next_wait_id + state.next_wait_id += 1 + + stream_key = state.runtime.stream_key(self.name) + state.runtime.register( + wait_id=wait_id, + stream_key=stream_key, + # Passed on registration rather than read back from the subscription + # later, because the runtime is what holds the quiescent set and + # reduces it with `min`. Leaving it out is not a smaller default -- + # it silently substitutes DEFAULT_IDLE_TIMEOUT for whatever + # `with_options` was given, so no configured value can ever reach + # the reduction and every set parks after one second. + idle_timeout=self.options.idle_timeout, + ) + # Registering a wait is not blocking on one. The quiescent snapshot is a + # request to Core to retain the Workflow Task and, once the idle timer + # expires, to park it; a subscription Workflow code has not begun + # iterating has no coroutine waiting on it, so including it would ask + # Core to hold a Workflow Task open -- and eventually park it -- for a + # wait nothing will ever resolve. The first `_await_readiness` is what + # enters the blocked state, and that transition is what the wait + # generation counts. + state.runtime.note_blocked(wait_id, False) + return ExternalStreamSubscription( + topic=self, wait_id=wait_id, stream_key=stream_key, state=state + ) + + +async def merge( + *subscriptions: ExternalStreamSubscription[Any], +) -> AsyncGenerator[tuple[ExternalStreamSubscription[Any], Any], None]: + """Iterates several subscriptions as one wait, in delivery order. + + Yields ``(subscription, value)`` rather than bare values: the streams may + carry different types, and a merged value whose origin had to be guessed + from its shape would be unusable for anything but logging. + + The subscriptions are already one wait set -- every subscription a Run holds + is -- so this adds no coordination of its own. What it adds is the ability to + *wait on all of them at once*: iterating them one at a time would block on + the first while records piled up on the second, and the idle timer covering + the set would then fire against a Workflow that was not actually idle. + + Each pass takes **at most one record from each subscription**, in ``wait_id`` + order, resuming after the subscription that last took one. All three of + those are load-bearing: + + - *In ``wait_id`` order*, which is what makes the interleaving reproduce. + Records that arrived in one batch across two streams have no inherent + order between them, so an order that depended on dict iteration, on + arrival time, or on which watcher happened to run first would replay + differently than it ran. ``wait_id`` comes from a per-Run counter in + ``subscribe()`` call order, so replay reconstructs the same total order + from the Workflow code itself, and the pass then depends only on which + waits have a record ready -- which replay reconstructs from the recorded + segments. + - *At most one record*, which is what makes it a merge rather than a + priority order. Draining one subscription's whole ready list first lets + the lowest ``wait_id`` spend the entire + :py:data:`MAX_RECORDS_PER_ACTIVATION` budget by itself, and the next + activation starts the same pass in the same order, so a continuously + backlogged first stream starves every later one forever. + - *Resuming after the last take*, which is what makes "at most one each" add + up to fairness across activations rather than only within a pass. The + budget covers the merged set, so a pass can be cut anywhere inside it, and + a pass that always restarted at the lowest wait id would be cut in the + same place every activation and re-privilege the same prefix forever. With + 257 ready subscriptions the 257th is never asked at all -- ``_fill`` + returns on the spent budget before it reaches the manager, so that wait is + not merely served nothing, it is never enquired after. With 100 the first + 56 take one record per activation more than the rest, and the gap grows + without bound. Rotating the start is what makes the skew between any two + continuously ready streams what it is claimed to be: a single record. + + A control record spends the subscription's turn: it is consumed, because it + occupies an offset inside a run, and the pass moves on. Filling one record + at a time is also what keeps the budget exact -- a fill that took the whole + remaining budget into one subscription's ready list would let the rest of + that list be consumed after the budget was spent, and would strand it there, + since the completion path re-arms readiness from the *manager's* buffers and + knows nothing about records already popped out of them. + + The recorded delivery schedule follows: alternating records across two + streams encode as one run per delivery, because a run is a maximal + consecutive stretch from a single wait. + + The delivery budget covers the whole set rather than each member, so a merge + over *n* never-empty streams still returns after + :py:data:`MAX_RECORDS_PER_ACTIVATION` records. A per-subscription budget + would let the activation run *n* times as long, which is the same deadlock + with a larger constant in front of it. + + **Ending a merge does not end the subscriptions it merged.** Stopping at a + value -- breaking out of the ``async for``, or ``aclose()`` on the generator + -- leaves every subscription open, at its own cursor, and consumable again + either singly or in another merge. Closing them is + :meth:`ExternalStreamSubscription.close`, and that is a separate decision, + because it is the one that drops what the buffer still holds without + consuming it. Nothing is left blocked in the meantime: the merged wait + unblocks every member when it ends, so a Workflow that stops consuming and + goes on to wait for something else is not holding a Workflow Task open for a + stream nobody is reading. + """ + ordered = sorted(subscriptions, key=lambda s: s.wait_id) + if not ordered: + raise ValueError("merge() needs at least one subscription") + if len({s.wait_id for s in ordered}) != len(ordered): + raise ValueError( + "merge() was given the same subscription twice; each one is a " + "separate wait with its own cursor, and merging one with itself " + "would deliver every record to it twice" + ) + + # The wait that last took a record, so the next pass resumes after it rather + # than restarting at the lowest wait id. + # + # Local to this generator, recorded nowhere, and that is what makes it + # replay-safe rather than merely convenient. Replay serves a drain from the + # **front** of the recorded segment and only while that front belongs to the + # asking wait, so a wait asked out of turn gets nothing and the record stays + # for whoever asks next. Every active wait is still asked exactly once per + # pass, so the yielded sequence is the recorded global order whatever + # position the pass starts at. Under replay the budget is unbounded, so + # passes are not cut where they were cut live and this cursor generally ends + # up somewhere else than it did -- which steers nothing but *later live* + # fairness, and no live schedule was ever recorded for History to contradict. + resume_after = 0 + while True: + # A closed subscription leaves the set rather than blocking it: it has + # no coroutine behind it, so including it in the wait would ask Core to + # retain the Workflow Task for a wait nothing can resolve. + active = [s for s in ordered if not s._finished] + if not active: + return + # The first wait past the last take, wrapping to the front when there is + # none -- which is also what happens when the wait the cursor named has + # since closed and left the set. + start = next((i for i, s in enumerate(active) if s.wait_id > resume_after), 0) + delivered_any = False + for step in range(len(active)): + subscription = active[(start + step) % len(active)] + subscription._fill(1) + record = subscription._peek() + if record is None: + continue + delivered_any = True + # Advanced only where a record was actually taken, control records + # included, because those spend the turn and the budget too. A wait + # that had nothing has not had its turn, and moving the cursor past + # it would cost it the turn it never got -- which is the starvation + # this exists to end, reintroduced from the other side. + resume_after = subscription.wait_id + if record.is_control: + subscription._commit(record) + continue + # Decoded before consumption is committed, for the same reason as in + # `_iterate`: a decode that raises must leave the record where a + # later pass can still find it. + value = subscription._decode(record) + subscription._commit(record) + yield subscription, value + if not delivered_any: + # Nothing anywhere: block on all of them at once. Whichever wait + # Core resolves first wakes this, and the next pass picks it up. + if await _await_any_readiness(active): + continue + + +async def _await_any_readiness( + subscriptions: Sequence[ExternalStreamSubscription[Any]], +) -> bool: + """Blocks until any one of the waits is resolved. Returns whether it skipped. + + **One scoped operation over the whole group**, and both halves of that follow + from it. + + Every wait is marked blocked while it runs, because the quiescent snapshot + must name the **complete** set the Workflow is waiting on: a set missing one + member would let Core park the Workflow Task while that member was still + live. + + And every wait leaves the blocked set when it ends, however it ends -- a + resolved future, the double-check below returning early, or cancellation. + Only one member can have a record in hand afterwards, and :meth:`_peek` clears + the flag only for that one, so a member left blocked here is one no coroutine + is awaiting: the quiescent snapshot names it, and Core retains and eventually + parks the Workflow Task on behalf of a group wait that is already over. The + next genuine group wait marks every member blocked again and advances its + wait generation, which is a new blocking epoch rather than a continuation of + this one. + """ + runtime = subscriptions[0]._state.runtime + assert runtime is not None + # Every member is checked before any is registered. A merge that refused + # half-way would leave the earlier waits registered and blocked with no + # coroutine behind them, which is the state that asks Core to retain a + # Workflow Task for nobody. + for subscription in subscriptions: + subscription._refuse_a_second_waiter() + futures: list[asyncio.Future[None]] = [] + try: + # Inside the `try`, so that registering the group is covered by the same + # cleanup that ends it. A raise part-way through would otherwise leave + # the members already registered blocked with their futures unreachable + # -- the very state the refusal loop above is careful not to create. + for subscription in subscriptions: + runtime.note_blocked(subscription.wait_id, True) + future = runtime.new_readiness_future() + runtime.register_pending(subscription.wait_id, future) + # Also on the subscription, so `close()` can resume a merge that is + # sitting on this wait: the runtime's map is keyed for the side that + # resolves readiness, and closing is neither that side nor this one. + subscription._pending_future = future + futures.append(future) + # The same last look the single-subscription path takes, for the same + # reason: a record buffered before these waits were registered had its + # readiness reported to nobody. And for the same reason as there, it + # goes through `_fill` and so finds nothing once the budget is spent -- + # otherwise a merge over a busy stream would resume immediately and the + # activation would never end. + # + # One record, matching the pass above: a fill that took the whole + # remaining budget here would put records into a ready list the budget + # can no longer pay for, and the completion path re-arms readiness from + # the manager's buffers, which no longer hold them. + for subscription in subscriptions: + subscription._fill(1) + if subscription._ready: + return True + await asyncio.wait(futures, return_when=asyncio.FIRST_COMPLETED) + return False + finally: + for subscription in subscriptions: + # Unblocking belongs here, beside the removal of the futures it is + # the counterpart of: a wait whose readiness future has just been + # discarded is a wait nothing can be waiting on. See this function's + # docstring for what leaving one blocked would ask of Core -- and + # note that a member the loop above never reached is already + # unblocked, so this is a no-op for it rather than a claim about it. + runtime.note_blocked(subscription.wait_id, False) + subscription._pending_future = None + runtime.discard_pending(subscription.wait_id) + for future in futures: + if not future.done(): + future.cancel() + + +class ExternalStreamSubscription(Generic[AnyType]): + """One subscription's async iterator over decoded values. + + **One consumer.** The cursor, the readiness future, and the blocked flag are + all the subscription's rather than an iterator's, so two coroutines waiting on + it at once is refused with + :class:`~temporalio.contrib.external_workflow_streams._errors.ConcurrentStreamConsumerError` + rather than served -- see :meth:`_refuse_a_second_waiter` for what sharing + them would do. Two consumers of the same *stream* is a supported shape and + the way to ask for it is a second ``subscribe()``: delivery is a broadcast + (ADR-021), so each wait gets every record and keeps its own cursor. + + Iterating the same subscription again after the previous consumer has stopped + is fine, and resumes where it left off. + """ + + def __init__( + self, + *, + topic: ExternalStreamTopic[AnyType], + wait_id: int, + stream_key: StreamKey, + state: _RunState, + ) -> None: + """Bind one wait ID and stream key to its per-Run state.""" + self._topic = topic + self._wait_id = wait_id + self._stream_key = stream_key + self._state = state + self._ready: list[StreamRecord] = [] + #: Set by :meth:`close`, and the only thing that ends iteration. + self._finished = False + #: The readiness future currently being awaited, if any. Held here as + #: well as on the runtime because :meth:`close` has to be able to resume + #: the coroutine sitting on it, and the runtime's map is keyed by wait + #: id for the *resolving* side rather than for this one. + self._pending_future: asyncio.Future[None] | None = None + + @property + def wait_id(self) -> int: + """The Run-local identity Core uses for this subscription.""" + return self._wait_id + + @property + def stream_key(self) -> StreamKey: + """The durable stream identity this subscription reads.""" + return self._stream_key + + @property + def idle_timeout(self) -> timedelta: + """This subscription's configured contribution to the wait-set timeout.""" + return self._topic.options.idle_timeout + + def __aiter__(self) -> AsyncIterator[AnyType]: + """A fresh generator over the *same* cursor, not an independent view. + + Which is why a second one running concurrently is refused rather than + interleaved: see :meth:`_refuse_a_second_waiter`. + """ + return self._iterate() + + async def _iterate(self) -> AsyncIterator[AnyType]: + while not self._finished: + # Re-filled before *every* record rather than once per batch. A + # record buffered while Workflow code was doing something else -- a + # timer, an activity, another stream -- has already had its readiness + # reported and consumed, so nothing will report it again. Blocking + # without looking would wait forever on a record that is already + # here. Filling is a buffer pop and costs nothing. + self._fill() + record = self._peek() + if record is not None: + if record.is_control: + self._commit(record) + continue + # Decoded first, committed second. Consumption is a claim that + # Workflow code *received* this record, and the claim is only + # true once a value exists: a decode that raises out of a + # mismatched converter never yields anything. Committing first + # makes the claim false in exactly that case -- the record + # leaves the ready list, the buffer it came from is already + # empty, and the consumption cursor a Continue-As-New successor + # resumes from has stepped over a record nothing ever saw. Left + # uncommitted it stays at the head of the ready list and the + # next `__anext__` retries it. + # + # Synchronous, so there is no longer a point *inside* decoding + # at which cancellation can land at all: the record either + # becomes a value or raises, and neither outcome can leave the + # ready list half-consumed. + value = self._decode(record) + self._commit(record) + yield value + continue + await self._await_readiness() + + def _fill(self, limit: int | None = None) -> None: + """Moves whatever is buffered into this subscription's ready list. + + Draining is a buffer pop and never touches the backend -- the record is + already here or it is not, and if it is not, only Core can say when to + look again. + + Bounded by the activation's remaining delivery budget, **and charged + against it here**. A drain that took the whole buffer would be handed + straight back by a producer that keeps refilling it, and this activation + would never return. Charging where the records move rather than where they + are later consumed is what makes the bound a reservation: a drain that + checked the budget and charged nothing left the same room visible to the + next subscription's drain, so two subscriptions consumed in two + independent coroutines took the whole budget each. + + ``limit`` bounds it further, for a caller that will consume fewer records + than the budget allows. :py:func:`merge` passes 1: it takes one record + per subscription per pass, and anything it pulled out of the manager's + buffer and did not take would be stranded there -- the completion path + re-arms readiness from the manager's buffers, so a record already popped + out of one has nothing left to announce it. + """ + if self._ready: + # Nothing to check the budget for: these records are already charged + # -- by the drain that took them, and again by the `begin_activation` + # of any later activation that inherits them, since each activation + # pays for what it may hand over. A check here would refuse a record + # this activation has already been billed for. + return + assert self._state.runtime is not None + budget = self._state.runtime.delivery_budget_remaining() + if limit is not None: + budget = min(budget, limit) + if budget <= 0: + # Nothing is taken even though records are sitting right here, so the + # caller blocks and the activation ends. The records are not lost: + # the completion path re-reports readiness for every buffer that is + # still non-empty, which is what brings the next activation in. + return + drained = self._state.runtime.drain(self._wait_id, budget) + for record in drained: + # Recorded for *every* record, control ones included: they occupy + # offsets inside a run, so a run's count includes them and their + # positions go in `control_positions`. Omitting them would make + # replay's range read find more records than the marker claims. + self._state.runtime.record_delivery(self._wait_id, record) + # Control records stay in the list rather than being filtered out here, + # so that consumption advances past them in order. A filtered control + # record would be neither consumed nor left behind, and the continuation + # cursor would step over it. + self._ready = list(drained) + + def _peek(self) -> StreamRecord | None: + """The next ready record, or ``None``. Commits nothing. + + Leaving this wait's blocked state is right here rather than at commit + time: a record is in hand, so no coroutine is waiting on this wait, and a + quiescent snapshot taken while the value is being decoded must not name + it. + """ + if not self._ready: + return None + assert self._state.runtime is not None + self._state.runtime.note_blocked(self._wait_id, False) + return self._ready[0] + + def _commit(self, record: StreamRecord) -> None: + """Pops the record and records that Workflow code now has it. + + Consumption is not delivery. A batch is delivered whole, but a Workflow + that stops iterating part-way through has consumed only its prefix, and + a successor Run resuming from the delivery cursor would step over the + rest. + + Called only once the record has actually become a value -- or, for a + control record, once it has been skipped, which is the whole of what + receiving one means. + + Spends no delivery budget: the drain that put this record here already + did. What it does tell the runtime is that the record has left the ready + list, which is what stops the next activation being charged for it. + """ + assert self._state.runtime is not None + assert self._ready and self._ready[0] is record + self._ready.pop(0) + self._state.runtime.record_consumption(self._wait_id, record) + + def _refuse_a_second_waiter(self) -> None: + """Refuses a second coroutine blocking on this subscription. + + One subscription, one consumer. Everything a blocked wait is found + through is single-slot -- :attr:`_pending_future` here, and the runtime's + pending map keyed by wait id -- so a second waiter does not queue behind + the first, it *replaces* it: readiness resolves only the newer future, its + ``finally`` removes the map entry, and the older one is left unreachable by + the readiness activation and by :meth:`close` alike. That coroutine is + stranded for the life of the Run, and the shared blocked flag can even say + the wait is no longer blocked while it is still sitting there. + + Refused here, at the point a *second waiter appears*, rather than at + ``__aiter__``. Two coroutines inside ``_iterate`` at once is only possible + while one of them is suspended, and the only suspension point in there is + this wait -- a generator suspended at its ``yield`` is not inside + ``_iterate`` at all, it is between ``__anext__`` calls. So this catches + every genuinely concurrent consumer, while a guard on the iterator would + also refuse the sequential shape that works today: taking some records, + breaking out of the loop, and coming back to the same subscription later. + A broken-out-of ``async for`` leaves its generator suspended rather than + closed, so the iterator-level guard cannot tell that shape from this one + (ADR-037). + """ + if self._pending_future is None: + return + raise ConcurrentStreamConsumerError( + f"external stream wait {self._wait_id} already has a coroutine " + "blocked on it. A subscription is a single consumer: it has one " + "cursor and one readiness future, so a second waiter would replace " + "the first and strand it permanently. Iterate one subscription from " + "one coroutine, or subscribe again -- a second subscription to the " + "same stream is its own wait with its own cursor, and delivery is a " + "broadcast." + ) + + async def _await_readiness(self) -> None: + """Blocks until the readiness activation resolves this wait. + + The future is resolved from the ``ResolveExternalStreamWaits`` branch of + the activation dispatch, which is the only thing that knows a record + arrived. + """ + assert self._state.runtime is not None + self._refuse_a_second_waiter() + # Entering the blocked state is what the wait generation counts, and is + # what later makes a readiness notification for *this* block + # distinguishable from one for a block already resolved. + self._state.runtime.note_blocked(self._wait_id, True) + future = self._state.runtime.new_readiness_future() + self._state.runtime.register_pending(self._wait_id, future) + self._pending_future = future + try: + # Look once more, now that this wait is registered. A record buffered + # between the last drain and this registration had its readiness + # reported while nothing was waiting for it, and no second + # notification is coming -- the watcher has already moved its + # prefetch cursor past it. Checking after registering is what closes + # that window: anything earlier is found here, anything later + # resolves the future. + # + # It goes through `_fill`, so it observes the delivery budget like + # every other drain. A refill that ignored the budget would hand this + # wait a record the instant the budget stopped it, and the activation + # would go right back to never returning -- the double-check would + # undo the cap rather than merely coexist with it. + self._fill() + if self._ready: + return + await future + except BaseException: + # Abandoned rather than resolved: cancellation, most often because + # Workflow code raced this stream against a timer and cancelled the + # loser. Nothing is awaiting this wait any more, so it must leave the + # blocked set -- a wait left in it is named by the quiescent + # snapshot, and Core would retain and eventually park the Workflow + # Task for a coroutine that no longer exists. + self._state.runtime.note_blocked(self._wait_id, False) + raise + finally: + self._pending_future = None + self._state.runtime.discard_pending(self._wait_id) + + def _decode(self, record: StreamRecord) -> AnyType: + """This record as a value of the topic's declared type. + + **Synchronous, and that is the point.** The Workflow thread runs one + half of ``DataConverter.decode`` -- ``from_payloads``, which needs the + topic's type and performs no I/O. The other half, external-payload + retrieval and the user's ``PayloadCodec``, already ran on the Worker's + loop before this record was buffered, exactly as it does for every other + payload an activation carries. Awaiting a codec here would perform real + I/O inside a deterministic event loop, synthesize Workflow commands out + of the codec's internals, and put arbitrary user work under the + 2-second deadlock timeout. + + Failures of *either* half surface from here, because here is where the + record would have become a value. A preparation error carried over from + the Worker's loop is raised at the delivery it belongs to rather than + where it happened: the watcher has no Workflow to tell, and a record + that is never delivered must not fail anything. + """ + assert self._state.runtime is not None + # By wait, not by Run: the asynchronous half ran under the stream key + # the record's own wait was bound to, and on replay that key comes from + # the marker rather than from this Run's identity. Passing the wait is + # what lets the two halves agree. + codec = self._state.runtime.codec_for(self._topic.value_type, self._wait_id) + prepared = getattr(record, "prepared_payload", None) + failed = getattr(record, "prepare_error", None) + try: + if failed is not None: + raise failed + if prepared is None: + prepared = codec.parse_unprepared(record.payload) + return codec.convert(prepared) + except StreamError: + # Already classified -- a storage failure reaching the payload store + # stays row one rather than being relabelled as the consumer's + # converter mismatch. + raise + except Exception as err: + # `range_validated=True` for both delivery paths, and for the same + # reason. Replay validated the recorded range before this record was + # prepared; live delivery read the record out of the backend and the + # provider guarantees a record's bytes cannot change once written + # (ADR-003). Either way the bytes are the bytes that were written, + # so what failed is the consumer's converter -- row three -- and not + # the stream (ADR-015). + raise classify_read_failure(range_validated=True, cause=err) from err + + def close(self) -> None: + """Ends this subscription: iteration stops and the wait goes away. + + Synchronous and idempotent, because the ordinary shape is a ``finally`` + that cannot know whether the iterator already ended. + + Three things happen, and each answers a distinct way an abandoned + subscription is visible: + + - the wait leaves the **blocked set**, so no quiescent snapshot names it + and Core is not asked to retain -- and eventually park -- the Workflow + Task for it; + - any coroutine sitting on this wait's readiness future is **resumed** + rather than cancelled, so ``_iterate`` reaches its loop condition and + the iteration ends with ``StopAsyncIteration`` instead of raising + ``CancelledError`` into Workflow code that merely closed a stream; + - records already drained but never handed over are **dropped without + being consumed**, so the consumption cursor stops short of them and a + Continue-As-New successor receives them. + + The Worker-side half -- stopping this wait's watcher and removing its + park intent -- is the manager's, and is reached through the runtime + rather than from here, because it has to be scheduled onto the Worker's + loop: this runs on the Workflow thread, where creating a task is + silently ineffective. + + What closing does **not** do is forget the subscription. Its recorded + binding is what replay needs to know what this wait was reading, and its + cursor is what stops a Continue-As-New successor restarting the stream + from the beginning. + """ + if self._finished: + return + self._finished = True + runtime = self._state.runtime + if runtime is None: + return + runtime.note_blocked(self._wait_id, False) + runtime.discard_pending(self._wait_id) + self._ready.clear() + future = self._pending_future + self._pending_future = None + if future is not None and not future.done(): + future.set_result(None) + runtime.unsubscribe(self._wait_id) + + +#: The entry point Workflow code uses. +external_stream = ExternalStreamOptions() diff --git a/temporalio/contrib/external_workflow_streams/_backend.py b/temporalio/contrib/external_workflow_streams/_backend.py new file mode 100644 index 0000000000..ca7e1ecdd3 --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_backend.py @@ -0,0 +1,459 @@ +"""The stream backend contract (P2). + +What a provider must implement to be registrable. The contract is small on +purpose -- every operation here exists because replay or parking cannot be made +correct without it. + +The obligations that are easy to get subtly wrong, and that the conformance +suite exists to catch: + +- **The range read is inclusive of both endpoints; the watch is exclusive.** + These are two different operations, not one with a flag. Replay issues an + inclusive read for a range the marker already names, and never asks the + backend what comes next. Implementing the inclusive read with exclusive + semantics is invisible until the first replay drops a record. +- **Offsets are compared by the provider's rule, not lexically.** Redis stream + ids compare as numeric ``(ms, seq)`` tuples, and string comparison is wrong + the moment the millisecond component changes width. +- **A cursor is a boundary, never the identity of a future record.** No + operation may require naming a record that does not exist yet, or a consumer + parked at the tail could not resume. +- **Append is idempotent on identity, not on key alone.** Reusing a + ``(session_id, sequence)`` with byte-identical content is a no-op returning + the original offset; reusing it with *different* bytes is an error, because + silently accepting it would let a retried producer overwrite history. +- **Records are immutable.** Mandatory, and checked at registration rather than + compensated for at runtime. +""" + +from __future__ import annotations + +import abc +import enum +from collections.abc import AsyncIterator +from dataclasses import dataclass +from datetime import timedelta +from typing import ClassVar, Final + +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + Cursor, + IdempotencyKey, + Offset, + StreamRecord, +) + +__all__ = [ + "AppendConflictError", + "ParkIntent", + "ParkIntentRemoval", + "StreamBackend", + "StreamDirection", + "StreamKey", +] + + +@dataclass(frozen=True) +class ParkIntent: + """One subscription's declaration that it is about to park. + + The key is ``(stream key, wait_id)``; everything here is the value. + """ + + wait_id: int + cursor: Cursor + """Where this subscription stopped consuming. What ``recheck`` reads past.""" + park_generation: int + """The quiescence generation being parked. A wake Signal names this.""" + run_id: str + """The Run that installed it -- part of the value so a new Run's intent + replaces its predecessor's for the same key rather than accumulating.""" + + def __post_init__(self) -> None: + """Reject the generation reserved for unparked wakes.""" + if self.park_generation < 1: + # 0 is the reserved unparked-wake sentinel, so a real park + # generation can never take it. + raise ValueError( + f"a park generation is a quiescence generation and starts at 1, " + f"got {self.park_generation}" + ) + + +@enum.unique +class StreamDirection(enum.Enum): + """Which endpoint role a physical stream serves. + + Direction is part of provider identity rather than a user-controlled prefix + on ``stream_name``. This permits an input and output topic to share a name + without sharing records, idempotency state, or coordination metadata. + """ + + INPUT = "input" + """Records produced externally and consumed by a Workflow.""" + + OUTPUT = "output" + """Records produced by a Workflow or Activity for an external client.""" + + +@dataclass(frozen=True) +class StreamKey: + """What identifies a stream, for the whole life of a Continue-As-New chain. + + ``first_execution_run_id`` rather than the current Run ID: it prevents + collisions after Workflow ID reuse while staying stable across the chain, + so a new Run continues the same stream rather than starting a fresh one. + """ + + namespace: str + workflow_id: str + first_execution_run_id: str + stream_name: str + direction: StreamDirection = StreamDirection.INPUT + """The endpoint direction, defaulting to the existing input protocol. + + The default preserves the meaning of every four-argument ``StreamKey`` + constructor used by existing input providers and applications. + """ + + def __post_init__(self) -> None: + """Require every component of the durable stream identity.""" + for field_name in ( + "namespace", + "workflow_id", + "first_execution_run_id", + "stream_name", + ): + if not getattr(self, field_name): + raise ValueError(f"a stream key needs a non-empty {field_name}") + if not isinstance(self.direction, StreamDirection): + raise TypeError( + "a stream key direction must be a StreamDirection, got " + f"{type(self.direction).__name__}" + ) + + def __str__(self) -> str: + """Return a slash-delimited diagnostic representation.""" + return ( + f"{self.namespace}/{self.workflow_id}/" + f"{self.first_execution_run_id}/{self.direction.value}/{self.stream_name}" + ) + + +class AppendConflictError(Exception): + """An idempotency key was reused with different bytes. + + Distinct from a successful idempotent retry, which returns the original + offset and appends nothing. Accepting this quietly would let a retried + producer attempt with changed input rewrite what a consumer may already + have delivered. + """ + + def __init__(self, key: IdempotencyKey) -> None: + """Create a conflict for the key reused with different content.""" + super().__init__( + f"idempotency key {key} was already used with different content; " + "an append is idempotent on identity, not on the key alone" + ) + self.key = key + + +#: How long a watch blocks before returning empty, when no timeout is given. +DEFAULT_WATCH_BLOCK: Final = timedelta(seconds=5) + + +class ParkIntentRemoval(enum.Enum): + """What a conditional park-intent removal found at the key. + + Three outcomes rather than a Boolean, because two of them mean the intent + the consumer named is gone and one means it never got out of the way, and + the consumer acts differently on each: a removal that happened ends a wake + suppression that may owe a record an announcement, while a live intent that + replaced it does not. + """ + + REMOVED = 1 + """The intent matched the named Run and park generation, and is gone.""" + + ABSENT = 2 + """Nothing was installed at the key. + + Reported for a key that is already clear -- including one this provider + itself cleared on a call whose reply never arrived. It is *not* the same as + a key some other intent holds. + """ + + MISMATCH = 3 + """A different Run or park generation is installed, and was left alone.""" + + +class StreamBackend(abc.ABC): + """A stream provider. + + Instances live **outside** the Workflow sandbox, on the Worker. Workflow + code never imports or holds the configured instance. + """ + + guarantees_immutability: ClassVar[bool | None] = None + """Whether this provider guarantees a record's bytes cannot change. + + ``None`` means *undeclared*, which is what an implementer who never + considered the question leaves it as. Registration requires ``True``, so + both "forgot" and "cannot" are rejected at Worker construction -- loudly, + before any Workflow can name the backend -- rather than at replay, quietly, + after data has already been consumed. + + The guarantee is what makes the four cheap range checks sufficient: given + it, the only damage replay has to detect is a record that is no longer + there. + """ + + provider_id: ClassVar[str] = "" + """Stable identifier recorded in every annotation header.""" + + provider_format_version: ClassVar[int] = 1 + """Bumped when this provider changes how it stores a record.""" + + # --- required operations ------------------------------------------------ + + @abc.abstractmethod + async def append(self, key: StreamKey, record: StreamRecord) -> StreamRecord: + """Appends one record, idempotently on ``(session_id, sequence)``. + + Returns: + The record as stored, with its assigned offset. For a repeat append + of byte-identical content this is the **original** offset and + nothing new is written. + + Raises: + AppendConflictError: The key was used before with different bytes. + """ + + @abc.abstractmethod + async def read_range( + self, key: StreamKey, first: Offset, last: Offset + ) -> list[StreamRecord]: + """Reads ``[first, last]`` -- **inclusive of both endpoints**. + + This is the replay read. The marker already names the range, so this + must never consult "what comes next"; it returns exactly what is + present in the closed interval, in increasing offset order, and says + nothing about what is missing. Detecting a missing record is the + caller's job, and it can only do it if this read is honest about what + it found. + """ + + @abc.abstractmethod + async def read_after( + self, + key: StreamKey, + after: Cursor, + *, + max_records: int, + block: timedelta | None = DEFAULT_WATCH_BLOCK, + ) -> list[StreamRecord]: + """Reads records **strictly after** ``after``, blocking up to ``block``. + + This is the live read. ``after`` is a boundary, so ``BEGINNING`` means + the start of the stream and ``AFTER(x)`` means "whatever follows the + record at ``x``" -- which must work whether or not such a record exists + yet, because that is exactly the state a consumer parked at the tail is + in. + + Returns an empty list when the block elapses with nothing new. Blocking + is a provider concern and never reaches the Workflow thread. + """ + + @abc.abstractmethod + def compare_offsets(self, left: Offset, right: Offset) -> int: + """Three-way comparison under **this provider's** ordering rule. + + Synchronous and pure: replay validation calls it once per adjacent pair + of a range, and an ordering that needed I/O would put backend latency + inside a validation loop. + """ + + # --- parking (P2b) ------------------------------------------------------ + + supports_leased_claims: ClassVar[bool] = False + """Whether :meth:`claim_park_generation` expires claims and permits takeover. + + An unleased claim introduces a failure mode with no recovery: a producer + that crashes between claiming a generation and sending its wake Signal + strands that generation, and every other producer concludes the wake is + already handled -- so the parked Workflow waits forever with data sitting + in the stream. + + A provider that cannot lease must declare ``False`` and expose + **observe-only** semantics instead: :meth:`claim_park_generation` always + grants, and every producer signals idempotently. That costs duplicate + Signals, which are harmless, rather than lost ones, which are not. + """ + + @abc.abstractmethod + async def install_park_intent(self, key: StreamKey, intent: ParkIntent) -> None: + """Records that one **subscription** intends to park. + + Keyed ``(stream key, wait_id)``, never by stream key alone: two + subscriptions in one Workflow to the same stream are two independent + waits, and a stream-keyed intent would have one overwrite the other -- + after which only one of them could ever be woken. + + The current Run ID is part of the intent's *value*, not its key, so a + new Run's intent deterministically replaces its predecessor's rather + than accumulating beside it. + """ + + @abc.abstractmethod + async def remove_park_intent(self, key: StreamKey, wait_id: int) -> None: + """Removes one subscription's intent. Idempotent.""" + + @abc.abstractmethod + async def remove_park_intent_if_matches( + self, + key: StreamKey, + wait_id: int, + *, + run_id: str, + park_generation: int, + ) -> ParkIntentRemoval: + """Atomically removes an intent only when its identity still matches. + + The comparison and removal must be one backend operation. A delayed + cleanup can overlap a Continue-As-New successor installing a new intent + at the same ``(stream key, wait_id)``; implementing this as + :meth:`park_intent` followed by :meth:`remove_park_intent` can delete + that successor's live park. + + All three outcomes must be reported distinctly. Collapsing + :attr:`ParkIntentRemoval.ABSENT` into + :attr:`ParkIntentRemoval.MISMATCH` tells the consumer that someone + else's park is in the way when in fact the key is clear -- which is what + a retry after a lost reply sees, the reply to a delete this provider + already performed. + """ + + @abc.abstractmethod + async def park_intent(self, key: StreamKey, wait_id: int) -> ParkIntent | None: + """The installed intent, if any.""" + + @abc.abstractmethod + async def recheck(self, key: StreamKey, wait_id: int) -> bool: + """Whether any record now sits past the installed intent's cursor. + + Called after every intent is installed, and it is this recheck -- not + the intent -- that closes the append/park race: a producer appends + before it observes the park generation, so an append is either seen + here or paired with a wake Signal. + """ + + @abc.abstractmethod + async def claim_park_generation( + self, + key: StreamKey, + wait_id: int, + park_generation: int, + *, + claimant: str, + lease: timedelta, + ) -> bool: + """Tries to take responsibility for waking ``park_generation``. + + Returns ``True`` when this caller now holds the claim -- including when + it is renewing a claim it already held, so a producer can extend rather + than losing it mid-flight. + + A provider that declares ``supports_leased_claims = False`` must always + return ``True``: observe-only semantics, where every producer signals. + """ + + @abc.abstractmethod + async def parked_wait_ids(self, key: StreamKey) -> list[int]: + """Every subscription with an installed intent on this stream. + + A producer knows the stream it published to and nothing about the + Workflow's subscriptions -- ``wait_id`` is allocated by a per-Run counter + inside ``subscribe()``, which no producer can see. The other five + parking operations all take a ``wait_id``, so without this one a producer + that has just appended has no way to address a wake at all. + + Enumeration rather than a new concept: intents are already keyed + ``(stream key, wait_id)``, and this is the ``wait_id`` half of that key + for one stream. + """ + + @abc.abstractmethod + async def current_park_generation(self, key: StreamKey, wait_id: int) -> int | None: + """The generation a producer should name in its wake Signal. + + ``None`` means no confirmed park is installed for this subscription, in + which case a producer sends an *unparked* wake -- generation 0 -- rather + than staying silent. + """ + + # --- provided on top of the required operations ------------------------- + + async def watch( + self, + key: StreamKey, + after: Cursor, + *, + max_records: int = 100, + block: timedelta | None = DEFAULT_WATCH_BLOCK, + ) -> AsyncIterator[StreamRecord]: + """Yields records forever, starting strictly after ``after``. + + A convenience over :meth:`read_after` for prefetch loops. The cursor + advances to ``AFTER(last delivered)``, so a resumption after any number + of empty blocks continues from where delivery stopped rather than from + where the stream currently is. + """ + cursor = after + while True: + batch = await self.read_after( + key, cursor, max_records=max_records, block=block + ) + for record in batch: + yield record + if batch: + last = batch[-1].offset + assert last is not None, "a stored record always has an offset" + cursor = AFTER(last) + + def is_before(self, left: Offset, right: Offset) -> bool: + """Whether ``left`` precedes ``right`` in this provider's order.""" + return self.compare_offsets(left, right) < 0 + + def strictly_increasing(self, offsets: list[Offset]) -> bool: + """Whether ``offsets`` ascend under this provider's ordering rule.""" + return all(self.compare_offsets(a, b) < 0 for a, b in zip(offsets, offsets[1:])) + + +def _validate_backend( # pyright: ignore[reportUnusedFunction] + backend: object, +) -> StreamBackend: + """Validate the backend before a Worker can consume through it.""" + if not isinstance(backend, StreamBackend): + raise TypeError( + f"external stream backend is a {type(backend).__name__}, which is not " + f"a {StreamBackend.__name__}" + ) + + declared = type(backend).guarantees_immutability + if declared is not True: + raise ValueError( + f"external stream backend {type(backend).__name__} declares " + f"guarantees_immutability = {declared!r}. Configuration requires True: " + "the provider must guarantee that a record's bytes cannot change " + "once written, because replay validates presence, count, order, and " + "control positions only. A provider that cannot make that guarantee " + "does not satisfy the backend contract." + ) + + if not type(backend).provider_id: + raise ValueError( + f"external stream backend {type(backend).__name__} declares no " + "provider_id; every replay annotation header records it, so replay " + "could not verify which provider wrote a marker" + ) + return backend diff --git a/temporalio/contrib/external_workflow_streams/_codec.py b/temporalio/contrib/external_workflow_streams/_codec.py new file mode 100644 index 0000000000..6e64ef2620 --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_codec.py @@ -0,0 +1,206 @@ +"""Payload encoding for stream records (P4). + +A record's ``payload`` bytes are a serialized +:py:class:`temporalio.api.common.v1.Payload`, not a bare value encoding. Keeping +the envelope means the payload's own ``encoding`` and ``messageType`` metadata +travel with it, so a consumer decodes exactly what the producer wrote rather +than having to be told out of band. + +Producer and consumer must use the **same** ``DataConverter``, including any +codec. A mismatch is detected on the consumer, and is classified as a decode +failure rather than as stream integrity loss -- the stream is fine; the +configuration is not. + +The one thing this module classifies itself is an **external payload store that +cannot be reached**. That is row one of the taxonomy, not row three: the bytes +in the stream are a reference, the value does not exist until the reference is +fetched, and a driver that cannot fetch it raises whatever its client raises. +Left unlabelled it reaches the consumer as an unclassified failure on a record +whose range validated, which the classification rule turns into a decode +failure -- sending an operator to change a converter during a storage outage. +Nothing further down can tell the two apart, so the label is applied at the call +that knows. + +Decoding on the consumer is **two halves, run in two different places**: + +- :py:meth:`StreamPayloadCodec.prepare`, on the **Worker's loop**: + external-payload retrieval and the user's ``PayloadCodec``. Arbitrary + asynchronous work, and none of it depends on the value's type. +- :py:meth:`StreamPayloadCodec.convert`, on the **Workflow thread**: + ``from_payloads`` with the topic's declared type. Synchronous, performs no + I/O, and is the only half that needs to know the type at all. + +That is the same split the Worker already applies to every other payload an +activation carries: ``decode_activation`` is awaited before ``activate()`` is +handed to the executor, and the payload converter runs inside it. A record whose +codec ran on the Workflow thread would perform real I/O inside a deterministic +event loop, turn the codec's awaits into Workflow commands, and put a +multi-second KMS round trip under a 2-second deadlock timeout. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Generic + +import temporalio.api.common.v1 +import temporalio.converter +from temporalio.contrib.external_workflow_streams._errors import ( + StreamError, + StreamStorageError, +) +from temporalio.types import AnyType + +__all__ = ["StreamPayloadCodec"] + + +@dataclass(frozen=True) +class StreamPayloadCodec(Generic[AnyType]): + """Encodes and decodes one topic's values. + + The type is carried by the topic -- ``topic("tokens", type=str)`` -- and is + passed to the converter as a decode hint, so a topic declared as ``str`` + yields ``str`` rather than whatever the default converter happens to infer. + """ + + data_converter: temporalio.converter.DataConverter + value_type: type[AnyType] | None = None + + async def encode(self, value: AnyType) -> bytes: + """Serializes one value to a record's payload bytes.""" + payloads = await self.data_converter.encode([value]) + if len(payloads) != 1: + # `DataConverter.encode` is documented as free to return fewer + # payloads than values. One value must still produce one payload, + # and a converter that collapses it would silently drop a record. + raise ValueError( + f"encoding one stream value produced {len(payloads)} payloads; " + "exactly one is required" + ) + return payloads[0].SerializeToString() + + async def decode(self, payload: bytes) -> AnyType: + """Parses a record's payload bytes back into a value. + + Both halves at once. Correct anywhere a coroutine may await arbitrary + work -- the producer side, and tests -- and **wrong on the Workflow + thread**, which is why the consumer side calls :meth:`prepare` and + :meth:`convert` separately instead. + + Raises whatever the converter or codec raises. Classification into the + failure taxonomy belongs to the caller, which is the only side that + knows whether the record's range validated first. + """ + return self.convert(await self.prepare(payload)) + + async def prepare(self, payload: bytes) -> temporalio.api.common.v1.Payload: + """The asynchronous half of decoding, and the half that has no type. + + ``DataConverter.decode`` is three steps: external-payload retrieval, the + user's ``PayloadCodec``, and the payload *converter*. The first two are + arbitrary asynchronous work -- a network fetch, a KMS round trip -- and + neither of them depends on the topic's declared type. They belong on the + Worker's loop, which is exactly where the Worker already awaits them for + every other payload an activation carries. + + Returns the payload as the payload converter will see it. What is left + is :meth:`convert`, which is synchronous and needs no I/O. + """ + parsed = temporalio.api.common.v1.Payload() + parsed.ParseFromString(payload) + # A deliberate reach into two private helpers rather than an oversight: + # this split is the whole point, and `DataConverter` exposes no public + # "everything except the payload converter" entry point. `decode()` + # itself is these two calls followed by `from_payloads`, so preparing + # here and converting in `convert` is byte-for-byte the same work in the + # same order -- only on two different threads. + try: + retrieved = await self.data_converter._external_retrieve_payload_sequence( + [parsed] + ) + except StreamError: + raise + except Exception as err: + # Row one of the taxonomy, and it has to be labelled *here*. An + # external-storage driver raises whatever its client raises -- a + # bare `ConnectionError` for an unreachable payload store -- and the + # converter does not wrap it. Everything downstream sees an + # unclassified exception on a record whose range validated, and the + # classification rule for that is row three: the operator is told to + # align a converter that is in fact fine while the payload store is + # the thing that is down. Nothing below this call can tell the two + # apart, because by then the only evidence is the exception type the + # driver chose. + # + # A payload store is as much a stream read as the stream backend is: + # the record's bytes are a reference, and the value does not exist + # until they are fetched. Transient by the same argument, and clears + # the same way. + raise StreamStorageError( + "an external stream record's payload could not be retrieved from " + f"external storage: {err}" + ) from err + decoded = await self.data_converter._decode_payload_sequence(retrieved) + if len(decoded) != 1: + raise ValueError( + f"preparing one stream record produced {len(decoded)} payloads; " + "exactly one is required" + ) + return decoded[0] + + def parse_unprepared(self, payload: bytes) -> temporalio.api.common.v1.Payload: + """The envelope parse alone, for a converter with nothing to prepare. + + A ``DataConverter`` with no codec and no external storage has an empty + asynchronous half, so a record that never passed through + :meth:`prepare` still converts correctly. One with either of them does + not, and the difference is not detectable from the bytes: a codec's + output is just another payload. Rather than return a plausible wrong + value, this refuses -- an unprepared record on the Workflow thread means + the record reached it by a path that does not prepare, which is a + routing defect and not a user's converter mismatch. + """ + # Asked of the two public members rather than through the converter's + # `_decode_payload_has_effect`, which is the same condition but is marked + # in its own source as a temporary shortcircuit to be removed. This + # refusal is a safety property; it should not stop working the day an + # unrelated cleanup lands upstream. + if ( + self.data_converter.payload_codec is not None + or self.data_converter.external_storage is not None + ): + raise RuntimeError( + "an external stream record reached the Workflow thread without " + "its DataConverter's asynchronous half having been applied. " + "Preparing it here would run a payload codec or an external " + "payload fetch inside activate(); every delivery path must " + "prepare records on the Worker's loop instead." + ) + parsed = temporalio.api.common.v1.Payload() + parsed.ParseFromString(payload) + return parsed + + def convert(self, prepared: temporalio.api.common.v1.Payload) -> AnyType: + """The synchronous half: the topic's declared type, applied. + + The only step that needs the type, and the only one a Workflow thread + may run: ``from_payloads`` performs no I/O and awaits nothing, which is + what every ordinary activation payload's conversion already does on this + same thread. + """ + type_hints: list[type] | None = ( + None if self.value_type is None else [self.value_type] + ) + values = self.data_converter.payload_converter.from_payloads( + [prepared], type_hints + ) + if len(values) != 1: + raise ValueError( + f"decoding one stream record produced {len(values)} values; " + "exactly one is required" + ) + return values[0] + + def with_type(self, value_type: type[Any] | None) -> StreamPayloadCodec[Any]: + """This codec bound to a different topic's declared type.""" + return StreamPayloadCodec(self.data_converter, value_type) diff --git a/temporalio/contrib/external_workflow_streams/_continuation.py b/temporalio/contrib/external_workflow_streams/_continuation.py new file mode 100644 index 0000000000..e721cee25c --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_continuation.py @@ -0,0 +1,161 @@ +"""The Continue-As-New cursor (P15). + +A stream spans a whole Continue-As-New chain, so a new Run must resume from the +position its predecessor committed -- on replay as well as live. + +**A cursor is never derived from mutable backend state, on any Run.** Reading the +current position from the backend, or from coordination state keyed by chain, +would give replay whatever the backend holds *now* rather than what the Run +originally started from, and two replays of one history could then diverge. The +position therefore travels in a reserved internal header on the Continue-As-New +command, persisted in the new Run's ``WorkflowExecutionStarted`` and restored +before any subscription is established (ADR-022). + +It populates the same annotation-header ``start_cursor`` field that a first +execution fills with ``BEGINNING``, so replay reads an explicit starting boundary +in **every** case -- including the case where the stream was empty for the +subscription's entire life. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import temporalio.api.common.v1 +from temporalio.contrib.external_workflow_streams._annotation import ( + AnnotationDecodeError, + _put_cursor, + _put_str, + _put_uvarint, + _Reader, +) +from temporalio.contrib.external_workflow_streams._record import Cursor + +__all__ = [ + "CONTINUATION_HEADER", + "Continuation", + "decode_continuation", + "encode_continuation", + "read_continuation_header", + "write_continuation_header", +] + +CONTINUATION_HEADER = "__temporal_external_stream_continuation" +"""Reserved, and in this feature's namespace rather than ``workflow_stream``'s. + +``temporalio.contrib.workflow_streams`` reserves ``__temporal_workflow_stream_*`` +for a different feature that coexists with this one (ADR-001). +""" + +CONTINUATION_ENCODING = b"binary/plain" +"""Serialized by this module, not by the user's ``DataConverter``. + +The header is written by one Run and read by the next, and a chain whose Runs +were deployed with different converter configuration would otherwise restart at +an unreadable cursor -- silently, since an unreadable cursor looks exactly like +no cursor at all. +""" + +_SCHEMA_VERSION = 1 +"""The only continuation schema this unreleased feature understands.""" + + +@dataclass(frozen=True) +class Continuation: + """Where each subscription had got to when the Run continued as new. + + Keyed by ``wait_id`` so two subscriptions to one stream restore + independently: they are separate waits with their own cursors, and a + stream-keyed continuation would restart one of them at the other's position. + + The three maps beside ``cursors`` are the **binding the cursor was produced + under**, and each is there because without it a cursor can be restored into + something it does not describe: + + - ``stream_names`` so a renumbered subscription is detectable. Restoring a + cursor onto a differently-numbered wait would resume one stream at + another's offset, which the backend would accept and no later check would + catch. + - ``provider_ids`` and ``provider_format_versions`` because a Worker can + be reconfigured to a different implementation. Marker replay checks both + before interpreting recorded offsets; the first live read in a successor + Run happens before that Run has written a marker that could perform the + check, so the continuation has to carry them itself. + """ + + cursors: dict[int, Cursor] + stream_names: dict[int, str] + provider_ids: dict[int, str] = field(default_factory=dict) + provider_format_versions: dict[int, int] = field(default_factory=dict) + + +def encode_continuation(continuation: Continuation) -> bytes: + """Encode current continuation state deterministically.""" + out = bytearray() + _put_uvarint(out, _SCHEMA_VERSION) + _put_uvarint(out, len(continuation.cursors)) + # Sorted, so the same state always encodes to the same bytes. The header is + # re-derived from live state every time the terminal command is built, and + # bytes that depended on dict insertion order would put a different cursor + # in front of the successor on a Workflow Task the server retried. + for wait_id in sorted(continuation.cursors): + _put_uvarint(out, wait_id) + _put_cursor(out, continuation.cursors[wait_id]) + _put_str(out, continuation.stream_names.get(wait_id, "")) + _put_str(out, continuation.provider_ids[wait_id]) + _put_uvarint(out, continuation.provider_format_versions[wait_id]) + return bytes(out) + + +def decode_continuation(raw: bytes) -> Continuation: + """Decode and validate a supported continuation schema.""" + reader = _Reader(raw) + version = reader.uvarint() + if version != _SCHEMA_VERSION: + raise AnnotationDecodeError( + f"the Continue-As-New cursor header is schema version {version}, but " + f"this Worker understands {_SCHEMA_VERSION}. The previous Run of this " + "chain was executed by a newer SDK." + ) + cursors: dict[int, Cursor] = {} + stream_names: dict[int, str] = {} + provider_ids: dict[int, str] = {} + provider_format_versions: dict[int, int] = {} + for _ in range(reader.uvarint()): + wait_id = reader.uvarint() + cursors[wait_id] = reader.cursor() + stream_names[wait_id] = reader.string() + provider_ids[wait_id] = reader.string() + provider_format_versions[wait_id] = reader.uvarint() + return Continuation( + cursors=cursors, + stream_names=stream_names, + provider_ids=provider_ids, + provider_format_versions=provider_format_versions, + ) + + +def write_continuation_header( + continuation: Continuation, +) -> temporalio.api.common.v1.Payload: + """Wrap encoded continuation state in its reserved Payload envelope.""" + return temporalio.api.common.v1.Payload( + metadata={"encoding": CONTINUATION_ENCODING}, + data=encode_continuation(continuation), + ) + + +def read_continuation_header( + headers: dict[str, temporalio.api.common.v1.Payload] | None, +) -> Continuation | None: + """The continuation a predecessor Run left, or ``None`` on a first execution. + + ``None`` is not a failure: it is what every chain's first Run sees, and the + subscription then starts at ``BEGINNING``. + """ + if not headers: + return None + payload = headers.get(CONTINUATION_HEADER) + if payload is None: + return None + return decode_continuation(payload.data) diff --git a/temporalio/contrib/external_workflow_streams/_errors.py b/temporalio/contrib/external_workflow_streams/_errors.py new file mode 100644 index 0000000000..c9450d2355 --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_errors.py @@ -0,0 +1,267 @@ +"""The external stream failure taxonomy (P18). + +Four outcomes, deliberately not collapsed, because each has a different operator +response. All four reach the server through the same channel -- a Workflow +activation completion is ``oneof status { Success | Failure }`` -- so what +distinguishes them is the **error type and the metric**, never the retry +behavior of the Workflow Task. + +**There is no protocol-level "non-retryable Workflow Task failure": the server +retries Workflow Task failures regardless of cause.** Alerting must therefore +key on the metrics here rather than on Workflow Task failure counts, which a +transient backend outage also increments. + +============================ ========================== ==================== +Condition Error type Operator response +============================ ========================== ==================== +Backend unreachable/erroring :class:`StreamStorageError` None -- clears when + the backend recovers +Recorded offset missing, :class:`StreamIntegrityError` Repair or restore +expired, reordered, or the backend, or +miscounted terminate the Run +Bytes intact but undecodable :class:`StreamDecodeError` Align the consumer's + converter with the + producer's +Annotation does not match ordinary nondeterminism Fix or version the +the subscriptions made Workflow code +============================ ========================== ==================== + +There is deliberately **no** ``workflow_failure_exception_types`` registration +here: integrity loss *blocks* a Workflow rather than terminating it (ADR-014). +Retention loss is usually an operational error and usually repairable, and a +blocked Workflow can be resumed after repair where a failed one cannot. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final + +import temporalio.common +import temporalio.exceptions + +__all__ = [ + "ConcurrentStreamConsumerError", + "ExternalStreamCapacityError", + "StreamDecodeError", + "StreamError", + "StreamIntegrityError", + "StreamStorageError", + "classify_read_failure", +] + + +class StreamError(Exception): + """Base class for every external stream failure this module classifies.""" + + +class StreamStorageError(StreamError): + """The backend was unreachable, timed out, or returned an error. + + Transient by definition: it clears when the backend recovers, so an + operator has nothing to do. This is the one row of the taxonomy that a + Workflow Task failure count alone describes adequately. + """ + + +class StreamIntegrityError(StreamError): + """A recorded range no longer reads back as it was written. + + Raised when an offset the annotation names is missing, when a range + contains the wrong number of records, when ordering is not strictly + increasing, or when control positions do not match. Those four checks are + the whole of replay validation, and they are sufficient precisely because + every provider guarantees a record's bytes cannot change (ADR-003) -- so + the only damage replay must detect is a record that is no longer there. + + **This must never resolve to an alternate stream result.** Substituting a + later record for a deleted one would hand Workflow code a different history + than the one its commands were derived from. + """ + + +class StreamDecodeError(StreamError): + """A record was present and intact but could not be decoded. + + Separated from :class:`StreamIntegrityError` because it is a configuration + error on the *consumer* -- a DataConverter or codec that does not match the + producer's -- and reporting it as integrity loss sends an operator to + restore a backend that was never damaged (ADR-015). + """ + + +class ExternalStreamCapacityError(temporalio.exceptions.ApplicationError): + """This Workflow's subscription set cannot be recorded in an annotation. + + Deliberately **not** one of the four taxonomy rows above. Those describe a + stream or a converter behaving unexpectedly at read time; this describes + Workflow code asking for more than the marker format can carry, and it is + known at ``subscribe()`` rather than at read time. The remedy is in the + Workflow: fewer streams, or shorter names. + + Non-retryable, and raised from ``subscribe()`` rather than from the + completion path, because both halves of that matter. A capacity limit + discovered while encoding a completion fails the *Workflow Task*, the server + retries Workflow Task failures regardless of cause, and the same subscription + set encodes to the same oversized header every time -- a Workflow stuck + forever with nothing durable to show why. Raised where the subscription is + made, it is an error in the Workflow's own code path, deterministic under + replay, and reported once. + """ + + def __init__(self, message: str) -> None: + """Create a deterministic, non-retryable capacity failure.""" + super().__init__( + message, type="ExternalStreamCapacityError", non_retryable=True + ) + + +class ConcurrentStreamConsumerError(temporalio.exceptions.ApplicationError): + """A second coroutine tried to wait on a subscription already being waited on. + + A subscription is **one consumer**. It has one cursor, one readiness future, + and one entry in the runtime's pending map, so a second coroutine blocking on + it replaces the first everywhere the first can be found: readiness resolves + the newer waiter, its ``finally`` clears the entry, and the older future + becomes unreachable by the readiness activation *and* by ``close()``. The + Workflow is then permanently stuck on a future nothing can resolve -- and + because the blocked-set state is shared too, Core may not even be retaining a + Workflow Task on that coroutine's behalf. + + Not one of the four taxonomy rows above, and for the same reason as + :class:`ExternalStreamCapacityError`: this is Workflow code asking for + something the design does not offer, not a stream or a converter misbehaving. + Non-retryable for the same reason too -- a second consumer is deterministic, + so a Workflow Task failure would be retried into the identical failure forever + with nothing durable to show why. + + Iterating a subscription again *after* the previous consumer has stopped is + not this: the check is on a live waiter rather than on the iterator, so the + ordinary shape of taking some records, doing something else, and coming back + to the same subscription keeps working (ADR-037). + """ + + def __init__(self, message: str) -> None: + """Create a deterministic, non-retryable consumer failure.""" + super().__init__( + message, type="ConcurrentStreamConsumerError", non_retryable=True + ) + + +#: Metric names. Operators are expected to alert on the integrity metric +#: specifically, since the transient row increments Workflow Task failures too. +METRIC_INTEGRITY: Final = "temporal_external_stream_integrity_failure" +METRIC_DECODE: Final = "temporal_external_stream_decode_failure" +METRIC_STORAGE: Final = "temporal_external_stream_storage_failure" +METRIC_SHUTDOWN_WAKE_FAILED: Final = "temporal_external_stream_shutdown_wake_failed" + + +@dataclass(frozen=True) +class StreamMetrics: + """The counters this taxonomy reports through. + + Each failure class gets its own counter rather than one counter with a + ``reason`` attribute, so an alert on integrity loss cannot be diluted by a + backend outage sharing the same series. + """ + + integrity: temporalio.common.MetricCounter + decode: temporalio.common.MetricCounter + storage: temporalio.common.MetricCounter + shutdown_wake_failed: temporalio.common.MetricCounter + + @staticmethod + def create(meter: temporalio.common.MetricMeter) -> StreamMetrics: + return StreamMetrics( + integrity=meter.create_counter( + METRIC_INTEGRITY, + "External stream replay found a recorded range that no longer reads back " + "as written. Repair the backend or terminate the Run; this does not clear " + "on its own.", + ), + decode=meter.create_counter( + METRIC_DECODE, + "External stream record bytes were intact but could not be decoded. The " + "stream is fine; the consumer's converter or codec does not match the " + "producer's.", + ), + storage=meter.create_counter( + METRIC_STORAGE, + "External stream backend was unreachable, timed out, or errored. " + "Transient; clears when the backend recovers.", + ), + shutdown_wake_failed=meter.create_counter( + METRIC_SHUTDOWN_WAKE_FAILED, + "A Worker shutdown wake Signal was still unacknowledged when the graceful " + "shutdown grace period expired.", + ), + ) + + def counter_for( + self, error_type_name: str + ) -> temporalio.common.MetricCounter | None: + """The counter for an error class *name*, or ``None`` for anything else. + + By name, because the side that reports these has only a name: a failure + raised on the Workflow thread is converted inside ``activate()``, and + what the Worker gets back is a completion whose application failure + carries the exception's class name. Keeping the mapping here rather than + at that call site is what stops the two from drifting apart -- the same + reason :meth:`record` exists for the side that still holds the + exception. + + ``None`` for every other failure, which is what keeps row four -- + ordinary nondeterminism -- and plain Workflow bugs out of these series. + """ + return { + StreamIntegrityError.__name__: self.integrity, + StreamDecodeError.__name__: self.decode, + StreamStorageError.__name__: self.storage, + }.get(error_type_name) + + def record(self, error: BaseException) -> None: + """Increments the counter matching ``error``'s class, if any.""" + if isinstance(error, StreamIntegrityError): + self.integrity.add(1) + elif isinstance(error, StreamDecodeError): + self.decode.add(1) + elif isinstance(error, StreamStorageError): + self.storage.add(1) + + +def classify_read_failure( + *, range_validated: bool, cause: BaseException +) -> StreamError: + """The mechanical classification rule, in one place. + + Unconditional, because structural immutability is required of every + provider: + + **If the range validated, the bytes are the bytes that were written**, + so any subsequent failure is a decode failure. Only a missing offset or + a range that fails validation is integrity loss. + + Args: + range_validated: Whether the four range checks passed for the range the + failing record belongs to. + cause: What actually went wrong, attached as the new error's cause. + + Returns: + The error to raise. Never raises on its own -- the caller decides + whether to raise, record a metric, or both. + """ + # A storage failure is about reaching the backend at all, so it is never + # reclassified by whether a range validated: there was no range to read. + if isinstance(cause, StreamStorageError): + return cause + + if range_validated: + err: StreamError = StreamDecodeError( + f"external stream record could not be decoded: {cause}" + ) + else: + err = StreamIntegrityError( + f"external stream recorded range failed validation: {cause}" + ) + err.__cause__ = cause + return err diff --git a/temporalio/contrib/external_workflow_streams/_manager.py b/temporalio/contrib/external_workflow_streams/_manager.py new file mode 100644 index 0000000000..86b521e8db --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_manager.py @@ -0,0 +1,2877 @@ +"""The per-Worker subscription and watcher manager (P8). + +**Any design in which ``_apply`` reads the backend is wrong by construction, not +merely slow.** ``_WorkflowInstanceImpl.activate()`` is synchronous, runs on a +thread-pool executor under a 2-second deadlock timeout, and drives a custom +deterministic event loop in which real network I/O cannot be awaited at all. + +So the manager lives out here, on the Worker's own asyncio loop, owns every +backend connection and watcher task, and communicates with the Workflow thread +through a **bounded, thread-safe buffer per subscription**: + +=========== ====================== ===================================== +Step Where What +=========== ====================== ===================================== +Register Workflow thread Records the wait. Non-blocking. +Prefetch Manager loop Reads ahead into the bounded buffer. +Readiness Manager loop Reports **only after** a record is + buffered -- never on a bare socket + event. +Drain Workflow thread Pops from the buffer. Nothing else. +=========== ====================== ===================================== + +That "only after buffered" rule is what makes the resulting activation +guaranteed non-blocking. Readiness for an unbuffered record would produce an +activation whose drain must block, reintroducing exactly the deadlock hazard +this structure exists to remove. + +Three cursors, and conflating them is how a speculative read becomes a durable +claim: + +=================== =================== =========== ================== +Cursor Advances on Owner Survives eviction +=================== =================== =========== ================== +``committed`` marker commit only the marker yes +``delivery`` hand-off to workflow the instance no +``prefetch`` buffering the manager no +=================== =================== =========== ================== + +``prefetch`` is speculative: reading a record is not consuming it, and consuming +it is not committing it. The manager may only move it *backwards* to +``committed``, never forwards past what a marker has committed. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +import threading +from collections import deque +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass, field, replace +from datetime import timedelta +from typing import Any + +import temporalio.converter +from temporalio.contrib.external_workflow_streams._annotation import StreamBinding +from temporalio.contrib.external_workflow_streams._backend import ( + ParkIntent, + ParkIntentRemoval, + StreamBackend, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._errors import ( + StreamError, + StreamStorageError, +) +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Cursor, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._replay import ( + ReplayPlan, + ReplaySegment, + build_replay_plan, +) +from temporalio.contrib.external_workflow_streams._wake import new_sender_identity + +__all__ = [ + "PreparedRecord", + "ReadinessResult", + "StreamSubscriptionManager", + "Subscription", +] + +logger = logging.getLogger(__name__) + + +@contextlib.contextmanager +def _as_storage_failure(what: str) -> Iterator[None]: + """Reports a backend failure on an activation path as row one. + + The park handshake is a backend transaction an activation is waiting on, and + a backend that is unreachable or erroring is the taxonomy's *transient* + row: nothing for an operator to do, and it clears when the backend + recovers. Left as whatever the provider raised, it reaches the server as an + anonymous Workflow Task failure with no cause and no counter, + indistinguishable from a bug in the Workflow's own code. + + Not used for the intent removal a resolve performs, which deliberately logs + rather than raises: that one is cleanup on the delivery path, and failing a + Workflow Task for it would trade a stale intent for a repeated Workflow + Task. + + Anything already classified passes through untouched, so an integrity + failure is never relabelled as a transient one. + """ + try: + yield + except asyncio.CancelledError: + raise + except StreamError: + raise + except Exception as err: + raise StreamStorageError(f"{what} failed: {err}") from err + + +DEFAULT_BUFFER_SIZE = 256 +"""How many records one subscription may hold ahead of the Workflow. + +Backpressure *is* this bound: a full buffer stops prefetch. It never drops a +record and never blocks the Workflow thread. +""" + +DEFAULT_WATCH_BLOCK = timedelta(seconds=5) + + +class ReadinessResult: + """The subset of Core's readiness answers the manager reacts to. + + Mirrors :py:class:`temporalio.bridge.worker.ExternalStreamReadyResult` but + is matched structurally, so the manager can be driven in tests without a + Core worker behind it. + """ + + ACCEPTED = "Accepted" + STALE = "Stale" + PARKED = "Parked" + NO_OPEN_WORKFLOW_TASK = "NoOpenWorkflowTask" + RUN_NOT_FOUND = "RunNotFound" + + +class RunStatus: + """The read-only answer to "what state is this Run's wait set in?". + + Mirrors :py:class:`temporalio.bridge.worker.ExternalStreamRunStatus` and is + matched structurally for the same reason :class:`ReadinessResult` is. + """ + + WFT_OPEN = "WftOpen" + PARKED = "Parked" + NO_OPEN_WORKFLOW_TASK = "NoOpenWorkflowTask" + RUN_NOT_FOUND = "RunNotFound" + + +READINESS_ATTEMPTS = 3 +"""How many times a failing readiness report is retried before a wake is owed. + +A raising notifier means the record is buffered and Core has not been told. That +is indistinguishable, from here, from a Run with no open Workflow Task -- so +after these attempts it is treated as exactly that. +""" + +READINESS_RETRY_DELAY = timedelta(milliseconds=200) + +STALE_REPORT_ATTEMPTS = 3 +"""How many times a stale readiness report is re-sent against a newer generation. + +Bounded rather than open-ended: a generation that keeps moving is a Workflow +consuming records happily, and owing a wake at the end of that costs one empty +Workflow Task, where giving up silently costs the record. +""" + +STALE_RETRY_DELAY = timedelta(milliseconds=50) + +SHUTDOWN_WAKE_ATTEMPTS = 3 +"""How many times one owed wake is attempted before it is given up on. + +More than one because a Worker shutting down is often shutting down *because* +something is unhealthy, so the first attempt is the one most likely to land in +the middle of it. Bounded because the alternative to giving up is holding +shutdown open, and the metric exists precisely so that giving up is visible. + +Governs the live path too, which owes wakes for exactly the same reason and has +a far worse backstop: a wake owed *there* is the only thing that will ever +produce a Workflow Task, so nothing re-attempts it -- see `_send_owed_wake`. +The name is the shutdown sweep's because that is where the retry started. +""" + +SHUTDOWN_WAKE_RETRY_DELAY = timedelta(milliseconds=200) +"""Short enough that three attempts fit comfortably inside the grace period.""" + +PARK_REMOVAL_ATTEMPTS = 3 +"""How many times an owed park-intent removal is retried in place. + +The cheap first line against a momentary backend blip, and *only* that: what +makes giving up here survivable is the ledger the failure is recorded in, which +outlives the Subscription and has its own retry loop. Bounded because these +inline retries sit on the close path and inside a registration, neither of +which may wait on a backend indefinitely. +""" + +PARK_REMOVAL_RETRY_DELAY = timedelta(milliseconds=100) +"""Shorter than the wake delays: this backs off *between* park-lock holds. + +The lock is a Run's park handshake serialization, and `prepare_park` runs inside +an activation under Core's deadlock timeout, so every millisecond spent backing +off is a millisecond that budget may have to absorb. +""" + +PARK_REMOVAL_MAX_RETRY_DELAY = timedelta(seconds=5) +"""Caps autonomous removal backoff while still retrying indefinitely.""" + +PARK_REMOVAL_CALL_TIMEOUT = timedelta(milliseconds=500) +"""How long park-intent work off the activation path may hold a Run's park lock. + +Bounds the *call*, which counting attempts cannot. A backend that hangs rather +than raises returns nothing to retry and raises nothing to catch, so the retry +loop, the reconciliation and the close each stop where they are -- and hold this +Run's park lock while they do, which the next park, resolve or eviction takes. + +The asymmetry is what makes giving up right: these holders retry, so an attempt +abandoned costs them a backoff, while the activation waiting behind them is +spending Core's deadlock timeout. Well inside that timeout, because waiting this +out is only the first thing the activation then has to do. + +Holds an activation takes *itself* are deliberately not bounded. Those are the +same backend exposure the park handshake already has -- `install_park_intent` +can hang exactly as a removal can -- so a bound there would move the wait rather +than remove it. + +Not a latency limit on the provider, though it becomes one for anything slower +than it: these are single-key operations, and a coordination backend that cannot +answer one in half a second cannot serve the park handshake either, which does +two of them per wait inside an activation. A deployment that far out is failing +Workflow Tasks before this timeout is what it notices. +""" + +DEFAULT_SHUTDOWN_GRACE = timedelta(seconds=10) +"""How long the sweep may hold shutdown open. + +Bounded on purpose: a Worker that could not reach the server would otherwise +hang on the way out, and a wake that has not been acknowledged by now is better +reported than waited on indefinitely. +""" + +DEFAULT_PROBE_GRACE = timedelta(seconds=2) +"""How long the probe phase may delay Core's stop-polling step. + +Much shorter than the sweep's grace, and for a different reason: the probe is +purely local -- one message on Core's own serialized input lane per Run -- so it +is either quick or wedged, and everything it delays is a Worker that has already +been asked to stop. A Run it does not reach is not lost; it is probed again by +the sweep, on whatever Core has left to say. +""" + + +def _status_value(status: Any) -> str: + return getattr(status, "value", status) + + +#: What the manager calls to tell Core a record is buffered. Returns one of the +#: five readiness results as a plain value (or an enum whose ``value`` is one). +ReadinessNotifier = Callable[[str, int, int], Awaitable[Any]] + +#: What the manager calls when local readiness could not be delivered. Filled in +#: by the producer wake-signal path (P14); until then a subscription simply +#: records that a wake was owed. +#: +#: Takes a :data:`WakeTarget` rather than a `Subscription`, because the wake a +#: retired stale intent owes outlives the subscription the record was buffered +#: on -- see :class:`_CleanupWake`. +WakeSender = Callable[["WakeTarget"], Awaitable[None]] + +#: The read-only Run-status probe (C4), returning one of the four run statuses. +RunStatusProbe = Callable[[str], Awaitable[Any]] + + +def _result_value(result: Any) -> str: + return getattr(result, "value", result) + + +@dataclass +class Subscription: + """One Workflow subscription, as the manager sees it.""" + + run_id: str + wait_id: int + stream_key: StreamKey + backend: StreamBackend + buffer_size: int = DEFAULT_BUFFER_SIZE + + committed_cursor: Cursor = BEGINNING + """Advances only when a marker commits. Reconstructed from History.""" + + delivery_cursor: Cursor = BEGINNING + """Advances when a record is handed to Workflow code. Rebuilt by replay.""" + + prefetch_cursor: Cursor = BEGINNING + """Advances on buffering. Speculative, and discarded outright on eviction.""" + + wait_generation: int = 0 + """Increments each time this wait re-enters the blocked state. + + Written by :meth:`note_wait_generation` from the Workflow thread and read by + the watcher on the Worker's loop, so both go through ``_lock``. Only the + runtime knows when a wait re-blocks, so nothing out here can derive it. + """ + + #: Records read ahead but not yet delivered. Appended by the manager loop, + #: popped by the Workflow thread, so every touch is under `_lock`. + _buffer: deque[StreamRecord] = field(default_factory=deque, repr=False) + _lock: threading.Lock = field(default_factory=threading.Lock, repr=False) + _has_room: asyncio.Event = field(default_factory=asyncio.Event, repr=False) + _watcher: asyncio.Task[None] | None = field(default=None, repr=False) + _cancelled: bool = field(default=False, repr=False) + #: Bumped whenever every speculative read is discarded. A `read_after` + #: already in flight started from the cursor being discarded, so its records + #: are speculative too; the watcher compares this across the await and drops + #: them rather than appending them on top of the reset -- which would put the + #: buffer straight back where it was and re-advance `prefetch_cursor`. + _prefetch_epoch: int = field(default=0, repr=False) + + #: Wakes owed because local readiness could not be delivered. Counted rather + #: than merely logged, so a test can tell "no wake was needed" from "a wake + #: was needed and dropped". + wakes_owed: int = 0 + + #: The sender's sequence number for the wake this subscription currently + #: owes, and what the unparked wake's request ID is derived from. + #: + #: Drawn from the *manager* rather than counted here, because this object + #: does not outlive its Run: an evicted Run that comes back gets a new + #: `Subscription` with `wakes_owed` back at zero, and a counter taken from + #: that would re-derive the request ID the previous incarnation's first + #: wake already used. The server deduplicates it, no Workflow Task is + #: created, and the Run waits on records it is already holding. The + #: derivation calls this "a per-sender monotonic counter", and the sender is + #: the manager. + wake_counter: int = 0 + + installed_park_generation: int | None = None + """The generation of the park intent this manager has in the backend, if any. + + The manager's mirror of one piece of backend state, kept so the invariant + *an intent exists only while that park is outstanding* is enforceable from + here: the manager is the only installer, so it is the only thing that can + know an intent is owed a removal. ``None`` means no park *this manager + confirmed* is outstanding, and the removal path is then free -- which + matters because a resolve activation is the ordinary live-delivery path, not + a rare one. + + It is deliberately **not** the only thing that can say a removal is owed. + An intent installed by a previous Worker is mirrored nowhere -- the mirror + went with the Worker -- and a removal that was attempted and failed is owed + after the Subscription holding this field has been dropped. Both live in the + manager's per-Run ledger instead; see + :meth:`StreamSubscriptionManager._remove_park_intent`. + """ + + def __post_init__(self) -> None: + """Mark a new subscription as ready to accept buffered records.""" + self._has_room.set() + + # --- the Workflow thread's half (synchronous, never blocks) ------------- + + @property + def buffered(self) -> int: + """Return the number of records currently buffered for delivery.""" + with self._lock: + return len(self._buffer) + + def drain(self, max_records: int | None = None) -> list[StreamRecord]: + """Pops buffered records. **The only thing ``_apply`` may call.** + + Bounded, non-blocking, and performs no I/O -- which is what makes the + activation it serves safe under the 2-second deadlock timeout. + """ + with self._lock: + limit = len(self._buffer) if max_records is None else max_records + popped = [ + self._buffer.popleft() for _ in range(min(limit, len(self._buffer))) + ] + if popped: + last = popped[-1].offset + assert last is not None + # Inside the lock with the pop, not after it: `reposition_to` + # runs on this same thread but between two activations, and a + # cursor advanced outside the lock could be written on top of + # the boundary a reposition had just committed. + self.delivery_cursor = AFTER(last) + return popped + + def note_wait_generation(self, generation: int) -> None: + """Takes the wait generation the runtime just moved to. + + Called from the Workflow thread the moment the wait re-enters the + blocked state. Without it this stays 0 for the life of the + subscription, every readiness report after the first block names a + generation Core has already left behind, and Core answers ``Stale`` -- + which the watcher treats as "re-probe later" while its prefetch cursor + is already past the record. The append is then never re-announced and + the Workflow is never woken. + """ + with self._lock: + self.wait_generation = generation + + def current_wait_generation(self) -> int: + """The generation to report readiness under, read on the Worker's loop.""" + with self._lock: + return self.wait_generation + + def blocked_cursor(self) -> Cursor: + """Where this subscription's deliveries stopped. + + The terminal's boundary. Fixed the moment the last activation returned, + so it is never refreshed against the backend -- doing that could name a + position replay must not reproduce. + """ + return self.delivery_cursor + + # --- the manager loop's half -------------------------------------------- + + def _append(self, records: list[StreamRecord], epoch: int) -> bool: + """Buffers a read's result unless a reposition has retracted it. + + The epoch is compared **here, under the lock that the reposition also + takes**, rather than by the watcher before it calls: repositioning is + synchronous on the Workflow thread (see + :meth:`StreamSubscriptionManager.reposition_to_committed`), so a check + made outside this lock could pass and then have the reposition land + before the append -- putting the retracted records straight back into + the buffer and re-advancing ``prefetch_cursor`` past them. + + Returns whether the records were buffered. ``False`` means they were + read from a position the marker already accounts for and the next pass + reads from the new cursor. + """ + with self._lock: + if self._prefetch_epoch != epoch: + return False + self._buffer.extend(records) + full = len(self._buffer) >= self.buffer_size + last = records[-1].offset + assert last is not None + self.prefetch_cursor = AFTER(last) + if full: + self._has_room.clear() + return True + + def _room(self) -> int: + with self._lock: + return max(0, self.buffer_size - len(self._buffer)) + + def note_drained(self) -> None: + """Tells the watcher there is room again. + + Called from the Workflow thread through the manager, which hops it onto + the loop -- an asyncio.Event is not thread-safe to set directly. + """ + self._has_room.set() + + def reset_to_committed(self) -> None: + """Discards every speculative read, as eviction and failure require. + + Reading is not consuming and consuming is not committing, so nothing + here was ever a claim. This is exactly why "no cursor advances unless + the marker commits" is safe to state. + """ + self._retract_to_committed() + self._has_room.set() + + def _retract_to_committed(self) -> None: + """The state half of :meth:`reset_to_committed`, safe on either thread. + + Every write is under ``_lock`` and nothing here touches an + ``asyncio`` primitive, which is what lets the Workflow thread call it + directly through :meth:`reposition_to`. Waking the watcher is the part + that must run on the manager's loop, so it stays with the callers. + """ + with self._lock: + self._retract_locked() + + def _retract_locked(self) -> None: + """Discards the speculative state. The caller holds ``_lock``.""" + self._buffer.clear() + self._prefetch_epoch += 1 + self.delivery_cursor = self.committed_cursor + self.prefetch_cursor = self.committed_cursor + + def commit(self, cursor: Cursor) -> None: + """Advances the committed cursor when a marker commits.""" + with self._lock: + self.committed_cursor = cursor + + def reposition_to(self, cursor: Cursor) -> None: + """Commits a marker's boundary and restarts every read from it. + + What replay needs afterwards. While replay was handing this Workflow the + records the marker recorded, the watcher was independently reading the + *same* records from the subscription's start cursor into this buffer -- + nothing had told it the marker exists. Leaving it there hands Workflow + code every replayed record a second time the moment the next live drain + happens. + + The boundary comes from the marker rather than from what the buffer + happens to hold, because the marker is the only durable statement of + where consumption reached. + + **Called synchronously from the Workflow thread**, and deliberately not + hopped onto the manager's loop. The drain that immediately follows + replay is on the Workflow thread too, so a reposition merely *posted* + to the loop leaves the retracted records in the buffer for it to hand + over a second time -- the reposition has to have happened by the time + this returns, not merely be scheduled. Nothing here touches an + ``asyncio`` primitive, and the epoch bump is what fences a read already + in flight; see :meth:`_append`. + """ + # One lock hold for both halves: a watcher append landing between them + # would be dropped by the epoch bump anyway, but a reader that saw the + # new committed cursor next to the old delivery cursor would see a state + # that never existed. + with self._lock: + self.committed_cursor = cursor + self._retract_locked() + + +class PreparedRecord(StreamRecord): + """A record whose payload has already had retrieval and codec applied. + + A subclass rather than a field on :py:class:`StreamRecord`, because being + prepared is not a property of a record -- it is a property of *this + Worker's* handling of one, and a record read by a producer, written by a + backend, or compared for idempotency has no such half-state. Everything that + reads a record reads the same fields; only the delivery path looks for the + two added here. + + ``payload`` is deliberately left as the backend's bytes. Replacing it would + make the record no longer equal to what the stream holds, which is what + idempotency comparison and integrity validation are expressed in. + """ + + #: The payload as the payload converter will see it, or ``None`` if + #: preparing it raised. + prepared_payload: Any = None + + #: What preparing raised, carried rather than raised on the Worker's loop. + #: The watcher has no Workflow Task to fail and no Workflow to tell, and a + #: record that is never delivered must not fail anything at all -- so the + #: error travels with the record and is raised by the delivery that would + #: have yielded its value. + prepare_error: BaseException | None = None + + def __eq__(self, other: object) -> bool: + """Equal to the record it was built from, and to any equal record. + + A frozen dataclass's generated ``__eq__`` compares classes exactly, so + without this a prepared record would be unequal to the identical + ``StreamRecord`` a caller built to compare against -- and being prepared + is a property of this Worker's handling, not of the record. Preparation + is deliberately not part of the comparison for the same reason. + """ + if isinstance(other, StreamRecord): + return self._fields(self) == self._fields(other) + return NotImplemented + + def __hash__(self) -> int: + """Hash the underlying record fields, excluding preparation state.""" + return hash(self._fields(self)) + + @staticmethod + def _fields(record: StreamRecord) -> tuple[Any, ...]: + return ( + record.kind, + record.payload, + record.producer_session_id, + record.sequence, + record.offset, + ) + + @classmethod + def of( + cls, + record: StreamRecord, + prepared: Any, + error: BaseException | None, + ) -> PreparedRecord: + """Create a prepared view of ``record`` with its decode outcome.""" + out = cls( + kind=record.kind, + payload=record.payload, + producer_session_id=record.producer_session_id, + sequence=record.sequence, + offset=record.offset, + ) + # `StreamRecord` is frozen, and these two are this subclass's own + # fields rather than dataclass fields, so they are set the same way a + # frozen dataclass sets anything. + object.__setattr__(out, "prepared_payload", prepared) + object.__setattr__(out, "prepare_error", error) + return out + + +@dataclass(frozen=True) +class _OwedRemoval: + """One park intent this manager knows is installed and no park sits behind. + + Not a Subscription field, and that is the whole point of it. Every removal + path this feature has reaches its intent *through* a Subscription -- the + resolve iterates the registered ones, the withdrawal walks the ones it just + installed, the close works on the one being dropped -- so a failure recorded + on the Subscription is a failure recorded on the very object the next step + throws away. Carrying the removal separately is what makes "still owed" + outlive the close, the eviction, and the Worker that installed nothing. + + The generation and Run ID are what the intent looked like when it was + recorded, kept so a delayed retry can tell it apart from a *different* + intent that has since taken the same key -- see + :meth:`StreamSubscriptionManager._drain_owed_removals`. + """ + + backend: StreamBackend + park_generation: int + run_id: str + + #: Whether a generation-0 handoff has already been made for this entry. Set + #: by the shutdown sweep, which wakes the subscriptions of a Run it is about + #: to tear down and composes those wakes as unparked precisely because this + #: entry says the installed intent is stale. The wake the cleanup would owe + #: afterwards is then the same wake, and sending it again would ask the + #: server for a second, empty Workflow Task. + #: + #: Out of the comparison on purpose: it is a note about a handoff, not part + #: of the identity of the removal, and `_owe_removal` treats an entry that + #: compares equal as debt it has already recorded. + announced: bool = field(default=False, compare=False) + + +@dataclass +class _CleanupWake: + """The wake a retired stale intent owes, with no Subscription left to owe it. + + A stale intent does not only leak, it *silences*: every wake sent while it + was installed named the generation behind it, and Core discards a non-zero + generation that is not the park it holds. Retiring the intent ends the + suppression and announces nothing, so the removal carries a wake obligation + with it -- and :meth:`StreamSubscriptionManager._reannounce_after_cleanup` + can only discharge that obligation through local readiness, which needs a + registered `Subscription` and a buffer to read. + + Eviction and close deliberately take that Subscription away while leaving + the ledger and its retry alive, which is the whole point of the ledger. So + the obligation is carried here instead: enough to compose the one wake that + is always correct after a confirmed removal -- the unparked one, which + carries no generation for Core to discard. It costs one empty Workflow Task + when no record arrived during the suppression window, and that is the + cheaper side of the trade: from out here, with no buffer left to consult, + "a record arrived" and "none did" look identical, and only one of them is + silent if guessed wrong. + + Shaped like the part of a `Subscription` a wake is composed from, because + the sender takes one target and does not care which of the two it is. + """ + + run_id: str + wait_id: int + stream_key: StreamKey + backend: StreamBackend + + #: Counted and drawn by `_count_owed_wake`, exactly as a subscription's is: + #: the request ID derives from it, so re-drawing between attempts would ask + #: for a second Workflow Task rather than re-send this one. + wakes_owed: int = 0 + wake_counter: int = 0 + + +#: What one owed wake is composed from: a `Subscription` while the wait is still +#: registered on this Worker, and a `_CleanupWake` once it is not. +WakeTarget = Subscription | _CleanupWake + + +@dataclass +class _PendingWake: + """One Run-wide wake cycle, correlated to its current send attempt.""" + + wait_id: int + attempt_started_at: int + acknowledged: bool = False + completed_terminal: bool | None = None + + +class StreamSubscriptionManager: + """Every subscription on one Worker, keyed by Run. + + Keying by Run ID is what keeps a stale Run from leaking connections: the + teardown path for eviction, Workflow Task failure, and shutdown is the same + one, and it is reached by Run. + """ + + def __init__( + self, + *, + backend: StreamBackend | None, + notify_ready: ReadinessNotifier, + send_wake: WakeSender | None = None, + run_status: Callable[[str], Awaitable[Any]] | None = None, + shutdown_wake_failed_metric: Callable[[Any], None] | None = None, + client_identity: str = "", + data_converter: temporalio.converter.DataConverter | None = None, + buffer_size: int = DEFAULT_BUFFER_SIZE, + watch_block: timedelta = DEFAULT_WATCH_BLOCK, + ) -> None: + """Create a manager for the external stream backend on one Worker.""" + self._backend = backend + #: The Worker's own converter, used for the **asynchronous half** of + #: decoding only -- retrieval and codec, never the payload converter, + #: which needs a type only Workflow code knows. Optional so a manager + #: can be constructed in isolation; a Worker always supplies it, and a + #: record that reaches the Workflow thread unprepared under a converter + #: that had something to do is refused there rather than decoded late. + self._data_converter = data_converter + self._notify_ready = notify_ready + self._send_wake = send_wake + self._run_status = run_status + self._metric = shutdown_wake_failed_metric + #: This Worker's sender identity for unparked wakes. Drawn once here + #: rather than taken from the client identity, because two Workers in + #: one process share a `Client`: with the client identity their first + #: unparked wakes derive the same request ID, the server deduplicates + #: the second, and the Run that the surviving Worker picked up never + #: gets a Workflow Task. Fixed for this manager's lifetime, so the + #: shutdown sweep's retry stays the same wake rather than a new one. + self.wake_sender_identity = new_sender_identity(client_identity) + #: The sequence the unparked wake counter is drawn from, monotonic for + #: this manager's lifetime -- which is what makes two wakes from one + #: sender distinct request IDs. Never reset: a repeat of any value this + #: sender has already used is a wake the server deduplicates away. + self._wake_sequence = 0 + #: Runs for which a wake is in flight or acknowledged, mapped to the + #: latest activation sequence when its send began. Only a later + #: activation can have been caused by that wake, so completion of a task + #: already in flight cannot prematurely release the gate. Failed tasks + #: are evicted and replayed before a successful completion, so keeping + #: this across eviction coalesces reconstructed readiness too. + self._pending_wake_runs: dict[str, _PendingWake] = {} + self._activation_sequences: dict[str, int] = {} + #: Wakes the shutdown sweep could not get acknowledged. Reported through + #: `external_stream_shutdown_wake_failed`; kept here so a test can tell + #: "no wake was needed" from "a wake was needed and lost". + self.shutdown_wake_failures = 0 + self._buffer_size = buffer_size + self._watch_block = watch_block + self._runs: dict[str, dict[int, Subscription]] = {} + #: Replay plans read and validated ahead of the delivering activation. + self._replay_plans: dict[str, ReplayPlan] = {} + # The loop the watchers run on, captured at construction because that is + # the Worker's loop. `register` is called from the *Workflow executor + # thread*, which has no loop of its own and must never touch this one + # directly. + self._loop = asyncio.get_event_loop() + self._shutting_down = False + #: Whether the sweep has already run, so `shutdown` can tell "nobody + #: swept" -- which it must then do itself -- from "already swept". + self._swept = False + #: Subscriptions the sweep has not yet decided anything about. Filled + #: when the sweep starts and drained as each Run's status is read; what + #: is left when the sweep stops -- cancelled by the grace period, or + #: stopped by a probe it could not complete -- is counted as a lost wake + #: rather than passed over, because a wake this Worker abandoned is + #: silent by nature and the counter is the only thing that says otherwise. + self._unaccounted: list[Subscription] = [] + #: What the probe phase heard, per Run. Recorded rather than acted on, + #: because the two halves of the sweep belong at different points of the + #: Worker's shutdown -- see `probe_runs`. + self._probed: dict[str, str] = {} + #: One Run's park-intent work, serialized. See `_park_lock`. + self._park_locks: dict[str, asyncio.Lock] = {} + #: Strong references to the in-flight registration-time reconciliations. + self._reconciliations: set[asyncio.Task[None]] = set() + #: Removals this manager decided on and did not get confirmed, per Run + #: and then per `(stream key, wait_id)`. The durable half of the intent + #: invariant: a removal that failed is a claim about the *backend*, and + #: recording it anywhere that a close or an eviction takes away is the + #: same as not recording it at all. Drained under `_park_lock`. + self._owed_removals: dict[str, dict[tuple[StreamKey, int], _OwedRemoval]] = {} + #: One autonomous retry loop per Run with outstanding removals. The + #: tasks are held strongly here because the ledger, rather than any + #: Subscription, owns their lifetime and can outlive cache eviction. + self._owed_removal_retries: dict[str, asyncio.Task[None]] = {} + #: Wakes a sleeping retry when another removal is added to its Run or + #: the last one is retired. A new debt should not wait behind an old + #: entry's maximum backoff, and an empty loop should stop promptly. + self._owed_removal_wakeups: dict[str, asyncio.Event] = {} + self._stopping_owed_removal_retries = False + #: In-flight cleanup wakes, and what each one is for. Strong references, + #: because the loop keeps only a weak one to a running task and this + #: one's whole job is a Signal; and a mapping rather than a set, because + #: a flush that runs out of grace has to be able to say *which* handoff + #: it abandoned. Deliberately not among the tasks + #: `_stop_background_tasks` cancels: these are not retries that can be + #: resumed later, they are the wakes themselves. + self._cleanup_wakes: dict[asyncio.Task[None], _CleanupWake] = {} + + # --- registration ------------------------------------------------------- + + def register( + self, + *, + run_id: str, + wait_id: int, + stream_key: StreamKey, + start_cursor: Cursor = BEGINNING, + ) -> Subscription: + """Registers a wait and starts its watcher. Non-blocking. + + Called from the Workflow thread, so it must not await: the watcher task + is scheduled onto the manager's loop rather than started here. + """ + backend = self._backend + if backend is None: + raise RuntimeError( + "external streams are not configured on this Worker; pass " + "external_stream_backend=... to the Worker" + ) + subscription = Subscription( + run_id=run_id, + wait_id=wait_id, + stream_key=stream_key, + backend=backend, + buffer_size=self._buffer_size, + committed_cursor=start_cursor, + delivery_cursor=start_cursor, + prefetch_cursor=start_cursor, + ) + replaced = self._runs.setdefault(run_id, {}).get(wait_id) + if replaced is not None: + # A wait re-registered under a key that already has one. The old + # subscription is unreachable from here on, and its watcher would go + # on polling the backend for the life of the process -- a leak that + # only shows up as a connection that never closes. + replaced._cancelled = True + self._loop.call_soon_threadsafe(self._cancel_watcher, replaced) + self._runs[run_id][wait_id] = subscription + # `create_task` is not thread-safe, and this runs on the Workflow + # executor thread. Scheduling the start onto the manager's loop is the + # difference between a watcher that runs and one that is silently never + # scheduled -- which looks exactly like a stream that never delivers. + self._loop.call_soon_threadsafe(self._start_watcher, subscription) + return subscription + + def _cancel_watcher(self, subscription: Subscription) -> None: + watcher = subscription._watcher + if watcher is not None and not watcher.done(): + watcher.cancel() + + def _start_watcher(self, subscription: Subscription) -> None: + """Starts a watcher, and reconciles the park state it inherited. + + Both on the manager's own loop, because both are `create_task` calls and + `register` runs on the Workflow executor thread. + """ + if subscription._cancelled or subscription._watcher is not None: + return + reconcile = self._loop.create_task(self._reconcile_inherited_park(subscription)) + # Held, because the loop keeps only a weak reference to a running task: + # an unreferenced one can be collected mid-await, and this one's whole + # job is a backend round trip. + self._reconciliations.add(reconcile) + reconcile.add_done_callback(self._reconciliations.discard) + subscription._watcher = self._loop.create_task(self._watch(subscription)) + + async def _reconcile_inherited_park(self, subscription: Subscription) -> None: + """Removes a park intent this Worker inherited rather than installed. + + `installed_park_generation` is a *mirror* of backend state, and it lives + on the Worker that installed the park. The intent it mirrors is durable: + it survives eviction, a Workflow Task that moved to another Worker, and + shutdown, all of which take the mirror with them. Removal keyed on the + mirror alone therefore cannot reach the one class of intent that most + needs reaching -- the intents whose installer is gone -- and no later + resolve can repair it, because the resolve looks at the same empty + mirror. + + Registration is where a Worker learns such an intent exists, and it is + also the moment its status is unambiguous. A subscription is registered + by user Workflow code running, and no user code runs inside a park + (`wft-lifecycle.md`), so an intent found here belongs to a park that is + over: the Core that confirmed it has either moved on or gone with the + Worker that held it. What leaving it costs is the invariant's whole + point -- `current_park_generation` keeps answering a generation Core has + discarded, every producer wake names that generation and Core discards + it as stale, and because a parked wake's request ID ignores sender + identity the second such wake is byte-identical to the first and the + server deduplicates it away. The Workflow then waits forever on a record + that is durably present. + + Best-effort in the sense that no failure here reaches a Workflow Task -- + this is the registration path of a Run that is already running, and a + momentary backend failure must not fail one. It is *not* best-effort in + the sense of one attempt, and neither half of the job may end with the + burst of cheap ones below. + + Once the intent has been *read* it is recorded in this Run's + owed-removal ledger, whose own retry makes backend recovery sufficient; + this loop hands the entry over rather than draining it twice. Until it + has been read there is nothing to record -- the ledger holds an + identity, and the identity is what could not be obtained -- so + discovery keeps retrying here, with backoff, for as long as this Worker + holds the subscription. Both halves answer the same argument: waiting + for "the next time this wait is registered" is a coincidence of + eviction, not a mechanism, and the Run that most needs the removal -- + one cached and blocked on something other than this stream -- is + precisely the one that never registers it again. + + Later parks, resolves, registrations and evictions remain eager fast + paths for whatever is in the ledger. + """ + run_id = subscription.run_id + key = (subscription.stream_key, subscription.wait_id) + attempt = 0 + delay = PARK_REMOVAL_RETRY_DELAY.total_seconds() + while True: + if subscription._cancelled: + # Eviction does not cancel this task -- it is held in + # `_reconciliations`, not on the Subscription -- and shutdown + # cancels it only after Run teardown, so this flag stops it in + # either path. Anything read by now is in the ledger and its + # autonomous retry. + # + # It is also what keeps this loop *safe* to run indefinitely. + # Removing whatever is installed at this key is justified by + # this Worker holding the Run and this wait being registered on + # it: any intent there belongs to a park that is over. Once the + # subscription is gone that justification is gone with it, and + # the intent this would read could be a park a *different* + # Worker is sitting in. + return + if await self._reconcile_inherited_park_once(subscription, attempt): + return + if key in self._owed_removals.get(run_id, {}): + # Discovery worked and the removal did not, which is the case + # the ledger exists for. It owns the retry from here; a second + # loop draining the same entry would only double the round + # trips. + logger.warning( + "Could not remove the inherited park intent for %s wait %s " + "in %s attempts; its autonomous owed-removal retry " + "continues in the background", + subscription.stream_key, + subscription.wait_id, + attempt + 1, + ) + return + attempt += 1 + if attempt == PARK_REMOVAL_ATTEMPTS: + # Past the cheap burst with nothing read, so there is no ledger + # entry to hand this to -- the ledger records an intent's + # identity, and that is precisely what could not be read. Giving + # up here is the liveness gap the ledger closed for removals, + # reopened for discovery: a backend down for the burst and + # healthy after it would leave the obsolete intent installed + # until this wait happens to be registered again, which for a + # Run that stays cached is never. + logger.warning( + "Could not read the park intent for %s wait %s in %s " + "attempts; this reconciliation keeps retrying it in the " + "background, with backoff, until it can say whether one is " + "installed", + subscription.stream_key, + subscription.wait_id, + PARK_REMOVAL_ATTEMPTS, + ) + if attempt >= PARK_REMOVAL_ATTEMPTS: + delay = min(PARK_REMOVAL_MAX_RETRY_DELAY.total_seconds(), delay * 2) + await asyncio.sleep(delay) + + async def _reconcile_inherited_park_once( + self, subscription: Subscription, attempt: int + ) -> bool: + """One read-and-remove pass. Returns whether nothing is left to do. + + The lock is taken per attempt rather than held across the backoff + between them, because `prepare_park` runs inside an activation under + Core's deadlock timeout: a reconciliation that kept this Run's lock + while it slept would spend that budget on cleanup for a park that is + already over. + + Bounding the hold itself is the same argument carried one step further. + This runs on a task of its own, not inside the activation, so the calls + below can hang without anything above them noticing -- and this Run's + park lock hangs with them, which is what the next park, resolve or + eviction takes, and what the owed-removal retry loop takes. An attempt + that runs out of time is treated as a failed attempt, because that is + what it is; the entry was recorded before the call, so nothing read is + lost by giving up on it. + """ + try: + async with self._park_lock(subscription.run_id): + return await asyncio.wait_for( + self._reconcile_inherited_park_locked(subscription, attempt), + PARK_REMOVAL_CALL_TIMEOUT.total_seconds(), + ) + except asyncio.TimeoutError: + logger.warning( + "Reconciling the inherited park intent for %s wait %s did not " + "finish within %s on attempt %s/%s; the park lock is released " + "and anything already read stays owed", + subscription.stream_key, + subscription.wait_id, + PARK_REMOVAL_CALL_TIMEOUT, + attempt + 1, + PARK_REMOVAL_ATTEMPTS, + ) + return False + + async def _reconcile_inherited_park_locked( + self, subscription: Subscription, attempt: int + ) -> bool: + """The pass itself, with this Run's park lock already held.""" + run_id = subscription.run_id + key = (subscription.stream_key, subscription.wait_id) + # Drained first, because a previous attempt of this same loop may + # already have read the intent and failed only to remove it. The + # ledger is that retry; re-reading would find exactly what it holds. + await self._drain_owed_removals(run_id) + if key in self._owed_removals.get(run_id, {}): + return False + if ( + subscription._cancelled + or subscription.installed_park_generation is not None + ): + # A park confirmed since this was scheduled is *this* manager's, + # and `resolve_park` owns it. Removing it here would take the + # intent out from under a park that really is outstanding. + return True + try: + inherited = await subscription.backend.park_intent( + subscription.stream_key, subscription.wait_id + ) + except asyncio.CancelledError: + raise + except Exception: + logger.warning( + "Reading the inherited park intent for %s wait %s failed on attempt %s", + subscription.stream_key, + subscription.wait_id, + attempt + 1, + exc_info=True, + ) + return False + if inherited is None: + # The ordinary case, and the reason this is a read before it + # is a write: a Run that never parked owes the backend + # nothing. + return True + if ( + subscription._cancelled + or subscription.installed_park_generation is not None + ): + # Asked again across the read, because this Run's state can + # have moved on entirely while it was in flight: evicted and + # picked up again, with a *new* park confirmed for the same + # key. Removing what was found before that would strand the + # park that replaced it. + return True + logger.info( + "Removing an inherited external stream park intent for %s " + "wait %s: park generation %s was confirmed by a Worker that " + "no longer holds this Run", + subscription.stream_key, + subscription.wait_id, + inherited.park_generation, + ) + # Owed from the moment it is known to exist, not from the moment a + # removal fails: this is the only mirror an inherited intent ever + # gets, and without it `_remove_park_intent` goes on short-circuiting + # on an empty `installed_park_generation` and every removal path -- + # the resolve, the withdrawal, the close -- stays disabled for it. + self._owe_removal( + run_id, + key, + _OwedRemoval( + backend=subscription.backend, + park_generation=inherited.park_generation, + run_id=inherited.run_id, + ), + ) + try: + outcome = await subscription.backend.remove_park_intent_if_matches( + subscription.stream_key, + subscription.wait_id, + run_id=inherited.run_id, + park_generation=inherited.park_generation, + ) + except asyncio.CancelledError: + raise + except Exception: + logger.warning( + "Removing the inherited park intent for %s wait %s failed on " + "attempt %s", + subscription.stream_key, + subscription.wait_id, + attempt + 1, + exc_info=True, + ) + return False + self._forget_owed_removal(run_id, key) + # This path is the *ordinary* one for an inherited intent -- the removal + # succeeding first time -- and it silences records exactly as a retried + # one does. `_start_watcher` runs this reconciliation and the watcher + # concurrently, so the watcher can read a record, find no open Workflow + # Task and send its wake naming the inherited generation while this is + # still in flight. Core discards that wake, and the watcher will not + # come back through here without a *new* append. + self._reannounce_after_cleanup(subscription, outcome) + return True + + def _owe_removal( + self, run_id: str, key: tuple[StreamKey, int], record: _OwedRemoval + ) -> None: + owed = self._owed_removals.setdefault(run_id, {}) + previous = owed.get(key) + owed[key] = record + retry = self._owed_removal_retries.get(run_id) + if retry is None or retry.done(): + if not self._stopping_owed_removal_retries: + self._owed_removal_wakeups.setdefault(run_id, asyncio.Event()) + retry = self._loop.create_task(self._retry_owed_removals(run_id)) + self._owed_removal_retries[run_id] = retry + elif previous != record: + # Do not make newly discovered debt wait behind the capped backoff + # of an older failing entry for this Run. + self._owed_removal_wakeups[run_id].set() + + def _forget_owed_removal(self, run_id: str, key: tuple[StreamKey, int]) -> None: + owed = self._owed_removals.get(run_id) + if owed is None: + return + owed.pop(key, None) + if not owed: + # Dropped rather than left empty, so a Run that owes nothing costs + # nothing to carry and `evict_run` can tell the two apart cheaply. + self._owed_removals.pop(run_id, None) + wakeup = self._owed_removal_wakeups.get(run_id) + if wakeup is not None: + wakeup.set() + + async def _retry_owed_removals(self, run_id: str) -> None: + """Retries one Run's ledger without needing another lifecycle event.""" + this_task = asyncio.current_task() + initial = PARK_REMOVAL_RETRY_DELAY.total_seconds() + delay = initial + maximum = PARK_REMOVAL_MAX_RETRY_DELAY.total_seconds() + wakeup = self._owed_removal_wakeups[run_id] + try: + while self._owed_removals.get(run_id): + try: + await asyncio.wait_for(wakeup.wait(), delay) + delay = initial + except asyncio.TimeoutError: + pass + wakeup.clear() + + before = len(self._owed_removals.get(run_id, {})) + await self._bounded_drain(run_id) + after = len(self._owed_removals.get(run_id, {})) + if after == 0: + return + delay = initial if after < before else min(maximum, delay * 2) + finally: + if self._owed_removal_retries.get(run_id) is this_task: + self._owed_removal_retries.pop(run_id, None) + self._owed_removal_wakeups.pop(run_id, None) + if run_id not in self._runs and not self._owed_removals.get(run_id): + self._park_locks.pop(run_id, None) + + async def _bounded_drain( + self, run_id: str, budget: timedelta = PARK_REMOVAL_CALL_TIMEOUT + ) -> None: + """Drains under this Run's park lock, and lets go of it either way. + + The retry loop and the last pass at shutdown are the drains that take + this lock with no activation above them and nothing else bounding them -- + the reconciliation's drain is bounded by the hold it sits inside, and + every other one is inside the activation whose budget a slow backend + spends regardless. A backend that hangs rather than raises would hold the + lock for as long as the call lasts, hence a timeout and not merely a + bounded number of attempts: attempts that never return are not bounded by + counting them. + + A timeout leaves the entry owed, which is what the loop above does with a + failure -- back off and come round again. + """ + try: + async with self._park_lock(run_id): + await asyncio.wait_for( + self._drain_owed_removals(run_id), budget.total_seconds() + ) + except asyncio.TimeoutError: + logger.warning( + "A background owed-removal drain for run %s did not finish " + "within %s; the park lock is released and it stays owed", + run_id, + budget, + ) + + async def _drain_owed_removals(self, run_id: str) -> None: + """Retries every removal this Run still owes. Never raises. + + A ledger entry is not "an intent exists"; it is a removal this manager + has already decided on and not had confirmed. That is what makes + draining safe from any holder of this Run's park lock rather than only + from the path that recorded it. The entry retains nothing -- no watcher, + no buffer, no connection beyond the backend the removal has to go + through. + + A Continue-As-New successor re-uses the stream key with wait ids that + start again at 1, so an entry a predecessor Run left could name a live + park's key. The backend therefore compares both the Run ID and park + generation as part of the delete. A read followed by an unconditional + removal would leave a cross-Run window in which this retry could delete + the successor's park. + """ + owed = self._owed_removals.get(run_id) + if not owed: + return + for key, record in list(owed.items()): + stream_key, wait_id = key + try: + outcome = await record.backend.remove_park_intent_if_matches( + stream_key, + wait_id, + run_id=record.run_id, + park_generation=record.park_generation, + ) + except asyncio.CancelledError: + raise + except Exception: + logger.warning( + "Retrying the owed park intent removal for %s wait %s " + "failed; it stays owed", + stream_key, + wait_id, + exc_info=True, + ) + continue + self._forget_owed_removal(run_id, key) + subscription = self.subscription(run_id, wait_id) + if subscription is not None and subscription.stream_key == stream_key: + # The mirror is only cleared once the backend agrees, here for + # the same reason `_remove_park_intent` does it in that order. + subscription.installed_park_generation = None + self._reannounce_after_cleanup(subscription, outcome) + else: + # No subscription to announce through, which is the *ordinary* + # case for this ledger rather than a corner of it: eviction and + # close both drop the Subscription and deliberately leave the + # entry and its retry alive, so by the time a recovered backend + # lets the removal through there is frequently nothing left + # holding the record the intent silenced. Announcing only when + # the wait happens to still be registered put the record-loss + # half of the suppression straight back. + self._wake_after_cleanup(record, stream_key, wait_id, outcome) + + def _wake_after_cleanup( + self, + record: _OwedRemoval, + stream_key: StreamKey, + wait_id: int, + outcome: ParkIntentRemoval, + ) -> None: + """Hands off a record a retired intent may have silenced, with no wait left. + + The Subscription-bound half of this is + :meth:`_reannounce_after_cleanup`, and it is the better one where it + applies: local readiness is free, and it fires only for a wait that is + actually holding records. Neither is knowable from here. The buffer went + with the Subscription, and whether a *remote* producer appended during + the suppression window was never visible from this Worker at all -- so + the choice is one unparked wake or silence, and silence is the failure + this whole path exists to prevent. + + ``MISMATCH`` is the one outcome that stays silent, for the reason + :meth:`_reannounce_after_cleanup` gives: an intent is still installed at + that key, so the suppression has not ended and it belongs to a park this + entry knows nothing about. + + Posted as a task rather than awaited, because every caller reaches the + drain holding this Run's park lock -- and one bounded by + `PARK_REMOVAL_CALL_TIMEOUT`, which is a budget for a single-key backend + call and not for a Signal with retries behind it. `shutdown` awaits what + is still in flight; see :meth:`_flush_cleanup_wakes`. + """ + if outcome is ParkIntentRemoval.MISMATCH: + return + if record.announced: + # The shutdown sweep already made this exact handoff -- an unparked + # wake for this wait, composed that way *because* this entry said the + # installed intent was stale. A second one is a second empty + # Workflow Task, not a second chance. + return + if self._send_wake is None: + return + target = _CleanupWake( + run_id=record.run_id, + wait_id=wait_id, + stream_key=stream_key, + backend=record.backend, + ) + task = self._loop.create_task(self._send_cleanup_wake(target)) + self._cleanup_wakes[task] = target + task.add_done_callback(self._cleanup_wakes.pop) + + async def _send_cleanup_wake(self, target: _CleanupWake) -> None: + """Sends one cleanup handoff. Never raises except on cancellation. + + Counted through `_count_owed_wake` like any other owed wake, so the + retries inside `_send_owed_wake` are the same wake rather than three + separate asks. + + A wake that never lands is reported, and while shutting down it is also + *counted*: the sweep's promise is that a handoff this Worker did not make + cannot be silent, and a removal confirmed during the final pass is as + much a part of that shutdown as a wake the sweep sent itself. On the live + path there is no counter with that meaning -- the metric is the shutdown + sweep's -- so the warning is what is left, and it says plainly that + nothing retries this. + """ + self._count_owed_wake(target) + if await self._send_owed_wake(target): + return + logger.warning( + "The wake owed for the record a retired park intent silenced was " + "not acknowledged for %s wait %s; nothing else announces it", + target.stream_key, + target.wait_id, + ) + if self._shutting_down: + self._record_shutdown_wake_failure(target) + + async def _flush_cleanup_wakes(self, budget: timedelta) -> None: + """Waits out the cleanup handoffs still in flight as the Worker exits. + + The last one of these is created by :meth:`_final_owed_removal_pass`, + which runs after everything else shutdown does, so without this the + process leaves with the Signal task merely scheduled -- which is the same + as not having sent it, and worse, because the failure counter would say a + clean shutdown. + + Bounded by what is left of the grace period, for the reason every other + step of shutdown is: cleanup may not hold the exit open. What the bound + cuts off is counted rather than dropped, because the wake it abandons is + silent by nature. + """ + pending = set(self._cleanup_wakes) + if not pending: + return + if budget.total_seconds() > 0: + _, pending = await asyncio.wait(pending, timeout=budget.total_seconds()) + for task in pending: + target = self._cleanup_wakes.get(task) + task.cancel() + if target is None: + continue + logger.warning( + "External stream shutdown abandoned the wake owed for a retired " + "park intent on %s wait %s", + target.stream_key, + target.wait_id, + ) + self._record_shutdown_wake_failure(target) + + def _reannounce_after_cleanup( + self, subscription: Subscription, outcome: ParkIntentRemoval + ) -> None: + """Announces again a record the retired intent had already silenced. + + A stale intent does not only leak. While it was installed every wake for + this stream named the generation behind it -- the producer reads + `current_park_generation`, and so does the Worker's own sender -- and + Core discards a non-zero generation that is not the park it is holding. + A record that arrived in that window was therefore announced to nobody, + and *removing the intent does not announce it*: `_report_ready` counted + its wake as sent, the watcher has moved `prefetch_cursor` past it and + returns on the next empty read, and `_send_owed_wake` re-sends only a + wake that **raised**. Without this, a Run can sit blocked on a record + this Worker is already holding until some unrelated event happens to + create a Workflow Task. + + Both outcomes that clear the key announce, and that is why the provider + contract does not answer with a Boolean. `ABSENT` is not "someone else's + intent is in the way", it is "the intent this entry named is gone" -- + which is what a retry sees after the reply to a delete that in fact + succeeded was lost, and what one cleanup owner sees after another + finished the job. Reading that as a mismatch would leave the record the + intent silenced silent for good. `MISMATCH` announces nothing: an intent + is still installed there, the suppression it causes has not ended, and it + belongs to a park this entry knows nothing about. + + Narrow otherwise. Only a wait still holding records, so a healthy park + costs no Workflow Task; and never during the sweep, which owns every wake + from that point and accounts for each one. Posted rather than awaited, + because callers hold the Run's park lock and `_report_ready` goes out to + Core. + + Every one of those narrowings reads the Subscription, which is why this + is only half the cleanup. Once the wait is closed or its Run evicted + there is no buffer to consult and no readiness channel to use, and the + obligation is discharged by :meth:`_wake_after_cleanup` instead. + """ + if outcome is ParkIntentRemoval.MISMATCH: + return + if subscription._cancelled or not subscription.buffered: + return + if self._shutting_down: + return + pending = self._pending_wake_runs.get(subscription.run_id) + if pending is not None and pending.wait_id == subscription.wait_id: + # The pending wake for this wait named the intent just retired, so + # Core discarded it. It cannot gate the unparked reannouncement + # whose purpose is to replace that silenced wake. + del self._pending_wake_runs[subscription.run_id] + self._loop.create_task(self._report_ready(subscription)) + + def _park_lock(self, run_id: str) -> asyncio.Lock: + """Serializes one Run's park-intent work on the manager's loop. + + The install/recheck handshake, the resolve, and the reconciliation above + all read-then-write the same `(stream key, wait_id)` objects, and the + reconciliation is scheduled from another thread, so their interleaving + is not otherwise constrained. Without this, a reconciliation that + overlapped a confirming park could remove the intent that park had just + installed -- an unwakeable Run produced by the very code that exists to + prevent one. + """ + lock = self._park_locks.get(run_id) + if lock is None: + lock = self._park_locks[run_id] = asyncio.Lock() + return lock + + def subscription(self, run_id: str, wait_id: int) -> Subscription | None: + """Return one registered subscription, if present.""" + return self._runs.get(run_id, {}).get(wait_id) + + def subscriptions(self, run_id: str) -> list[Subscription]: + """Return all subscriptions currently registered for a Run.""" + return list(self._runs.get(run_id, {}).values()) + + def runs_with_subscriptions(self) -> list[str]: + """Return the Run IDs that currently own at least one subscription.""" + return [run_id for run_id, subs in self._runs.items() if subs] + + # --- the Workflow thread's entry points --------------------------------- + + def drain( + self, run_id: str, wait_id: int, max_records: int | None = None + ) -> list[StreamRecord]: + """Pops from one subscription's buffer. Performs no I/O.""" + subscription = self.subscription(run_id, wait_id) + if subscription is None: + return [] + popped = subscription.drain(max_records) + if popped: + # Freeing room restarts prefetch. Hopped onto the loop because an + # asyncio.Event may not be set from another thread. + self._loop.call_soon_threadsafe(subscription.note_drained) + return popped + + def note_wait_generation(self, run_id: str, wait_id: int, generation: int) -> None: + """Hands one wait's current generation to the subscription reporting it. + + Called from the Workflow thread, and deliberately *not* hopped onto the + manager's loop: the next readiness report may be issued by a watcher + before a queued callback would run, and it would then name the stale + generation this exists to replace. The write is a single guarded field + assignment, so it is safe to make directly. + """ + subscription = self.subscription(run_id, wait_id) + if subscription is not None: + subscription.note_wait_generation(generation) + + def rearm_ready(self, run_id: str) -> None: + """Re-reports readiness for every buffer this Run left non-empty. + + The per-activation delivery budget stops an iterator with records still + buffered. Nothing else will announce them: readiness is reported once, + when the watcher buffers a record, and that watcher has long since moved + its prefetch cursor past these. Without this the Workflow blocks forever + on records already sitting in front of it -- and worse, it blocks + *quiescently*, so Core would start the idle timer and eventually park a + Workflow Task whose data had already arrived. + + Called from the Workflow thread at activation completion, so the work is + hopped onto the manager's loop: `create_task` is not thread-safe, and a + task created from the Workflow executor thread is silently never + scheduled -- indistinguishable from a stream that never delivers. + """ + self._loop.call_soon_threadsafe(self._rearm_ready, run_id) + + def _rearm_ready(self, run_id: str) -> None: + """The manager-loop half of :meth:`rearm_ready`.""" + for subscription in self.subscriptions(run_id): + if subscription._cancelled or not subscription.buffered: + continue + self._loop.create_task(self._report_ready(subscription)) + + def note_workflow_task_started(self, run_id: str) -> None: + """Records an activation boundary used to correlate an outstanding wake.""" + self._activation_sequences[run_id] = ( + self._activation_sequences.get(run_id, 0) + 1 + ) + + def note_workflow_task_completed(self, run_id: str, *, terminal: bool) -> None: + """Releases a coalesced wake after Core accepts a successful completion. + + Signal acknowledgement only says the wake entered History. The task it + creates can still fail or have its completion rejected; both paths are + replayed, and readiness reconstructed during that replay belongs behind + the same outstanding wake. The Worker's successful completion callback + is the first evidence that cycle is over. + + A non-terminal task may leave another generation buffered. Re-report it + after releasing the gate because a report queued by + :meth:`rearm_ready` while completion was in flight was deliberately + coalesced. A terminal task has no next activation and must not re-arm. + """ + pending = self._pending_wake_runs.get(run_id) + if pending is None: + return + # A completion for an activation already running when the send began is + # not evidence that the wake was processed. This exact interleaving is + # common in the closing window: readiness sends while the old task is + # reporting, then that old completion wins the race to Core. + if self._activation_sequences.get(run_id, 0) <= pending.attempt_started_at: + return + if not pending.acknowledged: + # The service call is still in flight. Its response decides whether + # this completion can belong to the attempt: on failure the next + # retry resets both the activation boundary and this observation. + pending.completed_terminal = terminal + return + self._release_pending_wake(run_id, pending, terminal=terminal) + + def _release_pending_wake( + self, run_id: str, pending: _PendingWake, *, terminal: bool + ) -> None: + """Ends ``pending`` if it is still this Run's current wake cycle.""" + if self._pending_wake_runs.get(run_id) is not pending: + return + del self._pending_wake_runs[run_id] + if not terminal: + self._rearm_ready(run_id) + + def _note_wake_attempt_started(self, target: WakeTarget) -> None: + """Moves a live wake's correlation boundary to this retry attempt.""" + if not isinstance(target, Subscription): + return + pending = self._pending_wake_runs.get(target.run_id) + if pending is None or pending.wait_id != target.wait_id: + return + pending.attempt_started_at = self._activation_sequences.get(target.run_id, 0) + pending.acknowledged = False + pending.completed_terminal = None + + def _note_wake_attempt_acknowledged(self, target: WakeTarget) -> None: + """Marks the current attempt delivered and applies any racing completion.""" + if not isinstance(target, Subscription): + return + pending = self._pending_wake_runs.get(target.run_id) + if pending is None or pending.wait_id != target.wait_id: + return + pending.acknowledged = True + if pending.completed_terminal is not None: + self._release_pending_wake( + target.run_id, + pending, + terminal=pending.completed_terminal, + ) + + def reposition_to_committed( + self, run_id: str, cursors: Mapping[int, Cursor] + ) -> None: + """Moves each named wait to the boundary a replayed marker committed. + + Called from the Workflow thread once replay has delivered a marker's + recorded ranges, and **completed before it returns** -- unlike + :meth:`rearm_ready`, which only has to happen eventually. The very next + thing the Workflow thread does is drain, and a reposition merely posted + to the manager's loop leaves the marker-covered records sitting in the + buffer for that drain to hand over a second time. Posting it and + returning made the fix depend on the manager loop winning a race that + nothing ordered; the observed symptom was a Workflow receiving + ``['alpha', 'alpha', 'beta']``. + + Only the watcher's wakeup is hopped, because that is the one part that + touches an ``asyncio`` primitive. The retraction itself is under the + subscription's lock and the epoch bump fences a read already in flight, + so a watcher appending concurrently is either cleared by this or + rejected by :meth:`Subscription._append`. + """ + repositioned: list[Subscription] = [] + for wait_id, cursor in cursors.items(): + subscription = self.subscription(run_id, wait_id) + if subscription is None or subscription._cancelled: + # A wait the marker records but this Run no longer holds. Replay + # itself reports that as nondeterminism; there is nothing to + # reposition here and nothing to say about it. + continue + subscription.reposition_to(cursor) + repositioned.append(subscription) + for subscription in repositioned: + # An `asyncio.Event` may not be set from another thread, so the one + # asynchronous consequence of the retraction -- a watcher parked on + # backpressure now having room -- is what gets hopped. + self._loop.call_soon_threadsafe(subscription.note_drained) + + def blocked_snapshot(self, run_id: str) -> dict[int, Cursor]: + """Where every active subscription's deliveries stopped. + + Read from manager state alone -- no provider call, no backend read. + The boundary is not "wherever the stream is now"; it is where this + Workflow Task's deliveries stopped, which is already fixed. + """ + return { + wait_id: sub.blocked_cursor() + for wait_id, sub in sorted(self._runs.get(run_id, {}).items()) + } + + # --- watchers ----------------------------------------------------------- + + async def _watch(self, subscription: Subscription) -> None: + """Prefetches into the buffer and reports readiness once buffered. + + Survives Workflow Task completion. Torn down only on cancellation, Run + eviction, or Worker shutdown -- because the window between tasks is + exactly when an append most needs someone watching for it. + """ + try: + while not subscription._cancelled: + room = subscription._room() + if room == 0: + # Backpressure. Stop reading; drop nothing, block nobody. + subscription._has_room.clear() + await subscription._has_room.wait() + continue + + # Captured before the read, so that a reposition landing while + # it is in flight is detected afterwards. The read started from + # a cursor that no longer describes this subscription, and + # appending its result would undo the reposition rather than + # merely race it. + with subscription._lock: + epoch = subscription._prefetch_epoch + try: + records = await subscription.backend.read_after( + subscription.stream_key, + subscription.prefetch_cursor, + max_records=room, + block=self._watch_block, + ) + except asyncio.CancelledError: + raise + except Exception: + # Watcher failures are retried within provider policy. The + # Workflow Task is not failed from out here -- it may not + # even exist. + logger.exception( + "External stream watcher failed for %s wait %s", + subscription.stream_key, + subscription.wait_id, + ) + await asyncio.sleep(0.05) + continue + + if not records: + continue + + # Prepared **before** buffering, so that what the Workflow + # thread pops is already a value waiting to be typed. This is + # the whole reason the manager holds a converter: a codec + # awaited on the Workflow thread performs I/O in a + # deterministic event loop. + # + # And before the epoch check rather than after it, because + # preparing awaits: a reposition landing while a user's codec + # runs would otherwise slip past a check that had already + # passed, and the retracted records would go into the buffer + # anyway. The check stays the last thing before the append. + prepared = await self._prepare( + [(subscription.stream_key, record) for record in records] + ) + + # The epoch is compared inside `_append`, under the lock the + # reposition also takes, because repositioning is synchronous on + # the Workflow thread: a check made out here could pass and then + # have the reposition land before the append. `False` means the + # records were read from a position a reposition has since + # retracted, so they are records the marker already accounts + # for; the next pass reads from the new cursor. + if not subscription._append(prepared, epoch): + continue + await self._report_ready(subscription) + except asyncio.CancelledError: + pass + + async def _prepare( + self, records: Sequence[tuple[StreamKey, StreamRecord]] + ) -> list[StreamRecord]: + """Runs the DataConverter's asynchronous half, here on the Worker's loop. + + Records arrive paired with the stream they came from, because this + manager serves every Run on the Worker and the converter has to be bound + to the Workflow each individual record belongs to. + + Every record is prepared, including ones the Workflow may never take: + this is the same bargain the Worker already makes for an activation's + payloads, which are decoded in full before the executor sees any of + them. The cost of preparing a record that is later discarded is one + codec call; the cost of *not* preparing it is that the Workflow thread + has to, and the Workflow thread cannot. + + **Nothing raises out of here.** A failure is carried on the record and + raised by the delivery that would have yielded its value, which is the + only point at which a Workflow exists to be told. Raising here would + kill the watcher for the Run -- taking every later record with it -- and + would report a failure for a record the Workflow might never have asked + for. + + Control records carry no payload by construction, so they pass through. + """ + if self._data_converter is None: + return [record for _, record in records] + from temporalio.contrib.external_workflow_streams._codec import ( + StreamPayloadCodec, + ) + + data_converter = self._data_converter + # One codec per Workflow, not one per manager. A converter bound at + # construction would be bound to nothing in particular: this manager + # outlives every Run on the Worker and prepares records for all of them + # at once, so the only correct context is the one the record's own + # stream carries. + # + # `workflow_id` rather than `run_id`, and both taken from the stream key + # rather than from a subscription: a stream spans the whole + # Continue-As-New chain, so a successor Run must decode its predecessor's + # records with the same key -- which is also why `decode_activation` + # keys on `workflow_id`. + # + # Memoized for this call only. Live batches are one stream, replay + # batches are one Workflow's waits, so the clone happens per batch + # rather than per record; holding the map on the manager instead would + # accumulate one entry per Workflow ID the Worker ever served. + codecs: dict[tuple[str, str], StreamPayloadCodec[Any]] = {} + + def codec_for(key: StreamKey) -> StreamPayloadCodec[Any]: + cached = codecs.get((key.namespace, key.workflow_id)) + if cached is None: + # No type: the type belongs to the topic, and the topic belongs + # to Workflow code. Nothing out here needs it, because nothing + # out here runs the payload converter. + cached = StreamPayloadCodec( + data_converter.with_context( + temporalio.converter.WorkflowSerializationContext( + namespace=key.namespace, + workflow_id=key.workflow_id, + ) + ), + None, + ) + codecs[(key.namespace, key.workflow_id)] = cached + return cached + + prepared: list[StreamRecord] = [] + for key, record in records: + if record.is_control: + prepared.append(record) + continue + try: + prepared.append( + PreparedRecord.of( + record, await codec_for(key).prepare(record.payload), None + ) + ) + except asyncio.CancelledError: + raise + except Exception as err: + logger.debug( + "Preparing external stream record at %s failed; the error " + "travels with the record", + record.offset, + exc_info=True, + ) + prepared.append(PreparedRecord.of(record, None, err)) + return prepared + + async def _report_ready(self, subscription: Subscription) -> None: + """Tells Core a record is buffered, and acts on which answer comes back. + + **Total, except for cancellation.** Nothing but `CancelledError` may + leave this method. The watcher calls it in its loop, so an exception + escaping ends that watcher for good -- the subscription stays registered, + its buffer keeps its records, and nothing ever announces them again. + `_rearm_ready` also launches it with `create_task`, where an exception + becomes a task result nobody retrieves and the failure is not even + logged. + """ + result = await self._notify_ready_with_retries(subscription) + + if result == ReadinessResult.ACCEPTED: + # Core will activate; the record is announced and this is done. + return + + if result == ReadinessResult.STALE: + # Core is holding a newer generation for this wait than the one just + # named, so the report was for a block that has already been + # resolved. Re-report against the current generation rather than + # returning: the watcher only calls back here after a *new* + # non-empty read, and `prefetch_cursor` is already past the record + # in the buffer, so nothing would announce it a second time. A + # record announced to nobody is a Workflow blocked forever on data + # it is already holding. + # + # `result` is **replaced** by what the retries ended on, not merely + # tested. The five answers differ in what happens to the watcher + # afterwards, and a retry can legitimately land on a different one + # than the first report did -- the Run can be evicted between the + # report that raced a generation change and the delayed retry that + # follows it. Keeping the original `Stale` there sent the owed wake + # (right) and then skipped the `RunNotFound` teardown (wrong), leaving + # a watcher, a buffer, a backend read loop and a Run-map entry alive + # for a Run this Worker no longer owns. + result = await self._retry_stale(subscription) + if result == ReadinessResult.ACCEPTED: + return + + # Everything left means local readiness could not be delivered. One + # Signal is owed for the Run, but only one until Core accepts the + # successful Workflow Task completion it caused. A second wait or a + # replayed copy of this one can report during the closing/failing window; + # sending each report creates another Workflow Task that can collide + # with the same completion and sustain the replay loop. + if subscription.run_id in self._pending_wake_runs: + return + self._pending_wake_runs[subscription.run_id] = _PendingWake( + wait_id=subscription.wait_id, + attempt_started_at=self._activation_sequences.get(subscription.run_id, 0), + ) + # + # Counted once, and then retried inside `_send_owed_wake`, because a + # single attempt here has nothing behind it. The watcher has already + # moved `prefetch_cursor` past the buffered record and returns on `if + # not records`, so it never comes back through here without a *new* + # append; `rearm_ready` needs the activation this lost wake was supposed + # to cause; and the idle timer only runs while a Workflow Task is + # retained, which `NoOpenWorkflowTask` says there is not. One failed + # attempt was a lost record. + self._count_owed_wake(subscription) + if not await self._send_owed_wake(subscription): + self._pending_wake_runs.pop(subscription.run_id, None) + logger.warning( + "External stream wake for %s wait %s was not acknowledged; it " + "stays owed and the shutdown sweep is the only backstop left", + subscription.stream_key, + subscription.wait_id, + ) + + if result == ReadinessResult.RUN_NOT_FOUND: + # The Run is gone from this Worker. Nothing here can serve it again. + # Dropped only now, after the retries above: this pop is what takes + # the subscription out of `_runs`, and the shutdown sweep iterates + # `_runs`, so a wake given up on before it has no backstop at all. + subscription._cancelled = True + self._runs.get(subscription.run_id, {}).pop(subscription.wait_id, None) + # PARKED and NO_OPEN_WORKFLOW_TASK both *keep* the watcher: the Run is + # still cached and this is the normal window between Workflow Tasks. + + async def _notify_ready_with_retries(self, subscription: Subscription) -> str: + """Reports readiness, retrying a failing call a bounded number of times. + + A raising notifier is a transport problem, not an answer. Letting it + escape kills the watcher; swallowing it and returning would claim the + record was announced. Retrying and then falling through to the + wake-owed branch treats an unreachable Core as what it is: local + readiness could not be delivered, which is the sixth case the five + answers do not name. + """ + for attempt in range(READINESS_ATTEMPTS): + try: + return _result_value( + await self._notify_ready( + subscription.run_id, + subscription.wait_id, + subscription.current_wait_generation(), + ) + ) + except asyncio.CancelledError: + raise + except Exception: + logger.warning( + "External stream readiness report %s/%s failed for %s wait %s", + attempt + 1, + READINESS_ATTEMPTS, + subscription.stream_key, + subscription.wait_id, + exc_info=True, + ) + if attempt + 1 < READINESS_ATTEMPTS: + await asyncio.sleep(READINESS_RETRY_DELAY.total_seconds()) + return ReadinessResult.NO_OPEN_WORKFLOW_TASK + + async def _retry_stale(self, subscription: Subscription) -> str: + """Re-reports a stale readiness against the generation Core now holds. + + Returns the result the retries ended on, so the caller can act on it. A + Boolean would answer only "was it announced", and the four + non-``Accepted`` answers are not interchangeable: they differ in what + happens to the watcher, and `RunNotFound` in particular requires the + watcher to be torn down. Discarding them left that teardown unreachable + from here. + + A generation moves when the wait re-enters the blocked state, which is + Workflow code coming back around to it -- so the report that raced it is + answered by the next one. Bounded, because a wait that keeps moving is a + Workflow consuming happily, and a wake owed at the end of that costs one + empty Workflow Task rather than a silent stall. + + `RunNotFound` ends the retries rather than using them up. It is the one + answer that cannot change back: the Run is gone from this Worker, so a + further report can only be answered the same way, and each attempt costs a + delay before the wake this record still needs. + """ + result = ReadinessResult.STALE + for _ in range(STALE_REPORT_ATTEMPTS): + await asyncio.sleep(STALE_RETRY_DELAY.total_seconds()) + if subscription._cancelled or not subscription.buffered: + # Nothing left to announce: either the wait is gone or an + # activation drained the buffer, which is the record having been + # delivered by the very block this report raced. Reported as + # `Accepted` because that is what the caller does with it -- stop, + # and owe no wake. + return ReadinessResult.ACCEPTED + result = await self._notify_ready_with_retries(subscription) + if result in (ReadinessResult.ACCEPTED, ReadinessResult.RUN_NOT_FOUND): + return result + return result + + # --- the runtime-only jobs' backend work (P19) -------------------------- + + async def prepare_park( + self, run_id: str, park_generation: int, blocked: Mapping[int, Cursor] + ) -> bool: + """Installs park intents, then rechecks every stream. + + Returns ``True`` if any stream became ready, which abandons this parking + generation. + + ``blocked`` **is the park set**, not a lookup table beside one: its keys + are the waits Core asked to park, taken from + ``PrepareExternalStreamPark.waits``, and its values are their cursor + boundaries. The manager's own registration list is a superset -- a + subscription that delivered a record and was not awaited again is + registered and not blocked -- and parking the superset is wrong in both + directions. A recheck for a wait outside the set finds records nothing + is waiting on and aborts a legitimate park, which then runs again on the + next idle timeout and aborts again; and an intent installed for a wait + outside the set is an intent with no park behind it, which is exactly + what `backend-contract.md` forbids leaving in a backend. + + The order is what closes the append/park race: a producer appends its + record *before* it observes the park generation, so an append is either + seen by the recheck below or paired with a wake Signal. Rechecking + before all the intents were installed would leave a window where it is + neither. + """ + async with self._park_lock(run_id): + with _as_storage_failure("installing external stream park intents"): + return await self._prepare_park(run_id, park_generation, blocked) + + async def _prepare_park( + self, run_id: str, park_generation: int, blocked: Mapping[int, Cursor] + ) -> bool: + """The handshake itself, under this Run's park lock.""" + # A park is the one moment this Run is certain to reach a backend, so it + # is the cheapest place to retire whatever a previous close or resolve + # failed to remove -- and it must happen *before* the installs below, + # which would otherwise re-key the entries this is meant to retire. + await self._drain_owed_removals(run_id) + subscriptions = [ + subscription + for subscription in self.subscriptions(run_id) + if subscription.wait_id in blocked + ] + missing = set(blocked) - {s.wait_id for s in subscriptions} + if missing: + # Core is parking a wait this manager no longer holds. Nothing can + # be installed for it and there is nothing to recheck; a producer on + # that stream finds no intent and sends the unparked wake, which + # Core accepts as a recheck request (ADR-023). + logger.warning( + "Core asked run %s to park waits %s, which this Worker does not " + "hold; they are left out of the park set", + run_id, + sorted(missing), + ) + + installed: list[Subscription] = [] + try: + for subscription in subscriptions: + await subscription.backend.install_park_intent( + subscription.stream_key, + ParkIntent( + wait_id=subscription.wait_id, + cursor=blocked[subscription.wait_id], + park_generation=park_generation, + run_id=run_id, + ), + ) + subscription.installed_park_generation = park_generation + # A fresh intent at this key supersedes anything owed for it: a + # removal recorded against the generation just overwritten would + # otherwise be retried against the park now sitting behind it. + self._forget_owed_removal( + run_id, (subscription.stream_key, subscription.wait_id) + ) + installed.append(subscription) + + became_ready = False + for subscription in subscriptions: + if await subscription.backend.recheck( + subscription.stream_key, subscription.wait_id + ): + became_ready = True + break + except BaseException: + # A park that failed part-way through is still a park every producer + # can see. The activation fails and Core confirms nothing, so those + # intents describe a park that does not exist -- and an eviction then + # takes the local mirror away while they stay. Rolling back is what + # keeps "all-or-nothing across the set" true of the failure path too. + # Removal failures are swallowed: the storage error that got here is + # the one worth reporting, and anything still installed stays owed + # in this Run's ledger for the next drain to retry. + # + # `BaseException`, not `Exception`, and not with the usual + # `except asyncio.CancelledError: raise` above it: cancellation is + # the *most* likely way a half-installed park is abandoned -- Core + # withdrawing the activation, the Worker shutting down -- and it + # derives from `BaseException`, so an `Exception` handler leaves + # exactly those intents behind. Deliberately not shielded: a single + # `Task.cancel()` still lets this await run to completion, whereas a + # shielded rollback would detach and land after `prepare_park` has + # released `_park_lock`, where it can strand a newer legitimate + # park's intent instead. + await self._withdraw_park(installed) + raise + + if became_ready: + # All-or-nothing: a park confirmed for a set with a ready member + # would lose that member's record until a producer happened to + # signal, so every intent installed above comes back out. + failures = await self._withdraw_park(installed) + if failures: + # Nothing may report `became_ready` while an intent for that + # generation is still installed. Failing the activation retries + # the whole Workflow Task, which commits no cursor and loses no + # record; the removal stays owed either way. + raise failures[0] + return became_ready + + async def _withdraw_park( + self, subscriptions: list[Subscription] + ) -> list[Exception]: + """Takes every intent one attempted park installed back out. + + Total: every subscription is attempted whatever the others do, because a + rollback that stopped at its first failure would leave behind exactly + the orphans it exists to prevent. The failures are returned rather than + raised so each caller can decide which error is the one worth reporting. + """ + failures: list[Exception] = [] + for subscription in subscriptions: + try: + await self._remove_park_intent(subscription) + except asyncio.CancelledError: + raise + except Exception as err: + logger.exception( + "Failed removing the park intent for %s wait %s while " + "withdrawing an unconfirmed park; its autonomous " + "owed-removal retry continues in the background", + subscription.stream_key, + subscription.wait_id, + ) + failures.append(err) + return failures + + async def resolve_park(self, run_id: str) -> None: + """Removes the intents of a park this Run is no longer sitting in. + + The other half of :meth:`prepare_park`, and the reason the invariant is + stated as *an intent exists only while that park is outstanding* rather + than as "an aborted park cleans up after itself". An aborted park was + only ever the cheaper half: a **confirmed** park ends too, when a wake + Signal or a fresh quiescent snapshot clears Core's ``park_generation``, + and Core's own state moving on is not something the backend can observe. + + A left-behind intent is not inert, because it is the answer + ``current_park_generation`` gives to everyone who asks: + + - a **producer** reads it to decide what its wake Signal names, and a + dead generation is precisely the claim Core is designed to discard as + stale -- so the producer appends, signals, and the Workflow is never + woken by it; + - the **shutdown sweep** reads it through the same call, and would send a + parked wake where P20 requires the unparked one (ADR-023). Worse than + useless there: a parked wake's request ID deliberately ignores sender + identity, so it comes out identical to the wake that already resolved + that generation and the server deduplicates it away. + + This is the half for parks *this* manager installed, plus whatever the + Run's owed-removal ledger still holds. An intent whose installer is gone + -- evicted, moved to another Worker, shut down -- leaves no mirror here + to remove it by, and is put into that ledger by + :meth:`_reconcile_inherited_park` when the wait is registered again. + + Driven by ``ResolveExternalStreamWaits``, which is Core telling this + Worker the wait set has moved on -- the one event that covers both ways a + confirmed park ends, and the aborted-in-Core-but-confirmed-in-lang case + the recheck cannot see either. + + Failure is logged rather than raised. The removal is cleanup, and turning + a momentary backend blip into a failed Workflow Task on the *delivery* + path would trade a stale intent for a repeated Workflow Task. It stays + in the ledger, so the next drain retries it. + """ + async with self._park_lock(run_id): + # Ahead of the registered waits, because the ledger is the only + # thing that still names a wait this Run has closed -- and a stale + # intent on a closed wait is not confined to it: it keeps + # `parked_wait_ids` non-empty, which suppresses the unparked wake + # for the whole stream and silences every live wait on it. + await self._drain_owed_removals(run_id) + for subscription in self.subscriptions(run_id): + try: + await self._remove_park_intent(subscription) + except Exception: + logger.exception( + "Failed removing the resolved park intent for %s wait %s; " + "its autonomous owed-removal retry continues in the " + "background", + subscription.stream_key, + subscription.wait_id, + ) + + async def _remove_park_intent(self, subscription: Subscription) -> None: + """Removes one installed intent, and forgets it only once it is gone. + + Two things can say an intent is installed, and either one is enough. + ``installed_park_generation`` is the mirror of a park *this* manager + confirmed. The ledger is a removal already decided on and not yet + confirmed -- including every inherited intent, which is mirrored nowhere + because the Worker that installed it is gone. Keying on the mirror alone + is what left `resolve_park`, `_withdraw_park` and `cancel` all silently + doing nothing for exactly the intents that most need reaching. + + No read guards this one, unlike :meth:`_drain_owed_removals`: it is + reached only with a Subscription of a Run this Worker still holds, under + that Run's park lock, so the successor Run whose intent a stale entry + could name does not exist yet. + """ + run_id = subscription.run_id + key = (subscription.stream_key, subscription.wait_id) + if key not in self._owed_removals.get(run_id, {}): + if subscription.installed_park_generation is None: + return + # Recorded before the call rather than after it fails: a removal + # that never comes back -- the backend raised, or this task was + # cancelled mid-await -- is owed either way, and the Subscription + # this was reached through is frequently dropped in the same breath. + self._owe_removal( + run_id, + key, + _OwedRemoval( + backend=subscription.backend, + park_generation=subscription.installed_park_generation, + run_id=run_id, + ), + ) + record = self._owed_removals[run_id][key] + await subscription.backend.remove_park_intent_if_matches( + subscription.stream_key, + subscription.wait_id, + run_id=record.run_id, + park_generation=record.park_generation, + ) + subscription.installed_park_generation = None + self._forget_owed_removal(run_id, key) + + async def prepare_replay(self, run_id: str, replay_annotation: bytes) -> ReplayPlan: + """Reads and validates every recorded range, before any delivery. + + Runs here rather than in ``_apply`` because an inclusive range read over + a whole recorded batch is exactly the multi-second operation the + Workflow thread's deadlock timeout cannot accommodate -- and it gets + worse the more records the marker recorded. + + Subscriptions may not exist yet: on replay the Workflow has not run far + enough to call ``subscribe()``. The annotation's own header carries the + stream key and configured provider identity for each wait, which is why + it records them. + """ + from temporalio.contrib.external_workflow_streams._annotation import ( + decode_annotation, + ) + + annotation = decode_annotation(replay_annotation) + backends: dict[int, StreamBackend] = {} + stream_keys: dict[int, StreamKey] = {} + for wait_id, binding in annotation.header.streams.items(): + stream_keys[wait_id] = binding.stream_key + resolved = self._replay_backend(wait_id, binding) + if resolved is not None: + backends[wait_id] = resolved + + plan = await build_replay_plan(replay_annotation, backends, stream_keys) + # Replay delivers through the same drain the live path does, so it has + # to arrive in the same condition: prepared. This is the only chance to + # do it -- by the time the segments are delivered, the Workflow thread + # is running, and the ranges have already been validated here, so a + # decode failure from now on is a converter mismatch rather than + # integrity loss (ADR-015). + for index, segment in enumerate(plan.segments): + plan.segments[index] = ReplaySegment( + tuple( + (wait_id, prepared) + for (wait_id, _), prepared in zip( + segment.deliveries, + # The annotation header is the only place a replayed + # record's stream is written down: the Workflow has not + # run far enough to have re-created the subscription + # that would otherwise carry it. + await self._prepare( + [ + (stream_keys[wait_id], record) + for wait_id, record in segment.deliveries + ] + ), + ) + ) + ) + self._replay_plans[run_id] = plan + return plan + + def _replay_backend( + self, wait_id: int, binding: StreamBinding + ) -> StreamBackend | None: + """The one backend a recorded wait may be read from. + + Returning ``None`` leaves the wait unresolved when no backend is + configured; :func:`build_replay_plan` then reports any recorded ranges + that cannot be read. + """ + backend = self._backend + if backend is None: + return None + + # Whether the Worker still carries the same implementation is a + # deployment question, not a Workflow one. It is raised before any read, + # so an incompatible implementation never interprets recorded offsets. + declared_id = type(backend).provider_id + if declared_id != binding.provider_id: + raise StreamStorageError( + f"external stream wait {wait_id} was recorded against provider " + f"{binding.provider_id!r}, but the backend configured on this " + f"Worker declares {declared_id!r}. This Worker cannot read what " + "that marker recorded; configure the recorded provider." + ) + declared_version = type(backend).provider_format_version + if declared_version != binding.provider_format_version: + raise StreamStorageError( + f"external stream wait {wait_id} was recorded by provider " + f"{binding.provider_id!r} format version " + f"{binding.provider_format_version}, but the configured backend " + f"implements format version " + f"{declared_version}. Reading it would interpret the recorded " + "offsets under a format they were not written in." + ) + return backend + + def take_replay_plan(self, run_id: str) -> ReplayPlan | None: + """The prepared plan, consumed once by the delivering activation.""" + return self._replay_plans.pop(run_id, None) + + # --- teardown ----------------------------------------------------------- + + def cancel_from_workflow_thread(self, run_id: str, wait_id: int) -> None: + """Schedules :meth:`cancel` from the Workflow thread. + + Workflow code closes a subscription inside the synchronous + ``activate()``, on the executor thread, where creating a task is not + merely unsafe but silently ineffective: the task is never scheduled, and + a watcher that was supposed to stop keeps running with nothing to say so. + """ + self._loop.call_soon_threadsafe( + lambda: self._loop.create_task(self.cancel(run_id, wait_id)) + ) + + async def cancel(self, run_id: str, wait_id: int) -> None: + """Cancels one subscription: remove its intent, drop its buffer, stop it. + + The subscription leaves the Run's map **first**, so a close always stops + the watcher. Keeping a cancelled subscription registered while a removal + is retried would resurrect the orphaned-watcher bug from the other + direction: it goes on prefetching into a buffer nothing can drain, for + the rest of the Worker's life. + + What makes dropping it safe is that the removal no longer lives *on* it. + It is attempted here, retried a bounded number of times, and whatever is + still owed after that stays in this Run's ledger -- which outlives the + subscription and retries autonomously. Parks, resolves, registrations + and eviction also drain it eagerly. Before that ledger existed there was + no retry at all: the resolve path iterates registered subscriptions, + eviction and the shutdown sweep remove no intents, and another wait's + park cannot touch a per-wait key, so an intent left behind here was left + behind for good. + + And it is not confined to the closed wait. A stale intent keeps + `parked_wait_ids` non-empty, which suppresses the unparked-wake fallback + for the **whole stream**: with no live wait parked, the producer sends + only the dead generation, Core discards it as stale, and dedup silences + every later publish -- so live waits across the Continue-As-New chain + lose their wakes too. + """ + subscription = self._runs.get(run_id, {}).pop(wait_id, None) + if subscription is None: + return + for attempt in range(PARK_REMOVAL_ATTEMPTS): + async with self._park_lock(run_id): + try: + # Bounded for the same reason the reconciliation's hold is: + # this runs on a task of its own, and a call that hangs + # rather than raises would hold this Run's park lock for as + # long as it lasts. A timeout arrives here as one more + # failed attempt, which is what it is. + await asyncio.wait_for( + self._remove_park_intent(subscription), + PARK_REMOVAL_CALL_TIMEOUT.total_seconds(), + ) + break + except Exception: + logger.warning( + "Removing the park intent for %s wait %s while closing " + "it failed on attempt %s/%s", + subscription.stream_key, + subscription.wait_id, + attempt + 1, + PARK_REMOVAL_ATTEMPTS, + exc_info=True, + ) + if attempt + 1 < PARK_REMOVAL_ATTEMPTS: + await asyncio.sleep(PARK_REMOVAL_RETRY_DELAY.total_seconds()) + else: + logger.warning( + "Could not remove the park intent for %s wait %s while closing " + "it; it stays in this Run's owed-removal ledger, and until a " + "drain retires it producers read a generation Core has " + "discarded and their wakes are ignored as stale", + subscription.stream_key, + subscription.wait_id, + ) + await self._stop(subscription) + + async def evict_run(self, run_id: str) -> None: + """Tears down every subscription for a Run. + + The same path serves eviction, Workflow Task failure mid-batch, and + shutdown, because all three leave exactly the same thing behind: + speculative reads that were never committed. + + Owed removals belong to the manager rather than its cache. Eviction + makes one eager attempt but leaves any failure in the autonomous retry + loop, so backend recovery is sufficient even if this Run never returns + to the Worker. Installed intents are deliberately untouched -- an + eviction is not the end of a park, and the intents of a park that really + is outstanding must survive it. What the removal *silenced* outlives this + too, and is handed off by :meth:`_wake_after_cleanup`: the subscriptions + dropped below are exactly what the Subscription-bound announcement needs. + + The eager attempt is skipped while shutting down, and that is not the + same as skipping cleanup. A retry loop stuck in a backend call holds this + Run's park lock, and a drain here would wait on it with no bound at all, + wedging the shutdown it is part of. The last attempt is made by + :meth:`_final_owed_removal_pass` instead, which runs after those loops + have been cancelled and awaited and is bounded by what is left of the + grace period. + """ + self._replay_plans.pop(run_id, None) + if self._owed_removals.get(run_id) and not self._shutting_down: + async with self._park_lock(run_id): + await self._drain_owed_removals(run_id) + for subscription in self._runs.pop(run_id, {}).values(): + await self._stop(subscription) + retry = self._owed_removal_retries.get(run_id) + if not self._owed_removals.get(run_id) and (retry is None or retry.done()): + self._park_locks.pop(run_id, None) + + async def probe_runs(self, *, grace: timedelta = DEFAULT_PROBE_GRACE) -> None: + """Records what state Core says each Run is in. The sweep's first half. + + Separated from the wakes because the two halves have opposite timing + requirements, and running them together means one of them is wrong: + + - the **probe** has to be asked while Core can still answer it. Its whole + value is telling ``WftOpen``, ``Parked`` and ``NoOpenWorkflowTask`` + apart, and all three are statements about a Run Core still holds. Core + keeps them only until the Worker's shutdown is initiated: an idle + cached Run has no pending work, so ``shutdown_done`` is satisfied by + the very first input after the shutdown token is cancelled and the + whole workflow-state lane ends. Every probe after that answers + ``RunNotFound`` -- which still owes a wake, so the sweep goes on + looking correct while its ``Parked`` branch (a Run that needs no wake + at all) and its ``WftOpen`` branch (a Run that must be left to C15b + rather than raced) have silently stopped being reachable. + - the **wakes** must not run there. Sending them before the pollers stop + would offer the Run to a task queue this Worker is still polling, which + is the opposite of a hand-off, and would hold the stop-polling step + open for as long as the server takes to acknowledge them. + + So this runs immediately *before* shutdown is initiated and does nothing + but ask and remember, and :meth:`sweep` acts on the answers afterwards. + The one thing that costs: a Run recorded as ``NoOpenWorkflowTask`` here + may still pick up a Workflow Task from an in-flight poll before the + pollers stop, and would then get both C15b's forced replacement and this + sweep's wake. That is one extra empty Workflow Task, which this design + permits, and it is the cheaper side of the trade -- the alternative is + never distinguishing the states at all. + + Bounded by ``grace`` because nothing may make stopping the pollers wait + indefinitely. Every Run left unprobed is simply probed again by the + sweep, where the answer is whatever Core has left to say. + """ + probe = self._run_status + if probe is None: + return + try: + await asyncio.wait_for(self._probe_runs(probe), grace.total_seconds()) + except asyncio.TimeoutError: + logger.warning( + "External stream shutdown probe did not finish within %s; the " + "Runs it did not reach are swept on whatever Core can still say", + grace, + ) + + async def _probe_runs(self, probe: RunStatusProbe) -> None: + for run_id in list(self._runs): + if not self._runs.get(run_id): + continue + try: + self._probed[run_id] = _status_value(await probe(run_id)) + except Exception: + logger.exception( + "External stream shutdown probe failed for run %s", run_id + ) + + async def sweep(self, *, grace: timedelta = DEFAULT_SHUTDOWN_GRACE) -> None: + """Sends the wakes the probe found owed. The sweep's second half. + + Separate from :meth:`shutdown` so teardown can stay where it belongs -- + after every activation has been answered, so per-Run teardown remains + driven by ``RemoveFromCache`` and a ``FinalizeExternalStreams`` in flight + is answered before the Run's state disappears. + + Never blocked past ``grace``. A wake that could not be acknowledged in + time is reported rather than dropped, and never counted as delivered. + """ + if self._swept: + return + self._swept = True + self._shutting_down = True + # Every subscription starts out unaccounted for, and each one leaves this + # set exactly once: when its Run's status says nothing is owed, or when + # its wake is acknowledged. Whatever is still in it when the sweep stops + # is a handoff this Worker did not make, and the whole point of the + # counter is that such a handoff cannot be silent. + self._unaccounted = [ + subscription + for run_id in list(self._runs) + for subscription in self._runs.get(run_id, {}).values() + ] + try: + await asyncio.wait_for(self._sweep(), grace.total_seconds()) + except asyncio.TimeoutError: + logger.warning( + "External stream shutdown sweep did not finish within %s; " + "tearing down anyway", + grace, + ) + finally: + # In a `finally`, and reached on the timeout path especially. + # `wait_for` cancels `_sweep()`, and cancellation lands wherever the + # sweep happened to be: inside a hanging Signal send, whose + # `CancelledError` `_send_owed_wake` deliberately re-raises, and + # before every subscription the serial loop had not reached yet. + # Accounting done only where the sweep managed to reach therefore + # reported a clean shutdown -- `shutdown_wake_failures == 0` -- for a + # Worker that had just abandoned every one of its handoffs. + self._account_unswept() + + async def shutdown(self, *, grace: timedelta = DEFAULT_SHUTDOWN_GRACE) -> None: + """Sweeps every Run that still holds subscriptions, then tears down. + + Per-Run teardown is normally eviction's job, and this is the backstop for + Runs eviction never reaches -- an idle cached Run gets no eviction + activation at shutdown at all, which is exactly the Run that most needs + the sweep: its records are buffered in a process that is about to exit, + and nothing else will ever tell the Workflow they arrived. + """ + deadline = self._loop.time() + max(0.0, grace.total_seconds()) + await self.sweep(grace=grace) + for run_id in list(self._runs): + await self.evict_run(run_id) + remaining = timedelta(seconds=max(0.0, deadline - self._loop.time())) + await self._stop_background_tasks(remaining) + await self._final_owed_removal_pass( + timedelta(seconds=max(0.0, deadline - self._loop.time())) + ) + # Last, because the pass above is what creates the final cleanup + # handoffs: a removal confirmed there ends a suppression window nothing + # else will ever announce, and its wake is posted as a task rather than + # awaited under the park lock. Leaving without waiting for those tasks is + # indistinguishable from never having sent them. + await self._flush_cleanup_wakes( + timedelta(seconds=max(0.0, deadline - self._loop.time())) + ) + + async def _final_owed_removal_pass(self, budget: timedelta) -> None: + """The last attempt at every removal still owed, and a count of what is left. + + Reached only after :meth:`_stop_background_tasks`, which is both what + makes this the last attempt and what makes it safe: the retry loops that + own these entries are cancelled and awaited by then, so nothing holds a + park lock behind this pass and nothing restarts a loop after it. + + Eviction is where this used to happen, and while shutting down it no + longer does -- a retry loop blocked in a backend call still held the lock + the eviction would have waited on. What was lost with it is the case this + exists for: an entry the loop failed once and is now sleeping out its + backoff, on a backend that has since recovered. Cancelling that sleep and + exiting would leave an intent installed that one call would have taken + out. + + Bounded by whatever is left of the grace the sweep was given, because + cleanup may not extend shutdown -- and silent about nothing, since an + intent still installed as the process exits suppresses the unparked wake + for whoever picks the Run up next. + """ + owed = list(self._owed_removals) + if not owed: + return + deadline = self._loop.time() + max(0.0, budget.total_seconds()) + for run_id in owed: + left = deadline - self._loop.time() + if left <= 0: + break + await self._bounded_drain( + run_id, min(PARK_REMOVAL_CALL_TIMEOUT, timedelta(seconds=left)) + ) + still = sum(len(entries) for entries in self._owed_removals.values()) + if still: + logger.warning( + "%s external stream park intent(s) are still installed as this " + "Worker exits, with no park behind them; the next Worker's " + "registration-time reconciliation is the only thing left that " + "removes them", + still, + ) + + async def _stop_background_tasks(self, grace: timedelta) -> None: + """Cancels manager-owned cleanup work without extending shutdown.""" + self._stopping_owed_removal_retries = True + tasks = set(self._reconciliations) + tasks.update(self._owed_removal_retries.values()) + if not tasks: + return + for task in tasks: + task.cancel() + # Give ordinary cancellation one turn even when the wake sweep consumed + # the whole grace period. Anything needing longer stays bounded by what + # remains rather than silently adding a second grace period to shutdown. + await asyncio.sleep(0) + pending = {task for task in tasks if not task.done()} + if pending and grace.total_seconds() > 0: + _, pending = await asyncio.wait(pending, timeout=grace.total_seconds()) + if pending: + logger.warning( + "%s external stream cleanup task(s) did not stop within %s", + len(pending), + grace, + ) + + async def _sweep(self) -> None: + """Acts on each Run's state and owes a wake only where one is owed. + + The probe is deliberately **not** the readiness call. Readiness means "a + record is buffered", so probing with it would assert something false and + manufacture a spurious Workflow Task on the way out of a Worker that is + shutting down. + + Nothing is counted as a failure in here beyond the wakes this actually + attempted. What was never reached is accounted for by + :meth:`_account_unswept`, which runs whether this returns or is cancelled + by the grace period. + """ + if self._run_status is None: + # No probe wired means there is no sweep on this manager at all -- + # the whole mechanism is defined in terms of what Core answers -- so + # there is no obligation to have failed to discharge. Resolved rather + # than counted: a metric that fires for a mechanism that was never + # configured tells an operator nothing about the deployment that has + # it. + self._resolve_unaccounted(self._unaccounted) + return + for run_id in list(self._runs): + subscriptions = list(self._runs.get(run_id, {}).values()) + if not subscriptions: + continue + status = self._probed.pop(run_id, None) + if status is None: + # Not probed while Core could still answer -- either nothing ran + # the probe phase, or this Run was cached from an in-flight poll + # after it. Ask anyway rather than guessing: the answer is + # whatever Core has left to say, and both answers it can still + # give owe a wake. + try: + status = _status_value(await self._run_status(run_id)) + except asyncio.CancelledError: + raise + except Exception: + # Left unaccounted for on purpose. A Run whose status cannot + # be read is a Run this Worker cannot say anything about, and + # "we could not tell" is not "nothing was owed": these + # subscriptions may each be holding a buffered record with + # nowhere to announce it. Sending a wake anyway would race a + # Workflow Task that may be open, so the honest outcome is to + # send nothing and let the counter say a handoff was lost. + logger.exception( + "External stream shutdown probe failed for run %s", run_id + ) + continue + + if status == RunStatus.WFT_OPEN: + # A Workflow Task is open and Core owns what happens to it + # (C15b). A wake here would race that transition and produce a + # second task for a Run already being attended to. + self._resolve_unaccounted(subscriptions) + continue + if status == RunStatus.PARKED: + # Already parked, so a producer's append will wake it through the + # ordinary path. Nothing is owed. + self._resolve_unaccounted(subscriptions) + continue + + # NoOpenWorkflowTask and RunNotFound both mean local readiness has + # nowhere to go. They differ in what happens to the Run afterwards, + # not in what is owed now. + for subscription in subscriptions: + await self._sweep_wake(subscription) + + async def _sweep_wake(self, subscription: Subscription) -> None: + """Sends one unparked wake, retrying within the grace period. + + Awaited rather than fired: an unacknowledged wake is the case this whole + sweep exists to prevent, and reporting shutdown as clean while a record + sits unannounced in a stream would make the wakeup-durability boundary + false. + + The grace period bounds the whole sweep, so the retries cannot extend + shutdown past it, and a wake that never lands is reported rather than + dropped -- which is the only thing that makes giving up acceptable. + + Counting the failure is deliberately **not** done here for the + cancellation case. The grace period expiring cancels this coroutine + wherever it is, `_send_owed_wake` re-raises `CancelledError` by design, + and no `except` here could both record the failure and leave the + cancellation intact for the subscriptions after this one -- which are not + reached either. The subscription stays in the unaccounted set instead and + :meth:`_account_unswept` counts it, which covers being cancelled and + being never visited with the same rule. + """ + self._count_owed_wake(subscription) + if await self._send_owed_wake(subscription): + self._note_swept_handoff(subscription) + self._resolve_unaccounted([subscription]) + else: + self._record_shutdown_wake_failure(subscription) + self._resolve_unaccounted([subscription]) + + def _note_swept_handoff(self, subscription: Subscription) -> None: + """Records that the sweep's wake also discharged an owed removal's handoff. + + Only where there is one. A stale intent on this wait means the wake just + sent was composed as unparked -- `wake_park_generation` saw the ledger + entry -- so it is the very generation-0 handoff the removal would owe + once it finally goes through, and `_final_owed_removal_pass` runs after + this. Without the note the two are one obligation sent twice, which + costs an empty Workflow Task; marked only on success, so a wake the + sweep could not get acknowledged leaves the handoff still owed. + """ + owed = self._owed_removals.get(subscription.run_id) + if not owed: + return + key = (subscription.stream_key, subscription.wait_id) + record = owed.get(key) + if record is not None and not record.announced: + owed[key] = replace(record, announced=True) + + def _resolve_unaccounted(self, subscriptions: Sequence[Subscription]) -> None: + """Marks these subscriptions as decided, however they were decided. + + Both outcomes are decisions: nothing was owed, or a wake was attempted + and its result recorded. What stays in the set is only what the sweep + never got to say anything about. + """ + decided = set(map(id, subscriptions)) + self._unaccounted = [ + subscription + for subscription in self._unaccounted + if id(subscription) not in decided + ] + + def _account_unswept(self) -> None: + """Counts every subscription the sweep never resolved as a lost handoff. + + The grace period is a bound on how long shutdown may be held open, not a + licence to stop counting: a Worker that abandons a wake and reports zero + failures makes the failure invisible in exactly the way this counter + exists to prevent. Both silent cases end up here -- the wake the + cancellation landed inside, and every subscription the serial loop never + reached. + """ + unaccounted, self._unaccounted = self._unaccounted, [] + for subscription in unaccounted: + if ( + self._runs.get(subscription.run_id, {}).get(subscription.wait_id) + is not subscription + ): + # Gone from the Run while the sweep ran, which the *live* readiness + # path does -- it drops a subscription on `RunNotFound`, and only + # after its own owed wake was acknowledged. `RunNotFound` is a + # likely answer during shutdown and the watchers keep running + # through the whole grace window, so this is an ordinary + # interleaving, not a corner. Counting it would report a lost + # handoff for one that was made, on the counter operators are told + # to alert on. + continue + logger.warning( + "External stream shutdown left a wake unresolved for %s wait %s", + subscription.stream_key, + subscription.wait_id, + ) + self._record_shutdown_wake_failure(subscription) + + async def wake_park_generation(self, target: WakeTarget) -> int | None: + """Which park one owed wake must name, or ``None`` for the unparked one. + + The Worker's sender asks this rather than reading + ``current_park_generation`` itself, because "installed" is not the same + as "live" and only the manager can tell the two apart. Two things make + the raw read wrong: + + - a `_CleanupWake` exists *because* the intent it names was just + retired. There is nothing left at that key that Core would accept, and + a generation another Worker has since installed belongs to a park this + handoff knows nothing about. + - a `Subscription` whose intent is in this Run's owed-removal ledger is + holding a generation this manager has **already decided to remove**. + Core discards a non-zero generation that is not the park it holds, so + naming it composes a wake that is ignored while the send reports + success -- which is exactly how the shutdown sweep came to count an + obsolete-generation Signal as a handoff it had made. + + Still a read and not a remembered value in the ordinary case: the park a + wake must name is whatever is installed *now*, and a generation cached + when the watcher started would name a park since abandoned and resolved. + """ + if isinstance(target, _CleanupWake): + return None + key = (target.stream_key, target.wait_id) + if key in self._owed_removals.get(target.run_id, {}): + return None + return await target.backend.current_park_generation( + target.stream_key, target.wait_id + ) + + def _count_owed_wake(self, target: WakeTarget) -> None: + """Counts one owed wake and draws the sequence number it is sent under. + + Both halves happen here, exactly once per wake, because + :meth:`_send_owed_wake` retries *the same* wake: re-drawing the counter + between attempts would derive a second request ID and ask the server for + a second Workflow Task rather than re-sending the one that may already + have arrived. + """ + target.wakes_owed += 1 + self._wake_sequence += 1 + target.wake_counter = self._wake_sequence + + async def _send_owed_wake(self, subscription: WakeTarget) -> bool: + """Sends the one wake ``wakes_owed`` already counts. Returns whether it landed. + + **The count belongs to the caller, and it is made exactly once, before + this is entered.** The wake's request ID is derived from it, so a retry + that re-counted would derive a *different* ID: the server would dedupe + nothing and answer with a second, empty Workflow Task, which is a + different thing from re-sending the wake that may in fact have arrived. + Every attempt below is therefore the same wake, which is precisely what + makes re-sending it safe. + + Retried because the common failure is a momentary one, and because both + callers have the same problem if it is not: on the live path nothing + re-attempts an owed wake at all, and on the shutdown path the process is + about to exit. + + Nothing but `CancelledError` leaves here. The live caller is reached + from the watcher loop, which an escaping exception would end for good -- + taking every later record on that subscription with it. + """ + if self._send_wake is None: + return False + for attempt in range(SHUTDOWN_WAKE_ATTEMPTS): + self._note_wake_attempt_started(subscription) + try: + await self._send_wake(subscription) + self._note_wake_attempt_acknowledged(subscription) + return True + except asyncio.CancelledError: + raise + except Exception: + # The wake sender reports failure by raising, because a wake + # counted as delivered when it was not is the failure this whole + # path exists to prevent. + logger.warning( + "External stream wake attempt %s/%s failed for %s wait %s", + attempt + 1, + SHUTDOWN_WAKE_ATTEMPTS, + subscription.stream_key, + subscription.wait_id, + exc_info=True, + ) + if attempt + 1 < SHUTDOWN_WAKE_ATTEMPTS: + await asyncio.sleep(SHUTDOWN_WAKE_RETRY_DELAY.total_seconds()) + return False + + def _record_shutdown_wake_failure(self, subscription: WakeTarget) -> None: + """Surfaces the wake that could not be acknowledged. + + Counted rather than only logged: a dropped wake is silent by nature -- + the Workflow simply waits, and nothing distinguishes that from a producer + having nothing to say. + """ + self.shutdown_wake_failures += 1 + if self._metric is not None: + self._metric(subscription) + + async def _stop(self, subscription: Subscription) -> None: + subscription._cancelled = True + subscription.reset_to_committed() + watcher = subscription._watcher + if watcher is not None and not watcher.done(): + watcher.cancel() + try: + await watcher + except asyncio.CancelledError: + pass diff --git a/temporalio/contrib/external_workflow_streams/_output_api.py b/temporalio/contrib/external_workflow_streams/_output_api.py new file mode 100644 index 0000000000..780291f61a --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_output_api.py @@ -0,0 +1,147 @@ +"""Workflow-facing publishers for external output streams. + +The publisher is deliberately a different type from the input subscription +handle. Workflow code only converts values into deterministic logical +payloads; the Worker applies codecs and performs backend I/O after the +activation has returned. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from datetime import timedelta +from typing import Any, Generic, Protocol, cast + +import temporalio.workflow +from temporalio.contrib.external_workflow_streams._api import _run_state +from temporalio.contrib.external_workflow_streams._record import RecordKind +from temporalio.types import AnyType + +__all__ = [ + "ExternalOutputStreamOptions", + "ExternalOutputStreamTopic", + "external_output_stream", +] + +DEFAULT_MAX_PUBLISH_LATENCY = timedelta(milliseconds=100) +DEFAULT_MAX_RECORDS_PER_BATCH = 256 +DEFAULT_MAX_LOGICAL_BYTES_PER_BATCH = 1024 * 1024 + + +class _OutputRuntime(Protocol): + """The deterministic half of the per-Run stream runtime.""" + + async def publish_output( + self, + *, + topic: str, + value: Any, + value_type: type | None, + kind: RecordKind, + max_publish_latency: timedelta, + max_records: int, + max_logical_bytes: int, + ) -> None: ... + + +@dataclass(frozen=True) +class ExternalOutputStreamOptions: + """Options inherited by every output topic created from this entry point.""" + + max_publish_latency: timedelta = DEFAULT_MAX_PUBLISH_LATENCY + max_records: int = DEFAULT_MAX_RECORDS_PER_BATCH + max_logical_bytes: int = DEFAULT_MAX_LOGICAL_BYTES_PER_BATCH + + def __post_init__(self) -> None: + """Validate direct construction as well as :meth:`with_options`.""" + if self.max_publish_latency <= timedelta(0): + raise ValueError("max_publish_latency must be positive") + if self.max_records < 1: + raise ValueError("max_records must be positive") + if self.max_logical_bytes < 1: + raise ValueError("max_logical_bytes must be positive") + + def with_options( + self, + *, + max_publish_latency: timedelta | None = None, + max_records: int | None = None, + max_logical_bytes: int | None = None, + ) -> ExternalOutputStreamOptions: + """Return a copy with validated latency and logical batch limits.""" + latency = ( + self.max_publish_latency + if max_publish_latency is None + else max_publish_latency + ) + records = self.max_records if max_records is None else max_records + logical_bytes = ( + self.max_logical_bytes if max_logical_bytes is None else max_logical_bytes + ) + if latency <= timedelta(0): + raise ValueError("max_publish_latency must be positive") + if records < 1: + raise ValueError("max_records must be positive") + if logical_bytes < 1: + raise ValueError("max_logical_bytes must be positive") + return replace( + self, + max_publish_latency=latency, + max_records=records, + max_logical_bytes=logical_bytes, + ) + + def topic( + self, name: str, *, type: type[AnyType] | None = None + ) -> ExternalOutputStreamTopic[AnyType]: + """Create a typed Workflow publisher for one output topic.""" + if not name: + raise ValueError("an output topic needs a non-empty name") + return ExternalOutputStreamTopic( + name=name, + value_type=type, + options=self, + ) + + +@dataclass(frozen=True) +class ExternalOutputStreamTopic(Generic[AnyType]): + """A Workflow-side topic with publishing operations and no subscription API.""" + + name: str + value_type: type[AnyType] | None + options: ExternalOutputStreamOptions + + async def publish(self, value: AnyType) -> None: + """Accept one value into the current Workflow Task's output batch.""" + await self._publish(value, RecordKind.DATA) + + async def finish(self) -> None: + """Append the ordered terminal record for this topic.""" + await self._publish(None, RecordKind.FINISH) + + async def _publish(self, value: Any, kind: RecordKind) -> None: + if temporalio.workflow.unsafe.is_read_only(): + raise temporalio.workflow.ReadOnlyContextError( + "While in read-only function, action attempted: publish external output" + ) + state = _run_state() + runtime = state.runtime + if runtime is None or not hasattr(runtime, "publish_output"): + raise RuntimeError( + "external output streams are not configured on this Worker; pass " + "external_stream_backend=... to the Worker" + ) + await cast(_OutputRuntime, cast(object, runtime)).publish_output( + topic=self.name, + value=value, + value_type=self.value_type, + kind=kind, + max_publish_latency=self.options.max_publish_latency, + max_records=self.options.max_records, + max_logical_bytes=self.options.max_logical_bytes, + ) + + +external_output_stream = ExternalOutputStreamOptions() +"""The default Workflow-facing external output stream entry point.""" diff --git a/temporalio/contrib/external_workflow_streams/_output_backend.py b/temporalio/contrib/external_workflow_streams/_output_backend.py new file mode 100644 index 0000000000..ec4b1f58ea --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_output_backend.py @@ -0,0 +1,324 @@ +"""Provider contract for Workflow-originated external output streams. + +This contract is separate from :class:`StreamBackend`. Input-only providers +remain valid: output staging adds a pending/committed state machine and an +ordering barrier that an ordinary append-only input stream does not promise. + +The unit of staging is one topic sub-batch from one Workflow Task attempt. A +Worker-minted ``stage_token`` identifies the attempt, while ``sub_batch_id`` +distinguishes topic batches covered by that token. Records are invisible while +pending. A reader stops at the first unresolved batch even when later committed +Activity records already have provider offsets. +""" + +from __future__ import annotations + +import abc +import enum +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import timedelta +from typing import ClassVar + +from temporalio.contrib.external_workflow_streams._backend import ( + DEFAULT_WATCH_BLOCK, + StreamDirection, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._record import ( + Cursor, + Offset, + RecordKind, + StreamRecord, +) + +__all__ = [ + "OutputReadResult", + "OutputStage", + "OutputStageConflictError", + "OutputStageManifest", + "OutputStageNotFoundError", + "OutputStageResolutionError", + "OutputStageStatus", + "OutputStreamBackend", + "OutputStreamRecord", + "PendingOutputBarrier", + "StagedOutputRecord", +] + + +@enum.unique +class OutputStageStatus(enum.Enum): + """The durable visibility state of a staged output sub-batch.""" + + PENDING = "pending" + COMMITTED = "committed" + ABORTED = "aborted" + + +@dataclass(frozen=True) +class OutputStageManifest: + """Immutable logical identity and reconciliation proof for one sub-batch. + + Encoded record bytes are deliberately absent. A payload codec may use + randomness, so an idempotent retry compares this pre-codec manifest and + preserves the first successfully staged encoded bytes. + """ + + stream_key: StreamKey + provider_id: str + provider_format_version: int + stage_token: str + run_id: str + history_floor_event_id: int + sub_batch_id: int + fingerprint_version: int + fingerprint: bytes + record_count: int + logical_byte_count: int + + def __post_init__(self) -> None: + """Reject identities that could not safely be reconciled.""" + if self.stream_key.direction is not StreamDirection.OUTPUT: + raise ValueError("an output stage manifest requires an OUTPUT stream key") + for field_name in ("provider_id", "stage_token", "run_id"): + if not getattr(self, field_name): + raise ValueError( + f"an output stage manifest needs a non-empty {field_name}" + ) + for field_name in ( + "history_floor_event_id", + "provider_format_version", + "fingerprint_version", + "record_count", + ): + if getattr(self, field_name) < 1: + raise ValueError( + f"an output stage manifest needs a positive {field_name}" + ) + for field_name in ( + "sub_batch_id", + "logical_byte_count", + ): + if getattr(self, field_name) < 0: + raise ValueError( + f"an output stage manifest needs a non-negative {field_name}" + ) + if len(self.fingerprint) != 32: + raise ValueError("an output stage SHA-256 fingerprint must be 32 bytes") + + +@dataclass(frozen=True) +class StagedOutputRecord: + """One encoded record offered as part of a pending Workflow sub-batch. + + ``publish_index`` is local to the topic sub-batch. The backend assigns the + provider offset, and a retry of the same manifest must return the originally + stored bytes rather than replacing them with this attempt's encoding. + """ + + publish_index: int + kind: RecordKind + payload: bytes + + def __post_init__(self) -> None: + """Validate the deterministic identity and control-record shape.""" + if self.publish_index < 0: + raise ValueError("an output publish index must be non-negative") + if self.kind.is_control and self.payload: + raise ValueError(f"a {self.kind.name} output record carries no payload") + + +@dataclass(frozen=True) +class OutputStreamRecord: + """One provider-placed output record returned to a client.""" + + kind: RecordKind + payload: bytes + offset: Offset + + def __post_init__(self) -> None: + """Keep terminal and fence records payload-free after storage.""" + if self.kind.is_control and self.payload: + raise ValueError(f"a {self.kind.name} output record carries no payload") + + +@dataclass(frozen=True) +class OutputStage: + """The backend's durable view of one Workflow output sub-batch.""" + + manifest: OutputStageManifest + records: tuple[OutputStreamRecord, ...] + status: OutputStageStatus + + def __post_init__(self) -> None: + """Require the placed records to agree with the sealed manifest.""" + if len(self.records) != self.manifest.record_count: + raise ValueError( + "an output stage record count does not match its manifest: " + f"{len(self.records)} != {self.manifest.record_count}" + ) + + +@dataclass(frozen=True) +class PendingOutputBarrier: + """The unresolved stage preventing a client from reading later offsets.""" + + manifest: OutputStageManifest + offset: Offset + """The first physical position occupied by this pending sub-batch.""" + + +@dataclass(frozen=True) +class OutputReadResult: + """A committed prefix and the unresolved barrier immediately after it.""" + + records: tuple[OutputStreamRecord, ...] + pending: PendingOutputBarrier | None = None + + +class OutputStageConflictError(Exception): + """A stage token/sub-batch identity was reused with another manifest.""" + + def __init__(self, manifest: OutputStageManifest) -> None: + """Capture the immutable stage identity that conflicted.""" + super().__init__( + f"output stage {manifest.stage_token}/{manifest.sub_batch_id} was " + "already sealed with a different logical manifest" + ) + self.manifest = manifest + + +class OutputStageNotFoundError(Exception): + """A commit or abort named a stage the provider does not hold.""" + + def __init__(self, manifest: OutputStageManifest) -> None: + """Capture the immutable stage identity that was absent.""" + super().__init__( + f"output stage {manifest.stage_token}/{manifest.sub_batch_id} was not found" + ) + self.manifest = manifest + + +class OutputStageResolutionError(Exception): + """A terminal stage transition tried to reverse an existing decision.""" + + def __init__( + self, + manifest: OutputStageManifest, + *, + current: OutputStageStatus, + requested: OutputStageStatus, + ) -> None: + """Capture the attempted reversal of a terminal stage decision.""" + super().__init__( + f"output stage {manifest.stage_token}/{manifest.sub_batch_id} is " + f"already {current.value} and cannot become {requested.value}" + ) + self.manifest = manifest + self.current = current + self.requested = requested + + +class OutputStreamBackend(abc.ABC): + """A provider supporting externally visible Workflow output. + + Providers may implement this contract, :class:`StreamBackend`, or both. + Implementing the input contract alone does not imply transactional output + staging support. + """ + + guarantees_immutability: ClassVar[bool | None] = None + provider_id: ClassVar[str] = "" + provider_format_version: ClassVar[int] = 1 + + @abc.abstractmethod + async def stage_output( + self, + manifest: OutputStageManifest, + records: Sequence[StagedOutputRecord], + ) -> OutputStage: + """Durably append and seal an immutable pending sub-batch. + + The records become a barrier at their assigned positions but are not + readable. Repeating the call with the same complete manifest is an + idempotent success returning the original records and offsets, even if + a randomized codec supplied different encoded bytes on the retry. + Reusing ``(stage_token, sub_batch_id)`` with any different manifest + raises :class:`OutputStageConflictError`. The provider validates that + record count and publish indexes match the manifest before writing any + part of the batch. + """ + + @abc.abstractmethod + async def commit_output(self, manifest: OutputStageManifest) -> OutputStage: + """Make this exact pending stage visible, idempotently. + + Repeating a committed transition returns the committed stage. An + unknown stage raises :class:`OutputStageNotFoundError`; an aborted stage + raises :class:`OutputStageResolutionError` and is never resurrected. + """ + + @abc.abstractmethod + async def abort_output(self, manifest: OutputStageManifest) -> OutputStage: + """Resolve this exact pending stage as skipped, idempotently. + + Aborted positions remain coordination boundaries but yield no records. + An unknown stage raises :class:`OutputStageNotFoundError`; a committed + stage raises :class:`OutputStageResolutionError` and is never hidden. + """ + + @abc.abstractmethod + async def output_stage(self, manifest: OutputStageManifest) -> OutputStage | None: + """Inspect the exact stage, including its reconciliation metadata. + + Returns ``None`` only when the identity has never been staged. A stage's + manifest, encoded records, offsets, and terminal status are immutable. + """ + + @abc.abstractmethod + async def append_output( + self, key: StreamKey, record: StreamRecord + ) -> OutputStreamRecord: + """Append an immediately committed Activity/external producer record. + + ``key`` must have direction ``OUTPUT``. The append is idempotent on the + input record's ``(session_id, sequence)`` and byte identity exactly as + :meth:`StreamBackend.append`: an identical retry returns the original + output record and offset; different bytes under the same key raise + :class:`~temporalio.contrib.external_workflow_streams._backend.AppendConflictError`. + It may not pass an already-positioned pending stage in client reads. + """ + + @abc.abstractmethod + async def read_output_after( + self, + key: StreamKey, + after: Cursor, + *, + max_records: int, + block: timedelta | None = DEFAULT_WATCH_BLOCK, + ) -> OutputReadResult: + """Read the committed prefix strictly after ``after``. + + Records are returned in provider order and limited to ``max_records``. + Reading stops before the first pending stage. If that barrier is the + next unresolved position, ``pending`` names it; no record after it may + be returned even when already committed. Aborted stages yield nothing + and do not remain barriers. An empty result with no pending barrier is + ordinary watch idleness. + """ + + @abc.abstractmethod + async def output_tail(self, key: StreamKey) -> Cursor: + """Return the end boundary of the currently readable committed prefix. + + The boundary never advances past a pending stage, because a caller that + resumes after the returned cursor must still see that stage if it later + commits. It may advance across aborted coordination positions. This is + a position boundary only and does not correlate output to an Update. + """ + + @abc.abstractmethod + def compare_offsets(self, left: Offset, right: Offset) -> int: + """Three-way comparison under this provider's output offset order.""" diff --git a/temporalio/contrib/external_workflow_streams/_output_client.py b/temporalio/contrib/external_workflow_streams/_output_client.py new file mode 100644 index 0000000000..e9d446606b --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_output_client.py @@ -0,0 +1,517 @@ +"""External clients for Workflow-originated output streams. + +Committed output is read directly from the configured provider. Temporal is +consulted only when the provider reports a pending stage at the head of the +topic: History is then the authority that proves whether that exact stage +committed or was discarded with its Workflow Task attempt. +""" + +from __future__ import annotations + +import asyncio +import hmac +from collections.abc import AsyncIterator +from dataclasses import dataclass +from enum import Enum +from typing import Any, Final, Generic, cast + +import temporalio.client +import temporalio.converter +import temporalio.service +from temporalio.contrib.external_workflow_streams._backend import ( + StreamDirection, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec +from temporalio.contrib.external_workflow_streams._errors import ( + StreamError, + StreamIntegrityError, + StreamStorageError, + classify_read_failure, +) +from temporalio.contrib.external_workflow_streams._output_backend import ( + OutputReadResult, + OutputStage, + OutputStageConflictError, + OutputStageManifest, + OutputStageNotFoundError, + OutputStageResolutionError, + OutputStageStatus, + OutputStreamBackend, + PendingOutputBarrier, +) +from temporalio.contrib.external_workflow_streams._producer import ( + ChainKeyMismatchError, + WorkflowChainKey, + _verify_chain_key, +) +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Cursor, + Offset, + RecordKind, +) +from temporalio.types import AnyType + +__all__ = [ + "ExternalOutputStreamClient", + "ExternalOutputStreamClientTopic", + "ExternalOutputStreamItem", +] + +_EXTERNAL_STREAM_MARKER_NAME: Final = "core_external_stream" +_EXTERNAL_STREAM_MARKER_DETAILS_KEY: Final = "external_stream" +_READ_BATCH_SIZE: Final = 100 +_PENDING_RECONCILIATION_BACKOFF_SECONDS: Final = 1.0 + + +@dataclass(frozen=True) +class ExternalOutputStreamItem(Generic[AnyType]): + """One decoded output value and its opaque provider resume offset.""" + + data: AnyType + offset: Offset + + +class ExternalOutputStreamClient: + """A verified client connection to one Workflow execution chain.""" + + def __init__( + self, + *, + backend: Any, + workflow: WorkflowChainKey, + client: temporalio.client.Client, + data_converter: temporalio.converter.DataConverter, + ) -> None: + """Create a reader after validating its output provider contract.""" + if not isinstance(backend, OutputStreamBackend): + raise TypeError("backend does not implement OutputStreamBackend") + if type(backend).guarantees_immutability is not True: + raise ValueError("an output backend must guarantee immutable record bytes") + if not type(backend).provider_id: + raise ValueError("an output backend must declare a provider_id") + self._backend = backend + self._workflow = workflow + self._client = client + self._data_converter = data_converter.with_context( + temporalio.converter.WorkflowSerializationContext( + namespace=workflow.namespace, + workflow_id=workflow.workflow_id, + ) + ) + + @classmethod + async def connect( + cls, + *, + backend: OutputStreamBackend, + workflow: WorkflowChainKey, + client: temporalio.client.Client, + data_converter: temporalio.converter.DataConverter | None = None, + ) -> ExternalOutputStreamClient: + """Verify the chain binding and create an external output reader. + + The Temporal client remains attached only for resolving a pending + staging barrier. Ordinary reads and :meth:`tail` are backend-only. + """ + try: + await _verify_chain_key(client, workflow) + except ChainKeyMismatchError: + raise + except StreamError: + raise + except Exception as err: + raise StreamStorageError( + f"could not verify external output Workflow chain: {err}" + ) from err + return cls( + backend=backend, + workflow=workflow, + client=client, + data_converter=data_converter or client.data_converter, + ) + + @property + def workflow(self) -> WorkflowChainKey: + """The Workflow chain whose output this client reads.""" + return self._workflow + + def topic( + self, name: str, *, type: type[AnyType] | None = None + ) -> ExternalOutputStreamClientTopic[AnyType]: + """Create a typed reader for one output topic.""" + if not name: + raise ValueError("an output topic needs a non-empty name") + return ExternalOutputStreamClientTopic( + backend=self._backend, + stream_key=self._workflow.stream_key( + name, direction=StreamDirection.OUTPUT + ), + codec=StreamPayloadCodec(self._data_converter, type), + client=self._client, + workflow=self._workflow, + ) + + +class ExternalOutputStreamClientTopic(Generic[AnyType]): + """A typed, resumable external reader for one output topic.""" + + def __init__( + self, + *, + backend: OutputStreamBackend, + stream_key: StreamKey, + codec: StreamPayloadCodec[AnyType], + client: temporalio.client.Client, + workflow: WorkflowChainKey, + ) -> None: + """Create a typed reader bound to one direction-isolated topic.""" + self._backend = backend + self._stream_key = stream_key + self._codec = codec + self._client = client + self._workflow = workflow + + async def tail(self) -> Cursor: + """Return the boundary after the readable committed prefix.""" + try: + return await self._backend.output_tail(self._stream_key) + except StreamError: + raise + except Exception as err: + raise StreamStorageError( + f"could not read external output tail for {self._stream_key}: {err}" + ) from err + + async def subscribe( + self, *, after: Cursor | Offset = BEGINNING + ) -> AsyncIterator[ExternalOutputStreamItem[AnyType]]: + """Yield committed values strictly after ``after``. + + A pending stage is a hard ordering barrier. The committed prefix before + it may be yielded, but the client never reads or yields a later offset + until History has positively committed or aborted the pending stage. + """ + candidate = cast(object, after) + if isinstance(candidate, Offset): + cursor = AFTER(candidate) + elif isinstance(candidate, Cursor): + cursor = candidate + else: + raise TypeError( + f"after must be a Cursor or Offset, got {type(candidate).__name__}" + ) + while True: + result = await self._read_after(cursor) + self._validate_read_result(cursor, result) + for record in result.records: + cursor = AFTER(record.offset) + if record.kind is RecordKind.FINISH: + return + if record.kind.is_control: + continue + try: + data = await self._codec.decode(record.payload) + except asyncio.CancelledError: + raise + except Exception as err: + raise classify_read_failure( + range_validated=True, cause=err + ) from err + yield ExternalOutputStreamItem(data=data, offset=record.offset) + + if result.pending is not None: + resolved = await self._resolve_pending(result.pending) + if not resolved: + await asyncio.sleep(_PENDING_RECONCILIATION_BACKOFF_SECONDS) + + async def _read_after(self, after: Cursor) -> OutputReadResult: + try: + return await self._backend.read_output_after( + self._stream_key, + after, + max_records=_READ_BATCH_SIZE, + ) + except StreamError: + raise + except Exception as err: + raise StreamStorageError( + f"could not read external output for {self._stream_key}: {err}" + ) from err + + def _validate_read_result(self, after: Cursor, result: OutputReadResult) -> None: + """Refuse provider output that crosses or misorders a pending barrier.""" + previous = after.offset + for record in result.records: + if ( + previous is not None + and self._backend.compare_offsets(previous, record.offset) >= 0 + ): + raise StreamIntegrityError( + "external output provider returned records outside the requested " + f"strictly-after order for {self._stream_key}" + ) + previous = record.offset + + pending = result.pending + if pending is None: + return + if pending.manifest.stream_key != self._stream_key: + raise StreamIntegrityError( + "external output provider returned a pending barrier for another " + f"topic while reading {self._stream_key}" + ) + if ( + pending.manifest.provider_id != type(self._backend).provider_id + or pending.manifest.provider_format_version + != type(self._backend).provider_format_version + ): + raise StreamIntegrityError( + "external output provider returned a pending barrier carrying a " + "different provider binding" + ) + if ( + previous is not None + and self._backend.compare_offsets(previous, pending.offset) >= 0 + ): + raise StreamIntegrityError( + "external output provider returned data at or beyond an unresolved " + f"pending barrier for {self._stream_key}" + ) + + async def _resolve_pending(self, pending: PendingOutputBarrier) -> bool: + """Resolve one pending head from the exact producing Run's History. + + Returns ``False`` only when History contains neither this token nor a + later durable task/Workflow boundary, so the stage must remain pending. + """ + decision = await self._history_decision(pending.manifest) + if decision is None: + return False + await _apply_output_stage_decision( + backend=self._backend, + manifest=pending.manifest, + decision=decision, + ) + return True + + async def _history_decision( + self, manifest: OutputStageManifest + ) -> _HistoryDecision | None: + return await _history_decision( + client=self._client, + workflow_id=self._workflow.workflow_id, + manifest=manifest, + ) + + @staticmethod + def _validate_resolved_stage( + manifest: OutputStageManifest, + stage: OutputStage, + *, + expected_status: OutputStageStatus, + ) -> None: + _validate_resolved_stage( + manifest, + stage, + expected_status=expected_status, + ) + + +class _HistoryDecision(Enum): + COMMIT = "commit" + ABORT = "abort" + + +async def _history_decision( + *, + client: temporalio.client.Client, + workflow_id: str, + manifest: OutputStageManifest, +) -> _HistoryDecision | None: + """Prove one staged token's outcome from its exact producing History.""" + try: + events = [ + event + async for event in client.get_workflow_handle( + workflow_id, + run_id=manifest.run_id, + ).fetch_history_events() + ] + except temporalio.service.RPCError as err: + if err.status in ( + temporalio.service.RPCStatusCode.NOT_FOUND, + temporalio.service.RPCStatusCode.DATA_LOSS, + ): + raise StreamIntegrityError( + "History required to resolve external output stage " + f"{manifest.stage_token!r} is no longer available: {err}" + ) from err + raise StreamStorageError( + "could not read Temporal History while resolving external output " + f"stage {manifest.stage_token!r}: {err}" + ) from err + except StreamError: + raise + except Exception as err: + raise StreamStorageError( + "could not read Temporal History while resolving external output " + f"stage {manifest.stage_token!r}: {err}" + ) from err + + floor = manifest.history_floor_event_id + if floor and not any(event.event_id == floor for event in events): + raise StreamIntegrityError( + "History no longer contains the exact floor event needed to resolve " + f"external output stage {manifest.stage_token!r}: {floor}" + ) + relevant = [event for event in events if event.event_id > floor] + + # A task's marker and closing event commit in one server transaction, but + # the closing event precedes command events in History. Search the complete + # response for the positive token proof before considering any boundary + # proof of absence. + for event in relevant: + if _event_has_output_stage_token(event, manifest.stage_token): + return _HistoryDecision.COMMIT + if any(_is_deciding_boundary(event) for event in relevant): + return _HistoryDecision.ABORT + return None + + +async def _apply_output_stage_decision( + *, + backend: OutputStreamBackend, + manifest: OutputStageManifest, + decision: _HistoryDecision, +) -> None: + """Apply a positively proven History decision idempotently.""" + try: + if decision is _HistoryDecision.COMMIT: + stage = await backend.commit_output(manifest) + expected_status = OutputStageStatus.COMMITTED + else: + stage = await backend.abort_output(manifest) + expected_status = OutputStageStatus.ABORTED + except ( + OutputStageConflictError, + OutputStageNotFoundError, + OutputStageResolutionError, + ) as err: + raise StreamIntegrityError( + "external output pending stage could not be reconciled with its " + f"History decision: {err}" + ) from err + except StreamError: + raise + except Exception as err: + raise StreamStorageError( + f"could not {decision.value} external output stage " + f"{manifest.stage_token}: {err}" + ) from err + _validate_resolved_stage( + manifest, + stage, + expected_status=expected_status, + ) + + +async def _reconcile_output_stage( # pyright: ignore[reportUnusedFunction] + *, + backend: OutputStreamBackend, + client: temporalio.client.Client, + workflow_id: str, + manifest: OutputStageManifest, +) -> bool: + """Resolve after reporting, leaving ambiguous outcomes safely pending.""" + decision = await _history_decision( + client=client, + workflow_id=workflow_id, + manifest=manifest, + ) + if decision is None: + return False + await _apply_output_stage_decision( + backend=backend, + manifest=manifest, + decision=decision, + ) + return True + + +def _validate_resolved_stage( + manifest: OutputStageManifest, + stage: OutputStage, + *, + expected_status: OutputStageStatus, +) -> None: + """Refuse a provider which resolved a different immutable stage.""" + if stage.manifest != manifest or stage.status is not expected_status: + raise StreamIntegrityError( + "external output provider returned a different stage or status " + "after applying a History reconciliation decision" + ) + + +def _is_deciding_boundary(event: Any) -> bool: + """Whether this event proves that the staged token can no longer appear.""" + return any( + event.HasField(field) + for field in ( + "workflow_task_completed_event_attributes", + "workflow_task_failed_event_attributes", + "workflow_task_timed_out_event_attributes", + "workflow_execution_completed_event_attributes", + "workflow_execution_failed_event_attributes", + "workflow_execution_timed_out_event_attributes", + "workflow_execution_canceled_event_attributes", + "workflow_execution_terminated_event_attributes", + "workflow_execution_continued_as_new_event_attributes", + ) + ) + + +def _event_has_output_stage_token(event: Any, stage_token: str) -> bool: + """Read output commit proofs from the shared external-stream marker.""" + if not event.HasField("marker_recorded_event_attributes"): + return False + attributes = event.marker_recorded_event_attributes + if attributes.marker_name != _EXTERNAL_STREAM_MARKER_NAME: + return False + + payloads = attributes.details.get(_EXTERNAL_STREAM_MARKER_DETAILS_KEY) + if payloads is None or len(payloads.payloads) != 1: + raise StreamIntegrityError( + "an external stream marker has no singular marker envelope" + ) + + from temporalio.bridge.proto.external_data import ExternalStreamMarkerData + + marker = ExternalStreamMarkerData() + try: + marker.ParseFromString(payloads.payloads[0].data) + except Exception as err: + raise StreamIntegrityError( + f"an external stream marker envelope could not be decoded: {err}" + ) from err + + output = _marker_output_commit(marker) + return output is not None and _stage_tokens_equal(output.stage_token, stage_token) + + +def _marker_output_commit(marker: Any) -> Any | None: + """Return the shared marker's output proof, when this task had output.""" + if "output" not in marker.DESCRIPTOR.fields_by_name: + raise StreamIntegrityError( + "this SDK's external stream marker schema cannot inspect output " + "commit proofs, so a pending stage cannot be resolved safely" + ) + return marker.output if marker.HasField("output") else None + + +def _stage_tokens_equal(recorded: str | bytes, expected: str) -> bool: + if isinstance(recorded, bytes): + return hmac.compare_digest(recorded, expected.encode()) + return hmac.compare_digest(recorded, expected) diff --git a/temporalio/contrib/external_workflow_streams/_output_codec.py b/temporalio/contrib/external_workflow_streams/_output_codec.py new file mode 100644 index 0000000000..859888d670 --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_output_codec.py @@ -0,0 +1,127 @@ +"""Canonical logical framing for Workflow-originated output. + +Workflow code runs only the deterministic payload converter. The resulting +``Payload`` is framed and fingerprinted before any asynchronous payload codec +or external-storage transform runs, so codec randomness and compression cannot +change replay identity or capacity boundaries. +""" + +from __future__ import annotations + +import hashlib +import struct +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Final + +import temporalio.api.common.v1 +from temporalio.contrib.external_workflow_streams._record import RecordKind + +__all__ = [ + "LOGICAL_FINGERPRINT_VERSION", + "LogicalOutputFingerprint", + "canonical_logical_record_frame", + "fingerprint_logical_frames", + "fingerprint_logical_records", +] + +LOGICAL_FINGERPRINT_VERSION: Final = 1 +"""The canonical frame and batch hash format implemented here.""" + +_LENGTH = struct.Struct(">Q") +"""Unsigned 64-bit network-order lengths used by fingerprint version 1.""" + + +@dataclass(frozen=True) +class LogicalOutputFingerprint: + """The replay identity and deterministic capacity cost of one batch.""" + + version: int + digest: bytes + logical_byte_count: int + record_count: int + + +def _length_prefixed(value: bytes) -> bytes: + if len(value) >= 1 << 64: + raise ValueError("a logical output frame field is too large to encode") + return _LENGTH.pack(len(value)) + value + + +def canonical_logical_record_frame( + topic: str, + kind: object, + payload: temporalio.api.common.v1.Payload | None, +) -> bytes: + """Build one fingerprint-version-1 logical record frame. + + The format is, in order: length-prefixed UTF-8 topic; length-prefixed + unsigned record kind; unsigned metadata-entry count; each metadata key and + value length-prefixed with keys sorted by their UTF-8 bytes; and + length-prefixed raw payload data. Control records use an empty metadata set + and empty data. The ``Payload`` protobuf itself is never serialized. + """ + if not topic: + raise ValueError("a logical output record needs a non-empty topic") + if not isinstance(kind, RecordKind): + raise TypeError( + f"a logical output record kind must be RecordKind, got {type(kind).__name__}" + ) + if kind is RecordKind.DATA and payload is None: + raise ValueError("a DATA logical output record needs a Payload") + if kind.is_control and payload is not None: + raise ValueError(f"a {kind.name} logical output record carries no Payload") + + topic_bytes = topic.encode("utf-8") + # RecordKind is an unsigned protocol number. Minimal big-endian encoding is + # unambiguous because the byte string itself is length-prefixed. + kind_value = int(kind) + kind_bytes = kind_value.to_bytes(max(1, (kind_value.bit_length() + 7) // 8), "big") + metadata = () if payload is None else payload.metadata.items() + sorted_metadata = sorted(metadata, key=lambda item: item[0].encode("utf-8")) + if len(sorted_metadata) >= 1 << 64: + raise ValueError("a logical output Payload has too many metadata entries") + + frame = bytearray() + frame.extend(_length_prefixed(topic_bytes)) + frame.extend(_length_prefixed(kind_bytes)) + frame.extend(_LENGTH.pack(len(sorted_metadata))) + for key, value in sorted_metadata: + frame.extend(_length_prefixed(key.encode("utf-8"))) + frame.extend(_length_prefixed(value)) + frame.extend(_length_prefixed(b"" if payload is None else payload.data)) + return bytes(frame) + + +def fingerprint_logical_frames( + frames: Iterable[bytes], +) -> LogicalOutputFingerprint: + """Hash ordered, length-prefixed canonical frames with SHA-256. + + ``logical_byte_count`` is the sum of the canonical frame lengths, excluding + the batch-level prefixes used only to prevent adjacent frames from being + repartitioned into the same hash input. + """ + digest = hashlib.sha256() + logical_byte_count = 0 + record_count = 0 + for frame in frames: + digest.update(_length_prefixed(frame)) + logical_byte_count += len(frame) + record_count += 1 + return LogicalOutputFingerprint( + version=LOGICAL_FINGERPRINT_VERSION, + digest=digest.digest(), + logical_byte_count=logical_byte_count, + record_count=record_count, + ) + + +def fingerprint_logical_records( + records: Iterable[tuple[str, RecordKind, temporalio.api.common.v1.Payload | None]], +) -> LogicalOutputFingerprint: + """Frame and fingerprint an ordered sequence of logical output records.""" + return fingerprint_logical_frames( + canonical_logical_record_frame(topic, kind, payload) + for topic, kind, payload in records + ) diff --git a/temporalio/contrib/external_workflow_streams/_output_continuation.py b/temporalio/contrib/external_workflow_streams/_output_continuation.py new file mode 100644 index 0000000000..3b75c76a3e --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_output_continuation.py @@ -0,0 +1,101 @@ +"""Must-understand Continue-As-New state for finished output topics.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import temporalio.api.common.v1 +from temporalio.contrib.external_workflow_streams._annotation import ( + AnnotationDecodeError, + _put_str, + _put_uvarint, + _Reader, +) + +__all__ = [ + "OUTPUT_CONTINUATION_HEADER", + "OutputContinuation", + "decode_output_continuation", + "encode_output_continuation", + "read_output_continuation_header", + "write_output_continuation_header", +] + +OUTPUT_CONTINUATION_HEADER = "__temporal_external_output_stream_continuation" +_OUTPUT_CONTINUATION_ENCODING = b"binary/plain" +_SCHEMA_VERSION = 1 + + +@dataclass(frozen=True) +class OutputContinuation: + """The output topics whose explicit terminal committed in this chain.""" + + finished_topics: frozenset[str] + provider_id: str + provider_format_version: int + + def __post_init__(self) -> None: + """Reject continuation state that cannot identify a durable provider.""" + if not self.finished_topics: + raise ValueError("output continuation state needs a finished topic") + if any(not topic for topic in self.finished_topics): + raise ValueError("an output continuation topic may not be empty") + if not self.provider_id: + raise ValueError("output continuation state needs a provider ID") + if self.provider_format_version < 1: + raise ValueError( + "output continuation state needs a positive format version" + ) + + +def encode_output_continuation(continuation: OutputContinuation) -> bytes: + """Encode finished topics and their provider binding deterministically.""" + out = bytearray() + _put_uvarint(out, _SCHEMA_VERSION) + _put_str(out, continuation.provider_id) + _put_uvarint(out, continuation.provider_format_version) + _put_uvarint(out, len(continuation.finished_topics)) + for topic in sorted(continuation.finished_topics): + _put_str(out, topic) + return bytes(out) + + +def decode_output_continuation(raw: bytes) -> OutputContinuation: + """Decode a supported output continuation or fail closed.""" + reader = _Reader(raw) + version = reader.uvarint() + if version != _SCHEMA_VERSION: + raise AnnotationDecodeError( + "the external output Continue-As-New header is schema version " + f"{version}, but this Worker understands {_SCHEMA_VERSION}" + ) + provider_id = reader.string() + provider_format_version = reader.uvarint() + topics = frozenset(reader.string() for _ in range(reader.uvarint())) + return OutputContinuation(topics, provider_id, provider_format_version) + + +def write_output_continuation_header( + continuation: OutputContinuation, +) -> temporalio.api.common.v1.Payload: + """Wrap output continuation bytes without using the user converter.""" + return temporalio.api.common.v1.Payload( + metadata={"encoding": _OUTPUT_CONTINUATION_ENCODING}, + data=encode_output_continuation(continuation), + ) + + +def read_output_continuation_header( + headers: dict[str, temporalio.api.common.v1.Payload] | None, +) -> OutputContinuation | None: + """Read predecessor output state, with absence meaning a first Run.""" + if not headers: + return None + payload = headers.get(OUTPUT_CONTINUATION_HEADER) + if payload is None: + return None + if payload.metadata.get("encoding") != _OUTPUT_CONTINUATION_ENCODING: + raise AnnotationDecodeError( + "the external output Continue-As-New header has an unsupported encoding" + ) + return decode_output_continuation(payload.data) diff --git a/temporalio/contrib/external_workflow_streams/_output_producer.py b/temporalio/contrib/external_workflow_streams/_output_producer.py new file mode 100644 index 0000000000..2e0800f723 --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_output_producer.py @@ -0,0 +1,301 @@ +"""Activity and external-process producers for output streams.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Generic + +import temporalio.converter +from temporalio.contrib.external_workflow_streams._backend import ( + AppendConflictError, + StreamDirection, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec +from temporalio.contrib.external_workflow_streams._output_backend import ( + OutputStreamBackend, + OutputStreamRecord, +) +from temporalio.contrib.external_workflow_streams._producer import ( + WorkflowChainKey, + _default_session_id, + _verify_chain_key, +) +from temporalio.contrib.external_workflow_streams._record import ( + RecordKind, + StreamRecord, +) +from temporalio.types import AnyType + +if TYPE_CHECKING: + import temporalio.client + +__all__ = [ + "ExternalOutputStreamProducer", + "ExternalOutputStreamProducerTopic", + "OutputAppendNotAcknowledgedError", +] + + +class OutputAppendNotAcknowledgedError(Exception): + """A direct output append may or may not have reached the provider. + + Re-run :meth:`ExternalOutputStreamProducerTopic.resolve_append` with the + exact record carried here. The provider's idempotency rule makes that safe; + calling ``publish`` again is not safe because it draws another sequence. + """ + + def __init__(self, *, stream_key: StreamKey, record: StreamRecord) -> None: + """Capture the exact record required for safe ambiguity recovery.""" + super().__init__( + "the output append did not report an outcome; resolve it with the " + "same record before publishing another value" + ) + self.stream_key = stream_key + self.record = record + + +@dataclass(eq=False) +class _AppendOperation: + sequence: int + settled: asyncio.Event = field(default_factory=asyncio.Event) + failure: BaseException | None = None + + +class ExternalOutputStreamProducer: + """A chain-bound producer whose records are committed immediately.""" + + def __init__( + self, + *, + backend: Any, + workflow: WorkflowChainKey, + data_converter: temporalio.converter.DataConverter, + session_id: str, + ) -> None: + """Create a producer after validating its provider and session identity.""" + if not isinstance(backend, OutputStreamBackend): + raise TypeError("backend does not implement OutputStreamBackend") + if type(backend).guarantees_immutability is not True: + raise ValueError("an output backend must guarantee immutable record bytes") + if not type(backend).provider_id: + raise ValueError("an output backend must declare a provider_id") + if not session_id: + raise ValueError("an output producer session ID may not be empty") + self._backend = backend + self._workflow = workflow + self._data_converter = data_converter.with_context( + temporalio.converter.WorkflowSerializationContext( + namespace=workflow.namespace, + workflow_id=workflow.workflow_id, + ) + ) + self._session_id = session_id + self._sequence = 0 + self._operations: dict[StreamKey, list[_AppendOperation]] = {} + self._unresolved: dict[StreamKey, list[StreamRecord]] = {} + self._closing: set[StreamKey] = set() + self._finished: set[StreamKey] = set() + + @classmethod + async def connect( + cls, + *, + backend: OutputStreamBackend, + workflow: WorkflowChainKey, + client: temporalio.client.Client, + data_converter: temporalio.converter.DataConverter | None = None, + session_id: str | None = None, + ) -> ExternalOutputStreamProducer: + """Verify the chain binding before allowing the first append.""" + await _verify_chain_key(client, workflow) + return cls( + backend=backend, + workflow=workflow, + data_converter=data_converter or client.data_converter, + session_id=session_id or _default_session_id(), + ) + + @property + def session_id(self) -> str: + """The retry-stable idempotency session.""" + return self._session_id + + @property + def workflow(self) -> WorkflowChainKey: + """The Workflow chain receiving these records.""" + return self._workflow + + def topic( + self, name: str, *, type: type[AnyType] | None = None + ) -> ExternalOutputStreamProducerTopic[AnyType]: + """Create a typed direct publisher for one output topic.""" + if not name: + raise ValueError("an output topic needs a non-empty name") + key = self._workflow.stream_key( + name, + direction=StreamDirection.OUTPUT, + ) + return ExternalOutputStreamProducerTopic( + producer=self, + stream_key=key, + codec=StreamPayloadCodec(self._data_converter, type), + ) + + def _next_sequence(self) -> int: + sequence = self._sequence + self._sequence += 1 + return sequence + + +class ExternalOutputStreamProducerTopic(Generic[AnyType]): + """One immediately committed output topic.""" + + def __init__( + self, + *, + producer: ExternalOutputStreamProducer, + stream_key: StreamKey, + codec: StreamPayloadCodec[AnyType], + ) -> None: + """Create a typed publisher bound to one output topic.""" + self._producer = producer + self._stream_key = stream_key + self._codec = codec + + @property + def stream_key(self) -> StreamKey: + """The direction-isolated physical stream identity.""" + return self._stream_key + + def _refuse_unresolved(self) -> None: + records = self._producer._unresolved.get(self._stream_key) + if records: + raise OutputAppendNotAcknowledgedError( + stream_key=self._stream_key, + record=records[0], + ) + + def _refuse_closed(self) -> None: + if self._stream_key in self._producer._finished: + raise RuntimeError("this output topic is already finished") + if self._stream_key in self._producer._closing: + raise RuntimeError("this output topic is already being finished") + + def _remember_unresolved(self, record: StreamRecord) -> None: + records = self._producer._unresolved.setdefault(self._stream_key, []) + if record not in records: + records.append(record) + + def _forget_unresolved(self, record: StreamRecord) -> None: + records = self._producer._unresolved[self._stream_key] + records.remove(record) + if not records: + del self._producer._unresolved[self._stream_key] + + async def _append(self, record: StreamRecord) -> OutputStreamRecord: + try: + return await self._producer._backend.append_output(self._stream_key, record) + except AppendConflictError: + raise + except asyncio.CancelledError: + self._remember_unresolved(record) + raise OutputAppendNotAcknowledgedError( + stream_key=self._stream_key, + record=record, + ) from None + except Exception as err: + # Provider-specific transport exceptions cannot reliably distinguish + # a refusal from a lost acknowledgement. Repeating the exact append + # is the only recovery safe in both cases. + self._remember_unresolved(record) + raise OutputAppendNotAcknowledgedError( + stream_key=self._stream_key, + record=record, + ) from err + + async def resolve_append(self, record: StreamRecord) -> OutputStreamRecord: + """Resolve an unknown append by repeating its exact identity and bytes.""" + outstanding = self._producer._unresolved.get(self._stream_key, ()) + if record not in outstanding: + raise ValueError("record is not this topic's unresolved output append") + try: + placed = await self._producer._backend.append_output( + self._stream_key, record + ) + except AppendConflictError: + self._forget_unresolved(record) + if record.kind is RecordKind.FINISH: + self._producer._closing.discard(self._stream_key) + raise + except (Exception, asyncio.CancelledError) as err: + raise OutputAppendNotAcknowledgedError( + stream_key=self._stream_key, + record=record, + ) from err + self._forget_unresolved(record) + if record.kind is RecordKind.FINISH: + self._producer._closing.discard(self._stream_key) + self._producer._finished.add(self._stream_key) + return placed + + async def publish(self, value: AnyType) -> OutputStreamRecord: + """Append one immediately committed data record.""" + self._refuse_unresolved() + self._refuse_closed() + sequence = self._producer._next_sequence() + operation = _AppendOperation(sequence) + operations = self._producer._operations.setdefault(self._stream_key, []) + operations.append(operation) + try: + record = StreamRecord( + kind=RecordKind.DATA, + payload=await self._codec.encode(value), + producer_session_id=self._producer.session_id, + sequence=sequence, + ) + placed = await self._append(record) + except BaseException as err: + operation.failure = err + operation.settled.set() + raise + else: + operation.settled.set() + return placed + finally: + operations.remove(operation) + + async def finish_writing(self) -> OutputStreamRecord: + """Append the ordered terminal after every earlier publish settles.""" + self._refuse_unresolved() + self._refuse_closed() + self._producer._closing.add(self._stream_key) + sequence = self._producer._next_sequence() + try: + preceding = tuple(self._producer._operations.get(self._stream_key, ())) + for operation in preceding: + await operation.settled.wait() + if operation.failure is not None: + raise RuntimeError( + f"output publish {operation.sequence} did not append; " + "refusing to place a terminal in front of it" + ) from operation.failure + record = StreamRecord( + kind=RecordKind.FINISH, + payload=b"", + producer_session_id=self._producer.session_id, + sequence=sequence, + ) + placed = await self._append(record) + except OutputAppendNotAcknowledgedError: + # The terminal may already be durable, so the topic stays closed to + # new publishes until resolve_append settles that exact record. + raise + except BaseException: + self._producer._closing.discard(self._stream_key) + raise + else: + self._producer._closing.discard(self._stream_key) + self._producer._finished.add(self._stream_key) + return placed diff --git a/temporalio/contrib/external_workflow_streams/_producer.py b/temporalio/contrib/external_workflow_streams/_producer.py new file mode 100644 index 0000000000..49ce2876fd --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_producer.py @@ -0,0 +1,1470 @@ +"""The producer side (P6a binding, P6 append). + +The producer handle is a **different type** from the Workflow-side one, and not +the same object passed across a process boundary. A Workflow handle is bound to +the running Workflow's identity and backend configuration; a producer +handle is constructed explicitly from credentials the producer holds. + +Every binding input is explicit, because none of them can be inferred: + +- **The Workflow chain key**, including the first execution Run ID. + ``temporalio.activity.Info`` exposes ``workflow_run_id`` but *not* the first + execution Run ID, so an Activity cannot derive the key -- the Workflow passes + it in and the producer verifies it by describing the Workflow before its + first append. Publishing under an unverified key is a configuration error, + not a silent no-op. +- **A backend connection**, **a Temporal client** (for the wake Signal), and + **the same DataConverter** the consuming Workflow uses. +- **A stable producer session ID**, which is what makes append idempotent under + Activity retry. Activities default it to something derived from the + Activity's own identity so a retried attempt reuses it. Plain processes must + supply one: a random default would look like it worked and duplicate every + record on the first retry. + +The **stream name appears exactly once**, in :meth:`ExternalStreamProducer.topic`. +``connect()`` takes the chain key and ``topic(name)`` completes it into the full +stream identity, so one connection serves several topics and no two arguments +can disagree about the name. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from datetime import timedelta +from typing import TYPE_CHECKING, Any, Generic + +import temporalio.activity +import temporalio.converter +from temporalio.contrib.external_workflow_streams._backend import ( + AppendConflictError, + StreamBackend, + StreamDirection, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec +from temporalio.contrib.external_workflow_streams._record import ( + Offset, + RecordKind, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._wake import ( + WakeRequest, + send_wake_signal, + wake_request_for, +) +from temporalio.types import AnyType + +DEFAULT_WAKE_CLAIM_LEASE = timedelta(seconds=30) +"""Long enough to cover a Signal round trip, short enough to recover from a crash. + +The lease bounds how long a claim taken by a producer that then died keeps +saying a wake is in flight. It is not what makes the wake safe -- expiry alone +schedules nobody to take over, which is why a producer that loses the claim +signals anyway rather than treating the claim as an acknowledgement. +""" + +if TYPE_CHECKING: + import temporalio.client + +__all__ = [ + "AppendNotAcknowledgedError", + "ChainKeyMismatchError", + "ExternalStreamProducer", + "ExternalStreamProducerTopic", + "PrecedingWriteFailedError", + "WakeNotAcknowledgedError", + "WorkflowChainKey", +] + + +class WakeNotAcknowledgedError(Exception): + """The record is durable but the wake step did not complete. + + Raised rather than swallowed, because the two halves fail differently and the + caller can only act on one of them: the append has already succeeded and must + not be retried, while the wake must be. Retrying the wake with the same + inputs is safe -- it derives the same request ID and the server deduplicates + it -- so this is a resumable state, not a lost record. + + **Cancellation after the append arrives this way too**, with + :attr:`cancelled` set. It leaves the identical state -- a durable record and + an unsent wake -- and a bare ``CancelledError`` there would be + indistinguishable from cancellation *before* the append: the caller would not + know whether the value still had to be published, and its obvious move, + calling ``publish()`` again, draws a fresh sequence number and appends a + second record (ADR-036). + """ + + def __init__( + self, + message: str, + *, + pending: list[WakeRequest], + restart: bool = False, + cancelled: bool = False, + ) -> None: + """Capture the wakes still owed after a durable append.""" + super().__init__(message) + self.pending = pending + """The wakes still owed, ready to be retried verbatim.""" + self.cancelled = cancelled + """Whether what ended the wake was cancellation rather than a failure. + + Set because the two are the same *state* -- durable record, unsent wake, + one recovery -- and a caller that wants to honour the cancellation after + recovering the wake still has to know it was asked to stop. The recovery + itself does not depend on this: ``pending`` and ``restart`` say what to do + whatever ended the attempt (ADR-036). + """ + self.restart = restart + """Whether the caller must call :meth:`.wake` again instead of retrying. + + The wake is three steps -- observe the parked set, claim the generation, + Signal -- and only the third produces the requests + :meth:`ProducerTopicHandle.retry_wake` re-sends. A failure in the first + two leaves nothing to re-send, so ``pending`` is empty and retrying it + would silently do nothing at all: the record would stay durable and + unannounced while the caller believed it had recovered. + + ``True`` therefore says "no wake was composed; compose one". Calling + :meth:`ProducerTopicHandle.wake` again is safe and is the whole recovery + -- it re-observes the parked set, and a parked wake's request ID is + derived from the generation rather than from the sender, so a wake some + other producer already sent deduplicates against it. + """ + self.offset: Offset | None = None + """Where the record landed, when raised from :meth:`publish`. + + Set because the two halves fail differently and the caller can only act + on one of them: the append succeeded and must not be retried, the wake + did not and must be. + """ + + +class AppendNotAcknowledgedError(Exception): + """The append's outcome is unknown: the record may or may not be durable. + + Raised when :meth:`StreamBackend.append` neither returns nor *refuses* -- a + cancellation delivered while it was in flight, a dropped connection, a + provider exception that says nothing about what the server did with the + write. A remote backend commits on its own side and only then answers, so an + append that did not answer is not an append that did not happen: the Redis + provider runs an atomic script server-side and receives its result in a + separate client-side step, and losing the client between the two leaves a + durable record whose offset nobody holds. + + Reading that window as failure breaks the producer's contract in both + directions at once, which is why it gets a type of its own. Retrying + ``publish()`` draws a **new** sequence number, and therefore a new + idempotency key, so a record that did land is appended a second time. Not + retrying can leave a durable record that no wake was ever sent for, which + strands a parked Workflow for its whole idle timeout. A bare + ``CancelledError`` or a bare ``ConnectionError`` supports neither choice, + because neither carries the identity of the record whose fate is in + question. + + So the record is carried out instead -- byte-identical, still holding its + ``(session_id, sequence)`` -- and + :meth:`ExternalStreamProducerTopic.resolve_append` re-appends *that* record. + Under the backend contract that one call is correct whichever way the first + attempt went: a repeat append of byte-identical content under a used key + writes nothing and returns the original offset, and a key the backend never + saw is appended now (ADR-038). + + Until it is resolved, the stream refuses further appends from this producer + with this same error, since that is the caller move that duplicates the + record. + + Distinct from :class:`WakeNotAcknowledgedError`, which is the *next* window + along: there the offset is known and only the wake is owed. Resolving this + one produces that offset, after which a failing wake raises the wake error in + the ordinary way. + """ + + def __init__( + self, + message: str, + *, + stream_key: StreamKey, + record: StreamRecord, + wake: bool, + lease: timedelta, + cancelled: bool = False, + ) -> None: + """Capture an append whose durable outcome must be resolved.""" + super().__init__(message) + self.stream_key = stream_key + """The stream the append was for. Where it must be settled. + + Carried because a record does not name its own stream and the backend's + idempotency scope does: `(session_id, sequence)` is unused on every + *other* stream, so the same record handed to another topic's + ``resolve_append`` would append a second copy there rather than + deduplicate. The recovery refuses that, and this is what a caller + holding several topics matches on. + """ + self.record = record + """The exact record whose fate is unknown. What ``resolve_append`` takes. + + Carried rather than described, because idempotency is on identity: a + payload the caller re-encodes into different bytes under the same key is + an ``AppendConflictError`` rather than the no-op the recovery depends on. + """ + self.wake = wake + """Whether the interrupted call was going to wake, so recovery can too.""" + self.lease = lease + """The claim lease the interrupted call was going to use.""" + self.cancelled = cancelled + """Whether cancellation ended any attempt to settle this operation. + + Reported rather than propagated, for the reason ADR-036 already gives + about the wake: a caller that wants to honour the cancellation re-raises + *after* resolving the append, and that is the only order that leaves + nothing owed. Sticky across recovery attempts: a later transport failure + must not erase a cancellation the caller still has to honour. + """ + + +class PrecedingWriteFailedError(Exception): + """A write fence was not appended, because an earlier write was not either. + + Raised from :meth:`ExternalStreamProducerTopic.finish_writing` when a + ``publish()`` invoked earlier on the same stream ended without a durable + record. The fence means every write in this producer session preceding it + has been appended, and a consumer that drains through one may park on that + assertion, so a fence written over a hole would tell it the batch is + complete when it is short a record. + + Reachable only from *concurrent* calls: a publish that failed before the + fence was invoked is already finished and the fence covers what it says it + covers. Nothing was appended for this fence, so the recovery is the caller's + choice between the two halves -- publish the failed value again, which draws + a new sequence number and lands ahead of a later fence, or accept the batch + without it. Either way ``finish_writing()`` again appends the fence: the + failed write is no longer outstanding. + + Distinct from :class:`AppendNotAcknowledgedError`, which the fence raises + instead when the earlier append's outcome is *unknown* rather than failed. + There the record may well be durable, the stream refuses appends until it is + settled, and the recovery is that operation's rather than this call's. + """ + + def __init__(self, message: str, *, stream_key: StreamKey, sequence: int) -> None: + """Identify the earlier failed write that makes a fence unsafe.""" + super().__init__(message) + self.stream_key = stream_key + """The stream whose fence was refused.""" + self.sequence = sequence + """The sequence number of the earlier write that did not land. + + Names the operation rather than describing it: the failure itself is the + ``__cause__``, and the caller already saw it raised from its own call. + """ + + +@dataclass(frozen=True) +class _UnresolvedAppend: + """One append whose outcome the producer never learned. + + The *operation*, not just its record. What the interrupted call owed is more + than the bytes: whether a wake was due, under which lease, and whether + cancellation is still to be honoured once the state is settled. Holding only + the record meant a later refusal had to invent those three from whatever call + happened to be refused, so a caller following the refusal's own instructions + could drop a wake the unresolved record required. + + Keyed by stream because idempotency is scoped per stream: the same + ``(session_id, sequence)`` is unused on every other one, so a record settled + against the wrong topic appends a second copy rather than deduplicating. + """ + + stream_key: StreamKey + record: StreamRecord + wake: bool + lease: timedelta + cancelled: bool + operation: _StreamOperation | None = None + """The append order entry whose outcome this recovery decides, if any. + + A fence has none: nothing waits behind one. A publish has exactly one, and + it is the *only* place a resolution can be reported back to a fence that + already captured that operation while its outcome was unknown -- see + :meth:`ExternalStreamProducerTopic.resolve_append`. Without it, "the record + is no longer unresolved" is all a waiting fence can see, and both a durable + resolution and an ``AppendConflictError`` produce that. + """ + + def error(self, message: str) -> AppendNotAcknowledgedError: + """Reports this operation's current canonical recovery.""" + return AppendNotAcknowledgedError( + message, + stream_key=self.stream_key, + record=self.record, + wake=self.wake, + lease=self.lease, + cancelled=self.cancelled, + ) + + +@dataclass(eq=False) +class _StreamOperation: + """One **publish** on a stream, from the moment it draws its sequence. + + The entry in the producer's per-stream append order. It exists so that + :meth:`ExternalStreamProducerTopic.finish_writing` can tell which earlier + calls on that stream have not reached the backend yet -- which nothing else + records. ``publish()`` draws its sequence *before* awaiting the codec, so a + publish that is still encoding is invisible to :meth:`_refuse_while_unresolved` + (it has no unresolved append) and to the backend (it has appended nothing). + + **A fence is not one of these.** The order exists to hold a fence behind the + *data writes* that precede it, which is the whole of what a fence asserts; + two fences make independent assertions about the publishes each of them came + after, so neither has to wait for the other. Putting them in the same ledger + made a fence that never reached the backend -- cancelled while waiting, or + refused -- look to a later fence like a data write that went missing, and a + valid fence was then refused with ``PrecedingWriteFailedError``. + + Compared by identity rather than by field, because two operations may hold + the same values and the order removes exactly the one that settled. + """ + + sequence: int + + settled: asyncio.Event = field(default_factory=asyncio.Event) + """Set once the append has an outcome -- durable, refused, or unknown.""" + + failure: BaseException | None = None + """Why no durable record came of it, or ``None`` if one did. + + Read only after :attr:`settled`, and only by a fence waiting behind this + operation. The caller of the failed call has already been raised at. + + Not final while it holds an :class:`AppendNotAcknowledgedError`: that is the + one outcome the producer does not yet know, and + :meth:`ExternalStreamProducerTopic.resolve_append` replaces it with the + durable or refused answer recovery learned. + """ + + +@dataclass(frozen=True) +class WorkflowChainKey: + """The Workflow chain a producer publishes to. + + A *chain* key, not a Run key: the stream spans the whole Continue-As-New + chain, and ``first_execution_run_id`` is what stays stable across it while + still preventing collisions after Workflow ID reuse. + """ + + namespace: str + workflow_id: str + first_execution_run_id: str + + def __post_init__(self) -> None: + """Require every component of the Workflow chain identity.""" + for field_name in ("namespace", "workflow_id", "first_execution_run_id"): + if not getattr(self, field_name): + raise ValueError( + f"a Workflow chain key needs a non-empty {field_name}; it is " + "passed in by the Workflow rather than inferred, because an " + "Activity cannot derive the first execution Run ID" + ) + + def stream_key( + self, + stream_name: str, + *, + direction: StreamDirection = StreamDirection.INPUT, + ) -> StreamKey: + """Return the durable key for one stream in this Workflow chain.""" + return StreamKey( + namespace=self.namespace, + workflow_id=self.workflow_id, + first_execution_run_id=self.first_execution_run_id, + stream_name=stream_name, + direction=direction, + ) + + +class ChainKeyMismatchError(Exception): + """The chain key does not describe the Workflow the server knows about. + + Loud on purpose. Publishing under a wrong key writes records into a stream + no consumer is watching, and the Workflow simply waits out its idle timeout + with nothing to show for it. + """ + + +def _default_session_id() -> str: + """A session ID stable across an Activity's retries, or a clear error. + + Derived from the Activity's identity and deliberately **not** its attempt + number: the whole point is that attempt 2 reuses attempt 1's key so its + re-appends are recognised as the same records. + """ + try: + info = temporalio.activity.info() + except RuntimeError: + raise ValueError( + "a producer session ID is required and has no default outside an " + "Activity. It is what makes append idempotent under retry, so a " + "random default would look correct and duplicate every record the " + "first time the producer was retried." + ) from None + return f"activity:{info.workflow_run_id}:{info.activity_id}" + + +class ExternalStreamProducer: + """A bound producer connection. Serves any number of topics.""" + + def __init__( + self, + *, + backend: StreamBackend, + workflow: WorkflowChainKey, + data_converter: temporalio.converter.DataConverter, + session_id: str, + client: temporalio.client.Client | None = None, + ) -> None: + """Bind a backend, converter, and producer session to one chain.""" + self._backend = backend + self._workflow = workflow + #: Bound to the Workflow these records are for, which is the same + #: context the consuming Worker decodes them under. The two sides share + #: one converter and must therefore share one context: a producer that + #: encrypts under no context while the consumer decrypts under a + #: Workflow-derived key is a mismatch that the append reports as + #: success and that only surfaces on the far side, as a decode failure + #: against a configuration that is in fact correct. + #: + #: A producer is bound to exactly one chain, so this is settled once + #: here rather than per topic. `workflow_id` and not a Run: the chain + #: key spans Continue-As-New, and a record written for one Run is read + #: by its successors. + self._data_converter = data_converter.with_context( + temporalio.converter.WorkflowSerializationContext( + namespace=workflow.namespace, + workflow_id=workflow.workflow_id, + ) + ) + self._session_id = session_id + self._client = client + self._sequence = 0 + self._wake_counter = 0 + #: Per stream, the appends that never reported an outcome. + #: + #: Kept because the operation *is* the recovery: only these exact bytes + #: under these exact `(session_id, sequence)` pairs re-append as a no-op + #: if the first attempt landed, and only what the interrupted call owed + #: says whether settling it still has a wake to send. A list rather than + #: a single slot because concurrent publishes to one stream are supported + #: and more than one of them can be interrupted; entries are added only + #: when an append fails to answer, so a publish that is merely still in + #: flight registers nothing and cannot block its own sibling. + #: + #: **Per producer instance, deliberately.** This is what binds recovery + #: to the object that still holds the session's sequence and wake + #: counters; a replacement producer built with the same session id has + #: both back at zero and recovers by re-running the same calls instead + #: (ADR-038). + self._unresolved: dict[StreamKey, list[_UnresolvedAppend]] = {} + #: Per stream, the operations that have drawn a sequence number and not + #: yet reached an append outcome. The ordered append coordinator a write + #: fence needs. + #: + #: A fence asserts that every preceding write in this session has been + #: appended, and nothing made that true: an earlier `publish()` could + #: still be inside its codec -- holding a lower sequence number, having + #: appended nothing, and unresolved nowhere -- while `finish_writing()` + #: appended the fence and sent the batch's only wake. A consumer then + #: parks on a fence whose data lands behind it, with the one wake that + #: would have moved it already spent. + #: + #: Keyed by stream because that is the scope of the fence's claim, and + #: held on the producer rather than the handle because `topic()` returns + #: a fresh handle per call: two handles for one name are one stream and + #: must join one order. Only the fence waits -- concurrent publishes have + #: no defined order between them to preserve, so ordering their appends + #: would serialize a batch for nothing. + self._appends: dict[StreamKey, list[_StreamOperation]] = {} + + @staticmethod + async def connect( + *, + backend: StreamBackend, + workflow: WorkflowChainKey, + client: temporalio.client.Client, + data_converter: temporalio.converter.DataConverter | None = None, + session_id: str | None = None, + ) -> ExternalStreamProducer: + """Binds a producer, verifying the chain key before it can append. + + Args: + backend: A provider instance. A plain process constructs one + directly; there is no Worker registry out here to name. + workflow: The chain key, passed in by the Workflow. + client: Used to verify the chain key, and later to send the wake + Signal. Required -- an unverified binding is a configuration + error, not a degraded mode. + data_converter: Must match the consuming Workflow's, including any + codec. Defaults to the client's. + session_id: Stable across retries. Defaults inside an Activity to a + value derived from the Activity's identity; required outside + one. + + Raises: + ChainKeyMismatchError: The server's first execution Run ID for this + Workflow ID is not the one given. + """ + await _verify_chain_key(client, workflow) + return ExternalStreamProducer( + backend=backend, + workflow=workflow, + data_converter=data_converter or client.data_converter, + session_id=session_id or _default_session_id(), + client=client, + ) + + @property + def session_id(self) -> str: + """The idempotency session shared by this producer's appends.""" + return self._session_id + + @property + def workflow(self) -> WorkflowChainKey: + """The Workflow chain receiving this producer's records.""" + return self._workflow + + def topic( + self, name: str, *, type: type[AnyType] | None = None + ) -> ExternalStreamProducerTopic[Any]: + """A handle for one stream. + + The only place the stream name appears on the producer side. + """ + if not name: + raise ValueError("a topic needs a non-empty name") + return ExternalStreamProducerTopic( + producer=self, + stream_key=self._workflow.stream_key(name), + codec=StreamPayloadCodec(self._data_converter, type), + ) + + def _next_wake_counter(self) -> int: + """Distinguishes this sender's unparked wakes from each other. + + An unparked wake names generation 0, which carries no attempt identity, + so without a per-sender counter every unparked wake from one sender would + derive the same request ID and the server would deduplicate all but the + first away. Held fixed across *retries* of one attempt -- it advances per + attempt, not per send. + """ + self._wake_counter += 1 + return self._wake_counter + + def _next_sequence(self) -> int: + """The next sequence number in this producer session. + + Per *connection*, not per topic: the pair ``(session_id, sequence)`` is + the idempotency key, and restarting the count for each topic would make + two topics' first records collide. + """ + sequence = self._sequence + self._sequence += 1 + return sequence + + +class ExternalStreamProducerTopic(Generic[AnyType]): + """One stream, from the producer's side. + + Has ``publish`` and ``finish_writing`` and no ``subscribe`` -- the two sides + are mirror images, not one type used twice. + """ + + def __init__( + self, + *, + producer: ExternalStreamProducer, + stream_key: StreamKey, + codec: StreamPayloadCodec[AnyType], + ) -> None: + """Bind a typed codec and stream key to a producer.""" + self._producer = producer + self._stream_key = stream_key + self._codec = codec + + @property + def stream_key(self) -> StreamKey: + """The durable identity of this producer topic.""" + return self._stream_key + + def _refuse_while_unresolved(self) -> None: + """Refuses a new append while an earlier one's outcome is unknown. + + This is the caller move that duplicates a record, so it is the one the + producer stops. A new ``publish()`` draws a fresh sequence number and + therefore a fresh idempotency key, so if the unresolved append did land, + the value is now in the stream twice under two keys the backend has no + way to relate. Refusing also keeps the recovery in order: the unresolved + record is re-appended before anything is appended behind it. + + The error it raises is **the outstanding operation's**, not this call's. + Reporting the refused call's ``wake`` and ``lease`` beside the older + call's record would make the error's own instructions wrong: a caller + following them settles the record with no wake, and a parked Workflow + sits on a durable record nobody announced. + + Checked at entry, so a publish already past this point when a sibling + becomes unresolved still completes -- concurrent publishes have no + defined order between them to preserve. + """ + outstanding = self._producer._unresolved.get(self._stream_key) + if not outstanding: + return + pending = outstanding[0] + more = ( + "" + if len(outstanding) == 1 + else f" ({len(outstanding)} appends on this stream are unsettled)" + ) + raise pending.error( + f"the append of {pending.record.idempotency_key} on stream " + f"{self._stream_key} has never been acknowledged, so this stream " + "will not take another append: a new one draws a fresh sequence " + "number, and if that record did land the value would be in the " + f"stream twice{more}. Settle it with resolve_append(`.record`), " + "which owes the wake this error reports rather than the one the " + "refused call asked for." + ) + + def _begin(self, sequence: int) -> _StreamOperation: + """Joins this publish to its stream's append order. + + Synchronous, and called in the same uninterrupted step as + :meth:`_refuse_while_unresolved` and the sequence draw. An operation + registered after its first await would be invisible to a fence that + began in between, which is the entire race this closes. + """ + outstanding = self._producer._appends.setdefault(self._stream_key, []) + operation = _StreamOperation(sequence=sequence) + outstanding.append(operation) + return operation + + def _preceding_publishes(self) -> tuple[_StreamOperation, ...]: + """The publishes a fence invoked now has to wait behind. + + Snapshotted at invocation rather than re-read at wait time, because a + fence's claim is about the calls that came *before* it. A publish invoked + behind the fence is outside that claim -- it takes its place in the + stream after the fence, exactly as a later publish does -- and a fence + that waited for one would be held by work the caller started after + asking for the fence. + + Synchronous, and taken in the same uninterrupted step as the sequence + draw, for the reason :meth:`_begin` gives. A fence does not join the + order it reads: see :class:`_StreamOperation`. + """ + return tuple(self._producer._appends.get(self._stream_key, ())) + + def _settle( + self, operation: _StreamOperation, failure: BaseException | None = None + ) -> None: + """Reports this call's append outcome and releases anything behind it. + + Called once the append has an outcome and **not** after the wake. The + fence asserts that preceding writes are *appended*, which a durable + record with a still-owed Signal already satisfies; holding the fence for + the wake as well would stall it on a state its own claim says nothing + about, and on the failure of one leave a caller unable to fence at all. + """ + operation.failure = failure + operation.settled.set() + outstanding = self._producer._appends.get(self._stream_key) + if outstanding is None: + return + outstanding[:] = [held for held in outstanding if held is not operation] + if not outstanding: + del self._producer._appends[self._stream_key] + + async def _await_preceding_appends( + self, preceding: tuple[_StreamOperation, ...] + ) -> None: + """Holds a fence until every earlier publish on this stream has landed. + + What makes :meth:`finish_writing`'s claim true rather than merely + documented. Each earlier operation is waited on rather than polled, and + it settles as soon as its append returns, so a fence with nothing + outstanding ahead of it does not yield at all. + + Every outcome is read *after* the last wait rather than as each one + settles, because an outcome read earlier can still change: an append + whose answer was lost settles as unknown, and the + :meth:`resolve_append` that turns that into a durable record or an + ``AppendConflictError`` can happen while this fence is still waiting on a + later publish. Reading in the loop took the stale unknown and passed over + a write that recovery had by then proved absent. + + Then, in order: + + - An append still *unresolved* refuses the fence with that operation's + canonical error -- the same refusal :meth:`_refuse_while_unresolved` + gives at entry, because the stream takes no further append until it is + settled (ADR-038). Reported ahead of any outright failure, since it is + the one that blocks the recovery for the other: republishing a failed + value is itself refused while an append is unsettled. + - An earlier *failure* is propagated as + :class:`PrecedingWriteFailedError` instead of being passed over: the + caller asked for that write before it asked for the fence, and a fence + appended over the hole tells a consumer the batch is complete when it + is short a record. + """ + for earlier in preceding: + await earlier.settled.wait() + # Re-checked here and not only at entry: an append can lose its answer at + # any point while this fence is held, including one from a publish that + # began behind it, and appending the fence in front of a record that may + # yet be settled into the stream is the ordering this method exists to + # refuse. No await separates it from the scan below, so neither reads a + # state the other has already moved past. + self._refuse_while_unresolved() + for earlier in preceding: + failure = earlier.failure + if failure is None: + continue + raise PrecedingWriteFailedError( + f"the write at sequence {earlier.sequence} on stream " + f"{self._stream_key} did not produce a durable record, so this " + "fence was not appended: a fence means every preceding write in " + "this producer session has been appended, and a consumer that " + "drains through one may park on that. Publish the value again " + "-- it draws a new sequence number and lands ahead of a later " + "fence -- or accept the batch without it; either way " + "finish_writing() again appends the fence, since the failed " + "write is no longer outstanding.", + stream_key=self._stream_key, + sequence=earlier.sequence, + ) from failure + + def _remember(self, pending: _UnresolvedAppend) -> _UnresolvedAppend: + """Makes ``pending`` the canonical state for this unsettled operation. + + A recovery may deliberately change the wake or lease. If its own append + then loses the response, the *recovery* is now the interrupted operation, + so retaining the older policy makes the next defaulted attempt contradict + the error it was handed. Cancellation is different: once delivered it is + still owed after every later failure, so it accumulates rather than being + replaced. The append order entry is the *first* attempt's, for the same + reason: it is the publish a fence captured, and no recovery of it creates + another. + """ + outstanding = self._producer._unresolved.setdefault(self._stream_key, []) + for index, held in enumerate(outstanding): + if held.record != pending.record: + continue + pending = _UnresolvedAppend( + stream_key=pending.stream_key, + record=pending.record, + wake=pending.wake, + lease=pending.lease, + cancelled=held.cancelled or pending.cancelled, + operation=held.operation, + ) + outstanding[index] = pending + return pending + outstanding.append(pending) + return pending + + def _forget(self, record: StreamRecord) -> None: + """Drops a settled append from the unresolved set. + + Matched on the idempotency key, where :meth:`_remember` and + :meth:`_outstanding` match on the whole record. Not a discrepancy that + is reachable: a record with a stored key but different bytes cannot be + in the set, because `_outstanding` refuses it with a ``ValueError`` + before the backend is touched and `publish` always draws an unused key. + If either of those ever stops holding, this drops both entries for the + key while the other two treat them as distinct, and one unsettled append + disappears without being settled. + """ + outstanding = self._producer._unresolved.get(self._stream_key) + if not outstanding: + return + remaining = [ + held + for held in outstanding + if held.record.idempotency_key != record.idempotency_key + ] + if remaining: + self._producer._unresolved[self._stream_key] = remaining + else: + del self._producer._unresolved[self._stream_key] + + async def _append( + self, + record: StreamRecord, + *, + wake: bool, + lease: timedelta, + operation: _StreamOperation | None, + ) -> StreamRecord: + """The append, with its acknowledgement window made explicit. + + Everything that is neither a return nor a contractual refusal leaves as + :class:`AppendNotAcknowledgedError`, because from here the two are + genuinely indistinguishable: the backend may have committed and lost the + answer. ``AppendConflictError`` is the exception, and the only one the + contract defines -- the key was used with *different* bytes, so this + record did not land and re-appending it would raise the identical error. + + ``KeyboardInterrupt`` and ``SystemExit`` are not converted either, for + the reason ADR-036 gives: the interpreter is going away and there is no + caller left to recover. + + ``operation`` is the append order entry this record belongs to, carried + onto the unresolved state so a recovery of it can report the outcome it + learns back to a fence holding that entry. ``None`` for a fence's own + append, which nothing waits behind. + """ + producer = self._producer + try: + placed = await producer._backend.append(self._stream_key, record) + except AppendConflictError: + self._forget(record) + raise + except (Exception, asyncio.CancelledError) as err: + pending = _UnresolvedAppend( + stream_key=self._stream_key, + record=record, + wake=wake, + lease=lease, + cancelled=isinstance(err, asyncio.CancelledError), + operation=operation, + ) + pending = self._remember(pending) + raise pending.error( + f"the append of {record.idempotency_key} did not report an " + f"outcome: {err!r}. Whether it landed is unknown -- a backend " + "commits before it answers -- so settle it with " + "resolve_append(`.record`) rather than by publishing again." + ) from err + assert placed.offset is not None + self._forget(record) + return placed + + async def _wake_for( + self, placed: StreamRecord, *, wake: bool, lease: timedelta + ) -> Offset: + """Steps 2 and 3 for a record now known to be durable.""" + assert placed.offset is not None + if wake: + try: + await self.wake(lease=lease) + except WakeNotAcknowledgedError as err: + # Re-raised carrying the offset: the caller has no other way to + # learn that the append it must *not* retry did succeed. + err.offset = placed.offset + raise + return placed.offset + + def _outstanding(self, record: StreamRecord) -> _UnresolvedAppend: + """This topic's unsettled append for ``record``, or why there is none. + + The lookup **is** the safety check, and it is three checks at once. It + binds the recovery to the stream, because a record does not name its own + stream and `(session_id, sequence)` is unused on every other one -- so a + record settled against the wrong topic appends a second copy of the value + rather than deduplicating, and leaves the real stream still blocked. It + binds the recovery to the exact bytes, because idempotency is on identity + and a re-encoded payload under the same key is an ``AppendConflictError`` + rather than the no-op this depends on. And it binds the recovery to the + producer *instance*, which is the one that still holds the session's + sequence and wake counters. + + Every failure is a ``ValueError`` naming which of the three it was, and + every one of them is raised before the backend is touched. + """ + if record.offset is not None: + raise ValueError( + f"record {record.idempotency_key} already carries offset " + f"{record.offset}, which means the backend acknowledged it. " + "There is nothing unresolved about it; if a wake is still owed, " + "that is what wake() and retry_wake() are for." + ) + producer = self._producer + for held in producer._unresolved.get(self._stream_key, ()): + if held.record == record: + return held + + if record.producer_session_id != producer.session_id: + raise ValueError( + f"record {record.idempotency_key} was written by producer " + f"session {record.producer_session_id!r}, not by this one " + f"({producer.session_id!r}). Only the session that drew the " + "sequence number can re-append under that key; from any other " + "session the same bytes are a different record." + ) + elsewhere = [ + key + for key, outstanding in producer._unresolved.items() + if any(held.record == record for held in outstanding) + ] + if elsewhere: + raise ValueError( + f"record {record.idempotency_key} has an unsettled append on " + f"stream {elsewhere[0]}, not on {self._stream_key}. Append " + "idempotency is scoped to the stream, so settling it here would " + "not deduplicate against the copy that may already be on the " + "other stream -- it would append the value a second time, on a " + "topic no consumer of it is watching, and leave the first " + "stream blocked. Settle it on the topic the error names." + ) + same_key = [ + held + for held in producer._unresolved.get(self._stream_key, ()) + if held.record.idempotency_key == record.idempotency_key + ] + if same_key: + raise ValueError( + f"the unsettled append under {record.idempotency_key} does not " + "hold these bytes. Idempotency is on identity, so re-appending " + "different content under that key is an AppendConflictError " + "rather than the no-op the recovery depends on. Pass the " + "`.record` the error carried, unmodified." + ) + raise ValueError( + f"this producer has no unsettled append under " + f"{record.idempotency_key} on stream {self._stream_key}, so there " + "is nothing here to settle. If the producer that made it is gone, " + "the recovery is not this call: rebuild the producer with the same " + "session id and re-run the same calls in the same order. That " + "re-derives the same sequence numbers and re-appends the same " + "bytes, which the backend deduplicates -- and unlike this call it " + "leaves the new session's own counters correct, where settling here " + "would let its next publish reuse a sequence number and its next " + "unparked wake reuse a request ID (ADR-038)." + ) + + async def resolve_append( + self, + record: StreamRecord, + *, + wake: bool | None = None, + lease: timedelta | None = None, + ) -> Offset: + """Settles an append :class:`AppendNotAcknowledgedError` left unknown. + + Takes the ``.record`` that error carried and re-appends **it**, which is + the whole of why this is not ``publish()`` again. The backend contract + makes the one call right for both possible histories: byte-identical + content under a used ``(session_id, sequence)`` writes nothing and + returns the original offset, and a key the backend never saw is appended + now. Either way the stream ends with exactly one copy of the record and + the caller ends with its offset. + + Must be called on the topic the append was for, from the producer that + made it. Both are checked before the backend is touched; see + :meth:`_outstanding` for what each one prevents. + + **This is where an unknown outcome becomes a known one**, so it is also + where a fence waiting behind that operation is told which one it became. + A durable record clears the operation's failure; an + ``AppendConflictError`` replaces it, because that is the contract's one + definite refusal and it says this record did not land. Leaving the + operation holding its original "unknown" and letting a fence infer + durability from the record no longer being unresolved read both the same + way, and a conflict then released a fence claiming a write that was + never in the stream. + + The wake runs afterwards on exactly the terms :meth:`publish` describes, + so a coordination or Signal failure here raises + :class:`WakeNotAcknowledgedError` carrying that offset. + + Safe to call repeatedly: if this attempt is itself interrupted, it raises + :class:`AppendNotAcknowledgedError` again with the same record, and the + next attempt is the same call. + + Args: + record: The ``.record`` the error carried, unmodified. + wake: Defaults to **what the interrupted call was going to do**, + rather than to ``True``. The recovery finishes that operation, so + inventing a wake policy for it is how a fence appended with + ``wake=False`` acquires a Signal, and how a record published with + ``wake=True`` loses one. Pass a value only to override + deliberately. + lease: Likewise defaults to the interrupted call's. + + Raises: + ValueError: The record is not this topic's outstanding append -- + wrong topic, wrong bytes, wrong session, or nothing outstanding + at all. Raised before any backend call. + + AppendConflictError: The key was used with different bytes, so this + record did not land and re-appending it cannot change that. The + operation is now definitively failed rather than unknown, and a + fence behind it raises :class:`PrecedingWriteFailedError` from + this. + """ + pending = self._outstanding(record) + wake = pending.wake if wake is None else wake + lease = pending.lease if lease is None else lease + try: + placed = await self._append( + record, wake=wake, lease=lease, operation=pending.operation + ) + except AppendConflictError as err: + self._resolve(pending.operation, err) + raise + # Only the append is reported: the fence's claim is about records being + # appended, which this one now is whatever the wake below does. + self._resolve(pending.operation, None) + return await self._wake_for(placed, wake=wake, lease=lease) + + def _resolve( + self, operation: _StreamOperation | None, failure: BaseException | None + ) -> None: + """Replaces an operation's unknown outcome with the recovered one. + + Already ``settled`` -- the call that lost the answer set that before it + raised -- so this only rewrites what a fence reads there. That is why + :meth:`_await_preceding_appends` reads every outcome after its last wait + rather than as each one settles. + """ + if operation is None: + return + operation.failure = failure + operation.settled.set() + + async def publish( + self, + value: AnyType, + *, + wake: bool = True, + lease: timedelta = DEFAULT_WAKE_CLAIM_LEASE, + ) -> Offset: + """Appends one record and completes only once its wake is acknowledged. + + Idempotent under Activity retry through ``(session_id, sequence)``: + re-appending byte-identical content is a no-op returning the original + offset, and the same key with *different* bytes is an error rather than + an overwrite. + + **Returning means the record is durable and a parked Workflow has been + told about it** -- told by *this* call, which sent the wake itself and + had it accepted. It never means that some other producer claimed + responsibility for sending one. That combination is what makes the + "durable producer" row of the wakeup-durability boundary true. An append + that lands but is never signalled leaves the Workflow parked on data + already sitting in the stream, and a ``publish()`` that returned success + there would report a delivery that never happened. + + The wake happens after the append and never instead of it: only + successfully appended records may trigger wakeup, since a wake for a + record that did not land produces a Workflow Task that finds nothing. + + Args: + wake: Set ``False`` to append without waking, then call :meth:`wake` + once for the batch. The wake is idempotent either way -- + every producer waking one generation derives the same request ID + and the server deduplicates -- so this saves round trips rather + than changing the outcome. The record is durable but **un-signalled** until + that call completes, which is the un-acknowledged state made + explicit rather than hidden. + + Raises: + WakeNotAcknowledgedError: The record is durable; the wake is not. + ``.offset`` says where the record landed and ``.pending`` carries + the wakes still owed, so the caller retries the half that failed + rather than re-appending the half that did not. **Every** failure + after the append arrives this way, including one from the + coordination steps that precede the Signal -- those set + ``.restart`` and leave ``.pending`` empty, meaning the recovery is + another :meth:`wake` rather than a :meth:`retry_wake`. This + matters because it is the only thing that tells the caller not to + retry ``publish()``: a second call draws a new sequence number and + therefore a new idempotency key, so it appends the record twice. + + **Cancellation delivered after the append also arrives this way**, + with ``.cancelled`` set, because it leaves the same state and has + the same recovery. A bare ``CancelledError`` there would be + indistinguishable from cancellation before the append -- it carries + no offset, no ``pending`` and no ``restart`` -- so the caller could + neither wake the record nor safely re-publish it (ADR-036). + Cancellation delivered *before* the append is reached -- while + the payload is still encoding -- still raises ``CancelledError``: + nothing was sent to the backend, so there is nothing to say about + a record. + + AppendNotAcknowledgedError: The append itself neither returned nor + refused, so whether the record is durable is **unknown**. This is + its own outcome and not a failure: a backend commits on its own + side before it answers, so a cancellation or a lost connection in + that window can leave a durable record with no offset in hand. + ``.record`` carries the exact bytes and identity to re-append, + and :meth:`resolve_append` is the only safe way forward -- + calling ``publish()`` again draws a new sequence number and + appends the value twice if the first one landed. Until it is + resolved, this stream refuses further appends with the same error + (ADR-038). + + AppendConflictError: The idempotency key was used before with + different bytes. Unlike the above this *is* a refusal: nothing + landed for this record and re-appending it cannot change that. + """ + # The sequence is drawn **before** the encode is awaited, and that + # ordering is the whole of the idempotency key's stability. The key is + # `(session_id, sequence)`, a retried Activity reuses the session id + # deliberately, and the payload codec is allowed to do real I/O -- an + # external payload store, a KMS round trip. Drawing the number after that + # await hands identities out in *encode-completion* order, so two + # concurrent publishes exchange sequence numbers whenever the store + # answers in the other order. On the same stream the backend then sees + # each stable key reused with different bytes and raises + # `AppendConflictError`; across topics, where deduplication is per stream + # key, the swap appends duplicates instead. Either way a valid concurrent + # Activity becomes permanently non-retryable on timing alone. + # + # Drawn here, identities follow *invocation* order, which is the order the + # retry re-runs the same calls in. An encode that raises leaves its number + # unused, which costs nothing: a gap in the sequence is not observable -- + # offsets come from the provider -- and the retry reuses the same number + # for the same call. + self._refuse_while_unresolved() + sequence = self._producer._next_sequence() + # Registered in the same step, and for the same reason the sequence is + # drawn in it: a concurrent `finish_writing()` has to be able to see this + # call before the encode below hands the event loop away, or its fence + # asserts durability for a record that has not been appended yet. + operation = self._begin(sequence) + try: + record = StreamRecord( + kind=RecordKind.DATA, + payload=await self._codec.encode(value), + producer_session_id=self._producer.session_id, + sequence=sequence, + ) + placed = await self._append( + record, wake=wake, lease=lease, operation=operation + ) + except (Exception, asyncio.CancelledError) as err: + # Settled with the failure rather than merely dropped: a fence behind + # this call is waiting on it, and what it does next depends on + # whether a record landed. Settled *before* the raise, so the two are + # one uninterrupted step. + self._settle(operation, err) + raise + self._settle(operation) + return await self._wake_for(placed, wake=wake, lease=lease) + + async def wake( + self, + *, + lease: timedelta = DEFAULT_WAKE_CLAIM_LEASE, + ) -> list[str]: + """Step 2 and 3 of the send sequence: observe or claim, then Signal. + + Only ever called **after** a successful append -- only successfully + appended records may trigger wakeup, since a wake for a record that did + not land produces a Workflow Task that finds nothing. + + For each subscription parked on this stream, the current generation is + claimed under a renewable lease -- and the Signal is then sent **whether + or not the claim was granted**. + + Losing the claim means another producer *intends* to send. It is not + evidence that one did: a lease permits takeover after it expires, but it + schedules nobody to take over, so a producer that crashed between + claiming and signalling strands the generation until some later producer + happens to append again. Staying silent there would leave a parked + Workflow on a record already sitting in the stream while this call + reported an acknowledged wake. The only recovery open to the caller was + to send the wake itself -- exactly what :meth:`retry_wake` does with a + pending request -- so this sends it now instead of reporting a failure + whose only fix is the same send. + + The duplicate that costs is the one that creates a second Workflow Task, + and this is not one: a **parked** wake's request ID is derived from the + generation and ignores sender identity, so racing producers issue + byte-identical requests and the server deduplicates them into a single + wake. A granted claim therefore saves a round trip, not a wakeup. A + caller appending a batch saves many more of them with ``wake=False`` + followed by one :meth:`wake`. + + A provider that cannot lease declares ``supports_leased_claims = False`` + and always grants, which is now the same behaviour every provider gets. + + Subscriptions with no installed intent get an *unparked* wake rather than + nothing (ADR-023). The Workflow may be cached with no open Workflow Task, + or evicted; neither is visible from out here, and staying silent in + either case loses the record until something else happens to wake the + Workflow. + + Returns: + The request ID of each Signal sent -- one per parked subscription, + or a single unparked wake when nothing on this stream is parked. + + Raises: + WakeNotAcknowledgedError: The wake did not complete. The record is + durable; the wake is not. A failed Signal fills ``.pending`` with + the wakes still owed, for :meth:`retry_wake`. A failure in the + observe or claim steps that precede it leaves ``.pending`` empty + and sets ``.restart``, because nothing was composed and calling + this method again is the recovery. **Cancellation is one of the + ways an attempt fails here**, not an exception to it: this method + is only ever reached after a durable append, so cancellation leaves + the same durable-but-unannounced record and gets the same + ``pending``/``restart`` recovery, with ``.cancelled`` set + (ADR-036). + """ + producer = self._producer + if producer._client is None: + raise RuntimeError( + "waking requires a Temporal client, and this producer was built " + "without one. Use ExternalStreamProducer.connect(), which " + "requires it." + ) + + backend = producer._backend + # The coordination steps are inside the same guarantee as the Signal, and + # every caller reaches here *after* a durable append. A provider outage in + # any of them used to escape as whatever the provider raised -- a bare + # `ConnectionError` -- which told the caller nothing about the record that + # had already landed. `publish()` catches only the durable-but- + # unacknowledged error, so the raw exception passed straight through it + # and lost the offset with it; retrying `publish()` then appended a + # *second* record, because the sequence had already advanced and the + # idempotency key with it. + try: + parked = await backend.parked_wait_ids(self._stream_key) + # An unparked wake still needs a wait id for the envelope; 0 is the + # "no particular subscription" value, and Python rechecks every active + # subscription on wakeup regardless of which one the Signal named. + targets: list[tuple[int, int | None]] = [ + ( + wait_id, + await backend.current_park_generation(self._stream_key, wait_id), + ) + for wait_id in parked + ] + if not targets: + targets = [(0, None)] + + requests: list[WakeRequest] = [] + for wait_id, generation in targets: + if generation is not None: + # Claimed, and then signalled whichever way the claim went. The + # claim is how a provider learns a wake is in flight and how an + # abandoned one becomes takeable after its lease, so it is still + # taken -- but its answer is not an acknowledgement, and the + # result is deliberately unused. See this method's docstring. + await backend.claim_park_generation( + self._stream_key, + wait_id, + generation, + claimant=producer.session_id, + lease=lease, + ) + requests.append( + wake_request_for( + producer.workflow, + stream_name=self._stream_key.stream_name, + wait_id=wait_id, + park_generation=generation, + sender_identity=producer.session_id, + wake_counter=( + producer._next_wake_counter() if generation is None else 0 + ), + ) + ) + except WakeNotAcknowledgedError: + raise + except (Exception, asyncio.CancelledError) as err: + # `pending` is empty and `restart` is set: nothing was composed, so + # there is nothing to re-send and `retry_wake` would be a no-op that + # looked like recovery. Calling `wake()` again is the recovery. + # + # Cancellation is in here rather than re-raised ahead of it. Every + # caller reaches this method after a durable append, so a + # `CancelledError` escaping bare says nothing about the record that + # already landed and cannot be told apart from cancellation before the + # append -- the one distinction the caller needs in order to know + # whether the value still has to be published (ADR-036). The + # cancellation is not lost: it is what `cancelled` reports. + raise WakeNotAcknowledgedError( + "the record was appended but its wake could not be composed: " + f"{err!r}. No Signal was sent. Call wake() again -- it re-observes " + "the parked set, and a parked wake's request ID is derived from " + "the generation, so a wake another producer already sent " + "deduplicates against it.", + pending=[], + restart=True, + cancelled=isinstance(err, asyncio.CancelledError), + ) from err + + sent: list[str] = [] + for index, request in enumerate(requests): + try: + sent.append( + await send_wake_signal( + producer._client, + request, + producer_session_id=producer.session_id, + ) + ) + except (Exception, asyncio.CancelledError) as err: + # Cancellation included, for the reason the coordination handler + # above gives: the record is durable either way, and here the + # wakes still owed are known, so the recovery is `retry_wake` with + # `pending` rather than a fresh `wake()`. + raise WakeNotAcknowledgedError( + f"the record was appended but its wake was not acknowledged: " + f"{err!r}. Retrying with the same request is safe -- it derives " + "the same request ID and the server deduplicates it.", + pending=requests[index:], + cancelled=isinstance(err, asyncio.CancelledError), + ) from err + return sent + + async def retry_wake(self, pending: list[WakeRequest]) -> list[str]: + """Re-sends the wakes a failed attempt still owed. + + Takes the requests verbatim rather than recomputing them: recomputing + would draw a fresh wake counter for an unparked wake, derive a different + request ID, and defeat the deduplication that makes the retry safe. + + Refuses an empty list rather than returning quietly. A + :class:`WakeNotAcknowledgedError` raised before any request was composed + carries no pending wakes and sets ``restart``; a caller that fed that + empty list to this method would get a successful-looking no-op while the + record stayed durable and unannounced. There is exactly one recovery from + that state and it is :meth:`wake`. + """ + if not pending: + raise ValueError( + "retry_wake() was given no wakes to re-send. A " + "WakeNotAcknowledgedError with an empty `pending` sets `restart`, " + "which means no Signal was composed and there is nothing to " + "re-send: call wake() again instead." + ) + producer = self._producer + assert producer._client is not None + sent: list[str] = [] + for index, request in enumerate(pending): + try: + sent.append( + await send_wake_signal( + producer._client, + request, + producer_session_id=producer.session_id, + ) + ) + except (Exception, asyncio.CancelledError) as err: + raise WakeNotAcknowledgedError( + f"the wake retry did not complete: {err!r}", + pending=pending[index:], + cancelled=isinstance(err, asyncio.CancelledError), + ) from err + return sent + + async def finish_writing( + self, + *, + wake: bool = True, + lease: timedelta = DEFAULT_WAKE_CLAIM_LEASE, + ) -> Offset: + """Appends an ordered write fence. **Does not close the stream.** + + The fence means only: every write in *this* producer session preceding + it has been appended. It asserts nothing about other producers, and a + later record from one does not violate it -- it simply wakes the + Workflow and consumption resumes. + + A fence on one stream alone does not park the Workflow Task either; the + task parks early only when every active subscription is immediately + parkable. + + **Waits for every earlier ``publish()`` on this stream to reach a durable + append**, across every handle ``topic()`` has returned for the name. The + claim is otherwise not one this method can make: ``publish()`` draws its + sequence before awaiting the payload codec, so a publish invoked first + can still be encoding when this one is called, and a fence that overtook + it would park a consumer in front of data that had not been written -- + having spent, if the publish was a ``wake=False`` member of a batch, the + one wake that would have moved it again. + + It does **not** wait for another concurrent fence. Each fence asserts + durability of the publishes it was invoked after, and neither of two + fences is inside the other's claim, so ordering them would only let one + that never reached the backend refuse the other. + + Raises: + PrecedingWriteFailedError: An earlier ``publish()`` on this stream, + still in flight when this call was made, ended with no durable + record -- including one whose outcome was unknown and which + :meth:`resolve_append` then proved absent. Nothing was appended + for the fence; see that class for what the two recoveries are. + + WakeNotAcknowledgedError: The fence is durable; the wake is not, on + exactly the terms :meth:`publish` describes -- cancellation after + the append included. Retrying ``finish_writing()`` appends a + *second* fence, for the same reason retrying ``publish()`` appends + a second record, and a fence that reads back twice is a producer + session that ended twice. + + AppendNotAcknowledgedError: The fence's append reported no outcome, + again on exactly :meth:`publish`'s terms. :meth:`resolve_append` + takes ``.record`` and settles it, leaving one fence whichever way + the interrupted attempt went. **Also raised for an earlier + append on this stream whose outcome is unknown**, carrying that + operation's record rather than a fence: the record may be durable, + so the fence may not go in front of it, and the stream takes no + further append until it is settled (ADR-038). + """ + # Drawn before the append for the same reason as in `publish`, though a + # fence encodes nothing and so cannot be reordered by a codec: what the + # two share is that the number belongs to the *call*, so a retry that + # makes the same calls in the same order derives the same keys. + self._refuse_while_unresolved() + sequence = self._producer._next_sequence() + # Read in the same uninterrupted step as the draw, and *not* registered + # alongside them: a fence is not a write anything else waits behind, so + # it reads the order without joining it (see `_StreamOperation`). + preceding = self._preceding_publishes() + # The sequence identity is this call's, but the *append* is ordered: the + # fence enters the backend only once every earlier publish on this stream + # has one. + await self._await_preceding_appends(preceding) + fence = StreamRecord( + kind=RecordKind.WRITE_FENCE, + payload=b"", + producer_session_id=self._producer.session_id, + sequence=sequence, + ) + placed = await self._append(fence, wake=wake, lease=lease, operation=None) + # A fence is the record most likely to find the Workflow parked -- it is + # what a producer appends when it has nothing more to say -- so an + # unsignalled one strands the Workflow for its whole idle timeout at + # exactly the moment it was waiting to be told. + return await self._wake_for(placed, wake=wake, lease=lease) + + +async def _verify_chain_key( + client: temporalio.client.Client, workflow: WorkflowChainKey +) -> None: + """Describes the Workflow and checks the first execution Run ID matches.""" + if client.namespace != workflow.namespace: + raise ChainKeyMismatchError( + f"the client is connected to namespace {client.namespace!r} but the " + f"chain key names {workflow.namespace!r}" + ) + + description = await client.get_workflow_handle(workflow.workflow_id).describe() + actual = description.raw_description.workflow_execution_info.first_run_id + if actual != workflow.first_execution_run_id: + raise ChainKeyMismatchError( + f"Workflow {workflow.workflow_id!r} in namespace " + f"{workflow.namespace!r} has first execution Run ID {actual!r}, not " + f"{workflow.first_execution_run_id!r}. Publishing under the wrong key " + "writes records into a stream no consumer is watching, so this is " + "refused rather than accepted silently." + ) diff --git a/temporalio/contrib/external_workflow_streams/_record.py b/temporalio/contrib/external_workflow_streams/_record.py new file mode 100644 index 0000000000..549c5c0b33 --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_record.py @@ -0,0 +1,301 @@ +"""The record model: offsets, cursor boundaries, and stream records (P1). + +Nothing here touches Temporal or a provider. These are the value types every +other piece of the feature is expressed in. +""" + +from __future__ import annotations + +import enum +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Final, Protocol + +__all__ = [ + "AFTER", + "BEGINNING", + "Cursor", + "IdempotencyKey", + "Offset", + "OffsetComparator", + "RecordKind", + "StreamRecord", +] + + +@enum.unique +class RecordKind(enum.IntEnum): + """What a record is, from the runtime's point of view. + + Data and control records share one offset sequence -- a fence occupies an + offset exactly as a data record does, and advances the cursor the same way. + They differ only in whether the runtime yields them to Workflow code. + """ + + DATA = 1 + """Carries a payload for Workflow code.""" + + WRITE_FENCE = 2 + """A producer session declaring its preceding writes are all appended. + + Consumed by the runtime, never yielded to Workflow code. It does not close + the stream and asserts nothing about other producers. + """ + + FINISH = 3 + """An output topic's ordered terminal record. + + Unlike a write fence, this closes the topic across its Continue-As-New + chain. It is explicit because Workflow failure or termination cannot be + relied on to execute cleanup code. + """ + + @property + def is_control(self) -> bool: + """Whether this record advances the cursor without yielding a value.""" + return self is not RecordKind.DATA + + +@dataclass(frozen=True, order=False) +class Offset: + """A provider's position token: opaque, stable, and serializable. + + Deliberately **not** comparable with ``<``. Offsets are totally ordered, + but by their provider's rule rather than lexically -- Redis stream IDs are + ``(milliseconds, sequence)`` tuples, and string comparison is wrong as soon + as the millisecond component changes width. Ordering goes through the + provider's ``compare_offsets``; see :class:`OffsetComparator`. + """ + + token: str + + def __post_init__(self) -> None: + """Reject an empty provider position token.""" + if not self.token: + raise ValueError("an offset token may not be empty") + + def __str__(self) -> str: + """Return the provider position token.""" + return self.token + + def serialize(self) -> str: + """Encode this offset as its opaque provider token.""" + return self.token + + @classmethod + def deserialize(cls, token: str) -> Offset: + """Decode an offset from its opaque provider token.""" + return cls(token) + + +class OffsetComparator(Protocol): + """A provider's total order over its own offsets. + + Returns a negative number, zero, or a positive number, in the manner of a + three-way comparison. + """ + + def __call__(self, left: Offset, right: Offset, /) -> int: + """Compare two offsets using the provider's total order.""" + ... + + +@dataclass(frozen=True, order=False) +class Cursor: + """A position **boundary**, never the identity of a record (ADR-002). + + Two forms only:: + + BEGINNING # the provider's beginning-of-stream boundary + AFTER(offset) # immediately following the record at ``offset`` + + ``AFTER(x)`` names a boundary whether or not a record after ``x`` exists + yet, which is what lets a consumer park at the tail without naming the id + of a record nobody has written. + """ + + offset: Offset | None = None + """``None`` is the beginning-of-stream boundary.""" + + @property + def is_beginning(self) -> bool: + """Whether this cursor is the beginning-of-stream boundary.""" + return self.offset is None + + def __str__(self) -> str: + """Return the cursor in the documented boundary grammar.""" + return "BEGINNING" if self.is_beginning else f"AFTER({self.offset})" + + def serialize(self) -> str: + """Round-trips through :meth:`deserialize`. + + The two forms are distinguished by a prefix rather than by emptiness so + a provider whose beginning sentinel is a real token (Redis uses + ``0-0``) cannot be confused for ``BEGINNING``. + """ + if self.is_beginning: + return "B" + assert self.offset is not None + return f"A{self.offset.serialize()}" + + @classmethod + def deserialize(cls, encoded: str) -> Cursor: + """Decode the beginning or exclusive-after cursor form.""" + if encoded == "B": + return BEGINNING + if encoded.startswith("A"): + return AFTER(Offset.deserialize(encoded[1:])) + raise ValueError(f"not a serialized cursor: {encoded!r}") + + def follows(self, other: Offset, compare: OffsetComparator) -> bool: + """Whether this boundary lies at or after the record at ``other``. + + ``BEGINNING`` follows nothing; ``AFTER(x)`` follows every offset up to + and including ``x``. + """ + if self.is_beginning: + return False + assert self.offset is not None + return compare(other, self.offset) <= 0 + + def excludes(self, candidate: Offset, compare: OffsetComparator) -> bool: + """Whether a record at ``candidate`` sits at or before this boundary. + + This is the predicate a watch applies: it reads records strictly after + the boundary, so anything this excludes has already been consumed. + """ + return self.follows(candidate, compare) + + +BEGINNING: Final[Cursor] = Cursor() +"""The beginning-of-stream boundary.""" + + +def AFTER(offset: Offset) -> Cursor: # noqa: N802 -- reads as the grammar's form + """The boundary immediately following the record at ``offset``.""" + return Cursor(offset) + + +@dataclass(frozen=True) +class IdempotencyKey: + """What makes an append idempotent under Activity retry. + + Idempotency is on ``(session_id, sequence)`` **and identity**: reusing a + key with byte-identical content is a no-op returning the original offset, + and reusing it with different bytes is an error (ADR-020). + """ + + session_id: str + sequence: int + + def __post_init__(self) -> None: + """Require a non-empty session and non-negative sequence.""" + if not self.session_id: + raise ValueError("a producer session id may not be empty") + if self.sequence < 0: + raise ValueError(f"sequence must be non-negative, got {self.sequence}") + + def __str__(self) -> str: + """Return the session and sequence as a diagnostic key.""" + return f"{self.session_id}/{self.sequence}" + + +#: Provider-neutral field names. A provider stores records as fields keyed by +#: these; the prefix keeps them from colliding with anything a provider or an +#: operator adds alongside. +_FIELD_KIND: Final = "__tes_kind" +_FIELD_PAYLOAD: Final = "__tes_payload" +_FIELD_SESSION: Final = "__tes_session" +_FIELD_SEQUENCE: Final = "__tes_seq" + +RESERVED_FIELDS: Final[frozenset[str]] = frozenset( + {_FIELD_KIND, _FIELD_PAYLOAD, _FIELD_SESSION, _FIELD_SEQUENCE} +) + + +@dataclass(frozen=True) +class StreamRecord: + """One appended record, data or control. + + ``offset`` is assigned by the provider at append time and is therefore + absent until then; the producer side builds a record without one and the + provider returns the placed record. + """ + + kind: RecordKind + payload: bytes + producer_session_id: str + sequence: int + offset: Offset | None = field(default=None) + + def __post_init__(self) -> None: + """Validate control payloads and producer sequence numbers.""" + if self.kind.is_control and self.payload: + raise ValueError(f"a {self.kind.name} record carries no payload") + if self.sequence < 0: + raise ValueError(f"sequence must be non-negative, got {self.sequence}") + + @property + def is_control(self) -> bool: + """Control records advance the cursor but are never yielded.""" + return self.kind.is_control + + @property + def idempotency_key(self) -> IdempotencyKey: + """The producer identity under which this record is appended.""" + return IdempotencyKey(self.producer_session_id, self.sequence) + + def placed_at(self, offset: Offset) -> StreamRecord: + """This record as the provider stored it.""" + return StreamRecord( + kind=self.kind, + payload=self.payload, + producer_session_id=self.producer_session_id, + sequence=self.sequence, + offset=offset, + ) + + def to_fields(self) -> dict[str, bytes]: + """The provider-neutral field encoding, excluding the offset. + + The offset is the provider's to assign, so it is not part of what gets + written -- which is also what makes two appends under one idempotency + key byte-comparable. + """ + return { + _FIELD_KIND: str(int(self.kind)).encode(), + _FIELD_PAYLOAD: self.payload, + _FIELD_SESSION: self.producer_session_id.encode(), + _FIELD_SEQUENCE: str(self.sequence).encode(), + } + + @classmethod + def from_fields( + cls, offset: Offset, fields: Mapping[str, bytes] | Mapping[bytes, bytes] + ) -> StreamRecord: + """Rebuild a record a provider stored, tolerating bytes-keyed maps. + + Redis hands back ``bytes`` keys unless the client decodes responses, + and a provider should not have to normalize before calling this. + """ + decoded: dict[str, bytes] = {} + for raw_key, raw_value in fields.items(): + key = raw_key.decode() if isinstance(raw_key, bytes) else raw_key + value = ( + raw_value if isinstance(raw_value, bytes) else str(raw_value).encode() + ) + decoded[key] = value + + missing = RESERVED_FIELDS - decoded.keys() + if missing: + raise ValueError( + f"record at {offset} is missing field(s): {sorted(missing)}" + ) + + return cls( + kind=RecordKind(int(decoded[_FIELD_KIND])), + payload=decoded[_FIELD_PAYLOAD], + producer_session_id=decoded[_FIELD_SESSION].decode(), + sequence=int(decoded[_FIELD_SEQUENCE]), + offset=offset, + ) diff --git a/temporalio/contrib/external_workflow_streams/_redis.py b/temporalio/contrib/external_workflow_streams/_redis.py new file mode 100644 index 0000000000..dd7cfceb1e --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_redis.py @@ -0,0 +1,876 @@ +"""The Redis Streams provider (P3). + +Redis Streams qualifies as a backend because an entry's fields cannot be +rewritten in place. It can be deleted by ``XDEL`` or removed by trimming -- +which is exactly the damage replay's four range checks detect -- but never +altered, which is what makes those four checks sufficient. + +The two reads are **different Redis commands**, not one with a flag: + +- ``XRANGE first last`` is **inclusive of both endpoints**, and is the only + command replay may use. The marker already names the range. +- ``XREAD BLOCK ... STREAMS key id`` returns entries **strictly after** ``id``, + which is exactly the exclusive-after semantics a live watch needs. It cannot + serve a replay read: it would silently drop the range's first record. + +``BEGINNING`` maps to the sentinel ``0-0`` for watching, so a consumer that has +drained to the tail holds ``AFTER()`` and never has to name the id of +a record nobody has written yet. + +Offsets are Redis ids, compared as numeric ``(milliseconds, sequence)`` pairs. +String comparison is wrong the moment the millisecond component changes width. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +from base64 import b64decode, b64encode +from binascii import Error as BinasciiError +from collections.abc import Sequence +from datetime import timedelta +from typing import Any, ClassVar, Final +from urllib.parse import quote + +from temporalio.contrib.external_workflow_streams._backend import ( + DEFAULT_WATCH_BLOCK, + AppendConflictError, + ParkIntent, + ParkIntentRemoval, + StreamBackend, + StreamDirection, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._errors import StreamIntegrityError +from temporalio.contrib.external_workflow_streams._output_backend import ( + OutputReadResult, + OutputStage, + OutputStageConflictError, + OutputStageManifest, + OutputStageNotFoundError, + OutputStageResolutionError, + OutputStageStatus, + OutputStreamBackend, + OutputStreamRecord, + PendingOutputBarrier, + StagedOutputRecord, +) +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Cursor, + Offset, + StreamRecord, +) + +__all__ = ["RedisStreamBackend"] + +DEFAULT_URL: Final = "redis://127.0.0.1:6379" +DEFAULT_KEY_PREFIX: Final = "temporal-external-stream" + +#: Redis' beginning-of-stream sentinel for `XREAD`. Not an offset: no record +#: ever has this id, which is why `BEGINNING` is a distinct cursor form rather +#: than `AFTER(Offset("0-0"))`. +_BEGINNING_SENTINEL: Final = "0-0" +_OUTPUT_STAGE_FIELD: Final = "__tes_output_stage" + + +def _parse(offset: Offset) -> tuple[int, int]: + """A Redis id as the ``(ms, seq)`` pair it actually is.""" + ms, _, seq = offset.token.partition("-") + return int(ms), int(seq or 0) + + +def _content_hash(record: StreamRecord) -> str: + """A stable digest of exactly what would be written. + + Excludes the offset, which the provider assigns -- so a re-append of + identical content hashes identically rather than looking like a conflict. + """ + digest = hashlib.sha256() + for name, value in sorted(record.to_fields().items()): + digest.update(name.encode()) + digest.update(b"\x00") + digest.update(value) + digest.update(b"\x00") + return digest.hexdigest() + + +#: Append-if-new-or-identical, atomically. +#: +#: Split across two commands without a script, a crash between `XADD` and the +#: idempotency write would leave a record no retry could recognise as its own, +#: so the retry would append a duplicate. +_APPEND_LUA: Final = """ +local existing = redis.call('HGET', KEYS[2], ARGV[1]) +if existing then + local sep = string.find(existing, '|') + local offset = string.sub(existing, 1, sep - 1) + local digest = string.sub(existing, sep + 1) + if digest == ARGV[2] then + return {'reused', offset} + end + return {'conflict', ''} +end +local id = redis.call('XADD', KEYS[1], '*', unpack(ARGV, 3)) +redis.call('HSET', KEYS[2], ARGV[1], id .. '|' .. ARGV[2]) +return {'appended', id} +""" + + +_STAGE_OUTPUT_LUA: Final = """ +local existing = redis.call('HGET', KEYS[2], ARGV[1]) +if existing then + if existing ~= ARGV[2] then + return {'conflict', '', ''} + end + return {'reused', redis.call('HGET', KEYS[3], ARGV[1]), redis.call('HGET', KEYS[4], ARGV[1])} +end +local count = tonumber(ARGV[3]) +local cursor = 4 +local offsets = {} +for i = 1, count do + local field_count = tonumber(ARGV[cursor]) + cursor = cursor + 1 + local fields = {} + for j = 1, field_count * 2 do + fields[j] = ARGV[cursor] + cursor = cursor + 1 + end + offsets[i] = redis.call('XADD', KEYS[1], '*', unpack(fields)) +end +local joined = table.concat(offsets, ',') +redis.call('HSET', KEYS[2], ARGV[1], ARGV[2]) +redis.call('HSET', KEYS[3], ARGV[1], 'pending') +redis.call('HSET', KEYS[4], ARGV[1], joined) +return {'staged', 'pending', joined} +""" + + +_RESOLVE_OUTPUT_LUA: Final = """ +local existing = redis.call('HGET', KEYS[1], ARGV[1]) +if not existing then + return {'missing', '', ''} +end +if existing ~= ARGV[2] then + return {'conflict', '', ''} +end +local current = redis.call('HGET', KEYS[2], ARGV[1]) +if current == ARGV[3] then + return {'resolved', current, redis.call('HGET', KEYS[3], ARGV[1])} +end +if current ~= 'pending' then + return {'reversed', current, redis.call('HGET', KEYS[3], ARGV[1])} +end +redis.call('HSET', KEYS[2], ARGV[1], ARGV[3]) +return {'resolved', ARGV[3], redis.call('HGET', KEYS[3], ARGV[1])} +""" + + +#: Compare-and-remove one park intent, atomically. The claim belongs to the +#: intent generation, so a successful removal retires both keys just like the +#: unconditional operation does. +#: +#: Returns 1 removed, 2 nothing installed, 0 someone else's intent. The absent +#: case is answered before the comparison, so a key this script already cleared +#: on a call whose reply was lost does not come back as a mismatch. +_REMOVE_PARK_INTENT_IF_MATCHES_LUA: Final = """ +local run_id = redis.call('HGET', KEYS[1], 'run_id') +if not run_id then + return 2 +end +local generation = redis.call('HGET', KEYS[1], 'generation') +if run_id == ARGV[1] and generation == ARGV[2] then + redis.call('DEL', KEYS[1], KEYS[2]) + return 1 +end +return 0 +""" + + +#: The script's return values, in the order it documents them. +_REMOVAL_OUTCOMES: Final = { + 0: ParkIntentRemoval.MISMATCH, + 1: ParkIntentRemoval.REMOVED, + 2: ParkIntentRemoval.ABSENT, +} + + +class RedisStreamBackend(StreamBackend, OutputStreamBackend): + """Redis Streams as an external workflow stream provider.""" + + guarantees_immutability: ClassVar[bool | None] = True + """`XADD` entries cannot be rewritten in place -- only deleted or trimmed.""" + + provider_id = "redis-streams" + provider_format_version = 1 + supports_leased_claims = True + + def __init__( + self, + *, + url: str = DEFAULT_URL, + client: Any | None = None, + key_prefix: str = DEFAULT_KEY_PREFIX, + ) -> None: + """Create a Redis-backed external workflow stream provider. + + Args: + url: Connection URL, used when ``client`` is not supplied. + client: An existing ``redis.asyncio.Redis``. It **must** have been + created with ``decode_responses=False`` -- payloads are + arbitrary bytes, and a decoding client would corrupt them. + key_prefix: Prepended to every key this backend creates, so one + Redis can serve several deployments. + """ + if client is None: + import redis.asyncio # imported lazily so redis stays optional + + client = redis.asyncio.from_url(url, decode_responses=False) + self._client = client + self._key_prefix = key_prefix + self._append_script = client.register_script(_APPEND_LUA) + self._stage_output_script = client.register_script(_STAGE_OUTPUT_LUA) + self._resolve_output_script = client.register_script(_RESOLVE_OUTPUT_LUA) + self._remove_park_intent_if_matches_script = client.register_script( + _REMOVE_PARK_INTENT_IF_MATCHES_LUA + ) + + # --- key layout --------------------------------------------------------- + + def stream_key(self, key: StreamKey) -> str: + """The physical Redis key for one stream identity. + + Injective: distinct :class:`StreamKey` values always render as distinct + Redis keys, because :func:`_escaped` removes the delimiter from every + component before they are joined. See that function for why joining the + raw fields is not merely untidy but a correctness failure. + + An identity with no reserved character in it renders byte-identically to + the delimiter-joined layout this replaced, and the feature is private and + unreleased, so there is nothing written under the old rendering worth + migrating and no compatibility path here to maintain. + """ + return f"{self._key_prefix}:{_escaped(key)}" + + def _idempotency_key(self, key: StreamKey) -> str: + return f"{self.stream_key(key)}:idem" + + def _intent_key(self, key: StreamKey, wait_id: int) -> str: + # `(stream key, wait_id)`, never the stream alone: two subscriptions to + # one stream are two independent waits. + return f"{self.stream_key(key)}:park:{wait_id}" + + def _claim_key(self, key: StreamKey, wait_id: int) -> str: + return f"{self.stream_key(key)}:claim:{wait_id}" + + def _output_manifest_key(self, key: StreamKey) -> str: + return f"{self.stream_key(key)}:output:manifest" + + def _output_status_key(self, key: StreamKey) -> str: + return f"{self.stream_key(key)}:output:status" + + def _output_offsets_key(self, key: StreamKey) -> str: + return f"{self.stream_key(key)}:output:offsets" + + # --- required operations ------------------------------------------------ + + async def append(self, key: StreamKey, record: StreamRecord) -> StreamRecord: + """Append a record idempotently and return it with its Redis offset.""" + fields = record.to_fields() + args: list[Any] = [ + str(record.idempotency_key).encode(), + _content_hash(record).encode(), + ] + for name, value in sorted(fields.items()): + args.append(name.encode()) + args.append(value) + + outcome, offset = await self._append_script( + keys=[self.stream_key(key), self._idempotency_key(key)], + args=args, + ) + if _text(outcome) == "conflict": + raise AppendConflictError(record.idempotency_key) + return record.placed_at(Offset(_text(offset))) + + async def read_range( + self, key: StreamKey, first: Offset, last: Offset + ) -> list[StreamRecord]: + """Read the inclusive offset range used during replay.""" + # redis-py's overload permits response modes that this client + # configuration cannot produce, so retain the runtime response shape. + entries: Any = await self._client.xrange( + self.stream_key(key), first.serialize(), last.serialize() + ) + return [_to_record(entry_id, fields) for entry_id, fields in entries] + + async def read_after( + self, + key: StreamKey, + after: Cursor, + *, + max_records: int, + block: timedelta | None = DEFAULT_WATCH_BLOCK, + ) -> list[StreamRecord]: + """Read up to ``max_records`` strictly after a live cursor.""" + start = ( + _BEGINNING_SENTINEL if after.is_beginning else after.offset.serialize() # type: ignore[union-attr] + ) + block_ms = None if block is None else int(block.total_seconds() * 1000) + if block_ms is not None and block_ms <= 0: + # `XREAD BLOCK 0` blocks *forever*, which is the opposite of what a + # zero timeout asks for, so a non-blocking read must omit BLOCK. + block_ms = None + + streams: Any = await self._client.xread( + {self.stream_key(key): start}, count=max_records, block=block_ms + ) + if not streams: + return [] + _, entries = streams[0] + return [_to_record(entry_id, fields) for entry_id, fields in entries] + + def compare_offsets(self, left: Offset, right: Offset) -> int: + """Compare Redis stream IDs numerically rather than lexically.""" + a, b = _parse(left), _parse(right) + return (a > b) - (a < b) + + # --- workflow-originated output --------------------------------------- + + async def stage_output( + self, + manifest: OutputStageManifest, + records: Sequence[StagedOutputRecord], + ) -> OutputStage: + """Atomically append an immutable pending output sub-batch.""" + self._validate_output_manifest(manifest) + if len(records) != manifest.record_count or any( + record.publish_index != index for index, record in enumerate(records) + ): + raise ValueError( + "staged output records must be contiguous from publish index zero " + "and match the manifest record count" + ) + stage_id = _output_stage_id(manifest) + encoded_manifest = _encode_output_manifest(manifest) + args: list[Any] = [stage_id, encoded_manifest, len(records)] + for record in records: + stream_record = StreamRecord( + kind=record.kind, + payload=record.payload, + producer_session_id=f"stage:{manifest.stage_token}", + sequence=record.publish_index, + ) + fields = { + **stream_record.to_fields(), + _OUTPUT_STAGE_FIELD: stage_id.encode(), + } + args.append(len(fields)) + for name, value in sorted(fields.items()): + args.extend((name.encode(), value)) + outcome, status, offsets = await self._stage_output_script( + keys=[ + self.stream_key(manifest.stream_key), + self._output_manifest_key(manifest.stream_key), + self._output_status_key(manifest.stream_key), + self._output_offsets_key(manifest.stream_key), + ], + args=args, + ) + if _text(outcome) == "conflict": + raise OutputStageConflictError(manifest) + return await self._output_stage_from_offsets( + manifest, + _output_stage_status(status), + _text(offsets), + ) + + async def commit_output(self, manifest: OutputStageManifest) -> OutputStage: + """Promote an exact pending stage, idempotently.""" + return await self._resolve_output(manifest, OutputStageStatus.COMMITTED) + + async def abort_output(self, manifest: OutputStageManifest) -> OutputStage: + """Skip an exact pending stage, idempotently.""" + return await self._resolve_output(manifest, OutputStageStatus.ABORTED) + + async def _resolve_output( + self, + manifest: OutputStageManifest, + requested: OutputStageStatus, + ) -> OutputStage: + self._validate_output_manifest(manifest) + outcome, current, offsets = await self._resolve_output_script( + keys=[ + self._output_manifest_key(manifest.stream_key), + self._output_status_key(manifest.stream_key), + self._output_offsets_key(manifest.stream_key), + ], + args=[ + _output_stage_id(manifest), + _encode_output_manifest(manifest), + requested.value, + ], + ) + result = _text(outcome) + if result == "missing": + raise OutputStageNotFoundError(manifest) + if result == "conflict": + raise OutputStageConflictError(manifest) + if result == "reversed": + raise OutputStageResolutionError( + manifest, + current=_output_stage_status(current), + requested=requested, + ) + return await self._output_stage_from_offsets( + manifest, + requested, + _text(offsets), + ) + + async def output_stage(self, manifest: OutputStageManifest) -> OutputStage | None: + """Inspect one exact staged sub-batch.""" + self._validate_output_manifest(manifest) + stage_id = _output_stage_id(manifest) + stored, status, offsets = await _gather_redis( + self._client.hget(self._output_manifest_key(manifest.stream_key), stage_id), + self._client.hget(self._output_status_key(manifest.stream_key), stage_id), + self._client.hget(self._output_offsets_key(manifest.stream_key), stage_id), + ) + if stored is None: + return None + if stored != _encode_output_manifest(manifest): + raise OutputStageConflictError(manifest) + if status is None or offsets is None: + raise StreamIntegrityError( + "an output stage is missing its status or offsets" + ) + return await self._output_stage_from_offsets( + manifest, + _output_stage_status(status), + _text(offsets), + ) + + async def append_output( + self, key: StreamKey, record: StreamRecord + ) -> OutputStreamRecord: + """Append an immediately committed Activity/external record.""" + self._require_output_key(key) + placed = await self.append(key, record) + assert placed.offset is not None + return OutputStreamRecord(placed.kind, placed.payload, placed.offset) + + async def read_output_after( + self, + key: StreamKey, + after: Cursor, + *, + max_records: int, + block: timedelta | None = DEFAULT_WATCH_BLOCK, + ) -> OutputReadResult: + """Return a committed prefix without crossing a pending stage.""" + self._require_output_key(key) + if max_records < 1: + raise ValueError("max_records must be positive") + cursor = after + first_read = True + committed: list[OutputStreamRecord] = [] + while len(committed) < max_records: + start = ( + _BEGINNING_SENTINEL + if cursor.is_beginning + else cursor.offset.serialize() # type: ignore[union-attr] + ) + block_for = block if first_read else timedelta(0) + block_ms = ( + None if block_for is None else int(block_for.total_seconds() * 1000) + ) + if block_ms is not None and block_ms <= 0: + block_ms = None + streams: Any = await self._client.xread( + {self.stream_key(key): start}, + count=max(64, max_records - len(committed)), + block=block_ms, + ) + first_read = False + if not streams: + break + _, entries = streams[0] + if not entries: + break + stage_ids = sorted( + { + _text(fields[_OUTPUT_STAGE_FIELD.encode()]) + for _, fields in entries + if _OUTPUT_STAGE_FIELD.encode() in fields + } + ) + statuses = ( + dict( + zip( + stage_ids, + await self._client.hmget( + self._output_status_key(key), list(stage_ids) + ), + ) + ) + if stage_ids + else {} + ) + for entry_id, fields in entries: + offset = Offset(_text(entry_id)) + cursor = AFTER(offset) + raw_stage_id = fields.get(_OUTPUT_STAGE_FIELD.encode()) + if raw_stage_id is not None: + stage_id = _text(raw_stage_id) + raw_status = statuses.get(stage_id) + if raw_status is None: + raise StreamIntegrityError( + f"output record at {offset} has no stage status" + ) + status = _output_stage_status(raw_status) + if status is OutputStageStatus.PENDING: + stored = await self._client.hget( + self._output_manifest_key(key), stage_id + ) + if stored is None: + raise StreamIntegrityError( + f"pending output record at {offset} has no manifest" + ) + return OutputReadResult( + tuple(committed), + PendingOutputBarrier( + _decode_output_manifest(key, _text(stored)), offset + ), + ) + if status is OutputStageStatus.ABORTED: + continue + record = _to_record(entry_id, fields) + assert record.offset is not None + committed.append( + OutputStreamRecord(record.kind, record.payload, record.offset) + ) + if len(committed) == max_records: + break + if len(entries) < max(64, max_records - len(committed)): + break + return OutputReadResult(tuple(committed)) + + async def output_tail(self, key: StreamKey) -> Cursor: + """Return the readable boundary without moving past pending output.""" + self._require_output_key(key) + cursor = BEGINNING + while True: + result = await self.read_output_after( + key, + cursor, + max_records=256, + block=timedelta(0), + ) + if result.records: + cursor = AFTER(result.records[-1].offset) + if result.pending is not None or not result.records: + return cursor + + def _validate_output_manifest(self, manifest: OutputStageManifest) -> None: + self._require_output_key(manifest.stream_key) + if manifest.provider_id != self.provider_id or ( + manifest.provider_format_version != self.provider_format_version + ): + raise ValueError("an output manifest names a different provider format") + + @staticmethod + def _require_output_key(key: StreamKey) -> None: + if key.direction is not StreamDirection.OUTPUT: + raise ValueError("an output operation requires an OUTPUT stream key") + + async def _output_stage_from_offsets( + self, + manifest: OutputStageManifest, + status: OutputStageStatus, + encoded_offsets: str, + ) -> OutputStage: + offsets = [Offset(value) for value in encoded_offsets.split(",") if value] + if len(offsets) != manifest.record_count: + raise StreamIntegrityError( + "an output stage's stored offsets do not match its manifest" + ) + entries: Any = await self._client.xrange( + self.stream_key(manifest.stream_key), + offsets[0].serialize(), + offsets[-1].serialize(), + ) + by_offset = {_text(entry_id): fields for entry_id, fields in entries} + records = [] + for offset in offsets: + fields = by_offset.get(offset.serialize()) + if fields is None: + raise StreamIntegrityError( + f"staged output record at {offset} is missing" + ) + record = StreamRecord.from_fields(offset, fields) + records.append(OutputStreamRecord(record.kind, record.payload, offset)) + return OutputStage(manifest, tuple(records), status) + + # --- parking (P3b) ------------------------------------------------------ + + async def install_park_intent(self, key: StreamKey, intent: ParkIntent) -> None: + """Persist the intent for one Workflow wait to remain parked.""" + await self._client.hset( + self._intent_key(key, intent.wait_id), + mapping={ + b"cursor": intent.cursor.serialize().encode(), + b"generation": str(intent.park_generation).encode(), + b"run_id": intent.run_id.encode(), + }, + ) + + async def remove_park_intent(self, key: StreamKey, wait_id: int) -> None: + """Remove a wait's park intent and any associated producer claim.""" + await self._client.delete( + self._intent_key(key, wait_id), self._claim_key(key, wait_id) + ) + + async def remove_park_intent_if_matches( + self, + key: StreamKey, + wait_id: int, + *, + run_id: str, + park_generation: int, + ) -> ParkIntentRemoval: + """Remove a park intent only when its Run and generation still match.""" + outcome = await self._remove_park_intent_if_matches_script( + keys=[self._intent_key(key, wait_id), self._claim_key(key, wait_id)], + args=[run_id, str(park_generation)], + ) + return _REMOVAL_OUTCOMES[int(outcome)] + + async def park_intent(self, key: StreamKey, wait_id: int) -> ParkIntent | None: + """Return the currently installed park intent for one wait, if any.""" + stored = await self._client.hgetall(self._intent_key(key, wait_id)) + if not stored: + return None + stored = {_text(name): value for name, value in stored.items()} + return ParkIntent( + wait_id=wait_id, + cursor=Cursor.deserialize(_text(stored["cursor"])), + park_generation=int(_text(stored["generation"])), + run_id=_text(stored["run_id"]), + ) + + async def recheck(self, key: StreamKey, wait_id: int) -> bool: + """Return whether a parked wait now has at least one available record.""" + intent = await self.park_intent(key, wait_id) + if intent is None: + return False + found = await self.read_after(key, intent.cursor, max_records=1, block=None) + return bool(found) + + async def claim_park_generation( + self, + key: StreamKey, + wait_id: int, + park_generation: int, + *, + claimant: str, + lease: timedelta, + ) -> bool: + """Acquire or renew the leased producer claim for one park generation.""" + claim_key = self._claim_key(key, wait_id) + value = f"{park_generation}|{claimant}".encode() + lease_ms = max(1, int(lease.total_seconds() * 1000)) + + # SET NX gives exclusion; the TTL is what makes a crashed producer's + # claim recoverable instead of stranding the generation forever. + if await self._client.set(claim_key, value, nx=True, px=lease_ms): + return True + + held = await self._client.get(claim_key) + if held is None: + # It expired between the SET and the GET -- take it. + return bool(await self._client.set(claim_key, value, nx=True, px=lease_ms)) + if held == value: + # Renewal by the holder, which must not read as contention. + await self._client.pexpire(claim_key, lease_ms) + return True + + held_generation = _text(held).split("|", 1)[0] + if held_generation != str(park_generation): + # A claim for a different generation says nothing about this one. + await self._client.set(claim_key, value, px=lease_ms) + return True + return False + + async def parked_wait_ids(self, key: StreamKey) -> list[int]: + """Return all wait IDs with installed park intents for a stream.""" + prefix = f"{self.stream_key(key)}:park:" + found = [] + # SCAN rather than KEYS: this runs on the producer's hot path after every + # append, and KEYS blocks the whole server for the length of the keyspace. + # + # The prefix is a *literal* here, so it is escaped before the trailing + # `*` makes it a pattern. `_escaped` already leaves no glob character in + # the identity half, but the operator-supplied `key_prefix` is not + # escaped, and a pattern that quietly widened would enumerate another + # stream's intents rather than fail. + async for name in self._client.scan_iter(match=f"{_as_glob_literal(prefix)}*"): + suffix = _text(name)[len(prefix) :] + if suffix.isdigit(): + found.append(int(suffix)) + return sorted(found) + + async def current_park_generation(self, key: StreamKey, wait_id: int) -> int | None: + """Return the currently parked generation for one wait, if any.""" + intent = await self.park_intent(key, wait_id) + return None if intent is None else intent.park_generation + + # --- lifecycle ---------------------------------------------------------- + + async def aclose(self) -> None: + """Close the underlying Redis client.""" + await self._client.aclose() + + async def delete_for_test(self, key: StreamKey, offset: Offset) -> None: + """Removes one record, standing in for trimming or retention expiry.""" + await self._client.xdel(self.stream_key(key), offset.serialize()) + + +#: What Redis' glob matcher treats as more than itself. `]` and `^` are special +#: only inside a class, but escaping them too costs nothing and keeps the rule +#: one line long. +_GLOB_METACHARACTERS: Final = frozenset("*?[]\\") + + +def _as_glob_literal(text: str) -> str: + """A literal string, made safe to embed in a `SCAN MATCH` pattern.""" + return "".join( + f"\\{character}" if character in _GLOB_METACHARACTERS else character + for character in text + ) + + +def _escaped(key: StreamKey) -> str: + r"""One stream identity as a single, unambiguous key component. + + Percent-encoded per field, then joined -- **not** joined raw. A Workflow ID + and a stream name are user-chosen strings in which `:` is an ordinary + character, so joining the raw fields is not injective: + + ("ns", "wf", r1, f"{r2}:tokens") and ("ns", f"wf:{r1}", r2, "tokens") + + both render as `ns:wf:r1:r2:tokens`. Two unrelated Workflows would then share + one stream, one idempotency hash, one park intent and one claim -- delivering + each other's records, and each concluding the other's claim had already taken + its wake. Same reasoning as `_wake.py`'s length-prefixed request-ID material, + applied to a key rather than to a digest. + + Percent-encoding rather than length prefixes because a key is read by humans: + an ordinary identity still renders verbatim in `redis-cli`, and only a field + that actually contains a delimiter pays for it. It buys one property the + length prefix does not -- the encoded form contains no `:`, `*`, `?`, `[` or + `\`, so the derived `:idem`, `:park:` and `:claim:` suffixes stay + unambiguous and `parked_wait_ids`' pattern cannot be widened by a stream name. + + Not reversible in practice, and not meant to be: `key_prefix` is + operator-supplied and unescaped, so only the identity half round-trips. + """ + components: tuple[str, ...] = ( + key.namespace, + key.workflow_id, + key.first_execution_run_id, + key.stream_name, + ) + if key.direction is StreamDirection.OUTPUT: + # Existing input records must remain addressable under their original + # four-component keys. Adding a reserved fifth-component layout only + # for output is still injective because every user component is escaped. + components = ( + key.namespace, + key.workflow_id, + key.first_execution_run_id, + key.direction.value, + key.stream_name, + ) + return ":".join(quote(component, safe="") for component in components) + + +def _output_stage_id(manifest: OutputStageManifest) -> str: + """An injective Redis hash field for one token/topic sub-batch.""" + return f"{len(manifest.stage_token)}:{manifest.stage_token}:{manifest.sub_batch_id}" + + +def _output_stage_status(value: Any) -> OutputStageStatus: + """Decode coordination metadata without downgrading corruption to outage.""" + try: + return OutputStageStatus(_text(value)) + except (TypeError, ValueError) as err: + raise StreamIntegrityError( + f"an output stage has an invalid stored status: {_text(value)!r}" + ) from err + + +def _encode_output_manifest(manifest: OutputStageManifest) -> bytes: + """Stable immutable metadata compared by the atomic stage script.""" + return json.dumps( + { + "provider_id": manifest.provider_id, + "provider_format_version": manifest.provider_format_version, + "stage_token": manifest.stage_token, + "run_id": manifest.run_id, + "history_floor_event_id": manifest.history_floor_event_id, + "sub_batch_id": manifest.sub_batch_id, + "fingerprint_version": manifest.fingerprint_version, + "fingerprint": b64encode(manifest.fingerprint).decode("ascii"), + "record_count": manifest.record_count, + "logical_byte_count": manifest.logical_byte_count, + }, + sort_keys=True, + separators=(",", ":"), + ).encode() + + +def _decode_output_manifest( + key: StreamKey, encoded: str | bytes +) -> OutputStageManifest: + """Rebuild reconciliation metadata stored beside a pending barrier.""" + try: + values = json.loads(encoded) + return OutputStageManifest( + stream_key=key, + provider_id=values["provider_id"], + provider_format_version=int(values["provider_format_version"]), + stage_token=values["stage_token"], + run_id=values["run_id"], + history_floor_event_id=int(values["history_floor_event_id"]), + sub_batch_id=int(values["sub_batch_id"]), + fingerprint_version=int(values["fingerprint_version"]), + fingerprint=b64decode(values["fingerprint"], validate=True), + record_count=int(values["record_count"]), + logical_byte_count=int(values["logical_byte_count"]), + ) + except ( + BinasciiError, + KeyError, + TypeError, + ValueError, + json.JSONDecodeError, + ) as err: + raise StreamIntegrityError( + "an output stage has invalid stored reconciliation metadata" + ) from err + + +async def _gather_redis(*awaitables: Any) -> tuple[Any, ...]: + """Retain precise tuple shape for the three independent hash reads.""" + return tuple(await asyncio.gather(*awaitables)) + + +def _text(value: Any) -> str: + return value.decode() if isinstance(value, (bytes, bytearray)) else str(value) + + +def _to_record(entry_id: Any, fields: Any) -> StreamRecord: + return StreamRecord.from_fields(Offset(_text(entry_id)), fields) diff --git a/temporalio/contrib/external_workflow_streams/_replay.py b/temporalio/contrib/external_workflow_streams/_replay.py new file mode 100644 index 0000000000..f136c9cbb2 --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_replay.py @@ -0,0 +1,264 @@ +"""The replay read path (P13). + +Replay reads the **exact ranges the marker recorded** and delivers them from +memory in the recorded order. It never asks the backend what comes next: the +answer is already in the annotation, and consulting current stream timing would +be reproducing something other than what happened. + +Three kinds of waiting exist, and only one of them occurs here: + +=============================================== ========================== +Waiting for *new* records -- watchers, idle Never. Core starts no +timer, park handshake timers and no watcher runs. +Reading *recorded* offsets Occurs. Ordinary blocking + provider I/O. +Blocking on backend latency or unavailability Not a determinism failure -- + a transient storage failure. +=============================================== ========================== + +**The four range checks are the whole of validation**, and they are sufficient +precisely because every provider guarantees a record's bytes cannot change +(ADR-003). Given that, the only damage replay has to detect is a record that is +no longer *there* -- and deletion, trimming, and retention expiry are all caught +by one of the four: + +1. both endpoints present, +2. the range contains exactly ``count`` records, +3. ordering is strictly increasing under the provider's comparator, +4. ``control_positions`` match. + +A first, middle, or last deletion each fails a different one of them. + +What these checks are **not** for is a marker that names a subscription this +Workflow never created. That is row four of the taxonomy -- ordinary +nondeterminism -- because nothing is wrong with the backend: the ranges are +exactly where they were written and the Workflow code changed underneath them. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import temporalio.workflow +from temporalio.contrib.external_workflow_streams._annotation import ( + Annotation, + Run, + decode_annotation, +) +from temporalio.contrib.external_workflow_streams._backend import ( + StreamBackend, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._errors import ( + StreamIntegrityError, + StreamStorageError, +) +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + Cursor, + StreamRecord, +) + +__all__ = ["ReplayPlan", "ReplaySegment", "build_replay_plan", "validate_run"] + + +@dataclass(frozen=True) +class ReplaySegment: + """One original activation's deliveries, in the order they happened. + + Segments are kept apart rather than concatenated because each one was a + separate event-loop drain. Collapsing them would reproduce the record order + while changing how many drains occurred, and ``wait_condition`` predicates + would then fire a different number of times (ADR-018). + """ + + deliveries: tuple[tuple[int, StreamRecord], ...] + + +@dataclass +class ReplayPlan: + """Everything one marker's replay needs, already read and validated.""" + + annotation: Annotation + segments: list[ReplaySegment] = field(default_factory=list) + + @property + def total_records(self) -> int: + """The number of records prepared across all replay segments.""" + return sum(len(s.deliveries) for s in self.segments) + + @property + def committed_boundaries(self) -> dict[int, Cursor]: + """Where this marker left each wait it recorded anything for. + + What live delivery has to resume from once replay has handed the + recorded ranges over. The terminal is the authority, and every marker + has one: Core asks for a terminal on the boundaries it decides, and the + completion path supplies one on the boundaries Python decides -- a + Workflow Task that ended carrying server-bound commands included + (ADR-008). + + The fallback to the last recorded delivery per wait is kept for a marker + that somehow arrives without one, and it is not an approximation: both + name the last record the Workflow Task handed over. What only the + terminal can say is where a wait that received *nothing* in its final + activation stopped. + + A wait the marker recorded nothing for is deliberately absent rather + than mapped to its start cursor. Nothing was delivered for it, so there + is nothing to resume past. + """ + boundaries: dict[int, Cursor] = {} + for segment in self.annotation.segments: + for run in segment.runs: + boundaries[run.wait_id] = AFTER(run.last_offset) + # The terminal wins wherever it exists: it is where deliveries stopped + # at the end of the Workflow Task, which can be past the last run for a + # wait whose final activation delivered nothing. + for wait_id, cursor in (self.annotation.terminal or {}).items(): + boundaries[wait_id] = cursor + return boundaries + + +def validate_run(run: Run, records: list[StreamRecord], backend: StreamBackend) -> None: + """Applies the four range checks. Raises :class:`StreamIntegrityError`. + + **Integrity loss must never resolve to an alternate stream result.** Every + check below reports what is wrong rather than repairing it, because + substituting a later record for a deleted one would hand Workflow code a + different history than the one its commands were derived from -- and that + would surface much later as an unrelated nondeterminism error. + """ + offsets = [r.offset for r in records] + if any(o is None for o in offsets): + raise StreamIntegrityError( + f"the provider returned a record with no offset for range " + f"[{run.first_offset}, {run.last_offset}]" + ) + + # 1. Both endpoints present. A first- or last-record deletion fails here, + # and nothing else would catch it: the count check alone cannot tell a + # missing endpoint from a missing interior record. + if not records or records[0].offset != run.first_offset: + raise StreamIntegrityError( + f"the record at {run.first_offset} is missing from the stream; the " + f"marker records it as the first of a run of {run.count}" + ) + if records[-1].offset != run.last_offset: + raise StreamIntegrityError( + f"the record at {run.last_offset} is missing from the stream; the " + f"marker records it as the last of a run of {run.count}" + ) + + # 2. Exact count. An interior deletion fails here. + if len(records) != run.count: + raise StreamIntegrityError( + f"range [{run.first_offset}, {run.last_offset}] contains " + f"{len(records)} record(s) but the marker records {run.count}; " + "records have been deleted, trimmed, or have expired" + ) + + # 3. Strictly increasing under the *provider's* comparator, never lexically. + if not backend.strictly_increasing(offsets): # type: ignore[arg-type] + raise StreamIntegrityError( + f"range [{run.first_offset}, {run.last_offset}] did not read back in " + "strictly increasing offset order" + ) + + # 4. Control positions. A control record replaced by a data record at the + # same position would otherwise be delivered to Workflow code, which never + # saw it the first time. + actual = tuple(i for i, record in enumerate(records) if record.is_control) + if actual != tuple(run.control_positions): + raise StreamIntegrityError( + f"range [{run.first_offset}, {run.last_offset}] has control records at " + f"{actual} but the marker records them at {tuple(run.control_positions)}" + ) + + +async def build_replay_plan( + annotation_bytes: bytes, + backends: dict[int, StreamBackend], + stream_keys: dict[int, StreamKey], +) -> ReplayPlan: + """Reads and validates every recorded range, up front. + + All of it before any delivery, deliberately: a range that fails validation + must stop replay before Workflow code has seen anything from the marker, + rather than part-way through with some records already delivered. + + Args: + annotation_bytes: The marker's opaque annotation. + backends: The provider for each ``wait_id`` in the annotation. + stream_keys: The stream each ``wait_id`` was subscribed to. + + Raises: + StreamStorageError: The backend was unreachable or errored. Transient; + it clears when the backend recovers. + StreamIntegrityError: A recorded range no longer reads back as written. + Keeps failing until an operator repairs the backend. + """ + annotation = decode_annotation(annotation_bytes) + plan = ReplayPlan(annotation=annotation) + + # One read per run rather than one per record: replay I/O cost is a function + # of the consumed range, not of how the batch was segmented. + for segment in annotation.segments: + deliveries: list[tuple[int, StreamRecord]] = [] + for run in segment.runs: + backend = backends.get(run.wait_id) + key = stream_keys.get(run.wait_id) + if key is not None and backend is None: + # The wait is in the annotation, so the Workflow did create it. + # The deployment is missing the store needed to read it. + raise StreamStorageError( + f"the marker records external stream wait {run.wait_id} " + "but no external stream backend is configured on this " + "Worker; pass external_stream_backend=... to the Worker" + ) + if backend is None or key is None: + # Row four of the failure taxonomy, and **not** integrity loss. + # Nothing is wrong with the backend: the recorded ranges are + # exactly where they were written, and the Workflow code has + # changed underneath them. Reporting this as integrity loss + # would send an operator to repair a backend that is fine, when + # the fix is to version the Workflow code. + # + # Reported rather than skipped either way, because skipping + # would silently deliver a different stream result. + raise temporalio.workflow.NondeterminismError( + f"the marker records external stream wait {run.wait_id}, " + "which this Workflow did not create. A subscribe() call was " + "inserted, removed, or reordered, which renumbers every " + "later wait; gate the change behind workflow.patched() " + "exactly as an inserted timer would be." + ) + records = await _read_range(backend, key, run) + validate_run(run, records, backend) + deliveries.extend((run.wait_id, record) for record in records) + plan.segments.append(ReplaySegment(tuple(deliveries))) + + return plan + + +async def _read_range( + backend: StreamBackend, key: StreamKey, run: Run +) -> list[StreamRecord]: + """The inclusive read, with provider failures classified as transient. + + ``read_range`` and not ``read_after``: the range is closed on both ends and + the marker already names it. An exclusive read here would silently drop each + run's first record, and nothing before the first replay would notice. + """ + try: + return await backend.read_range(key, run.first_offset, run.last_offset) + except StreamStorageError: + raise + except Exception as err: + # Not integrity loss: nothing has been shown to be missing, only + # unreadable right now. An operator sent to repair a backend that was + # merely unreachable would find nothing wrong with it. + raise StreamStorageError( + f"could not read recorded range [{run.first_offset}, {run.last_offset}] " + f"from {key}: {err}" + ) from err diff --git a/temporalio/contrib/external_workflow_streams/_runtime.py b/temporalio/contrib/external_workflow_streams/_runtime.py new file mode 100644 index 0000000000..d3e0914454 --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_runtime.py @@ -0,0 +1,2110 @@ +"""The per-Workflow-instance runtime handle (P10a, P10b). + +One of these is created per Run and handed to Workflow code as an opaque handle. +It is the only thing that crosses the sandbox boundary, and it exposes only +stream operations. No provider instance is reachable through it. + +It also builds the **observation delta** -- the thing that makes replay possible +-- because it is the only component that sees deliveries in the order Workflow +code actually received them. Neither the manager (which sees buffering order) +nor Core (which is annotation-blind) can reconstruct that. + +Two rules here are easy to get subtly wrong: + +- **Emission is not conditional on records having been consumed.** A + subscription's first observation, an activation that drained nothing, and the + boundary an activation returned on *all* produce deltas. A subscription to an + empty stream must still record its provider, stream key, and start cursor, or + replay has no starting point at all. +- **Control records are recorded even though they are never yielded.** They + occupy offsets inside a run, so a run's ``count`` includes them and + ``control_positions`` marks which relative indices they were. Leaving them out + would make every replay range read find more records than the marker claims. +""" + +from __future__ import annotations + +import asyncio +import secrets +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import timedelta +from typing import Any + +import temporalio.api.common.v1 +import temporalio.converter +import temporalio.workflow +from temporalio.contrib.external_workflow_streams._annotation import ( + MAX_ANNOTATION_BYTES, + AnnotationAccumulator, + AnnotationHeader, + Run, + Segment, + SegmentEndReason, + StreamBinding, + encode_bindings, + encode_header, + encode_segment, + encode_terminal, + encoded_run_size, + encoded_segment_size, +) +from temporalio.contrib.external_workflow_streams._api import ( + MAX_RECORDS_PER_ACTIVATION, +) +from temporalio.contrib.external_workflow_streams._backend import ( + StreamBackend, + StreamDirection, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec +from temporalio.contrib.external_workflow_streams._continuation import Continuation +from temporalio.contrib.external_workflow_streams._errors import ( + ExternalStreamCapacityError, + StreamIntegrityError, + StreamStorageError, +) +from temporalio.contrib.external_workflow_streams._manager import ( + StreamSubscriptionManager, +) +from temporalio.contrib.external_workflow_streams._output_backend import ( + OutputStageConflictError, + OutputStageManifest, + OutputStageStatus, + OutputStreamBackend, + StagedOutputRecord, +) +from temporalio.contrib.external_workflow_streams._output_codec import ( + LOGICAL_FINGERPRINT_VERSION, + canonical_logical_record_frame, + fingerprint_logical_frames, +) +from temporalio.contrib.external_workflow_streams._output_continuation import ( + OutputContinuation, +) +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Cursor, + RecordKind, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._replay import ReplayPlan + +__all__ = ["QuiescentWait", "StagedWorkflowOutput", "WorkflowStreamRuntime"] + +_RUN_COST_FLOOR = 64 +"""The per-record annotation cost assumed until a real run has been measured. + +Pessimistic on purpose, and only ever a *starting* price: the largest run the +current annotation has actually encoded replaces it as soon as there is one, and +that measurement survives the segment it was taken in. A floor alone is not a +bound -- a provider chooses how long its offsets are, and a run costs two of +them -- which is why :data:`_SEGMENT_SPILL_CAP` exists behind it. + +Well above what a typical provider's two offsets cost (a Redis stream ID is +fifteen bytes), so in the ordinary case this bounds nothing: the per-activation +record cap is reached long first. +""" + +_SEGMENT_SPILL_CAP = 4096 +"""The largest margin a segment frame may overrun the affordability line into. + +The runtime prices a record it has not seen yet, so it can misprice one -- and a +segment frame cannot be moved to the next annotation, because its deliveries +happened in this Workflow Task and that task's marker is where replay has to find +them. The margin is what makes a misprice a *rollover* instead of a refused +frame: the segment spills into it, the same completion asks Core to end the +Workflow Task, and the terminal is still reserved behind it. + +Capped rather than a plain fraction so that a large budget does not hold back +kilobytes it will never need; floored by the fraction so that a small budget -- +which only tests use -- still has a margin at all. +""" + +_REPLAY_UNBOUNDED = 2**63 - 1 +"""The budget reported while a recorded segment is being delivered. + +A number rather than ``None`` so that every caller keeps one code path: a +segment is finite and already recorded, so "as many as the segment holds" and +"no limit" are the same answer. +""" + +_OUTPUT_MANIFEST_SCHEMA_VERSION = 1 +_MAX_OUTPUT_MANIFEST_BYTES = 64 * 1024 +"""Hard ceiling for output metadata placed in the compact History marker.""" + + +@dataclass(frozen=True) +class QuiescentWait: + """One member of the complete blocked snapshot Core is asked to retain for.""" + + wait_id: int + generation: int + immediately_parkable: bool + + +@dataclass(frozen=True) +class _LogicalOutputRecord: + topic: str + kind: RecordKind + payload: temporalio.api.common.v1.Payload | None + frame: bytes + + +@dataclass(frozen=True) +class StagedOutputTopic: + """One topic sub-batch staged under a Workflow Task's shared token.""" + + manifest: OutputStageManifest + finished: bool + + +@dataclass(frozen=True) +class StagedWorkflowOutput: + """Compact marker material for one durably staged Workflow output batch.""" + + schema_version: int + fingerprint_version: int + stage_token: str + history_floor_event_id: int + run_id: str + provider_id: str + provider_format_version: int + topics: tuple[StagedOutputTopic, ...] + segment_record_counts: tuple[tuple[int, ...], ...] + + +@dataclass +class _SubscriptionState: + """What the runtime tracks per subscription, on the Workflow thread.""" + + wait_id: int + stream_key: StreamKey + start_cursor: Cursor + idle_timeout: timedelta + + delivery_cursor: Cursor = BEGINNING + """How far records have been handed to the *subscription's* buffer. + + What the annotation records, and what replay must reproduce. + """ + consumption_cursor: Cursor = BEGINNING + """How far Workflow code has actually taken records. + + Trails :attr:`delivery_cursor` by whatever is still sitting unread in the + buffer. That buffer dies with the Run, so this -- not delivery -- is what a + successor Run must resume from. + """ + generation: int = 0 + """Increments each time this wait re-enters the blocked state.""" + + announced: bool = False + """Whether the *current* annotation has carried this subscription's binding. + + The first observation must carry provider identity, stream key, and start + cursor even if no record was ever delivered -- otherwise replay of a + subscription to an empty stream has nowhere to begin. + + Read rather than merely set, because ``register`` accepts a subscription at + any activation of a retained Workflow Task and the header frame for that + task has usually already gone to Core. A wait that joined after it is bound + by its own bindings frame, and this is what says which waits still need + one. + """ + + fence_reached: bool = False + """Set by a write fence, cleared by any later record. + + A fence means only that *this* producer session's preceding writes are all + appended. A later record does not violate it; it simply clears it. + """ + + blocked: bool = True + """Whether Workflow code is currently waiting on this subscription.""" + + ready_records: int = 0 + """Records delivered to this subscription's ready list and not yet consumed. + + What separates :attr:`delivery_cursor` from :attr:`consumption_cursor`, as a + count rather than as two boundaries -- until a close drops the ready list, + which zeroes this while leaving both cursors where they were. Kept + because the per-activation delivery budget has to be charged for records + *already in Workflow code's hands* at the start of an activation as well as + for the ones it delivers during it: a ready list carried across an activation + boundary is consumed with no drain and therefore no budget check, and n + subscriptions each carrying one would otherwise let one activation run for n + times the cap. + """ + + closed: bool = False + """Whether Workflow code has ended this subscription. + + A closed wait keeps its binding and its cursor -- replay and a successor Run + both still need them -- but can never be blocked again, so it can never put + Core back in the position of retaining a Workflow Task for a wait nothing is + awaiting. + """ + + +class WorkflowStreamRuntime: + """The handle Workflow code holds, and the observation delta's author.""" + + def __init__( + self, + *, + manager: StreamSubscriptionManager, + backend: StreamBackend | None, + run_id: str, + namespace: str, + workflow_id: str, + first_execution_run_id: str, + data_converter: temporalio.converter.DataConverter, + default_idle_timeout: timedelta, + max_annotation_bytes: int = MAX_ANNOTATION_BYTES, + continuation: Continuation | None = None, + output_continuation: OutputContinuation | None = None, + ) -> None: + """Create the per-Run runtime and restore any continuation state.""" + self._manager = manager + self._backend = backend + self._run_id = run_id + self._namespace = namespace + self._workflow_id = workflow_id + self._first_execution_run_id = first_execution_run_id + #: Already bound to this Workflow's `WorkflowSerializationContext` by + #: the Worker that built this runtime, so `codec_for` hands Workflow + #: code a converter carrying the same context every other payload in the + #: activation was converted with. Bound out there rather than in here + #: because `with_context` runs user code, and this object lives on the + #: far side of the sandbox boundary. + self._data_converter = data_converter + self._default_idle_timeout = default_idle_timeout + self._max_annotation_bytes = max_annotation_bytes + #: The margin a segment frame may overrun the affordability line into. + #: See :data:`_SEGMENT_SPILL_CAP`. + self._spill_bytes = min(_SEGMENT_SPILL_CAP, max(64, max_annotation_bytes // 8)) + #: What the predecessor Run committed, or None on a first execution. + #: Restored from History before any subscription is established, never + #: read from the backend -- a cursor derived from mutable backend state + #: would give replay whatever the stream holds now (ADR-022). + self._continuation = continuation + self._subscriptions: dict[int, _SubscriptionState] = {} + #: Where the *current* annotation begins, per wait. Captured when the + #: annotation begins rather than when its header is first needed -- + #: lazily reading the delivery cursor would let a record delivered + #: before the first emission slip in front of the start cursor, and + #: replay of that marker would then never deliver it. + self._annotation_start: dict[int, Cursor] = {} + self._accumulator: AnnotationAccumulator | None = None + #: Whether this Workflow Task's annotation has already been closed with + #: its terminal. Distinct from "no accumulator": one means nothing has + #: been written yet and a terminal must create the header for it, the + #: other that the annotation is finished and a second terminal would + #: append a whole second annotation to the marker. + self._annotation_closed = False + self._pending_deltas: list[bytes] = [] + #: Runs recorded since the current segment opened, in delivery order. + self._runs: list[Run] = [] + #: What each of those runs costs encoded, parallel to `_runs`. Kept + #: because the open segment's frame size is what decides whether another + #: record can be *delivered* at all: a segment frame that no longer fits + #: the annotation cannot be deferred to the next one -- its deliveries + #: happened in this Workflow Task and that task's marker is where replay + #: has to find them -- so the only place left to act is before the record + #: that would grow it is handed over. Measured, never estimated: a run's + #: cost is two provider-supplied offset strings whose length this side + #: does not choose. + self._run_sizes: list[int] = [] + #: The largest run *this annotation* has encoded, in bytes. Unlike + #: `_run_sizes` this survives `close_segment`, which is the whole point: + #: a record is priced before it is delivered, and pricing the first record + #: of each activation at the bare floor -- as a per-segment maximum + #: does -- hands over a record the closing segment then cannot record. + self._max_run_bytes = 0 + #: How many segments this annotation already holds. Zero means the next + #: record would be its first, which is the one case delivery is never + #: refused for: a fresh annotation is the most room there will ever be, so + #: refusing there rolls over to an annotation that refuses identically. + self._segments_in_annotation = 0 + #: Set when a subscription is registered or a record delivered, so an + #: activation that changed nothing at all emits nothing. + self._observed_this_activation = False + #: `wait_id -> Future`, awaited by Workflow code and resolved by the + #: readiness activation. It lives here rather than on either side alone + #: because the two halves are in different modules and a second map + #: would mean the side that resolves is never the side that registered. + self._pending: dict[int, asyncio.Future[None]] = {} + #: Non-``None`` only while a recorded segment is being delivered. + self._replay_ready: list[tuple[int, StreamRecord]] | None = None + #: The bindings of the marker currently being replayed. Non-``None`` + #: only for the length of one replay job, which is what makes a + #: registration made during it checkable against what was recorded. + self._replay_bindings: dict[int, StreamBinding] | None = None + #: `wait_id -> the converter a *recorded* wait's records convert with`. + #: Installed by the Worker before a replay job reaches this thread; see + #: :meth:`install_replay_converters` for why it exists and why it is not + #: torn down when the replay ends. + self._replay_converters: dict[int, temporalio.converter.DataConverter] = {} + #: Records this activation has put into Workflow code's hands, whether by + #: delivering them or by starting with them already in a ready list. + #: Counted here rather than per subscription because the cap is an + #: *activation* budget: `merge()` consumes from several subscriptions + #: inside one activation, and a per-subscription counter would let n + #: streams run n times as long. + self._delivered_this_activation = 0 + + # Workflow-originated output is logical Payload data until the Worker + # stages it after the activation. Keeping it here makes backend handles + # and codecs unreachable from Workflow code while still letting replay + # compare the exact deterministic pre-codec identity. + self._output_records: list[_LogicalOutputRecord] = [] + self._output_topics: list[str] = [] + if output_continuation is not None: + if not isinstance(backend, OutputStreamBackend): + raise StreamStorageError( + "Workflow History contains external output continuation state, " + "but the configured backend does not support output streams" + ) + if ( + output_continuation.provider_id != type(backend).provider_id + or output_continuation.provider_format_version + != type(backend).provider_format_version + ): + raise StreamStorageError( + "the configured output provider does not match the provider " + "recorded by the preceding Run" + ) + self._output_finished: set[str] = set( + () if output_continuation is None else output_continuation.finished_topics + ) + self._output_segments: list[dict[str, int]] = [] + self._output_history_floor_event_id: int | None = None + self._output_max_records: int | None = None + self._output_max_logical_bytes: int | None = None + self._output_max_publish_latency: timedelta | None = None + self._output_logical_bytes = 0 + self._output_rollover_requested = False + self._output_capacity_waiters: list[asyncio.Future[None]] = [] + self._output_stage_token: str | None = None + self._staged_output: StagedWorkflowOutput | None = None + self._output_replay_manifest: Any = None + self._output_replay_finished_before: set[str] | None = None + + # --- the per-activation delivery budget --------------------------------- + + def begin_activation(self, history_floor_event_id: int | None = None) -> None: + """Resets the delivery budget. Called once per activation. + + The budget is per activation because that is the unit the deadlock + timeout applies to: what must be bounded is how long one ``activate()`` + call can run, not how much a Run receives over its life. + + **Reset to what is already in Workflow code's hands, not to zero.** A + batch is delivered whole and consumed one record at a time, so an + activation that stops iterating part-way through leaves the rest in the + subscription's ready list -- where the *next* activation consumes it with + no drain, and so with no budget check of any kind. Zeroing here would make + that carried-over remainder free, and it accumulates: one subscription can + drain a full batch, consume one record and block elsewhere on every + activation in turn, so n subscriptions arrive at an activation holding + roughly n times the cap between them and hand all of it over in one + `activate()` call. Starting the count at the carry-over makes what an + activation may hand over -- carried-over plus newly delivered -- exactly + the cap, whatever the schedule. + """ + self._delivered_this_activation = sum( + state.ready_records for state in self._subscriptions.values() + ) + if history_floor_event_id is not None: + if ( + self._output_history_floor_event_id is not None + and self._output_history_floor_event_id != history_floor_event_id + and self._output_records + ): + raise RuntimeError( + "a new Workflow Task began before the preceding output batch " + "was staged" + ) + if self._output_history_floor_event_id != history_floor_event_id: + self._output_segments = [] + self._output_history_floor_event_id = history_floor_event_id + for waiter in self._output_capacity_waiters: + if not waiter.done(): + waiter.set_result(None) + self._output_capacity_waiters = [] + self._output_segments.append({}) + + def delivery_budget_remaining(self) -> int: + """How many more records this activation may hand to Workflow code. + + Two budgets, and the smaller wins. The record cap bounds how long one + ``activate()`` call can run; the **annotation budget** bounds what the + marker for this Workflow Task can record. The second belongs here for the + same reason as the first: a record handed to Workflow code has to be + recorded, a segment frame that no longer fits cannot be moved to the next + annotation -- its deliveries happened in *this* Workflow Task and that + task's marker is where replay must find them -- and there is no third + option. So the budget is spent before the record is delivered rather than + checked after the segment is built. + + Unbounded during replay. Delivery then comes from the recorded segments + rather than a live producer, so it is already finite, the recorded + boundaries already say how many records each activation received -- + re-cutting them here would deliver a different schedule than the one in + History -- and nothing is being written to a new annotation at all. + """ + if self._replay_ready is not None: + return _REPLAY_UNBOUNDED + return min( + max(0, MAX_RECORDS_PER_ACTIVATION - self._delivered_this_activation), + self._annotation_records_affordable(), + ) + + def delivery_budget_exhausted(self) -> bool: + """Whether this activation stopped delivering because of a budget. + + The completion path asks, because records left buffered by a budget have + no readiness notification coming: the watcher moved its prefetch cursor + past them when it buffered them. Their readiness has to be re-reported or + the Workflow waits forever on records already in front of it. + + Either budget counts. The annotation one leaves records buffered in + exactly the same way, and its rollover ends the Workflow Task rather than + the activation -- so the successor task has to be told the buffer is not + empty just as the next activation would have been. + + **Conservative rather than exact**, in both of the ways the count can + reach the cap: an activation whose last drain happened to empty the buffer + did not stop *because* of the budget, and an activation pre-charged for a + carried-over ready list may not have tried to deliver at all. Both + over-report, and over-reporting is the safe direction -- a re-arm for an + empty buffer is skipped by the manager, and a wait wrongly withheld from + immediate parkability is retained for its idle timeout instead of parked. + Under-reporting loses records: nothing else announces what a budget left + behind. + """ + return self._replay_ready is None and ( + self._delivered_this_activation >= MAX_RECORDS_PER_ACTIVATION + or self.annotation_budget_exhausted + ) + + # --- the annotation byte budget (ADR-007) -------------------------------- + + @property + def annotation_budget_exhausted(self) -> bool: + """Whether the annotation can no longer afford another record. + + Read by the completion path, which turns it into a rollover request, and + by :meth:`_segment_end_reason`, which records it as the reason this + segment ended. Not a failure: approaching the budget is a runtime event + and rollover is the mechanism for it (ADR-007). + + Deliberately **not** conditioned on this activation having observed + anything. That is a property of the activation and this is a property of + the annotation, and conflating the two wedges the Workflow: the completion + path reads this *after* ``take_observation_delta`` -- which it must, since + that call is what closes the crossing segment -- and that call clears the + observed flag. The rollover would then go unrequested, the next activation + would begin against the same full annotation, and + :meth:`delivery_budget_remaining` would hand it a budget of zero. Nothing + delivered means nothing observed, which means no rollover, forever. + """ + return self._replay_ready is None and self._annotation_records_affordable() <= 0 + + def _annotation_records_affordable(self) -> int: + """How many more records this activation's segment can afford to record. + + Bytes converted into records by the most expensive run **this annotation** + has encoded, floored at :data:`_RUN_COST_FLOOR`. Annotation-wide and not + per-segment, which is the difference between a price and a guess: the open + segment is emptied at the end of every activation, so a per-segment + maximum prices the first record of *every* activation at the bare floor, + and a real run costs two provider-chosen offset strings. Delivering on + that price hands over a record the closing segment then cannot record. + + Still only a price, never a proof: nothing here has seen the offsets of + the record it is pricing. What makes the arithmetic safe is that a + misprice spills into :data:`_SEGMENT_SPILL_CAP` and becomes a rollover, + and that the reserve behind the margin keeps the terminal affordable + regardless. + + **One record is always affordable while the annotation holds no segment.** + A fresh annotation is the most room there will ever be, so refusing there + would roll over to an annotation that refuses identically -- a Workflow + that delivers nothing, observes nothing, and therefore never even asks for + the rollover that was supposed to save it. + """ + headroom = self._annotation_headroom() - self._segment_bytes() + price = max(_RUN_COST_FLOOR, self._max_run_bytes) + if headroom >= price: + return headroom // price + if self._segments_in_annotation == 0 and not self._runs: + return 1 + return 0 + + def _annotation_headroom(self) -> int: + """Bytes this annotation still has for frames that are not its closers. + + The reserve is re-priced here rather than read off the accumulator, which + holds whatever it was told last. A terminal entry costs one byte while a + wait sits at ``BEGINNING`` and more once it has a cursor, so the figure + the accumulator is carrying goes stale on the first delivery -- and a + stale reserve on this path is an *under*-reserve, which is the direction + that matters. + + When there is no accumulator yet, the first observation is about to create + one and its header comes out of the same budget, so the header is priced + here rather than discovered to be unaffordable after the fact. + """ + size = ( + self._accumulator.size + if self._accumulator is not None + else len(encode_header(self._header_preview())) + ) + return max( + 0, + self._max_annotation_bytes + - size + - self._reserve_bytes() + - self._spill_bytes, + ) + + def _segment_bytes(self) -> int: + """What the open segment would cost as a frame, right now.""" + if not self._run_sizes: + return 0 + return encoded_segment_size(self._run_sizes) + + def _reserve_bytes(self) -> int: + """What closing this annotation will cost: the terminal, plus bindings. + + Held back rather than checked, because neither frame may ever be refused: + both record something that already happened, and an annotation Core + writes without a terminal is durable and cannot be decoded past the frame + after it. See :attr:`AnnotationAccumulator.reserved`. + """ + reserve = len( + encode_terminal( + { + wait_id: state.delivery_cursor + for wait_id, state in sorted(self._subscriptions.items()) + } + ) + ) + late = { + wait_id: self._binding(state) + for wait_id, state in sorted(self._subscriptions.items()) + if not state.announced + } + if late and self._accumulator is not None: + # Only with an accumulator: without one the header is about to carry + # every one of these, and `_annotation_headroom` prices it there. + reserve += len(encode_bindings(late)) + return reserve + + def _update_reserve(self) -> None: + """Re-prices the closing frames on the accumulator that holds them.""" + if self._accumulator is not None: + self._accumulator.reserve(self._reserve_bytes(), spill=self._spill_bytes) + + def _check_segment_recordable(self) -> None: + """Refuses a segment that has grown past even the spill margin. + + The last line, and the one thing the arithmetic above cannot rule out: a + record is priced before its offsets are seen, so a provider whose offsets + are far longer than anything measured can make one run cost more than the + whole budget has left. Rollover does not help -- a fresh annotation still + has to carry this run -- so the boundary is genuinely unrecordable. + + Raised **here**, where the record has been drained but not yet handed to + Workflow code, and as the same non-retryable capacity error + ``subscribe()`` raises. Two things follow, and both are the point: + ``AnnotationBudgetExceeded`` is not what surfaces, so the message names + the provider's offsets rather than an internal byte budget; and the + Workflow fails rather than its Workflow Task, so the encoding that cannot + fit is not retried forever (ADR-007). + """ + # Headroom plus the margin: what a segment may still spend, counting the + # header and every earlier frame this annotation already holds. Measured + # against the reserve alone it would miss the header entirely, and the + # header is the largest frame most annotations carry. + emittable = self._annotation_headroom() + self._spill_bytes + segment = self._segment_bytes() + if segment <= emittable: + return + largest = max(self._run_sizes) if self._run_sizes else 0 + raise ExternalStreamCapacityError( + f"this Workflow Task's external stream deliveries no longer fit the " + f"replay annotation: the segment recording them needs {segment} bytes " + f"against {emittable} available, and its largest single run costs " + f"{largest}. A run is two provider-supplied offsets, so this means the " + "backend's offsets are far longer than the marker format is sized for. " + "Rolling the Workflow Task over cannot help -- the next annotation has " + "to carry the same run. Use a backend with shorter offsets, or " + "subscribe to fewer streams from one Workflow." + ) + + def _check_annotation_capacity(self, wait_id: int) -> None: + """Refuses a subscription set no annotation could ever record. + + A header is one indivisible frame, and a binding carries four + caller-chosen strings -- namespace, Workflow ID, first-execution Run ID, + stream name -- plus the provider id. Enough + subscriptions, or long enough valid names, and the header alone is larger + than the whole budget. Nothing downstream can recover from that: rollover + writes a *fresh* header, so the next annotation is the same size and the + one after that too, and the Workflow Task fails identically on every + retry with no marker ever written. ADR-007 exists to keep the budget from + producing exactly that. + + So the capacity question is asked where the answer is still actionable: + inside the Workflow's ``subscribe()`` call. It is deterministic -- the + same subscriptions in the same order give the same answer -- so replay + reproduces the refusal rather than diverging on it. + + Priced against an **empty** annotation, and against everything such an + annotation is nonetheless obliged to carry: its header, its terminal, one + segment frame -- an activation that drained and observed nothing still + encodes one, and it is meaningful (ADR-018) -- and the spill margin a + mispriced record overruns into. Leaving any of those out accepts a + subscription set that clears the check and then cannot encode its very + first completion, which is the failure this exists to prevent rather than + to relocate. + """ + floor = ( + len(encode_header(self._header_preview())) + + len( + encode_terminal( + { + other: state.delivery_cursor + for other, state in sorted(self._subscriptions.items()) + } + ) + ) + + len(encode_segment(Segment((), SegmentEndReason.NO_DATA_AVAILABLE))) + + self._spill_bytes + ) + if floor <= self._max_annotation_bytes: + return + raise ExternalStreamCapacityError( + f"subscribing external stream wait {wait_id} would make this " + f"Workflow's replay annotation need {floor} bytes before recording a " + f"single record -- its header, its terminal, one segment frame, and " + f"the {self._spill_bytes}-byte margin -- past the " + f"{self._max_annotation_bytes}-byte budget. " + "Rolling the Workflow Task over cannot help: every annotation begins " + "with a header of this size. Subscribe to fewer streams from one " + "Workflow, or shorten the stream and provider names." + ) + + def rearm_readiness(self) -> None: + """Re-reports readiness for every buffer this Run left non-empty. + + Called on the Workflow thread, so the hop onto the manager's loop happens + inside the manager -- the same reason ``register`` hops. + """ + self._manager.rearm_ready(self._run_id) + + # --- the ExternalStreamRuntime protocol --------------------------------- + + def stream_key(self, stream_name: str) -> StreamKey: + """Build one stream's durable key in this Workflow chain.""" + return StreamKey( + namespace=self._namespace, + workflow_id=self._workflow_id, + first_execution_run_id=self._first_execution_run_id, + stream_name=stream_name, + ) + + # --- Workflow-originated output --------------------------------------- + + async def publish_output( + self, + *, + topic: str, + value: Any, + value_type: type | None, + kind: RecordKind, + max_publish_latency: timedelta, + max_records: int, + max_logical_bytes: int, + ) -> None: + """Convert and buffer one deterministic logical output record. + + This method runs on the Workflow thread. It deliberately performs only + the synchronous payload-converter step; codecs, external payload + storage, and the output provider run later on the Worker's loop. + """ + del value_type # A decode hint for readers; encoding depends on the value. + if self._output_replay_manifest is None and not isinstance( + self._backend, OutputStreamBackend + ): + raise RuntimeError( + "the configured external stream backend does not support output staging" + ) + if topic in self._output_finished: + raise temporalio.workflow.NondeterminismError( + f"external output topic {topic!r} was already finished" + ) + + payload: temporalio.api.common.v1.Payload | None + if kind is RecordKind.DATA: + payloads = self._data_converter.payload_converter.to_payloads([value]) + if len(payloads) != 1: + raise ValueError( + f"converting one output value produced {len(payloads)} Payloads; " + "exactly one is required" + ) + payload = payloads[0] + else: + payload = None + frame = canonical_logical_record_frame(topic, kind, payload) + replaying = self._output_replay_manifest is not None + if not replaying and len(frame) > max_logical_bytes: + raise ExternalStreamCapacityError( + f"one record for output topic {topic!r} has {len(frame)} logical " + f"bytes, exceeding the configured maximum {max_logical_bytes}" + ) + + while not replaying: + effective_records = min( + max_records, + self._output_max_records or max_records, + ) + effective_bytes = min( + max_logical_bytes, + self._output_max_logical_bytes or max_logical_bytes, + ) + if not self._output_records or ( + len(self._output_records) + 1 <= effective_records + and self._output_logical_bytes + len(frame) <= effective_bytes + ): + break + # Awaiting is the deterministic backpressure point. The activation + # returns with the current batch, the Worker stages it, and Core's + # forced replacement task resolves this future from begin_activation. + self._output_rollover_requested = True + waiter = self.new_readiness_future() + self._output_capacity_waiters.append(waiter) + await waiter + + if topic not in self._output_topics: + self._output_topics.append(topic) + if not self._output_segments: + self._output_segments.append({}) + self._output_segments[-1][topic] = self._output_segments[-1].get(topic, 0) + 1 + self._output_records.append(_LogicalOutputRecord(topic, kind, payload, frame)) + self._output_logical_bytes += len(frame) + # Replay validates the batch History recorded. It must not retain the + # currently configured capacity or latency as policy for the next live + # batch: those options may have changed since this marker was written, + # and neither is active while replaying. + if not replaying: + self._output_max_records = min( + max_records, + self._output_max_records or max_records, + ) + self._output_max_logical_bytes = min( + max_logical_bytes, + self._output_max_logical_bytes or max_logical_bytes, + ) + self._output_max_publish_latency = min( + max_publish_latency, + self._output_max_publish_latency or max_publish_latency, + ) + if kind is RecordKind.FINISH: + self._output_finished.add(topic) + + @property + def has_output(self) -> bool: + """Whether this Workflow Task has logical output awaiting staging.""" + return bool(self._output_records) + + @property + def output_rollover_requested(self) -> bool: + """Whether capacity backpressure is waiting for a replacement task.""" + return self._output_rollover_requested + + @property + def output_max_publish_latency(self) -> timedelta | None: + """The minimum latency configured by non-empty topics in this batch.""" + return self._output_max_publish_latency + + async def stage_output(self, history_floor_event_id: int) -> StagedWorkflowOutput: + """Apply outbound transforms and durably stage the current batch.""" + if history_floor_event_id < 1: + raise RuntimeError( + "Core did not provide the exact predecessor of this Workflow " + "Task's Scheduled event; refusing to stage output" + ) + if not self._output_records: + raise RuntimeError("there is no Workflow output to stage") + if self._output_history_floor_event_id not in (None, history_floor_event_id): + raise RuntimeError("the output batch belongs to a different Workflow Task") + backend = self._backend + if not isinstance(backend, OutputStreamBackend): + raise RuntimeError("the configured backend does not support output staging") + + token = self._output_stage_token + if token is None: + token = secrets.token_urlsafe(24) + self._output_stage_token = token + + records_by_topic = { + topic: [record for record in self._output_records if record.topic == topic] + for topic in self._output_topics + } + topic_manifests: list[StagedOutputTopic] = [] + for sub_batch_id, topic in enumerate(self._output_topics): + logical_records = records_by_topic[topic] + fingerprint = fingerprint_logical_frames( + record.frame for record in logical_records + ) + key = StreamKey( + namespace=self._namespace, + workflow_id=self._workflow_id, + first_execution_run_id=self._first_execution_run_id, + stream_name=topic, + direction=StreamDirection.OUTPUT, + ) + manifest = OutputStageManifest( + stream_key=key, + provider_id=type(backend).provider_id, + provider_format_version=type(backend).provider_format_version, + stage_token=token, + run_id=self._run_id, + history_floor_event_id=history_floor_event_id, + sub_batch_id=sub_batch_id, + fingerprint_version=fingerprint.version, + fingerprint=fingerprint.digest, + record_count=fingerprint.record_count, + logical_byte_count=fingerprint.logical_byte_count, + ) + topic_manifests.append( + StagedOutputTopic( + manifest, + any(record.kind is RecordKind.FINISH for record in logical_records), + ) + ) + + staged = StagedWorkflowOutput( + schema_version=_OUTPUT_MANIFEST_SCHEMA_VERSION, + fingerprint_version=LOGICAL_FINGERPRINT_VERSION, + stage_token=token, + history_floor_event_id=history_floor_event_id, + run_id=self._run_id, + provider_id=type(backend).provider_id, + provider_format_version=type(backend).provider_format_version, + topics=tuple(topic_manifests), + segment_record_counts=tuple( + tuple(segment.get(topic, 0) for topic in self._output_topics) + for segment in self._output_segments + ), + ) + # Bound the exact protobuf that will enter History before running a + # payload codec, external payload store, or provider operation. Topic + # names and a many-activation segment schedule can otherwise make a + # payload-free marker exceed the server event budget after external + # side effects have already happened. + from temporalio.bridge.proto.external_data import ExternalOutputStreamManifest + + marker_manifest = ExternalOutputStreamManifest( + schema_version=staged.schema_version, + fingerprint_version=staged.fingerprint_version, + stage_token=staged.stage_token, + history_floor_event_id=staged.history_floor_event_id, + run_id=staged.run_id, + provider_id=staged.provider_id, + provider_format_version=staged.provider_format_version, + ) + for staged_topic in staged.topics: + source = staged_topic.manifest + marker_topic = marker_manifest.topics.add() + marker_topic.topic = source.stream_key.stream_name + marker_topic.record_count = source.record_count + marker_topic.logical_byte_count = source.logical_byte_count + marker_topic.logical_fingerprint = source.fingerprint + marker_topic.finished = staged_topic.finished + for counts in staged.segment_record_counts: + marker_manifest.segments.add().record_counts_by_topic.extend(counts) + manifest_bytes = marker_manifest.ByteSize() + if manifest_bytes > _MAX_OUTPUT_MANIFEST_BYTES: + raise ExternalStreamCapacityError( + "this Workflow Task's external output manifest needs " + f"{manifest_bytes} bytes, exceeding the hard " + f"{_MAX_OUTPUT_MANIFEST_BYTES}-byte marker budget; shorten output " + "topic names or publish from fewer activation segments" + ) + + # Save the immutable attempt before the first provider await. A lost + # acknowledgement retries the same token and manifest; encoded bytes may + # change under a randomized codec, and the provider retains the first. + if self._staged_output is None: + self._staged_output = staged + elif self._staged_output != staged: + raise RuntimeError( + "one stage token was reused for different logical output" + ) + + encoded_by_topic: dict[str, list[StagedOutputRecord]] = {} + for topic in self._output_topics: + encoded_records: list[StagedOutputRecord] = [] + for publish_index, record in enumerate(records_by_topic[topic]): + if record.payload is None: + encoded = b"" + else: + copied = temporalio.api.common.v1.Payload() + copied.CopyFrom(record.payload) + transformed = await self._data_converter._encode_payload_sequence( + [copied] + ) + transformed = ( + await self._data_converter._external_store_payload_sequence( + transformed + ) + ) + if len(transformed) != 1: + raise ValueError( + "transforming one output Payload did not produce exactly one" + ) + encoded = transformed[0].SerializeToString() + encoded_records.append( + StagedOutputRecord(publish_index, record.kind, encoded) + ) + encoded_by_topic[topic] = encoded_records + + for topic in self._output_topics: + topic_stage = next( + value + for value in staged.topics + if value.manifest.stream_key.stream_name == topic + ) + try: + provider_stage = await backend.stage_output( + topic_stage.manifest, encoded_by_topic[topic] + ) + except OutputStageConflictError as err: + raise StreamIntegrityError( + "the output provider rejected an idempotent stage retry with " + f"a conflicting manifest for topic {topic!r}" + ) from err + except (StreamIntegrityError, StreamStorageError): + raise + except Exception as err: + raise StreamStorageError( + f"failed staging external output topic {topic!r}" + ) from err + if ( + provider_stage.manifest != topic_stage.manifest + or provider_stage.status is not OutputStageStatus.PENDING + ): + raise StreamIntegrityError( + "the output provider did not retain the exact stage as pending " + f"for topic {topic!r}" + ) + return staged + + def output_stage_recorded(self) -> None: + """Start the next batch after its compact commit command was emitted.""" + self._output_records = [] + self._output_topics = [] + self._output_segments = [] + self._output_logical_bytes = 0 + self._output_max_records = None + self._output_max_logical_bytes = None + self._output_max_publish_latency = None + self._output_rollover_requested = False + self._output_stage_token = None + self._staged_output = None + + def begin_output_replay(self, manifest: Any) -> None: + """Install the marker's output expectations without touching a backend.""" + if self._output_replay_manifest is not None: + raise RuntimeError("an external output replay is already in progress") + if manifest.schema_version != _OUTPUT_MANIFEST_SCHEMA_VERSION: + raise RuntimeError( + "external output marker schema version " + f"{manifest.schema_version} is unsupported" + ) + if manifest.fingerprint_version != LOGICAL_FINGERPRINT_VERSION: + raise RuntimeError( + "external output logical fingerprint version " + f"{manifest.fingerprint_version} is unsupported" + ) + self._output_replay_manifest = manifest + self._output_replay_finished_before = set(self._output_finished) + self._output_records = [] + self._output_topics = [] + self._output_segments = [] + self._output_logical_bytes = 0 + self._output_max_records = None + self._output_max_logical_bytes = None + self._output_max_publish_latency = None + self._output_rollover_requested = False + self._output_stage_token = None + + def begin_output_replay_segment(self) -> None: + """Attach output expectations to the shared input/output drain frame.""" + self._output_segments.append({}) + + def verify_output_replay(self) -> None: + """Match replayed publish calls against the compact logical manifest.""" + expected = self._output_replay_manifest + if expected is None: + return + actual_topics = [] + for topic in self._output_topics: + records = [ + record for record in self._output_records if record.topic == topic + ] + fingerprint = fingerprint_logical_frames(record.frame for record in records) + actual_topics.append( + ( + topic, + fingerprint.record_count, + fingerprint.logical_byte_count, + fingerprint.digest, + any(record.kind is RecordKind.FINISH for record in records), + ) + ) + expected_topics = [ + ( + topic.topic, + topic.record_count, + topic.logical_byte_count, + bytes(topic.logical_fingerprint), + topic.finished, + ) + for topic in expected.topics + ] + expected_segments = [ + tuple(segment.record_counts_by_topic) for segment in expected.segments + ] + actual_segments = [ + tuple(segment.get(topic, 0) for topic in self._output_topics) + for segment in self._output_segments + ] + if actual_topics != expected_topics or actual_segments != expected_segments: + raise temporalio.workflow.NondeterminismError( + "Workflow output publish calls do not match the external output " + "manifest recorded in History" + ) + self._output_replay_manifest = None + self._output_replay_finished_before = None + self._output_records = [] + self._output_topics = [] + self._output_segments = [] + self._output_logical_bytes = 0 + self._output_max_records = None + self._output_max_logical_bytes = None + self._output_max_publish_latency = None + self._output_rollover_requested = False + + def abandon_output_replay(self) -> None: + """Drop a failed marker's expectations without validating partial output.""" + if self._output_replay_manifest is None: + return + if self._output_replay_finished_before is not None: + self._output_finished = self._output_replay_finished_before + self._output_replay_manifest = None + self._output_replay_finished_before = None + self._output_records = [] + self._output_topics = [] + self._output_segments = [] + self._output_logical_bytes = 0 + self._output_max_records = None + self._output_max_logical_bytes = None + self._output_max_publish_latency = None + self._output_rollover_requested = False + + def register( + self, + *, + wait_id: int, + stream_key: StreamKey, + idle_timeout: timedelta | None = None, + start_cursor: Cursor | None = None, + ) -> None: + """Registers a wait with the Worker's manager. Non-blocking, no I/O. + + The start cursor defaults to what the predecessor Run committed for this + ``wait_id``, so a chain resumes where it left off without the Workflow + code saying anything about it -- and a first execution gets ``BEGINNING`` + from the same path. + + Raises: + ExternalStreamCapacityError: This subscription set cannot be recorded + in an annotation at all -- see + :meth:`_check_annotation_capacity`. Raised *here*, out of the + Workflow's own ``subscribe()`` call, because that is the only + point at which the answer is still "do not make this + subscription" rather than "this Workflow Task cannot be + completed". + temporalio.workflow.NondeterminismError: This wait's stream is not + the one the predecessor Run recorded -- see + :meth:`restored_start`. + StreamStorageError: The configured backend no longer uses the + provider the predecessor read through. + """ + if self._backend is None: + raise RuntimeError( + "external streams are not configured on this Worker; pass " + "external_stream_backend=... to the Worker" + ) + if start_cursor is None: + start_cursor = self.restored_start(wait_id, stream_key.stream_name) + if self._replay_bindings is not None: + # A subscription made while a marker is being replayed -- which is + # every subscription, on the activation that both starts the + # Workflow and replays its first marker. Checked before the state + # exists, so a mismatch is reported before a single recorded record + # can be handed through the wrong subscription. + binding = self._replay_bindings.get(wait_id) + if binding is not None: + self._verify_binding(wait_id, stream_key, binding) + state = _SubscriptionState( + wait_id=wait_id, + stream_key=stream_key, + start_cursor=start_cursor, + delivery_cursor=start_cursor, + consumption_cursor=start_cursor, + idle_timeout=idle_timeout or self._default_idle_timeout, + ) + self._subscriptions[wait_id] = state + # A subscription created part-way through an annotation begins at its + # own start cursor, not at wherever the others happen to be. + self._annotation_start[wait_id] = start_cursor + try: + self._check_annotation_capacity(wait_id) + except Exception: + # Rolled back so the refusal leaves no half-registered wait behind: + # a state entry with no manager registration would reach the next + # annotation's header as a binding for a wait nothing is watching. + del self._subscriptions[wait_id] + del self._annotation_start[wait_id] + raise + self._update_reserve() + self._manager.register( + run_id=self._run_id, + wait_id=wait_id, + stream_key=stream_key, + start_cursor=start_cursor, + ) + # A registration alone is replay-visible: it is what puts the stream in + # the annotation header, without which replay cannot start. + self._observed_this_activation = True + + def drain(self, wait_id: int, max_records: int | None = None) -> list[StreamRecord]: + """Pops buffered records. Performs no I/O and never blocks. + + During replay the records come from the segment currently being + delivered rather than from the live buffer. ``_apply`` cannot tell the + difference, which is the point: replay and live delivery run the same + Workflow code down the same path. + """ + if self._replay_ready is not None: + # Deliberately *not* the live buffer, even if the manager happens to + # hold something: a watcher that ran before the Run was evicted may + # have prefetched past where the marker stops, and delivering that + # would replay records this Workflow Task never saw. + # + # Taken from the **front**, and only while the front belongs to this + # wait. A segment is the recorded global order across every stream -- + # the order `record_delivery` saw, which is the order Workflow code + # took records in -- so a drain that searched past a record belonging + # to another wait would hand this one a record that came *after* it + # live. A segment recorded as (wait 2, wait 1) replays through + # `merge` as (wait 1, wait 2) under that reading, because `merge` + # asks in `wait_id` order and the search obliges every time. + # + # Live, that drain would simply have found nothing: the record was + # not in this wait's buffer yet. Returning nothing here is the same + # answer, and it leaves the record where the drain that recorded it + # will find it. + taken: list[StreamRecord] = [] + while ( + self._replay_ready + and self._replay_ready[0][0] == wait_id + and (max_records is None or len(taken) < max_records) + ): + taken.append(self._replay_ready.pop(0)[1]) + return taken + return self._manager.drain(self._run_id, wait_id, max_records) + + # --- replay delivery ---------------------------------------------------- + + def take_replay_plan(self) -> ReplayPlan | None: + """The plan prepared for this Run, if a replay job is being delivered.""" + return self._manager.take_replay_plan(self._run_id) + + def install_replay_converters( + self, converters: Mapping[int, temporalio.converter.DataConverter] + ) -> None: + """The context a recorded wait's records must be *converted* with. + + Called on the Worker's loop, before the replay job reaches this thread, + because the converters are produced by ``with_context`` -- user code + that clones the component converters, and that this side of the sandbox + boundary must not run. See ``_handle_external_stream_jobs``. + + The asynchronous half of decoding is already bound this way: the manager + prepares a replayed record under the *annotation's* stream key rather + than under anything the running Run knows, because the Workflow has not + run far enough to have re-created the subscription. Leaving the + synchronous half on the runtime's own converter converts one record + under two different Workflow identities, which is not hypothetical: an + offline ``Replayer`` runs under ``ReplayNamespace`` by default, so a + converter keyed on the namespace decodes a valid history to a different + value or refuses it outright. The same accommodation + :meth:`_verify_binding` makes for a replay harness's namespace, applied + to conversion. + + Merged rather than replaced, and deliberately **not** cleared when the + replay ends. Delivery is not consumption: a batch drained during replay + can be consumed over several later activations, and the record's context + is a property of where it was recorded rather than of the marker that + happened to still be open. Keeping the entry costs nothing on a live + Worker -- a wait's stream key is the Run's own identity, so the recorded + context *is* the runtime's -- and the two only differ under a replay + harness, which delivers nothing live to mis-convert. + """ + self._replay_converters.update(converters) + + def begin_replay(self, bindings: Mapping[int, StreamBinding]) -> None: + """Holds the marker's bindings open, and checks the ones already made. + + Delivery joins a recorded run to a subscription by ``wait_id`` alone, + and an integer is not an identity: it says nothing about *what* the wait + was subscribed to. Without this the same wait number pointing at a + different stream quietly delivers the recorded stream's bytes through + the new subscription -- the failure taxonomy's row four turned into a + silently different stream result, which is the one outcome the design + says must never happen. + + Checked here as well as in :meth:`register` because the two moments are + different Workflow Tasks' worth of history: a marker replayed after the + Workflow has already run past its ``subscribe()`` calls finds the + subscriptions in place, and one replayed in the same activation that + starts the Workflow finds none of them yet. + """ + self._replay_bindings = dict(bindings) + for wait_id, state in sorted(self._subscriptions.items()): + binding = self._replay_bindings.get(wait_id) + if binding is not None: + self._verify_binding(wait_id, state.stream_key, binding) + + def _verify_binding( + self, + wait_id: int, + stream_key: StreamKey, + binding: StreamBinding, + ) -> None: + """Row four, and deliberately not integrity loss. + + The stream compared here was chosen by Workflow code, so a difference + means the code moved, not that anything is wrong with the backend. + + Only the **stream name** is compared, not the whole key. The other three + components -- namespace, Workflow id, first execution Run id -- are the + Run's identity rather than anything the code chose, and a replay harness + legitimately supplies its own: `Replayer` runs under `ReplayNamespace`, + so comparing the full key would report every replayed history as + nondeterministic. The key is still *recorded* whole, because replay has + to read the ranges it names. + + The start cursor is **not** compared. It is the position the wait stood + at when that Workflow Task opened, not a property of the subscription, + so for every marker after a Run's first it legitimately differs from the + cursor the subscription was registered with. What guards the cursor + across Runs is ADR-022's continuation check in :meth:`restored_start`. + """ + if stream_key.stream_name != binding.stream_key.stream_name: + raise temporalio.workflow.NondeterminismError( + f"the marker records external stream wait {wait_id} on stream " + f"{binding.stream_key.stream_name!r}, but this Workflow " + f"subscribes it to {stream_key.stream_name!r}. A subscribe() " + "call was inserted, removed, or reordered, which renumbers " + "every later wait; gate the change behind workflow.patched() " + "exactly as an inserted timer would be." + ) + + def begin_replay_segment( + self, deliveries: Sequence[tuple[int, StreamRecord]] + ) -> None: + """Makes one recorded segment the only thing a drain can see.""" + self._verify_replay_consumed() + self._replay_ready = list(deliveries) + + def verify_replay_consumed(self) -> None: + """Everything the marker recorded must have happened again. + + Two things, checked once the last segment has been delivered. + + A record in a run was handed to Workflow code during the activation the + run was recorded in, so a replay that leaves one behind is running + different code. The common way to get here is a removed ``subscribe()`` + call: its wait is never registered, nothing ever drains it, and the + marker's records for it would otherwise be discarded in silence -- the + Workflow reaching its next command having consumed less than History + says it consumed. + + And every wait the marker *bound* must have been recreated, which is a + strictly larger claim: a binding is written for a subscription whether + or not anything was ever delivered through it. The first observation has + to carry provider identity, stream key, and start cursor even for a + stream that stayed quiet for the whole Workflow Task, so a binding with + no runs behind it is the ordinary shape of a quiet subscription rather + than an exotic one -- and it is invisible to every check that reasons + from deliveries. + """ + self._verify_replay_consumed() + # The one check that reads in the other direction. Every other one walks + # what the code *did*: :meth:`begin_replay` iterates the subscriptions + # that exist and verifies only those the marker also names, + # :meth:`register` verifies a wait only when it is registered at all, + # and :meth:`_verify_replay_consumed` has nothing to report unless + # records were left undelivered. A recorded wait that the code no longer + # creates and that the marker holds no records for is therefore reached + # by none of them. + # + # Two removals escape without this, and the second is the worse one: + # + # - the **last** ``subscribe()`` removed, on a stream that was quiet. + # Nothing renumbers, no records go undelivered, and the replay is + # accepted although the Workflow now holds one subscription fewer than + # History says it did -- with its own live reads never made. + # - a **middle** ``subscribe()`` removed where the later waits name the + # same stream and backend. Every survivor renumbers down by one, so + # the binding comparison compares equal for all of them, and the + # records the marker recorded for wait *k* are consumed by what was + # subscription *k+1*: a different cursor, a different consumer, and + # nothing left over for the delivery check to notice. + # + # Deliberately **not** symmetric, and deliberately not in + # :meth:`_verify_replay_consumed`, which :meth:`begin_replay_segment` + # also calls before each segment. "every bound wait was registered" is + # the invariant; the converse is not, because replay runs the Workflow + # forward past the Workflow Task the marker covers, and the + # subscriptions it makes there belong to the *next* marker's header. An + # added ``subscribe()`` at the end is a supported change and must stay + # one. + missing = sorted(set(self._replay_bindings or {}) - set(self._subscriptions)) + if missing: + # Row four, not integrity loss: the recorded ranges are exactly + # where they were written, and it is the Workflow code that moved. + raise temporalio.workflow.NondeterminismError( + f"the marker records external stream wait(s) {missing} that this " + "Workflow never created. A subscribe() call was removed or " + "reordered, which renumbers every later wait and can hand one " + "wait's recorded records to another wait's subscription; gate " + "the change behind workflow.patched() exactly as a removed " + "timer would be." + ) + + def _verify_replay_consumed(self) -> None: + if not self._replay_ready: + return + wait_ids = sorted({wait_id for wait_id, _ in self._replay_ready}) + raise temporalio.workflow.NondeterminismError( + f"the marker records {len(self._replay_ready)} delivery(s) for " + f"external stream wait(s) {wait_ids} that this Workflow never took. " + "A subscribe() call was removed, reordered, or is no longer " + "consumed, which leaves recorded records undelivered; gate the " + "change behind workflow.patched() exactly as a removed timer would " + "be." + ) + + def reposition_after_replay(self, boundaries: Mapping[int, Cursor]) -> None: + """Moves the manager's cursors to what the replayed marker committed. + + Replay delivers from the annotation, not from the manager's buffer -- + but the watcher has been filling that buffer from the subscription's + start cursor the whole time, because nothing out there knows a marker + for these records exists. Without this the first live drain after a + replay hands Workflow code every replayed record a second time; the + symptom is a Workflow that received ``['alpha', 'alpha', 'beta']``. + + The boundaries come from the marker rather than from wherever the + replay's deliveries happened to stop -- the marker is History's own + statement of what was committed -- and only the waits it recorded + something for are named. + + Performed **synchronously**, before this returns: the drain that + follows replay runs on this same thread, so a reposition merely posted + to the manager's loop would still have the marker-covered records in the + buffer when that drain reaches it. Only the watcher's wakeup is hopped; + see :meth:`StreamSubscriptionManager.reposition_to_committed`. + """ + self._manager.reposition_to_committed(self._run_id, boundaries) + + def end_replay(self) -> None: + """Hands drains back to the live buffer. + + Called from a ``finally``: a partial replay that left this set would + make every later drain on this Run return nothing at all, turning one + marker's failure into a Workflow that silently never receives again. + """ + self._replay_ready = None + self._replay_bindings = None + + def codec_for( + self, value_type: type | None, wait_id: int + ) -> StreamPayloadCodec[Any]: + """Return the wait's typed codec under its recorded serialization context.""" + # The recorded context wins wherever one was installed, so that both + # halves of a replayed record's decoding see the same Workflow identity. + # Everything else -- every live record, and every wait no marker bound -- + # converts with this Run's own converter, which is the same object the + # activation's other payloads used. + converter = self._replay_converters.get(wait_id, self._data_converter) + return StreamPayloadCodec(converter, value_type) + + def new_readiness_future(self) -> asyncio.Future[None]: + """A future belonging to the Workflow's own deterministic event loop.""" + return asyncio.get_event_loop().create_future() + + def register_pending(self, wait_id: int, future: asyncio.Future[None]) -> None: + """Register the future readiness will resolve for one blocked wait.""" + self._pending[wait_id] = future + + def discard_pending(self, wait_id: int) -> None: + """Forget a wait that is no longer awaiting readiness.""" + self._pending.pop(wait_id, None) + + def resolve_all_pending(self) -> int: + """Resumes every waiting subscription. Returns how many were resumed. + + *Every* one, not only those a readiness job named: the job's hints are + hints, not an exhaustive availability claim, so each subscription is + left to find out for itself whether anything arrived. Resolving only the + hinted ones would strand a record whose notification was coalesced away. + """ + resumed = 0 + for wait_id in sorted(self._pending): + future = self._pending.pop(wait_id, None) + if future is not None and not future.done(): + future.set_result(None) + resumed += 1 + return resumed + + def clear_pending(self) -> None: + """Drops every waiting future, as eviction requires.""" + self._pending.clear() + + # --- recording what Workflow code observed ------------------------------- + + def record_delivery(self, wait_id: int, record: StreamRecord) -> None: + """Notes one record reaching the runtime, in observed global order. + + Called for **every** record including control records: they occupy + offsets inside a run, so a run's count includes them and their relative + indices go in ``control_positions``. Omitting them would make replay's + range read find more records than the marker claims and fail as + integrity loss. + + **Also where the activation's delivery budget is spent.** Delivery is the + moment a record leaves the manager's buffer for a subscription's private + ready list, and that is what has to be charged rather than the later + consumption of it, for two reasons that point the same way. The budget is + a *reservation*: a drain that checked the budget and charged nothing let + the next subscription's drain see the same room and take it again, so two + subscriptions consumed independently -- no `merge()` involved -- delivered + twice the cap in one activation, and n of them n times it. And delivery is + what the annotation records, so a cap charged at consumption bounded a + different quantity than the segment it is supposed to bound: the recorded + segment could hold more records than the cap replay will divide it by. + """ + state = self._subscriptions.get(wait_id) + if self._replay_ready is None: + # Counted before the guards below so that a record whose subscription + # has already gone still costs its budget: the cap has to bound the + # activation whatever the bookkeeping says. + # + # Replayed records are deliberately not counted, not merely not + # capped. Counting them would leave the budget spent for any live + # delivery later in the same activation, and would make the completion + # re-arm readiness on a purely replayed Workflow Task. + self._delivered_this_activation += 1 + if state is None or record.offset is None: + return + + state.delivery_cursor = AFTER(record.offset) + state.ready_records += 1 + state.fence_reached = record.is_control + + if self._replay_ready is not None: + # Replay re-delivers records the marker already recorded. Advancing + # the cursors is right -- the runtime has to end up in the state the + # live run was in -- but accumulating them into a *new* annotation is + # not: Core would be asked to write a second marker for observations + # already in History, and the command would be matched against the + # very event it was read from. + return + + self._observed_this_activation = True + + # Extend the open run when this is another consecutive delivery from the + # same stream. A run is *maximal*, which is what makes a single-stream + # batch of 100,000 records cost one run rather than 100,000. + if self._runs and self._runs[-1].wait_id == wait_id: + previous = self._runs[-1] + positions = previous.control_positions + if record.is_control: + positions = (*positions, previous.count) + self._runs[-1] = Run( + wait_id=wait_id, + first_offset=previous.first_offset, + last_offset=record.offset, + count=previous.count + 1, + control_positions=positions, + ) + # Only the run that changed is re-measured. Extending one is cheap; + # re-encoding the whole segment once per record would not be. + self._run_sizes[-1] = encoded_run_size(self._runs[-1]) + self._max_run_bytes = max(self._max_run_bytes, self._run_sizes[-1]) + self._check_segment_recordable() + return + + self._runs.append( + Run( + wait_id=wait_id, + first_offset=record.offset, + last_offset=record.offset, + count=1, + control_positions=(0,) if record.is_control else (), + ) + ) + self._run_sizes.append(encoded_run_size(self._runs[-1])) + self._max_run_bytes = max(self._max_run_bytes, self._run_sizes[-1]) + self._check_segment_recordable() + + def unsubscribe(self, wait_id: int) -> None: + """Ends a wait, and tells the Worker to stop serving it. + + The subscription's state is **kept**, not removed. Two things are built + from it after a wait ends: the annotation header, whose binding replay + needs in order to know what the closed wait was reading -- without it a + replay of unchanged code fails as though the Workflow never created that + wait -- and the Continue-As-New continuation, whose cursor is what stops + a successor Run restarting that stream from the beginning. Closing + changes exactly one thing about the state: the wait can never be blocked + again. + """ + state = self._subscriptions.get(wait_id) + if state is None or state.closed: + return + state.closed = True + state.blocked = False + # Whatever was drained but never handed over is dropped rather than + # consumed -- that is what leaves the consumption cursor short of it, so a + # Continue-As-New successor receives it. It also stops being carry-over the + # next activation's budget has to pay for: those records are gone, and a + # count left standing would shrink every later activation's budget by the + # size of a ready list nobody can consume. + state.ready_records = 0 + self._manager.cancel_from_workflow_thread(self._run_id, wait_id) + + def note_blocked(self, wait_id: int, blocked: bool) -> None: + """Records whether Workflow code is waiting on this subscription.""" + state = self._subscriptions.get(wait_id) + if state is None or (state.closed and blocked): + # A closed wait may leave the blocked set but never re-enter it: the + # coroutine that was awaiting it is gone, so retaining a Workflow + # Task for it would be retaining for nobody. + return + if blocked and not state.blocked: + # Re-entering the blocked state is what the wait generation counts, + # and it is what makes a readiness notification for the previous + # block recognisable as stale. + state.generation += 1 + # Pushed to the manager immediately. The manager is what reports + # readiness to Core, and Core compares the generation it is given + # against the one this runtime put in the quiescent snapshot. A + # manager left at the generation it registered with reports every + # block after the first under a stale number, Core answers `Stale`, + # and the watcher -- whose prefetch cursor is already past the + # record -- never re-announces it. The Workflow then waits forever + # on an append that did arrive. + self._manager.note_wait_generation(self._run_id, wait_id, state.generation) + state.blocked = blocked + + # --- the observation delta (P10b) ---------------------------------------- + + def close_segment(self, reason: SegmentEndReason | None = None) -> None: + """Closes this activation's segment, whether or not it saw anything. + + An empty segment is meaningful: an activation that drained and found + nothing still ran one event-loop drain, and replay must reproduce that + drain or ``wait_condition`` predicates fire a different number of times. + """ + if not self._observed_this_activation: + return + if reason is None: + reason = self._segment_end_reason() + accumulator = self._ensure_accumulator() + self._pending_deltas.append( + accumulator.add_segment(Segment(tuple(self._runs), reason)) + ) + self._runs = [] + self._run_sizes = [] + self._segments_in_annotation += 1 + + def _segment_end_reason(self) -> SegmentEndReason: + """Why this activation stopped delivering. + + Decided here rather than passed in, because the reason has to be right + whether the caller remembered it or not: it is durable, and a reader of + the annotation has nothing else to tell it why the activation ended. + + ``BATCH_LIMIT`` outranks the two availability reasons. Both of them + assert that nothing more was available -- ``FENCE_REACHED`` additionally + that the set is immediately parkable -- and both are simply false when the + runtime stopped with records still sitting in the buffer. Recording + ``NO_DATA_AVAILABLE`` for a budget cut-off would put a claim in History + that the stream ran dry when it did not. + + ``BUDGET_ROLLOVER`` outranks even that, because it says something the + others do not: the batch continues in the *following marker*. A rollover + recorded as ``BATCH_LIMIT`` would leave replay treating this segment as + the end of the consumption. + """ + if self.annotation_budget_exhausted: + # Ranked first: it is the only one of the four that says the batch + # continues in the *next marker*, which is what tells replay to + # reassemble across the rollover boundary rather than treat this + # segment as the end of the consumption. + return SegmentEndReason.BUDGET_ROLLOVER + if self.delivery_budget_exhausted(): + return SegmentEndReason.BATCH_LIMIT + if self._subscriptions and all( + s.fence_reached for s in self._subscriptions.values() + ): + return SegmentEndReason.FENCE_REACHED + return SegmentEndReason.NO_DATA_AVAILABLE + + def take_observation_delta(self) -> bytes | None: + """The bytes to put on this completion's `WorkflowStreamProgress`. + + ``None`` means nothing replay-visible changed, which is the only case + where a completion legitimately carries no progress command -- and a + replay delivery is exactly that case, since everything it delivered is + already recorded in the marker being replayed. + """ + if self._replay_ready is not None: + return None + self.close_segment() + if not self._pending_deltas: + return None + delta = b"".join(self._pending_deltas) + self._pending_deltas = [] + self._observed_this_activation = False + return delta + + def add_terminal(self) -> bytes: + """Encodes the blocked snapshot that closes this Workflow Task. + + Read from in-memory state alone: the boundary is not "wherever the + stream is now", it is where this task's deliveries stopped, which was + fixed the moment the last activation returned. Refreshing it against the + backend could name a position replay must not reproduce. + + Returns **no bytes at all** when this Workflow Task's annotation is + already closed and nothing has been accumulated since. Core can decide a + boundary for a task whose last completion had already ended one -- a + rollover deadline or a shutdown landing on a completion that left no + subscription blocked -- and answering with a second header and terminal + would append a whole second annotation to the one Core is about to + write. The marker would then decode as far as the first terminal and + fail on the frame after it. + """ + if self._annotation_closed and self._accumulator is None: + return b"" + accumulator = self._ensure_accumulator() + terminal = accumulator.add_terminal( + { + wait_id: state.delivery_cursor + for wait_id, state in sorted(self._subscriptions.items()) + } + ) + # Anything still pending goes out with the terminal. Usually there is + # nothing -- each activation flushes its own delta -- but a Workflow Task + # that observed *nothing* creates its accumulator right here, and the + # header rides that creation. Returning the terminal alone would produce + # an annotation that begins at a terminal frame and cannot be decoded. + delta = b"".join([*self._pending_deltas, terminal]) + + # The annotation ends at the terminal. A later activation on this Run -- + # the Continue-As-New's own, say -- opens a fresh annotation rather than + # appending past it: Core writes one marker per finalized annotation, and + # a segment recorded after the terminal could never be read back. + self.start_new_annotation() + self._annotation_closed = True + return delta + + @property + def annotation_started(self) -> bool: + """Whether Core is holding accumulated bytes for the current annotation. + + True from the moment the header is emitted, which is what makes it the + right question to ask before closing an annotation with its terminal: an + annotation nobody has begun has nothing to terminate, and asking for a + terminal anyway would create a header and a terminal for a Workflow Task + that never touched a stream. + """ + return self._accumulator is not None + + @property + def request_rollover(self) -> bool: + """Whether Core should be asked to end this Workflow Task's annotation. + + Two conditions, and the second is not redundant. The high-water mark is a + *fraction* of the budget, and it only becomes true once a frame that + crossed it has been emitted -- but an indivisible frame can be larger than + the fraction that was left, so waiting for the mark alone is not a bound. + The second condition is the runtime having stopped delivering because the + annotation could no longer afford to record another record: that one is + true in the same activation, before anything has overflowed, and it is + what makes "an annotation can never exceed the budget" (ADR-007) hold for + a frame of any size. + """ + if self.annotation_budget_exhausted: + return True + return self._accumulator is not None and self._accumulator.request_rollover + + def start_new_annotation(self) -> None: + """Begins a fresh annotation, as the next Workflow Task requires. + + Its header records the *current* cursors rather than the original start + cursors, so consumption continues uninterrupted across a rollover or any + other Workflow Task boundary. + """ + self._accumulator = None + self._annotation_closed = False + self._pending_deltas = [] + self._runs = [] + self._run_sizes = [] + self._max_run_bytes = 0 + self._segments_in_annotation = 0 + self._observed_this_activation = False + self._annotation_start = { + wait_id: state.delivery_cursor + for wait_id, state in self._subscriptions.items() + } + + def record_consumption(self, wait_id: int, record: StreamRecord) -> None: + """Notes that Workflow code has been handed this record. + + Distinct from delivery, which advances by whole drained batches. The + difference is exactly the records sitting in a subscription's buffer that + the Workflow never asked for -- which die with the Run, and which a + successor must therefore still receive. + + Not where the delivery budget is spent -- :meth:`record_delivery` is, + because the budget has to be reserved by the drain that moves a whole + batch into a ready list rather than charged one record at a time + afterwards. What this does spend is the *carry-over* the next + :meth:`begin_activation` starts its count from: a record consumed here is + one the next activation no longer has in hand. + """ + state = self._subscriptions.get(wait_id) + if state is None or record.offset is None: + return + state.consumption_cursor = AFTER(record.offset) + state.ready_records = max(0, state.ready_records - 1) + + def continuation(self) -> Continuation: + """Where each subscription had got to, for the successor Run. + + The **consumption** cursor, which is the only one of the four positions + a successor may resume from. + + Not prefetch, which is speculative. Not delivery: a batch is delivered + whole, but a Workflow that stops iterating part-way through has taken + only its prefix, and the buffer holding the difference dies with this + Run -- so a successor starting from delivery would step silently over + records nothing ever showed to Workflow code. Not the committed cursor + either: by the time this is read the terminal command is being built, + and the observation delta covering these deliveries commits on that same + path, so a continuation taken from the last *marker* would restart the + successor at a stale cursor and lose the final segment. + + Carried with each cursor is the whole binding it is a position in -- the + stream and configured provider identity -- taken from :meth:`_binding`, + the same place the annotation header takes it from. An offset means + nothing outside the store that produced it. + """ + bindings = { + wait_id: self._binding(state) + for wait_id, state in self._subscriptions.items() + } + return Continuation( + cursors={ + wait_id: state.consumption_cursor + for wait_id, state in self._subscriptions.items() + }, + stream_names={ + wait_id: binding.stream_key.stream_name + for wait_id, binding in bindings.items() + }, + provider_ids={ + wait_id: binding.provider_id for wait_id, binding in bindings.items() + }, + provider_format_versions={ + wait_id: binding.provider_format_version + for wait_id, binding in bindings.items() + }, + ) + + def output_continuation(self) -> OutputContinuation | None: + """Finished output topics that a successor Run must not reopen.""" + if not self._output_finished: + return None + backend = self._backend + if not isinstance(backend, OutputStreamBackend): + raise RuntimeError( + "finished output topics have no configured output provider" + ) + return OutputContinuation( + finished_topics=frozenset(self._output_finished), + provider_id=type(backend).provider_id, + provider_format_version=type(backend).provider_format_version, + ) + + def restored_start(self, wait_id: int, stream_name: str) -> Cursor: + """The start cursor for a subscription, from the predecessor Run if any. + + ``BEGINNING`` on a first execution -- the same field, filled the same + way, so replay reads an explicit boundary in either case. + + The whole binding is compared before the cursor is handed back. The + stream is Workflow code, so a change is nondeterminism; the provider is + Worker configuration, so a change is a storage failure. Both are raised + before the cursor reaches the manager. + """ + if self._continuation is None: + return BEGINNING + restored = self._continuation.cursors.get(wait_id) + if restored is None: + # A subscription the predecessor did not have. Safe: adding one on a + # path the chain has not reached yet is the supported change. + return BEGINNING + # Row four of the taxonomy, not integrity loss: the cursor is exactly + # what the predecessor committed, and it is the Workflow code that moved. + recorded = self._continuation.stream_names.get(wait_id, "") + if recorded and recorded != stream_name: + raise temporalio.workflow.NondeterminismError( + f"the predecessor Run recorded external stream wait {wait_id} " + f"on stream {recorded!r}, but this Run subscribes it to " + f"{stream_name!r}. A subscribe() call was inserted, removed, or " + "reordered, which renumbers every later wait; gate the change " + "behind workflow.patched() exactly as an inserted timer would be." + ) + self._verify_restored_provider(wait_id) + return restored + + def _verify_restored_provider(self, wait_id: int) -> None: + """Whether the configured backend is what produced the cursor. + + A deployment question rather than a Workflow one, and classified exactly + as :meth:`StreamSubscriptionManager._replay_backend` classifies it: the + Workflow is unchanged and neither store is damaged, so it is a storage + failure -- retried, and it clears when a Worker carrying the recorded + implementation picks the task up. + + Marker replay makes this check for every recorded range it reads. A + successor Run's *first* read happens before it has written a marker that + could make it, which is why the continuation carries the provider + identity at all. + + The backend contract reserves no format-version value, so zero is + compared exactly like every other version. + """ + assert self._continuation is not None + recorded_id = self._continuation.provider_ids[wait_id] + backend = self._backend + if backend is None: + raise RuntimeError( + "external streams are not configured on this Worker; pass " + "external_stream_backend=... to the Worker" + ) + declared_id = type(backend).provider_id + if declared_id != recorded_id: + raise StreamStorageError( + f"external stream wait {wait_id} was continued from a Run that " + f"read it through provider {recorded_id!r}, but the backend " + f"configured on this Worker declares {declared_id!r}. The " + "restored cursor is a position in the store that produced it; " + "configure the recorded provider." + ) + declared_version = type(backend).provider_format_version + recorded_version = self._continuation.provider_format_versions[wait_id] + if recorded_version != declared_version: + raise StreamStorageError( + f"external stream wait {wait_id} was continued from a Run that " + f"read it through provider {recorded_id!r} format version " + f"{recorded_version}, but the configured backend implements " + f"format version {declared_version}. " + "Resuming would interpret the restored cursor under a format it " + "was not written in." + ) + + def _ensure_accumulator(self) -> AnnotationAccumulator: + if self._accumulator is None: + self._accumulator = AnnotationAccumulator( + self._header(), max_bytes=self._max_annotation_bytes + ) + # The header has to ride the *first delta*. Core accumulates by byte + # append and never parses, so anything the runtime holds but does + # not emit simply never reaches the marker -- and a marker whose + # annotation starts at a segment frame cannot be decoded at all. + self._pending_deltas.append(self._accumulator.accumulated()) + else: + self._announce_late_subscriptions() + # Re-priced after any late binding went out, so what is held back is what + # is still owed rather than what was owed when the accumulator was made. + self._update_reserve() + return self._accumulator + + def _announce_late_subscriptions(self) -> None: + """Binds every subscription the emitted header does not already carry. + + ``register`` accepts a subscription at any activation of a retained + Workflow Task, and only the first of those activations gets to write the + header -- Core appends the deltas it is handed and never rewrites what it + already holds, so the header cannot be extended in place. A wait that + joined later is bound by its own frame instead. + + Called from :meth:`_ensure_accumulator`, which is on the path of both + things that can follow a registration: the segment that records the + wait's first run, and the terminal that records where it stopped. The + binding therefore always precedes both, and no wait can reach the marker + as a run or a terminal entry with no stream key, no backend, and no start + cursor -- which replay reports as a wait "this Workflow did not create" + even when the code is unchanged. + """ + assert self._accumulator is not None + late = { + wait_id: state + for wait_id, state in sorted(self._subscriptions.items()) + if not state.announced + } + if not late: + return + for state in late.values(): + state.announced = True + self._pending_deltas.append( + self._accumulator.add_bindings( + {wait_id: self._binding(state) for wait_id, state in late.items()} + ) + ) + + def _header(self) -> AnnotationHeader: + """One binding per subscription, each naming **its own** backend. + + Marks every subscription announced, which is also how the flag is reset + for the *next* annotation: a header is written from scratch whenever an + accumulator is created, so whatever exists when it is carries its + binding there rather than in a bindings frame. + + The provider identity is read from the backend this subscription is + actually registered against rather than from a single annotation-wide + one. The API lets every topic name a different backend, so a shared + label would be right for at most one wait and would send replay to the + wrong store for all the others. + """ + for state in self._subscriptions.values(): + state.announced = True + return AnnotationHeader( + streams={ + wait_id: self._binding(state) + for wait_id, state in sorted(self._subscriptions.items()) + }, + ) + + def _header_preview(self) -> AnnotationHeader: + """The header a fresh annotation would carry, without claiming it. + + :meth:`_header` marks every subscription announced, which is right when + the header is actually being emitted and wrong when the question is only + what it would cost. + """ + return AnnotationHeader( + streams={ + wait_id: self._binding(state) + for wait_id, state in sorted(self._subscriptions.items()) + }, + ) + + def _binding(self, state: _SubscriptionState) -> StreamBinding: + backend = self._backend + if backend is None: + raise RuntimeError( + "external streams are not configured on this Worker; pass " + "external_stream_backend=... to the Worker" + ) + return StreamBinding( + stream_key=state.stream_key, + start_cursor=self._annotation_start.get(state.wait_id, state.start_cursor), + provider_id=type(backend).provider_id, + provider_format_version=type(backend).provider_format_version, + ) + + # --- quiescence (P10a) ---------------------------------------------------- + + def quiescent_snapshot(self) -> list[QuiescentWait] | None: + """The **complete** set of waits Workflow code is blocked on. + + ``None`` when nothing is blocked, which is the signal not to ask for + retention at all. A *partial* set would be worse than none: it would let + one idle stream park a Workflow Task another stream is still driving. + """ + blocked = [s for s in self._subscriptions.values() if s.blocked] + if not blocked: + return None + # A set the delivery budget stopped is blocked but not quiescent: records + # are still sitting in the local buffer, and the only reason nobody is + # reading them is that this activation ran out of budget. Calling any of + # it immediately parkable would ask Core to park a Workflow Task whose + # data has already arrived at the Worker. + exhausted = self.delivery_budget_exhausted() + return [ + QuiescentWait( + wait_id=s.wait_id, + generation=s.generation, + immediately_parkable=s.fence_reached and not exhausted, + ) + for s in sorted(blocked, key=lambda s: s.wait_id) + ] + + def effective_idle_timeout(self) -> timedelta: + """The quiescent set's timeouts reduced to one value. + + The reduction is ``min``, applied in ``wait_id`` order over the blocked + set and nothing else, so the result is deterministic and reproduces on + replay. It has to reduce at all because the timeout is a property of the + *set*: subscriptions configured differently still share one timer. + """ + blocked = sorted( + (s for s in self._subscriptions.values() if s.blocked), + key=lambda s: s.wait_id, + ) + if not blocked: + return self._default_idle_timeout + reduced = blocked[0].idle_timeout + for state in blocked[1:]: + reduced = min(reduced, state.idle_timeout) + return reduced + + # --- teardown ------------------------------------------------------------- + + def subscriptions(self) -> list[int]: + """Return the registered wait IDs in deterministic order.""" + return sorted(self._subscriptions) + + def blocked_snapshot(self) -> dict[int, Cursor]: + """Return each registered wait's latest delivery boundary.""" + return { + wait_id: state.delivery_cursor + for wait_id, state in sorted(self._subscriptions.items()) + } diff --git a/temporalio/contrib/external_workflow_streams/_wake.py b/temporalio/contrib/external_workflow_streams/_wake.py new file mode 100644 index 0000000000..a083068b1b --- /dev/null +++ b/temporalio/contrib/external_workflow_streams/_wake.py @@ -0,0 +1,260 @@ +"""The producer wake-signal path (P14). + +The server-visible wakeup, used whenever no open Workflow Task can accept local +readiness. It carries no stream payload -- only enough identity for Core to +decide whether it is a wake, a stale claim, or someone else's chain. + +Two properties are load-bearing and neither is available from the public Signal +API, which is why this path does not reuse it: + +**Codec bypass.** The envelope goes out through a raw ``SignalWorkflowExecution`` +built with the protocol's own serialization rather than the user's +``DataConverter`` (ADR-025). Core is the component that must *read* this Signal, +and Core has no access to a user codec -- a codec that encrypts payloads would +make the envelope unreadable to the only reader that matters. Nothing user-owned +is in it, so bypassing the codec leaks nothing. + +**Stable request ID.** The Temporal ``request_id`` is derived deterministically +from the wake's identity, so a producer retrying after an ambiguous failure sends +the identical request and the server deduplicates it. The public Signal path +generates a fresh UUID per attempt, which would turn every retry into a second +wake. +""" + +from __future__ import annotations + +import hashlib +import uuid +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import temporalio.api.common.v1 +import temporalio.api.workflowservice.v1 +from temporalio.bridge.proto.external_stream.external_stream_pb2 import WakeSignal + +if TYPE_CHECKING: + import temporalio.client + from temporalio.contrib.external_workflow_streams._producer import WorkflowChainKey + +__all__ = [ + "WakeRequest", + "new_sender_identity", + "send_wake_signal", + "wake_request_id", +] + +WAKE_SIGNAL_NAME = "__temporal_external_stream_wake" +"""Fixed, and versioned by the envelope rather than by the name. + +Distinct from every ``__temporal_workflow_stream_*`` name already reserved by +``temporalio.contrib.workflow_streams`` (ADR-001), which is a different feature +that coexists with this one. +""" + +WAKE_SIGNAL_MESSAGE_TYPE = "coresdk.external_stream.WakeSignal" +WAKE_SIGNAL_ENCODING = b"binary/protobuf" +WAKE_SIGNAL_ENVELOPE_VERSION = 1 + +UNPARKED_WAKE_GENERATION = 0 +"""No confirmed park is known; ask for a Workflow Task anyway (ADR-023). + +Not silence. The runtime rechecks every active subscription on wakeup regardless +of which stream the Signal named, so an unnecessary unparked wake costs at most +one empty Workflow Task -- while staying silent because no park was observed +would lose the record until something else happened to wake the Workflow. +""" + +_REQUEST_ID_NAMESPACE = uuid.UUID("6f1f8a1e-6f26-5f3f-9f2b-4b0a2c7d8e10") +"""A fixed namespace, so the derivation is stable across processes and versions. + +Deriving a UUID rather than passing the tuple through keeps the request ID +inside the length the server accepts while staying a pure function of its +inputs. +""" + + +def new_sender_identity(client_identity: str = "") -> str: + """One sender instance's identity for unparked wakes. + + Unique per *sender*, not per client identity. Two Workers in one process + share a ``Client`` and therefore one client identity, so deriving from that + alone gives their first unparked wakes byte-identical request IDs -- the + server deduplicates the second, no Workflow Task is created, and the Run + stalls. That is precisely the loss the sender identity exists to prevent. + + Drawn once and held for the sender's lifetime, which is what "fixed across + retries of that one attempt" requires: a value redrawn per attempt would + turn the shutdown sweep's in-grace retry into a second wake. + + Random, which is safe here: the wake Signal is sent from the Worker's own + event loop and never from Workflow code, so it is not replay-visible and + cannot affect determinism. The client identity is kept as a prefix so a + request ID stays traceable to a Worker in server-side logs. + """ + return f"{client_identity}#{uuid.uuid4()}" + + +@dataclass(frozen=True) +class WakeRequest: + """One wake attempt's complete identity. + + Everything the request ID is derived from lives here, so "the same attempt" + is a property of a value rather than of how carefully a caller re-passed six + arguments. + """ + + namespace: str + workflow_id: str + first_execution_run_id: str + stream_name: str + wait_id: int + park_generation: int + sender_identity: str = "" + wake_counter: int = 0 + + @property + def is_unparked(self) -> bool: + """Whether this wake requests a task outside a confirmed park.""" + return self.park_generation == UNPARKED_WAKE_GENERATION + + def __post_init__(self) -> None: + """Validate the identity required for server-side deduplication.""" + if self.is_unparked and not self.sender_identity: + raise ValueError( + "an unparked wake needs a sender identity. Generation 0 carries " + "no attempt identity of its own, so without it two Workers " + "shutting down at different times would derive the same request " + "ID and the server would deduplicate the second wake away -- " + "turning a correct retry mechanism into silent loss." + ) + + +def wake_request_id(request: WakeRequest) -> str: + """The deterministic ``request_id`` for one wake attempt. + + A **parked** wake is identified by its generation: a generation is woken + once, so every producer that races to wake it derives the same ID and the + server keeps one. + + An **unparked** wake has no generation to identify it, so the derivation + additionally includes the sender's identity -- unique per sender *instance*, + see :py:func:`new_sender_identity` -- and a per-sender monotonic counter, + both held fixed across retries of that one attempt. Two senders' unparked + wakes must stay distinct even though neither knows about the other. + """ + parts = [ + request.namespace, + request.workflow_id, + request.first_execution_run_id, + request.stream_name, + str(request.wait_id), + str(request.park_generation), + ] + if request.is_unparked: + parts += [request.sender_identity, str(request.wake_counter)] + # Length-prefixed rather than delimiter-joined: a stream literally named + # "a\x00b" must not collide with the pair ("a", "b"). + material = b"".join( + len(encoded).to_bytes(4, "big") + encoded + for encoded in (part.encode("utf-8") for part in parts) + ) + return str(uuid.uuid5(_REQUEST_ID_NAMESPACE, hashlib.sha256(material).hexdigest())) + + +def wake_envelope(request: WakeRequest, *, producer_session_id: str = "") -> bytes: + """The Signal's single argument, serialized at the protocol level.""" + return WakeSignal( + envelope_version=WAKE_SIGNAL_ENVELOPE_VERSION, + stream_name=request.stream_name, + wait_id=request.wait_id, + park_generation=request.park_generation, + first_execution_run_id=request.first_execution_run_id, + producer_session_id=producer_session_id, + ).SerializeToString() + + +def build_signal_request( + request: WakeRequest, + *, + identity: str, + producer_session_id: str = "", +) -> temporalio.api.workflowservice.v1.SignalWorkflowExecutionRequest: + """The raw request, built without touching the user's ``DataConverter``. + + Addressed to the Workflow ID with **no Run ID**, so it always lands on the + current Run of the chain -- a wake sent while a Continue-As-New is in flight + must reach the successor, not fail against a Run that has already closed. + ``first_execution_run_id`` inside the envelope is what lets Core tell the + chain apart from a reused Workflow ID. + """ + return temporalio.api.workflowservice.v1.SignalWorkflowExecutionRequest( + namespace=request.namespace, + workflow_execution=temporalio.api.common.v1.WorkflowExecution( + workflow_id=request.workflow_id + ), + signal_name=WAKE_SIGNAL_NAME, + input=temporalio.api.common.v1.Payloads( + payloads=[ + temporalio.api.common.v1.Payload( + metadata={ + "encoding": WAKE_SIGNAL_ENCODING, + "messageType": WAKE_SIGNAL_MESSAGE_TYPE.encode(), + }, + data=wake_envelope( + request, producer_session_id=producer_session_id + ), + ) + ] + ), + identity=identity, + request_id=wake_request_id(request), + ) + + +async def send_wake_signal( + client: temporalio.client.Client, + request: WakeRequest, + *, + producer_session_id: str = "", +) -> str: + """Sends one wake and returns the request ID it was sent under. + + Retrying with the same :class:`WakeRequest` is safe and is the intended + recovery from an ambiguous failure: the server deduplicates on the request + ID, so the second attempt resolves the first rather than waking twice. + """ + signal_request = build_signal_request( + request, + identity=client.service_client.config.identity, + producer_session_id=producer_session_id, + ) + await client.workflow_service.signal_workflow_execution(signal_request) + return signal_request.request_id + + +def wake_request_for( + workflow: WorkflowChainKey, + *, + stream_name: str, + wait_id: int, + park_generation: int | None, + sender_identity: str, + wake_counter: int = 0, +) -> WakeRequest: + """Builds the request from a chain key and an observed generation. + + ``park_generation=None`` -- no confirmed park for this subscription -- + becomes an unparked wake rather than no wake at all. + """ + return WakeRequest( + namespace=workflow.namespace, + workflow_id=workflow.workflow_id, + first_execution_run_id=workflow.first_execution_run_id, + stream_name=stream_name, + wait_id=wait_id, + park_generation=( + UNPARKED_WAKE_GENERATION if park_generation is None else park_generation + ), + sender_identity=sender_identity, + wake_counter=wake_counter, + ) diff --git a/temporalio/worker/_replayer.py b/temporalio/worker/_replayer.py index c7ae8d59d0..1387face74 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -8,6 +8,7 @@ from collections.abc import AsyncIterator, Mapping, Sequence from contextlib import AbstractAsyncContextManager, asynccontextmanager from dataclasses import dataclass +from typing import Any from typing_extensions import TypedDict @@ -57,6 +58,7 @@ def __init__( runtime: temporalio.runtime.Runtime | None = None, disable_safe_workflow_eviction: bool = False, header_codec_behavior: HeaderCodecBehavior = HeaderCodecBehavior.NO_CODEC, + external_stream_backend: Any | None = None, stream_provider: temporalio.streams.StreamProvider | None = None, ) -> None: """Create a replayer to replay workflows from history. @@ -70,6 +72,13 @@ def __init__( will default to a new thread pool executor with no max_workers set that will be shared across all replay calls and never explicitly shut down. Users are encouraged to provide their own if needing more control. + + A history containing external workflow stream markers can only be + replayed with the same backend configured, because replay re-reads the + recorded ranges from the provider. Without them the replay fails the + same way the original Workflow would have with no backends configured -- + which is correct, but is a configuration error rather than a finding + about the history. """ self._config = ReplayerConfig( workflows=list(workflows), @@ -88,6 +97,7 @@ def __init__( runtime=runtime, disable_safe_workflow_eviction=disable_safe_workflow_eviction, header_codec_behavior=header_codec_behavior, + external_stream_backend=external_stream_backend, stream_provider=stream_provider, ) self._initial_config = self._config.copy() @@ -282,6 +292,7 @@ def on_eviction_hook( ), should_enforce_versioning_behavior=False, assert_local_activity_valid=lambda a: None, + external_stream_backend=self._config.get("external_stream_backend"), encode_headers=self._config.get( "header_codec_behavior", HeaderCodecBehavior.NO_CODEC ) @@ -437,6 +448,7 @@ class ReplayerConfig(TypedDict, total=False): runtime: temporalio.runtime.Runtime | None disable_safe_workflow_eviction: bool header_codec_behavior: HeaderCodecBehavior + external_stream_backend: Any | None stream_provider: temporalio.streams.StreamProvider | None diff --git a/temporalio/worker/_worker.py b/temporalio/worker/_worker.py index fade69c4d5..04dfdd562a 100644 --- a/temporalio/worker/_worker.py +++ b/temporalio/worker/_worker.py @@ -153,6 +153,7 @@ def __init__( ), disable_payload_error_limit: bool = False, max_workflow_task_external_storage_concurrency: int = _DEFAULT_WORKFLOW_TASK_EXTERNAL_STORAGE_CONCURRENCY, + external_stream_backend: Any | None = None, stream_provider: temporalio.streams.StreamProvider | None = None, ) -> None: """Create a worker to process workflows and/or activities. @@ -279,6 +280,16 @@ def __init__( graceful_shutdown_timeout: Amount of time after shutdown is called that activities are given to complete before their tasks are cancelled. + external_stream_backend: External Workflow Stream provider instance. + The provider must declare ``guarantees_immutability = True``; + Worker construction fails otherwise. Input subscriptions use + ``StreamBackend``. Workflow-originated output additionally + requires the provider to implement ``OutputStreamBackend``; + input-only providers remain supported. + + .. warning:: + This option and the External Workflow Streams feature are + experimental. workflow_failure_exception_types: The types of exceptions that, if a workflow-thrown exception extends, will cause the workflow/update to fail instead of suspending the workflow via @@ -398,6 +409,7 @@ def __init__( nexus_task_poller_behavior=nexus_task_poller_behavior, disable_payload_error_limit=disable_payload_error_limit, max_workflow_task_external_storage_concurrency=max_workflow_task_external_storage_concurrency, + external_stream_backend=external_stream_backend, stream_provider=stream_provider, ) @@ -424,6 +436,20 @@ def _init_from_config(self, client: temporalio.client.Client, config: WorkerConf Client is safe to take separately since it can't be modified by worker plugins. """ self._config = config + + # The external stream backend is validated here rather than lazily so + # that a provider which cannot guarantee record immutability fails + # Worker construction, before any Workflow can consume through it. + self._external_stream_backend = config.get("external_stream_backend") + if self._external_stream_backend is not None: + from temporalio.contrib.external_workflow_streams._backend import ( + _validate_backend, + ) + + self._external_stream_backend = _validate_backend( + self._external_stream_backend + ) + if not ( config.get("activities") or config.get("nexus_service_handlers") @@ -551,6 +577,7 @@ def check_activity(activity: str): self._workflow_worker = _WorkflowWorker( bridge_worker=lambda: self._bridge_worker, + client=client, namespace=config["client"].namespace, # type: ignore[reportTypedDictNotRequiredAccess] task_queue=config["task_queue"], # type: ignore[reportTypedDictNotRequiredAccess] workflows=workflows, @@ -578,6 +605,7 @@ def check_activity(activity: str): encode_headers=client_config["header_codec_behavior"] != HeaderCodecBehavior.NO_CODEC, max_workflow_task_external_storage_concurrency=max_workflow_task_external_storage_concurrency, + external_stream_backend=self._external_stream_backend, stream_provider=stream_provider, ) @@ -882,6 +910,20 @@ async def raise_on_shutdown(): f"Beginning worker shutdown, will wait {graceful_timeout} before cancelling activities" ) + # Ask Core what state each streaming Run is in, before shutdown is + # initiated. This is the first half of P20's sweep and it has to be here: + # an idle cached Run has no pending work, so Core's workflow-state lane + # finishes on the first input after the shutdown token is cancelled, and + # every Run-status probe from then on answers RunNotFound. That answer + # owes a wake just as NoOpenWorkflowTask does, so a sweep that probes + # later still sends its wake and still looks correct -- while Parked (no + # wake needed) and WftOpen (C15b's, not ours to race) stop being + # reachable answers at all. Only the asking happens here; the wakes go + # out below, once the pollers have stopped, because offering a Run to a + # task queue this Worker is still polling is not a hand-off. + if self._workflow_worker: + await self._workflow_worker.probe_external_stream_runs() + # Initiate core worker shutdown self._bridge_worker.initiate_shutdown() @@ -915,6 +957,23 @@ async def raise_on_shutdown(): if self._nexus_worker: await self._nexus_worker.wait_all_completed() + # Send the wakes the probe above found owed, then tear down external + # stream subscriptions. This is the second half of P20's obligation and + # it happens here rather than inside the workflow worker's poll loop + # because both shutdown paths have to reach it: the poll drain above only + # replaces a worker task whose run() raised, so a clean shutdown would + # otherwise leave every Run registered, its watchers running, and its + # owed wake Signals unsent. + # + # After the wait above, so per-Run teardown stays driven by + # RemoveFromCache: every activation, eviction included, has been + # answered by now, so a FinalizeExternalStreams in flight was answered + # before the manager's state for that Run disappears. Before + # finalize_shutdown, because a Run the probe phase never reached is + # probed here and that still needs the bridge worker. + if self._workflow_worker: + await self._workflow_worker.shutdown_external_streams() + # Do final shutdown try: await self._bridge_worker.finalize_shutdown() @@ -1034,6 +1093,7 @@ class WorkerConfig(TypedDict, total=False): nexus_task_poller_behavior: PollerBehavior disable_payload_error_limit: bool max_workflow_task_external_storage_concurrency: int + external_stream_backend: Any | None stream_provider: temporalio.streams.StreamProvider | None diff --git a/temporalio/worker/_workflow.py b/temporalio/worker/_workflow.py index f4001b32e9..5b3cf9a021 100644 --- a/temporalio/worker/_workflow.py +++ b/temporalio/worker/_workflow.py @@ -16,8 +16,10 @@ from typing import Any, cast import temporalio.api.common.v1 +import temporalio.api.enums.v1 import temporalio.bridge.proto.common import temporalio.bridge.proto.workflow_activation +import temporalio.bridge.proto.workflow_commands import temporalio.bridge.proto.workflow_completion import temporalio.bridge.runtime import temporalio.bridge.worker @@ -49,6 +51,7 @@ WorkflowInstance, WorkflowInstanceDetails, WorkflowRunner, + _is_workflow_terminal_command, _WorkflowExternFunctions, _WorkflowLogicFlag, ) @@ -128,6 +131,54 @@ def _set_external_storage_metrics( target.driver_names.extend(sorted(metrics.driver_names)) +def _try_buffer_external_output( + completion: temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion, + runtime: Any, +) -> bool: + """Tell Core to retain a live output batch when this completion can wait. + + A quiescent input snapshot is the ordinary proof that Core can keep this + Workflow Task open. A park recheck that became ready is the other safe + retained transition: the attempted park lost to readiness and Core resumes + the same task. Any user/server command, a confirmed park, or capacity + backpressure takes the ordinary staging path instead. + """ + if runtime.output_rollover_requested: + return False + max_publish_latency = runtime.output_max_publish_latency + if max_publish_latency is None: + return False + + commands = completion.successful.commands + variants = [command.WhichOneof("variant") for command in commands] + became_ready = ( + len(commands) == 1 + and commands[0].HasField("external_stream_park_result") + and commands[0].external_stream_park_result.HasField("became_ready") + ) + quiescent = "workflow_stream_quiescent" in variants and all( + variant in ("workflow_stream_progress", "workflow_stream_quiescent") + for variant in variants + ) + if not became_ready and not quiescent: + return False + + command = temporalio.bridge.proto.workflow_commands.WorkflowCommand() + command.workflow_output_stream_buffered.max_publish_latency.FromTimedelta( + max_publish_latency + ) + # Keep input progress first and the quiescent snapshot last. This mirrors + # output commits and lets Core observe the cursor delta before retaining the + # task whose output deadline it is about to arm. + insert_at = 0 + while insert_at < len(commands) and commands[insert_at].HasField( + "workflow_stream_progress" + ): + insert_at += 1 + commands.insert(insert_at, command) + return True + + class _WorkflowWorker: # type:ignore[reportUnusedClass] def __init__( self, @@ -157,6 +208,8 @@ def __init__( encode_headers: bool, max_workflow_task_external_storage_concurrency: int, default_workflow_logic_flags: frozenset[_WorkflowLogicFlag] | None = None, + external_stream_backend: Any | None = None, + client: Any = None, stream_provider: temporalio.streams.StreamProvider | None = None, ) -> None: # Debug mode is enabled if specified or if the TEMPORAL_DEBUG env var is truthy @@ -222,9 +275,45 @@ def __init__( # itself, after every user interceptor has done its own setup. self._interceptor_classes.append(_StreamHooksInterceptor) + # External Workflow Streams. The manager is per-Worker and owns the + # backend connection and watcher tasks; it is created lazily on the + # Worker's own event loop because that is the loop the watchers must run + # on, and __init__ is not necessarily called from it. + # The runtime exists even when no backend is configured. Recorded stream + # state comes from History, so replay annotations and continuation + # headers must be decoded and validated independently of whether this + # Worker can create a new live subscription. Workflow code still sees + # the feature as unconfigured through `external_streams_configured` in + # its instance details below. + self._external_stream_backend = external_stream_backend + self._external_streams_configured = external_stream_backend is not None + #: Held only to send the reserved wake Signal, which is a raw service + #: call rather than anything the bridge can do -- Core cannot signal a + #: Workflow on this Worker's behalf. + self._client = client + # The taxonomy's counters, created from its own module rather than + # here: P18 names them, documents them, and knows which error class + # belongs to which one. A counter created ad hoc here is a second + # definition of a name an operator alerts on. + from temporalio.contrib.external_workflow_streams._errors import StreamMetrics + + self._stream_metrics = StreamMetrics.create(metric_meter) + self._external_stream_manager: Any = None + self._workflow_failure_exception_types = workflow_failure_exception_types self._patch_activation_callback = patch_activation_callback self._running_workflows: dict[str, _RunningWorkflow] = {} + #: Per-Run stream runtimes, held here rather than read off the instance. + #: A sandboxed Workflow's `instance` is a proxy that exposes only the + #: `WorkflowInstance` protocol, so reaching through it for the runtime + #: silently found nothing -- and the jobs that must never reach + #: `activate()` quietly went to `_apply` instead. + self._external_stream_runtimes: dict[str, Any] = {} + # Stages survive activations within the cached Run because Core may + # defer the server completion behind a local activity. A marker can be + # committed by that later activation even though it staged no new + # output of its own. + self._pending_external_output_stages: dict[str, list[Any]] = {} self._disable_eager_activity_execution = disable_eager_activity_execution self._on_eviction_hook = on_eviction_hook self._disable_safe_eviction = disable_safe_eviction @@ -338,6 +427,58 @@ async def drain_poll_queue(self) -> None: except PollShutdownError: return + async def probe_external_stream_runs(self) -> None: + """Asks Core what state each streaming Run is in. P20's sweep, first half. + + Called by the Worker immediately *before* Core's shutdown is initiated, + because that is the last moment the answer exists. An idle cached Run has + no pending work, so Core's ``shutdown_done`` is satisfied by the first + input after the shutdown token is cancelled and the workflow-state lane + ends there; every probe afterwards answers ``RunNotFound``. Since + ``RunNotFound`` owes a wake just as ``NoOpenWorkflowTask`` does, a sweep + that probes too late still sends its wake and still looks right, while + the two answers that mean *don't* send one -- ``Parked``, which needs no + wake, and ``WftOpen``, which belongs to C15b -- can no longer occur. + + Only the asking happens here. The wakes are sent from + :py:meth:`shutdown_external_streams`, after the pollers have stopped: + offering the Run to a task queue this Worker is still polling would be + the opposite of a hand-off. + + The manager bounds this with its own short grace period, so a wedged + Core cannot delay the stop-polling step. + """ + if self._external_stream_manager is not None: + await self._external_stream_manager.probe_runs() + + async def shutdown_external_streams(self) -> None: + """Sends the owed wakes and tears the manager down. P20's second half. + + Called by the Worker once every activation has been dealt with, on + *both* shutdown paths. It deliberately does not live in + :py:meth:`drain_poll_queue`, which the Worker substitutes only for a + worker task whose ``run()`` raised: wiring it there means a clean + ``Worker.shutdown()`` never sweeps at all, leaving Runs registered, + watchers running, and buffers and backend connections open in a process + that is about to exit. + + Here, and not with the probe, because per-Run teardown is driven by + ``RemoveFromCache`` and nothing else: a ``FinalizeExternalStreams`` in + flight has to be answered before the manager's state for that Run + disappears. + + It is not folded into eviction either, because an *idle cached Run + receives no eviction activation at shutdown at all* -- ``shutdown_done`` + treats a Run with no pending work as finished -- and that is exactly the + Run that most needs a wake: its records are buffered here and nothing + else will ever tell the Workflow they arrived. + + The manager bounds the sweep with its own grace period, so this cannot + hold shutdown open indefinitely. + """ + if self._external_stream_manager is not None: + await self._external_stream_manager.shutdown() + async def _activate_inline_for_debug( self, loop: asyncio.AbstractEventLoop, @@ -388,12 +529,16 @@ async def _handle_activation( await self._handle_cache_eviction(act, cache_remove_job) return + if self._external_stream_manager is not None: + self._external_stream_manager.note_workflow_task_started(act.run_id) + # Build default success completion (e.g. remove-job-only activations) completion = ( temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion() ) completion.successful.SetInParent() workflow = None + workflow_id: str | None = None data_converter = self._data_converter download_metrics = temporalio.converter._extstore.StorageOperationMetrics() try: @@ -454,7 +599,19 @@ async def _handle_activation( ) self._running_workflows[act.run_id] = workflow - if self._debug_mode: + # Two of the four stream jobs are *themselves* backend operations, + # and a third has to be prepared by one. Routing them through the + # synchronous `activate()` would put a multi-second transaction + # inside a call running under a 2-second deadlock timeout -- failing + # the Workflow Task for a perfectly healthy backend, and getting + # worse the more records replay must validate. + # + # So they are partitioned here, in the async layer that already + # awaits `decode_activation` before handing anything to the executor. + handled = await self._handle_external_stream_jobs(act, workflow) + if handled is not None: + completion = handled + elif self._debug_mode: # Inline on the main thread so pdb / breakpoint() can read # stdin. The loop blocks during the activation — that's the # intended single-stepping semantic. @@ -492,6 +649,19 @@ async def _handle_activation( workflow.deadlocked_activation_task = activate_task raise deadlock_exc from None + output_runtime = self._external_stream_runtimes.get(act.run_id) + if ( + output_runtime is not None + and completion.HasField("successful") + and not act.is_replaying + and output_runtime.has_output + ): + await self._stage_or_buffer_external_output( + act, + completion, + output_runtime, + ) + except Exception as err: if isinstance(err, _DeadlockError): err.swap_traceback() @@ -539,6 +709,26 @@ async def _handle_activation( f"Failed converting activation exception: {inner_err}" ) + # One place for every failed completion, however it was reached. An + # external stream failure arrives two ways -- raised out here by a + # runtime-only job, or raised on the Workflow thread by a delivery and + # already turned into a failure by `activate()` -- and both must carry + # the same cause and increment the same counter. + if self._external_stream_backend is not None and completion.HasField("failed"): + try: + self._note_external_stream_failure(completion) + except Exception: + # Reporting a failure may not *become* one. This runs outside + # the block that turns an exception into a failed completion, + # so anything raised here would escape with the completion + # unsent -- turning a Workflow Task failure the server can see + # into a Workflow Task timeout it cannot explain. + logger.exception( + "Failed classifying an external stream failure on workflow " + "with run ID %s", + act.run_id, + ) + completion.run_id = act.run_id # Encode completion @@ -590,6 +780,29 @@ async def _handle_activation( logger.exception( "Failed completing activation on workflow with run ID %s", act.run_id ) + else: + if workflow_id is not None: + await self._promote_external_output( + run_id=act.run_id, + workflow_id=workflow_id, + ) + # A wake accepted by the service stays outstanding until the task it + # caused completes successfully. Failed/rejected tasks replay, and + # readiness rebuilt during that replay must remain coalesced behind + # the original wake. Only Core accepting this completion proves the + # cycle ended and permits another buffered generation to wake. + if self._external_stream_manager is not None and completion.HasField( + "successful" + ): + self._external_stream_manager.note_workflow_task_completed( + act.run_id, + terminal=any( + map( + _is_workflow_terminal_command, + completion.successful.commands, + ) + ), + ) async def _handle_cache_eviction( self, @@ -693,6 +906,16 @@ async def _handle_cache_eviction( # Remove from map and send completion if act.run_id in self._running_workflows: del self._running_workflows[act.run_id] + # Per-Run stream teardown is driven by `RemoveFromCache` and nothing + # else, so a `FinalizeExternalStreams` in flight is always answered + # before the manager's state for the Run disappears. It is done here + # rather than inside the instance because `disable_safe_eviction` + # skips the instance's eviction job entirely, and a Run whose + # watchers outlived it would keep a backend connection open forever. + self._external_stream_runtimes.pop(act.run_id, None) + self._pending_external_output_stages.pop(act.run_id, None) + if self._external_stream_manager is not None: + await self._external_stream_manager.evict_run(act.run_id) try: await self._bridge_worker().complete_workflow_activation( temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion( @@ -786,6 +1009,8 @@ def _create_workflow_instance( ) # Create instance from details + runtime = self._create_external_stream_runtime(act, init) + self._external_stream_runtimes[act.run_id] = runtime det = WorkflowInstanceDetails( # Sharing the underlying converter could let another workflow's state # affect conversion and replay, so each instance needs a fresh one. @@ -802,6 +1027,8 @@ def _create_workflow_instance( last_completion_result=init.last_completion_result, last_failure=last_failure, default_workflow_logic_flags=frozenset(self._default_workflow_logic_flags), + external_stream_runtime=runtime, + external_streams_configured=self._external_streams_configured, stream_provider=self._stream_provider, ) if defn.sandboxed: @@ -809,6 +1036,596 @@ def _create_workflow_instance( else: return self._unsandboxed_workflow_runner.create_instance(det) + async def _handle_external_stream_jobs( + self, + act: temporalio.bridge.proto.workflow_activation.WorkflowActivation, + _workflow: _RunningWorkflow, + ) -> ( + temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion | None + ): + """Answers the runtime-only stream jobs without calling ``activate()``. + + Returns the synthesized completion when this activation was entirely + one of those jobs, or ``None`` to let the normal path run. + + ``PrepareExternalStreamPark`` and ``FinalizeExternalStreams`` run no + user Workflow code at all -- they cannot resolve futures and their only + legal answers are their own result command or an activation failure. + ``ReplayExternalStreams`` is *prepared* here rather than handled: its + recorded ranges are read and validated into the buffers, then the job + passes through so ``_apply`` delivers from memory exactly as it does for + a live resolve. + """ + runtime = self._external_stream_runtimes.get(act.run_id) + if runtime is None: + if any(j.HasField("replay_external_streams") for j in act.jobs): + raise RuntimeError( + "received an external stream replay job without the per-Run " + "validation runtime" + ) + return None + + if any(j.HasField("resolve_external_stream_waits") for j in act.jobs): + # Core is telling this Worker the wait set has moved on, which is + # also the only notice that a *confirmed* park is over: a wake Signal + # and a fresh quiescent snapshot both clear Core's `park_generation`, + # and neither is visible in the backend. The intent installed for + # that park comes out here, before the Workflow resumes, so nothing + # can read a generation that no longer exists -- not a producer + # choosing what its wake names, and not this Worker's own shutdown + # sweep. The job itself passes through to `_apply` unchanged. + await self._stream_manager().resolve_park(act.run_id) + + park = next( + ( + j.prepare_external_stream_park + for j in act.jobs + if j.HasField("prepare_external_stream_park") + ), + None, + ) + finalize = next( + ( + j.finalize_external_streams + for j in act.jobs + if j.HasField("finalize_external_streams") + ), + None, + ) + replay = next( + ( + j.replay_external_streams + for j in act.jobs + if j.HasField("replay_external_streams") + ), + None, + ) + + if replay is not None: + # Prepared, not handled: fill and validate every recorded range + # before the job reaches the Workflow thread. A transient backend + # error or an integrity violation surfaces from here through the + # activation-failure path rather than as a deadlock timeout, which + # would misattribute a storage problem to the Workflow's own code. + plan = None + if replay.replay_annotation: + plan = await self._stream_manager().prepare_replay( + act.run_id, replay.replay_annotation + ) + # The other half of the same record's decoding. `prepare_replay` + # bound the codec to the stream key the *marker* recorded; the + # converter that runs inside `activate()` would otherwise stay bound + # to this Run's identity, and an offline `Replayer` supplies its own + # namespace for that -- one record, two Workflow identities. Bound + # out here for the reason `_create_external_stream_runtime` gives: + # `with_context` clones the user's component converters, and the + # runtime crosses into the Workflow sandbox. + if plan is not None: + runtime.install_replay_converters(self._replay_stream_converters(plan)) + return None + + if park is None and finalize is None: + return None + + completion = ( + temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion() + ) + completion.successful.SetInParent() + + if finalize is not None: + # **No backend work at all** (ADR-010). The terminal is read from + # the runtime's own in-memory blocked snapshot: the boundary is not + # "wherever the stream is now", it is where this Workflow Task's + # deliveries stopped, which was fixed the moment the last activation + # returned. Refreshing it against the backend would be actively + # wrong -- it could name a position replay must not reproduce. + command = completion.successful.commands.add() + command.external_stream_finalized.quiescence_generation = ( + finalize.quiescence_generation + ) + command.external_stream_finalized.final_observation_delta = ( + runtime.add_terminal() + ) + return completion + + assert park is not None + # The set to park is **Core's**, not this Worker's registration list. + # `park.waits` is the complete blocked snapshot Core is holding the + # Workflow Task for; `blocked_snapshot()` is every subscription the + # runtime has registered, which is a superset -- a subscription that + # delivered a record and was not awaited again is registered and not + # blocked. Parking the superset rechecks a wait nothing is waiting on, + # so its records abort a park that was entirely legitimate and the + # handshake repeats on every idle timeout; and it installs an intent for + # a wait Core is not parking, which is an intent with no park behind it. + # The runtime supplies only the cursor boundary for each of Core's + # waits. + snapshot = runtime.blocked_snapshot() + became_ready = await self._stream_manager().prepare_park( + act.run_id, + park.quiescence_generation, + { + wait.wait_id: snapshot[wait.wait_id] + for wait in park.waits + if wait.wait_id in snapshot + }, + ) + command = completion.successful.commands.add() + command.external_stream_park_result.quiescence_generation = ( + park.quiescence_generation + ) + if became_ready: + # A recheck found records, so this parking generation is abandoned + # and Core issues a normal resolve activation next -- rather than + # running user code from inside the park path. + command.external_stream_park_result.became_ready.SetInParent() + else: + command.external_stream_park_result.confirmed.SetInParent() + command.external_stream_park_result.final_observation_delta = ( + runtime.add_terminal() + ) + return completion + + async def _stage_or_buffer_external_output( + self, + act: temporalio.bridge.proto.workflow_activation.WorkflowActivation, + completion: temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion, + runtime: Any, + ) -> None: + """Retain a safe batch, or stage it before reporting this completion.""" + if _try_buffer_external_output(completion, runtime): + return + staged_output = await self._stage_external_output(act, completion, runtime) + pending = self._pending_external_output_stages.setdefault(act.run_id, []) + for topic in staged_output.topics: + if topic.manifest not in pending: + pending.append(topic.manifest) + + async def _stage_external_output( + self, + act: temporalio.bridge.proto.workflow_activation.WorkflowActivation, + completion: temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion, + runtime: Any, + ) -> Any: + """Stage Workflow output before its compact marker command reaches Core.""" + staged = await runtime.stage_output(act.history_floor_event_id) + command = temporalio.bridge.proto.workflow_commands.WorkflowCommand() + output = command.workflow_output_stream_commit + output.request_rollover = runtime.output_rollover_requested + manifest = output.manifest + manifest.schema_version = staged.schema_version + manifest.fingerprint_version = staged.fingerprint_version + manifest.stage_token = staged.stage_token + manifest.history_floor_event_id = staged.history_floor_event_id + manifest.run_id = staged.run_id + manifest.provider_id = staged.provider_id + manifest.provider_format_version = staged.provider_format_version + for staged_topic in staged.topics: + source = staged_topic.manifest + topic = manifest.topics.add() + topic.topic = source.stream_key.stream_name + topic.record_count = source.record_count + topic.logical_byte_count = source.logical_byte_count + topic.logical_fingerprint = source.fingerprint + topic.finished = staged_topic.finished + for counts in staged.segment_record_counts: + segment = manifest.segments.add() + segment.record_counts_by_topic.extend(counts) + + # Input progress remains first. It validates consumed input before any + # command that could depend on it, while this internal commit sits ahead + # of the user's server-bound commands. + commands = completion.successful.commands + insert_at = 0 + while insert_at < len(commands) and commands[insert_at].HasField( + "workflow_stream_progress" + ): + insert_at += 1 + commands.insert(insert_at, command) + runtime.output_stage_recorded() + return staged + + async def _promote_external_output( + self, + *, + run_id: str, + workflow_id: str, + ) -> None: + """Promote a staged batch only after History proves its marker. + + Core's completion call has already made its server RPC by this point, + but it deliberately absorbs report errors in order to drive eviction + and task retry. A successful bridge await therefore is not itself a + commit proof. Reusing the cold-client History predicate keeps definite + rejection and ambiguous transport outcomes safely pending/aborted while + making the healthy path visible without waiting for a reader. + """ + backend = self._external_stream_backend + client = self._client + if backend is None or client is None: + return + + from temporalio.contrib.external_workflow_streams._output_client import ( + _reconcile_output_stage, + ) + + pending = self._pending_external_output_stages.get(run_id) + if not pending: + return + unresolved: list[Any] = [] + for manifest in pending: + try: + resolved = await _reconcile_output_stage( + backend=backend, + client=client, + workflow_id=workflow_id, + manifest=manifest, + ) + if not resolved: + unresolved.append(manifest) + except asyncio.CancelledError: + raise + except Exception as err: + # Promotion is an optimization over the durable pending + # barrier. A cold client repeats this exact reconciliation, so + # a transient History/backend failure must not retroactively + # turn an already-reported Workflow Task into an SDK failure. + self._stream_metrics.record(err) + logger.warning( + "Could not opportunistically reconcile external output stage " + "%s for Workflow %s; it remains pending for a client", + manifest.stage_token, + workflow_id, + exc_info=True, + ) + unresolved.append(manifest) + if unresolved: + self._pending_external_output_stages[run_id] = unresolved + else: + self._pending_external_output_stages.pop(run_id, None) + + def _replay_stream_converters( + self, plan: Any + ) -> dict[int, temporalio.converter.DataConverter]: + """One converter per wait the marker binds, in the recorded context. + + From the annotation's header, which is the only place a replayed + record's stream is written down -- the Workflow has not run far enough + to have re-created the subscription that would otherwise carry it. + + Bound from the Worker's own converter rather than from the runtime's, + which is already bound to this Run: a second `with_context` over the + first would ask a user's component converter to rebind itself, and + nothing in the protocol promises that composes. + + Memoized per ``(namespace, workflow_id)`` exactly as the manager's own + preparation is, so a marker binding several waits of one Workflow clones + the component converters once rather than once per wait. + """ + bound: dict[tuple[str, str], temporalio.converter.DataConverter] = {} + converters: dict[int, temporalio.converter.DataConverter] = {} + for wait_id, binding in plan.annotation.header.streams.items(): + key = binding.stream_key + cached = bound.get((key.namespace, key.workflow_id)) + if cached is None: + cached = self._data_converter.with_context( + temporalio.converter.WorkflowSerializationContext( + namespace=key.namespace, + workflow_id=key.workflow_id, + ) + ) + bound[(key.namespace, key.workflow_id)] = cached + converters[wait_id] = cached + return converters + + def _stream_manager(self) -> Any: + """The Worker's subscription manager, created on first use. + + Lazily, because it captures the running event loop -- the one its + watcher tasks must run on -- and ``__init__`` is not necessarily called + from that loop. + """ + if self._external_stream_manager is None: + from temporalio.contrib.external_workflow_streams._manager import ( + StreamSubscriptionManager, + ) + + self._external_stream_manager = StreamSubscriptionManager( + backend=self._external_stream_backend, + # The Worker's converter, for the **asynchronous half** of + # decoding a record: external-payload retrieval and the user's + # PayloadCodec. Both are arbitrary asynchronous work and neither + # needs the topic's declared type, so both belong out here on + # this loop -- the same place `decode_activation` awaits them + # for every other payload an activation carries. Without this + # the Workflow thread awaits them inside `activate()`, which + # performs I/O in a deterministic event loop and puts a user + # codec under the 2-second deadlock timeout. + data_converter=self._data_converter, + notify_ready=self._bridge_worker().notify_external_stream_ready, + # A Worker whose Run cannot take local readiness owes the same + # Signal a producer owes. Without this the record sits buffered + # while the Workflow waits for a task that nothing will create: + # parked, cached-with-no-open-task, and evicted all look the + # same from here, and all three are answered the same way. + send_wake=self._send_external_stream_wake, + # Read-only, and deliberately not the readiness call: readiness + # asserts a buffered record, so probing with it on the way out + # would manufacture a Workflow Task for a Run that had nothing + # waiting. + run_status=self._bridge_worker().external_stream_run_status, + shutdown_wake_failed_metric=self._record_shutdown_wake_failed, + # Only a prefix on this Worker's sender identity, which the + # manager makes unique per instance. Passed so a wake request ID + # stays traceable to a client in server-side logs; it is + # deliberately not the identity itself, since two Workers + # sharing one Client share this string. + client_identity=( + self._client.service_client.config.identity + if self._client is not None + else "" + ), + ) + return self._external_stream_manager + + def _note_external_stream_failure( + self, + completion: temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion, + ) -> None: + """Applies the external stream failure taxonomy to a failed completion. + + Three of the taxonomy's four rows are Workflow Task failures that differ + from every other Workflow Task failure -- and from each other -- only in + **the error type and the metric**. The server retries a failed Workflow + Task regardless of cause, so nothing about the retry distinguishes a + backend outage that will clear on its own from integrity loss that needs + an operator, or from a converter mismatch that needs a code change. + This is what makes those rows tellable apart: + + - ``force_cause`` is set to the external-storage cause, which is what + separates all three from an ordinary Workflow bug in server-side + Workflow Task failure reporting; + - the matching counter, and only the matching counter, is incremented, + so an alert on integrity loss is not diluted by a backend outage. + + Row four -- the annotation not matching the subscriptions Workflow code + creates -- is deliberately absent: it is ordinary nondeterminism, gets + no stream cause and no stream counter, and is fixed by versioning the + Workflow rather than by touching the backend. + + The failure is inspected rather than the exception, because half of + these failures never exist as an exception out here: a decode that + raises on the Workflow thread is converted inside ``activate()``, and + what comes back is a completion. The application failure type is the + exception's class name, and the chain is walked because the raising + frame may have wrapped it. + """ + failure = completion.failed.failure + counter = None + while True: + counter = self._stream_metrics.counter_for( + failure.application_failure_info.type + ) + if counter is not None or not failure.HasField("cause"): + break + failure = failure.cause + if counter is None: + return + + completion.failed.force_cause = temporalio.api.enums.v1.WorkflowTaskFailedCause.WORKFLOW_TASK_FAILED_CAUSE_EXTERNAL_STORAGE_FAILURE + counter.add(1) + + def _record_shutdown_wake_failed(self, subscription: Any) -> None: + """Counts a shutdown wake that could not be acknowledged. + + A dropped wake is silent by nature -- the Workflow simply waits, and + nothing distinguishes that from a producer having nothing to say -- so it + gets a metric rather than only a log line. + """ + self._stream_metrics.shutdown_wake_failed.add( + 1, + { + "namespace": self._namespace, + "task_queue": self._task_queue, + "stream_name": subscription.stream_key.stream_name, + }, + ) + + async def _send_external_stream_wake(self, subscription: Any) -> None: + """Sends the reserved wake Signal for a subscription that owes one. + + Addressed to the Workflow ID with no Run ID, so it lands on the current + Run of the chain -- which may already be a successor by the time this + runs. + + A failure is logged **and re-raised to the manager that called this**. + There is no Workflow Task to fail out here, so the exception is not a + way of reporting anything to a Workflow -- it is how the caller learns + the wake was not acknowledged. The shutdown sweep is the caller that + acts on it: it retries within its grace period and counts the wake on + ``external_stream_shutdown_wake_failed`` when the retries run out. + Returning normally instead makes an unacknowledged wake + indistinguishable from a delivered one, which ends that retry loop after + one attempt and reports a clean shutdown that lost a record. + + The request ID is derived from the wake's identity, so a retry is the + same wake rather than a second one -- which is what makes re-sending an + attempt that may in fact have arrived safe. + + The manager's live watcher path is the other caller, and it must guard + this call itself: an exception escaping ``_report_ready`` ends the + watcher task for good. + + ``subscription`` is whatever the manager composes a wake from, which is + not always a subscription: a stale park intent retired after its wait was + closed or its Run evicted still owes the wake it silenced, and the + manager carries that obligation on an object of its own. + """ + from temporalio.contrib.external_workflow_streams._wake import ( + WakeRequest, + send_wake_signal, + ) + + if self._client is None: + # Raised for the same reason a failed send is: the wake is owed and + # will not be sent, and a caller told nothing counts it as + # delivered. + raise RuntimeError( + "an external stream wake Signal is owed but this Worker has no " + "client to send it with; the Workflow would wait out its idle " + "timeout instead" + ) + + key = subscription.stream_key + # Asked of the manager rather than read straight from the backend. The + # park this wake must name is whatever is installed *now* -- a generation + # cached when the watcher started would name a park since abandoned and + # resolved -- but "installed" is not "live", and only the manager's + # owed-removal ledger can tell an intent Core is still parked on from one + # this Worker has already decided to remove. Naming the latter sends a + # Signal Core discards while reporting success, which is how the shutdown + # sweep came to count an obsolete generation as a handoff it had made. + generation = ( + await self._stream_manager().wake_park_generation(subscription) or 0 + ) + try: + await send_wake_signal( + self._client, + WakeRequest( + namespace=key.namespace, + workflow_id=key.workflow_id, + first_execution_run_id=key.first_execution_run_id, + stream_name=key.stream_name, + wait_id=subscription.wait_id, + park_generation=generation, + # This Worker's own identity, not the client's: two Workers + # in one process share a Client, and a shared identity would + # give their unparked wakes the same request ID for the + # server to deduplicate -- losing the second Worker's wake. + sender_identity=self._stream_manager().wake_sender_identity, + # Each wake cycle is a separate ask, not a retry of the last + # completed one. The manager coalesces reports until Core + # accepts that cycle's successful task completion. + wake_counter=subscription.wake_counter, + ), + ) + except Exception: + logger.exception( + "Failed sending external stream wake Signal for %s wait %s", + key, + subscription.wait_id, + ) + raise + + def _create_external_stream_runtime( + self, + act: temporalio.bridge.proto.workflow_activation.WorkflowActivation, + init: temporalio.bridge.proto.workflow_activation.InitializeWorkflow, + ) -> Any: + """The per-Run handle Workflow code reaches the manager through. + + Created even when no backends are registered because recorded state is + authoritative: continuation headers and replay annotations still need + to be decoded, and quiet marker bindings still need their reverse + nondeterminism check. The instance installs this handle into Workflow + code only when backends were configured, preserving ``subscribe()``'s + explicit error naming the Worker option. + """ + from temporalio.contrib.external_workflow_streams._api import ( + DEFAULT_IDLE_TIMEOUT, + ) + from temporalio.contrib.external_workflow_streams._continuation import ( + read_continuation_header, + ) + from temporalio.contrib.external_workflow_streams._output_continuation import ( + read_output_continuation_header, + ) + from temporalio.contrib.external_workflow_streams._runtime import ( + WorkflowStreamRuntime, + ) + + # Decode before consulting configuration. The reserved header is + # must-understand History state, so an unsupported encoding cannot turn + # into an ordinary first execution on a Worker whose current Workflow + # code no longer calls subscribe(). A valid non-empty continuation also + # needs its recorded backend configuration even if current code removed + # every stream call; otherwise the successor silently discards the + # cursor and binding its predecessor committed. + continuation = read_continuation_header(dict(init.headers)) + output_continuation = read_output_continuation_header(dict(init.headers)) + if ( + continuation is not None + and continuation.cursors + or output_continuation is not None + ) and not self._external_streams_configured: + raise RuntimeError( + "Workflow History contains external stream continuation state, " + "but external streams are not configured on this Worker; pass " + "external_stream_backend=... to the Worker" + ) + + return WorkflowStreamRuntime( + manager=self._stream_manager(), + backend=self._external_stream_backend, + run_id=act.run_id, + namespace=self._namespace, + workflow_id=init.workflow_id, + # The *chain* key, not this Run: the stream spans the whole + # Continue-As-New chain, so a new Run continues the same stream + # rather than starting a fresh one. + first_execution_run_id=init.first_execution_run_id, + # Bound to the consuming Workflow, the same way `decode_activation` + # binds every other payload this activation carries. A converter + # that derives a key from the Workflow it serves would otherwise get + # the right context for the Workflow's own argument and no context + # at all for a stream record delivered in the very same activation. + # + # Bound *here* rather than inside the runtime's `codec_for`: the + # runtime crosses into the Workflow sandbox, and `with_context` runs + # user code to clone the component converters -- work that belongs + # on this side of the boundary, and that a per-Run handle need do + # only once. + # + # `with_context` and not `_with_contexts`: the store context names + # the Workflow as the payload's *storer*, and a stream record is + # stored by its producer, not by this Run. `with_context` also + # returns the converter unchanged unless a component implements + # `WithSerializationContext`, so the default converter is untouched. + data_converter=self._data_converter.with_context( + temporalio.converter.WorkflowSerializationContext( + namespace=self._namespace, + workflow_id=init.workflow_id, + ) + ), + default_idle_timeout=DEFAULT_IDLE_TIMEOUT, + # Read here, before the Workflow object exists and therefore before + # any subscribe() call: a start cursor restored after a subscription + # was established would already have been overwritten by BEGINNING + # (ADR-022). + continuation=continuation, + output_continuation=output_continuation, + ) + def nondeterminism_as_workflow_fail(self) -> bool: return any( issubclass(temporalio.workflow.NondeterminismError, typ) diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py index 58a0dcfae9..5dea6ebcb5 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -107,6 +107,18 @@ LOG_IGNORE_DURING_DELETE = False +def _is_workflow_terminal_command( + command: temporalio.bridge.proto.workflow_commands.workflow_commands_pb2.WorkflowCommand, +) -> bool: + """Whether a command ends the current Workflow Run.""" + return ( + command.HasField("complete_workflow_execution") + or command.HasField("continue_as_new_workflow_execution") + or command.HasField("fail_workflow_execution") + or command.HasField("cancel_workflow_execution") + ) + + async def _shield_await(fut: asyncio.Future[Any]) -> Any: """Await a future without cancelling it if the awaiting task is cancelled. @@ -193,6 +205,20 @@ class WorkflowInstanceDetails: default_workflow_logic_flags: frozenset[_WorkflowLogicFlag] = field( default_factory=lambda: _DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS ) + external_stream_runtime: Any = None + """Opaque per-Run handle to the Worker's External Workflow Stream manager. + + Crosses into the sandbox by reference, like the rest of these details, which + is the whole point: the manager owns the Worker's backend connections and + watcher tasks, and a copy re-created inside the sandbox would watch nothing. + Workflow code only ever sees this handle -- never a provider instance. + + Present for validation even when no ``external_stream_backend`` was + configured on the Worker, because recorded state must be handled based on + History rather than current Worker configuration. + """ + external_streams_configured: bool = False + """Whether Workflow code may use the runtime for new subscriptions.""" stream_provider: temporalio.streams.StreamProvider | None = None @@ -308,6 +334,17 @@ def __init__(self, det: WorkflowInstanceDetails) -> None: ) self._extern_functions = det.extern_functions + self._external_stream_runtime = det.external_stream_runtime + self._external_streams_configured = det.external_streams_configured + self._pending_replay_finish: Any = None + """A marker replay whose last segment the activation's own drain serves. + + Set by ``_apply_replay_external_streams`` and consumed by + ``_finish_replay_external_streams``, which the activation calls once that + drain has run. Held here rather than on the runtime because what it + defers is a step of the *activation*, not of the stream runtime. + """ + self._pending_output_replay_finish = False self._disable_eager_activity_execution = det.disable_eager_activity_execution self._worker_level_failure_exception_types = ( det.worker_level_failure_exception_types @@ -486,6 +523,22 @@ def activate( ) self._time_ns = act.timestamp.ToNanoseconds() self._is_replaying = act.is_replaying + if self._external_stream_runtime is not None and not any( + job.HasField("remove_from_cache") for job in act.jobs + ): + # Re-arms the per-activation delivery budget. It has to be reset here + # rather than anywhere later: a producer that keeps a subscription's + # buffer non-empty would otherwise keep the iterator fed for the + # whole of this call, and this call runs on a thread-pool executor + # under a 2-second deadlock timeout that every retry would hit again. + # + # A cache-removal activation is not a Workflow Task. In particular, + # it must still be able to tear down a runtime whose output staging + # failed: treating the eviction's (empty) history floor as a new task + # would reject the still-uncommitted batch and poison eviction. + self._external_stream_runtime.begin_activation( + getattr(act, "history_floor_event_id", None) + ) self._current_thread_id = threading.get_ident() self._current_internal_flags = act.available_internal_flags self._single_batch_activation = self._workflow_logic_flag_enabled( @@ -520,30 +573,63 @@ def activate( if start_job: self._workflow_input = self._make_workflow_input(start_job) - if self._single_batch_activation: - # Applying every job before giving workflow tasks a chance to - # run prevents their order in the activation from hiding state - # that arrived in the same workflow task. - for job_set in job_sets: - for job in job_set: - # Let errors bubble out of these to the caller to fail the task - self._apply(job) - if any(job_sets): - self._run_once(check_conditions=bool(job_sets[1] or job_sets[2])) + try: + if self._single_batch_activation: + # Applying every job before giving workflow tasks a chance to + # run prevents their order in the activation from hiding state + # that arrived in the same workflow task. + for job_set in job_sets: + for job in job_set: + # Let errors bubble out of these to the caller to fail the task + self._apply(job) + if any(job_sets): + self._run_once( + check_conditions=bool(job_sets[1] or job_sets[2]) + ) + else: + # Preserve the legacy scheduling order for histories which do + # not contain the single-batch workflow logic flag. + for index, job_set in enumerate(job_sets): + if not job_set: + continue + for job in job_set: + # Let errors bubble out of these to the caller to fail the task + self._apply(job) + + # Run one iteration of the loop. We do not allow conditions to + # be checked in patch jobs (first index) or query jobs (last + # index). + self._run_once(check_conditions=index == 1 or index == 2) + except BaseException: + # An error is already on its way out, so the replay is *abandoned* + # rather than closed. Closing runs `verify_replay_consumed`, which + # raises whenever a recorded delivery is still armed -- which it is, + # since the drain that would have taken it is the one that just + # failed -- and an exception raised while another is propagating + # **replaces** it. A user's failure, or an unrelated nondeterminism + # error, would reach the completion path as a nondeterminism error + # blaming a `subscribe()` call nobody touched: the wrong diagnosis, + # and the wrong classification with it, since a workflow-failing + # error would be retried as a task failure instead. + # + # Replay mode is still left, because a Run stuck in it has every + # later drain return nothing at all. Only the checks and the + # cursor move are skipped -- and the reposition must be skipped + # here in any case: an activation that failed committed nothing. + self._abandon_replay_external_streams() + raise else: - # Preserve the legacy scheduling order for histories which do - # not contain the single-batch workflow logic flag. - for index, job_set in enumerate(job_sets): - if not job_set: - continue - for job in job_set: - # Let errors bubble out of these to the caller to fail the task - self._apply(job) - - # Run one iteration of the loop. We do not allow conditions to - # be checked in patch jobs (first index) or query jobs (last - # index). - self._run_once(check_conditions=index == 1 or index == 2) + # A replay marker's last recorded segment is drained by the + # activation's *own* drain above, not by the replay driver, so + # what closes the replay -- the consumed check, the cursor + # reposition, and leaving replay mode -- can only happen here. + self._finish_replay_external_streams() + + # Detected *after* the drain, not during it: `_run_once` already + # drains `self._ready` until empty, so "no coroutine is runnable" is + # its post-condition and quiescence is a registry check rather than + # an event-loop change. + self._emit_external_stream_commands() except Exception as err: # We want some errors during activation, like those that can happen # during payload conversion, to be able to fail the workflow not the @@ -608,17 +694,12 @@ def activate( ) self._current_completion.failed.failure.application_failure_info.SetInParent() - def is_completion( - command: temporalio.bridge.proto.workflow_commands.workflow_commands_pb2.WorkflowCommand, - ): - return ( - command.HasField("complete_workflow_execution") - or command.HasField("continue_as_new_workflow_execution") - or command.HasField("fail_workflow_execution") - or command.HasField("cancel_workflow_execution") + if any( + map( + _is_workflow_terminal_command, + self._current_completion.successful.commands, ) - - if any(map(is_completion, self._current_completion.successful.commands)): + ): self._warn_if_unfinished_handlers() return self._current_completion @@ -640,6 +721,10 @@ def _apply( self._apply_remove_from_cache(job.remove_from_cache) elif job.HasField("resolve_activity"): self._apply_resolve_activity(job.resolve_activity) + elif job.HasField("resolve_external_stream_waits"): + self._apply_resolve_external_stream_waits(job.resolve_external_stream_waits) + elif job.HasField("replay_external_streams"): + self._apply_replay_external_streams(job.replay_external_streams) elif job.HasField("resolve_child_workflow_execution"): self._apply_resolve_child_workflow_execution( job.resolve_child_workflow_execution @@ -819,6 +904,236 @@ def _apply_fire_timer( if handle: self._ready.append(handle) + def _apply_resolve_external_stream_waits( + self, + job: temporalio.bridge.proto.workflow_activation.ResolveExternalStreamWaits, + ) -> None: + """Resumes every wait that has something buffered. + + **Performs no I/O.** Everything this touches is already in memory: the + manager reported readiness only once a record was buffered, so the drain + that follows is bounded and cannot block. Reading the backend here would + put a multi-second transaction inside a synchronous activation running + under a 2-second deadlock timeout. + + The job's ``ready_hints`` are hints, not an exhaustive availability + claim, so *every* active wait is resolved and left to find out for + itself whether anything arrived. Resolving only the hinted ones would + strand a record whose readiness notification was coalesced away. + """ + del job # The hints are deliberately unused; see above. + if self._external_stream_runtime is not None: + self._external_stream_runtime.resolve_all_pending() + + def _apply_replay_external_streams( + self, + job: temporalio.bridge.proto.workflow_activation.ReplayExternalStreams, + ) -> None: + """Delivers a recorded marker's observations, from memory only. + + The recorded ranges were read and validated before this activation was + handed to the Workflow thread, so this performs no I/O at all. + + The live run's *k* activations become *k* drains inside this one + activation. Segments are walked in order with one event-loop drain each, + rather than delivered all at once, because reproducing the record order + while changing how many drains occurred would make ``wait_condition`` + predicates fire a different number of times than they did live + (ADR-018). + + *k*, not *k + 1*: the activation runs a drain of its own for the job set + this job arrived in, so only the first *k - 1* segments are drained here + and the last is left to that one. Closing the replay therefore also + moves to after it, in :meth:`_finish_replay_external_streams`. + + This is safe with respect to Workflow time: every segment of a marker + belongs to one Workflow Task, so ``workflow.now()`` is constant across + them in both directions. + """ + runtime = self._external_stream_runtime + if runtime is None: + raise RuntimeError( + "received an external stream replay job without the per-Run " + "validation runtime" + ) + has_field = getattr(job, "HasField", None) + try: + has_output = bool(has_field and has_field("output")) + except ValueError: + # Compatible with a generated binding from before the additive + # output field (and with input-only replay driver test doubles). + has_output = False + output = cast(Any, job).output if has_output else None + output_segments = output.segments if output is not None else () + if has_output: + runtime.begin_output_replay(output) + self._pending_output_replay_finish = True + plan = runtime.take_replay_plan() + if plan is None: + if has_output: + # An output-only marker deliberately has no input ReplayPlan. + # Its segment list is still the shared activation schedule: run + # the first k - 1 drains here and leave the last to activate()'s + # ordinary trailing drain, exactly as the input-driven path + # below does. Treating all output-only markers as one segment + # collapsed retained multi-activation Workflow Tasks on replay. + for _segment in output_segments[:-1]: + runtime.begin_output_replay_segment() + runtime.resolve_all_pending() + self._run_once(check_conditions=True) + if output_segments: + runtime.begin_output_replay_segment() + runtime.resolve_all_pending() + return + + # Nothing was prepared -- the input replay job reached the Workflow + # thread without its ranges being read, which would mean delivering + # from a buffer that live watching filled. Better to deliver nothing + # than to deliver something replay did not record. + runtime.resolve_all_pending() + return + + deferred = False + closed = False + try: + # The bindings first, and before any delivery. A recorded run is + # joined to a subscription by `wait_id`, which is an integer and + # therefore says nothing about what that wait was subscribed to; + # without checking the binding, wait 1 moved from one stream to + # another delivers the first stream's recorded bytes through the + # second's subscription rather than failing as nondeterminism. + runtime.begin_replay(plan.annotation.header.streams) + # Every segment but the **last** is drained here. The last one is + # armed and left to the drain the activation runs for the job set + # this job arrived in, because that drain happens whatever this + # method does: a driver that drained all *k* segments itself + # produced *k + 1* drains for the *k* the marker records, and + # ADR-018 requires exactly *k*. Leaving the last one to the + # activation is also what the live run did -- there each + # activation's single trailing drain served the records that + # activation had just been handed. + for segment in plan.segments[:-1]: + if has_output: + runtime.begin_output_replay_segment() + runtime.begin_replay_segment(list(segment.deliveries)) + runtime.resolve_all_pending() + self._run_once(check_conditions=True) + # Reached only once every segment but the last has been drained. + if plan.segments: + if has_output: + runtime.begin_output_replay_segment() + runtime.begin_replay_segment(list(plan.segments[-1].deliveries)) + runtime.resolve_all_pending() + # Closing the replay here would end replay mode before the drain + # that delivers this final segment, so it is handed to + # `_finish_replay_external_streams` -- which the activation calls + # once that drain has run, and which is what leaves replay mode + # from then on. + self._pending_replay_finish = plan + deferred = True + else: + if has_output: + runtime.begin_output_replay_segment() + # A marker that recorded no activation of its own -- a Workflow + # Task that bound a wait and blocked with the stream never + # delivering. **Nothing is deferred, and that is not asymmetry + # for its own sake.** With no recorded segment for the + # activation's drain to serve, that drain is a *live* one: + # records that arrived while this Run was evicted are already in + # the buffer and it would hand them over. Repositioning after it + # would then retract exactly what it had just delivered -- cursor + # back to the marker's boundary, buffer cleared -- and the + # watcher would re-read and re-deliver records Workflow code + # already had. Closing first retracts the buffer *before* the + # drain, so the drain finds nothing and the watcher re-reads from + # the boundary the marker committed. The records are not lost; + # they arrive on the activation the re-read announces. + self._pending_replay_finish = plan + closed = True + if has_output: + # Input has no recorded drain and must reposition before the + # activation's live drain. Output, however, is produced by + # that drain and is validated after it. + self._pending_output_replay_finish = False + self._finish_replay_external_streams() + self._pending_output_replay_finish = True + else: + self._finish_replay_external_streams() + finally: + if not deferred and not closed: + # The walk raised part-way. Leaving replay mode set would make + # every later drain on this Run return nothing at all, turning + # one marker's failure into a Workflow that silently never + # receives again -- and the checks the close performs must not + # run here, where they would replace the error that got us here + # with one of their own. + self._pending_replay_finish = None + self._pending_output_replay_finish = False + if has_output: + runtime.abandon_output_replay() + runtime.end_replay() + + def _abandon_replay_external_streams(self) -> None: + """Leaves replay mode without running any of the checks a close runs. + + For the path where an activation is already failing. See the caller. + """ + self._pending_replay_finish = None + self._pending_output_replay_finish = False + runtime = self._external_stream_runtime + if runtime is not None: + runtime.abandon_output_replay() + runtime.end_replay() + + def _finish_replay_external_streams(self) -> None: + """Closes the replay whose last segment the activation's drain served. + + Called from the activation's ``finally`` so that a drain which raised + still leaves replay mode, for the reason + :meth:`_apply_replay_external_streams` gives. + """ + plan = self._pending_replay_finish + output_pending = getattr(self, "_pending_output_replay_finish", False) + if plan is None and not output_pending: + return + self._pending_replay_finish = None + self._pending_output_replay_finish = False + runtime = self._external_stream_runtime + if runtime is None: + return + try: + # Nothing recorded may be left over. Each of these records was handed + # to Workflow code in the activation its run was recorded in, so a + # replay that ends holding one has run code that consumes less than + # History says was consumed -- what a removed `subscribe()` call + # looks like from here. + if plan is not None: + runtime.verify_replay_consumed() + # The manager knew nothing about this marker while replay was + # running: its watcher has been reading the very same records from + # the subscription's start cursor into the live buffer. The next live + # drain would hand them over again -- observed end-to-end as one + # marker's record delivered twice. Moving the cursors to what the + # marker committed is what makes live delivery resume *after* those + # records rather than in front of them. + # + # After every segment's drain, not inside the walk: a replay that + # raised part-way committed nothing, and advancing a committed cursor + # for records this Workflow may never have received would lose them + # outright. + if plan is not None: + runtime.reposition_after_replay(plan.committed_boundaries) + if output_pending: + runtime.verify_output_replay() + finally: + if output_pending: + # Successful verification already cleared the manifest, so + # this is a no-op on success and cleanup of partial output on + # every validation/input-reposition failure path. + runtime.abandon_output_replay() + if plan is not None: + runtime.end_replay() + def _apply_query_workflow( self, job: temporalio.bridge.proto.workflow_activation.QueryWorkflow ) -> None: @@ -890,6 +1205,12 @@ def _apply_remove_from_cache( # We consider eviction to be under replay so that certain code like # logging that avoids replaying doesn't run during eviction either self._is_replaying = True + # Drop the stream readiness futures before cancelling tasks. Nothing + # will ever resolve them now, and a coroutine cancelled while awaiting + # one would otherwise be woken by the cancellation and then find its + # future still registered. + if self._external_stream_runtime is not None: + self._external_stream_runtime.clear_pending() # Cancel everything for task in self._tasks: task.cancel() @@ -2340,6 +2661,277 @@ def _add_command(self) -> temporalio.bridge.proto.workflow_commands.WorkflowComm self._assert_not_read_only("add command") return self._current_completion.successful.commands.add() + def _attach_external_stream_continuation(self, command: Any) -> None: + """Puts each subscription's position on the Continue-As-New command. + + The successor Run restores it from its own ``WorkflowExecutionStarted`` + rather than reading the backend, so replay sees the boundary the Run + actually started from instead of wherever the stream has since got to + (ADR-022). + + Written when the command is created and again once the activation has + quiesced -- see :meth:`_refresh_external_stream_continuations`, which is + what makes the second write the one that counts. + + Nothing is attached when the Run held no subscriptions -- an empty header + on every Continue-As-New in every Workflow would be pure overhead. + """ + runtime = self._external_stream_runtime + if runtime is None: + return + continuation = runtime.continuation() + if continuation.cursors: + from temporalio.contrib.external_workflow_streams._continuation import ( + CONTINUATION_HEADER, + write_continuation_header, + ) + + command.headers[CONTINUATION_HEADER].CopyFrom( + write_continuation_header(continuation) + ) + output_continuation = runtime.output_continuation() + if output_continuation is not None: + from temporalio.contrib.external_workflow_streams._output_continuation import ( + OUTPUT_CONTINUATION_HEADER, + write_output_continuation_header, + ) + + command.headers[OUTPUT_CONTINUATION_HEADER].CopyFrom( + write_output_continuation_header(output_continuation) + ) + + def _refresh_external_stream_continuations(self) -> None: + """Re-takes the continuation snapshot now the activation has quiesced. + + Creating the Continue-As-New command does not end the activation. + :meth:`_run_once` drains ``self._ready`` until it is empty and then + re-checks conditions, so a stream consumer scheduled -- or unblocked -- + before the terminal command still gets its turn afterwards, and every + record it takes calls ``record_consumption`` after the header was + serialized. Those consumptions do reach the predecessor's final stream + marker, which is closed on the way out; a header left at the earlier + boundary would therefore describe a different boundary than History + does, and the successor would be handed the record a second time. + + No Workflow code runs from here on, so this is the first point at which + the consumption boundary is stable for the activation. + + Every Continue-As-New command is refreshed rather than the last one, + because two top-level functions can each raise one and only Core decides + which terminal survives. + + **Not gated behind an internal flag.** Core matches a Continue-As-New + command to its ``WorkflowExecutionContinuedAsNew`` event by command type + alone and never compares the command's headers against the recorded ones + (``continue_as_new_workflow_state_machine.rs``), so a replay that + regenerates the header at the later boundary cannot disagree with a + History written at the earlier one. What the successor resumes from is + the copy in its own ``WorkflowExecutionStarted``, which is durable and + which replaying the predecessor does not rewrite -- so a chain already in + flight keeps the boundary it started with. + """ + for command in self._current_completion.successful.commands: + if command.HasField("continue_as_new_workflow_execution"): + self._attach_external_stream_continuation( + command.continue_as_new_workflow_execution + ) + + def _emit_external_stream_commands(self) -> None: + """Answers the two independent questions every activation return poses. + + 1. **Did replay-visible stream state change?** If so a + ``WorkflowStreamProgress`` carries the observation delta, on *every* + completion path. This is what commits the cursor boundary, and it is + not conditional on why the Workflow Task ended: if a consumed record + influenced a command that lands in History while the consumption + itself is never marked, replay re-delivers that record while the + command it produced is already durable. + 2. **Should the Workflow Task be retained?** If so a + ``WorkflowStreamQuiescent`` asks Core to hold it open. It carries no + annotation data; the two questions are genuinely separate. + + The progress command is emitted **first**, because it must precede every + command whose value could depend on the consumed data -- on replay that + is what guarantees a record is validated before the command derived from + it is matched. + + The two are separate in *time* as well, which is why replay answers only + the second: the annotation is already in History, but the wait set the + quiescent command registers is per-Worker runtime state that a replayed + Run has to rebuild. + + Quiescence is also what a Continue-As-New header has been waiting for, so + :meth:`_refresh_external_stream_continuations` runs from here too -- the + boundary it records and the boundary the delta below commits have to be + the same one. + """ + runtime = self._external_stream_runtime + if runtime is None or self._deleting: + return + + # First, because the boundary a Continue-As-New header has to carry is + # the one this activation is about to commit, and it only stops moving + # here. + self._refresh_external_stream_continuations() + + # Read before anything is added below, so "what the Workflow itself + # produced this activation" stays answerable. A terminal command also + # determines whether buffered readiness can ever be useful again. + commands = self._current_completion.successful.commands + produced_commands = len(commands) > 0 + terminal = any(map(_is_workflow_terminal_command, commands)) + + # Records still buffered when a non-terminal activation ends have no readiness + # notification coming: the watcher moved its prefetch cursor past them + # when it buffered them, and it only reports again after a *new* non-empty + # read. Re-reporting is what brings the next activation in, and it is the + # only thing that can. + # + # **Asked of every completion, not only of one a delivery budget + # stopped.** The budget is the obvious way to end an activation with a + # full buffer and it is not the only one. Readiness accepted while an + # activation was open is kept by Core only if the quiescent snapshot that + # follows reports the *same* wait generation: same generation means the + # Workflow never saw the record, a bumped one means it drained and + # re-blocked. So a record that arrives after this activation's last drain + # is accepted at generation G, the wait re-blocks to G+1 on the way out, + # and the snapshot legitimately drops a readiness that now refers to a + # resolved block. Nothing else announces it, and the Workflow waits + # forever on data it is already holding -- until an unrelated later event. + # Gating this on the budget left that hole open, and it is a race, so it + # showed up as an occasional stall rather than a reproducible one. + # + # Cheap where it is not needed: the manager skips every subscription whose + # buffer is empty, which is the ordinary case, and a redundant report is + # answered `Accepted` or `Stale` and retried against the current + # generation. + # + # A terminal completion is the exception. This Run cannot consume a + # later activation, and re-arming after its terminal command reports the + # buffer into the Workflow's closing window. Core answers + # `NoOpenWorkflowTask`, which turns that report into a follow-up Signal + # that can race the terminal report and make the same task replay. + # + # Before every early return below: the waits involved are marked blocked, + # so they are in the quiescent snapshot, and a snapshot is what lets Core + # start the idle timer and eventually park. Parking a Workflow Task whose + # records are already in the local buffer would be wrong, and this is what + # makes Core resolve instead of park. + if not terminal: + runtime.rearm_readiness() + + # The snapshot goes out whether or not this completion retains the task. + # A completion carrying a timer, activity, child workflow, or signal must + # be reported so the server can act on it, so it cannot ask for + # retention -- but the subscriptions stay registered, and the wake Signal + # covers the window that leaves. + # + # "Stay registered" is the load-bearing phrase, and it is Core that + # decides how to read this: it registers the waits and, when the + # completion cannot be retained for, arms no timer. Withholding the + # snapshot here instead is what left a Workflow whose very first block + # rode such a completion registered nowhere, unresumable by any wake -- + # a deadlock in ordinary user code. + snapshot = runtime.quiescent_snapshot() + + retaining = bool(snapshot) and not produced_commands + + if self._is_replaying: + # Every marker for a replayed Workflow Task is already in History, + # so there is nothing here for Core to write. Re-deriving the + # annotation is not merely redundant: the command would be matched + # against the very event it was read from, and Core would see a + # marker arriving for a machine already resolved from lookahead. + # + # What lang accumulated while replaying is dropped rather than + # carried, because otherwise it rides the *next* completion -- which + # is live, and would write a second marker for observations the + # first one already recorded. A registration alone is enough to + # trigger that, so it happens even for a stream that never delivered + # anything. + # + # The *quiescent* command below is emitted all the same, because it + # is not about the annotation at all: it is the only thing that + # registers a wait set with Core, and Core's wait set is runtime + # state that a replayed Run rebuilds from nothing. A Run handed over + # and replayed that reported no snapshot would finish replay with an + # empty wait set -- every later readiness answered as though the Run + # had no subscriptions, every wake Signal marking nothing ready, and + # every Workflow Task it created completing with no activation in + # it. That Run is unresumable, with its records sitting in the + # stream. + # + # The cost is that Core starts its idle timer here too, and on + # replay it should not run one at all: a replay slower than the idle + # timeout can therefore queue a park handshake between two replay + # activations. That is a rare timing hazard against a certain and + # total failure, so it is the right trade to make from this side -- + # but removing it needs Core either to hold its timers while + # `replaying` or to register the wait set from the marker lookahead + # it already reads. + runtime.start_new_annotation() + else: + parts: list[bytes] = [] + # The delta first. `take_observation_delta` is what *closes* this + # activation's segment, so a rollover decision read before it is a + # decision about the annotation as it stood one activation ago: the + # segment that crossed the high-water mark went out with + # `request_rollover = false`, and the following activation was then + # free to add another frame and overflow before Core had ever been + # asked to roll over. Reading it after also picks up the runtime + # having stopped delivering because the annotation could afford no + # more, which is decided while this same segment is being closed. + delta = runtime.take_observation_delta() + if delta is not None: + parts.append(delta) + # And before the terminal is added: closing the annotation starts a + # fresh one, whose accumulator is nowhere near the high-water mark. + request_rollover = runtime.request_rollover + # Both of these end the Workflow Task. Retention is refused outright + # for the first; for the second Core takes the rollover as + # authoritative *over* a retention request, because the annotation + # is the thing that has to stop growing. + if runtime.annotation_started and (not retaining or request_rollover): + # This completion ends the Workflow Task, so Core writes the + # marker for everything accumulated and clears it. Two things + # follow, and neither is optional. + # + # The terminal has to ride *this* delta. Core is + # annotation-blind and never manufactures one, and it asks for + # one only on the boundaries it decides itself -- so on a + # completion Python decided, a marker whose annotation has no + # terminal is what gets written, and that is durable and wrong + # (ADR-008). + # + # And the next annotation has to begin from a fresh header. Core + # accumulates by byte append, so an accumulator that kept the + # header it already emitted would start the *next* marker at + # whatever frame came first -- the observed failure is an + # annotation beginning with a terminal frame, read back as + # "schema version 3". + parts.append(runtime.add_terminal()) + + if parts: + progress = temporalio.bridge.proto.workflow_commands.WorkflowCommand() + progress.workflow_stream_progress.observation_delta = b"".join(parts) + progress.workflow_stream_progress.request_rollover = request_rollover + # Inserted at the front rather than appended: the progress + # command must precede every command whose value could depend on + # the consumed data. + commands.insert(0, progress) + + if not snapshot: + return + + quiescent = self._add_command().workflow_stream_quiescent + quiescent.quiescence_generation = 0 + quiescent.idle_timeout.FromTimedelta(runtime.effective_idle_timeout()) + for wait in snapshot: + entry = quiescent.waits.add() + entry.wait_id = wait.wait_id + entry.generation = wait.generation + entry.immediately_parkable = wait.immediately_parkable + def _workflow_logic_flag_enabled(self, flag: _WorkflowLogicFlag) -> bool: if flag in self._current_internal_flags: return True @@ -2654,6 +3246,20 @@ def _instantiate_workflow_object(self) -> Any: else: workflow_instance = self._defn.cls() + # Hand the workflow object its External Stream handle. It goes on the + # object rather than in a module global because per-Run state must share + # the Run's lifetime exactly -- a global would outlive an evicted Run and + # hand its wait ids to the next one. + if ( + self._external_stream_runtime is not None + and self._external_streams_configured + ): + from temporalio.contrib.external_workflow_streams._api import ( + _install_runtime, + ) + + _install_runtime(workflow_instance, self._external_stream_runtime) + if self._defn.versioning_behavior: self._versioning_behavior = self._defn.versioning_behavior # If there's a dynamic config function, call it now after we've instantiated the object @@ -3932,6 +4538,7 @@ def _apply_command(self) -> None: v.backoff_start_interval.FromTimedelta(self._input.backoff_start_interval) if self._input.headers: temporalio.common._apply_headers(self._input.headers, v.headers) + self._instance._attach_external_stream_continuation(v) if self._input.retry_policy: self._input.retry_policy.apply_to_proto(v.retry_policy) if memo_payloads: diff --git a/temporalio/worker/workflow_sandbox/_restrictions.py b/temporalio/worker/workflow_sandbox/_restrictions.py index 3e6e0d6ba6..147c8c7e8e 100644 --- a/temporalio/worker/workflow_sandbox/_restrictions.py +++ b/temporalio/worker/workflow_sandbox/_restrictions.py @@ -464,6 +464,13 @@ def with_child_unrestricted(self, *child_path: str) -> SandboxMatcher: SandboxRestrictions.passthrough_modules_minimum = { "grpc", + # The External Workflow Streams subscription manager. It owns the Worker's + # backend connections, its watcher tasks, and the buffers the Workflow + # thread pops from -- all of which live on the Worker's own asyncio loop. + # Re-importing it inside the sandbox would give the Workflow a second, + # empty manager watching nothing, so only an opaque handle to the real one + # may cross the boundary. + "temporalio.contrib.external_workflow_streams._manager", # Due to some side-effecting calls made on import, these need to be # allowed "pathlib", @@ -779,6 +786,34 @@ def _public_callables(parent: Any, *, exclude: set[str] = set()) -> set[str]: "zipfile": SandboxMatcher( children={"ZipFile": SandboxMatcher(use={"extract", "extractall"})} ), + "temporalio": SandboxMatcher( + children={ + "contrib": SandboxMatcher( + children={ + # External stream provider instances hold connections + # and credentials and live on the Worker, outside the + # sandbox. Workflow code uses the backend configured on + # its Worker; it never constructs or imports one, and a + # provider reached from in here would be a second + # instance that no watcher owns. + "external_workflow_streams": SandboxMatcher( + children={ + "_redis": SandboxMatcher( + match_self=True, + leaf_message=( + "External stream providers may not be " + "imported from Workflow code. Register " + "the provider on the Worker with " + "external_stream_backend=... instead." + ), + ), + }, + access={"RedisStreamBackend"}, + ), + } + ), + } + ), "zoneinfo": SandboxMatcher( children={ "ZoneInfo": SandboxMatcher( diff --git a/tests/contrib/external_workflow_streams/__init__.py b/tests/contrib/external_workflow_streams/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/contrib/external_workflow_streams/conformance.py b/tests/contrib/external_workflow_streams/conformance.py new file mode 100644 index 0000000000..92b4489106 --- /dev/null +++ b/tests/contrib/external_workflow_streams/conformance.py @@ -0,0 +1,853 @@ +"""The backend conformance suite (P2). + +**The suite is the deliverable, not the interface.** Each check below encodes +one obligation from the backend contract and fails with a message naming the +obligation, so a provider author reads what they broke rather than a diff of +two lists. + +Checks are plain async functions taking ``(backend, key)`` and raising +``AssertionError``. That shape is what lets two different callers use them: a +provider's own test parametrizes over :data:`CONFORMANCE_CHECKS` and expects +every one to pass, while the suite's own tests call a single check against a +deliberately-broken stub and assert it fails *for the right reason*. A suite +that cannot be shown to fail is not evidence of anything. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from datetime import timedelta + +from temporalio.contrib.external_workflow_streams._backend import ( + AppendConflictError, + ParkIntent, + ParkIntentRemoval, + StreamBackend, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Offset, + RecordKind, + StreamRecord, +) + +Check = Callable[[StreamBackend, StreamKey], Awaitable[None]] + +NO_BLOCK = timedelta(0) + + +def _data(session: str, seq: int, payload: bytes) -> StreamRecord: + return StreamRecord(RecordKind.DATA, payload, session, seq) + + +def _fence(session: str, seq: int) -> StreamRecord: + return StreamRecord(RecordKind.WRITE_FENCE, b"", session, seq) + + +async def _append_all( + backend: StreamBackend, key: StreamKey, payloads: list[bytes], session: str = "s" +) -> list[StreamRecord]: + return [ + await backend.append(key, _data(session, i, p)) for i, p in enumerate(payloads) + ] + + +# --- ordering --------------------------------------------------------------- + + +async def check_offsets_are_totally_ordered( + backend: StreamBackend, key: StreamKey +) -> None: + """Appended records must ascend under the provider's own comparator.""" + placed = await _append_all(backend, key, [b"a", b"b", b"c"]) + offsets = [r.offset for r in placed] + + assert all(o is not None for o in offsets), ( + "append must return the record with its assigned offset" + ) + assert backend.strictly_increasing(offsets), ( # type: ignore[arg-type] + f"appended offsets must strictly increase under compare_offsets, got {offsets}" + ) + + +async def check_offsets_are_not_compared_lexically( + backend: StreamBackend, key: StreamKey +) -> None: + """Ordering must survive a width change in the offset's leading component. + + This is the check that catches string comparison, and it is invisible + without it: ``"9-0" < "10-0"`` is true numerically and false lexically, so + a lexical comparator is correct for nine appends out of ten and then + silently reorders a range. + """ + early, late = await _append_at_widening_boundary(backend, key) + + assert backend.compare_offsets(early, late) < 0, ( + f"offset {early} was appended before {late}, so compare_offsets must " + "order it first; comparing offsets lexically fails here" + ) + assert backend.compare_offsets(late, early) > 0, ( + "compare_offsets must be antisymmetric" + ) + assert backend.compare_offsets(early, early) == 0, ( + "compare_offsets must report an offset as equal to itself" + ) + + +async def _append_at_widening_boundary( + backend: StreamBackend, key: StreamKey +) -> tuple[Offset, Offset]: + """Two offsets whose leading components differ in digit width. + + A provider that assigns offsets from a clock it exposes (the in-memory + reference) is driven across the boundary directly. Anything else -- a real + Redis -- gets whatever it assigns, and the check still holds because the + comparator must order *its own* offsets correctly either way. + """ + now = getattr(backend, "now_ms", None) + if now is not None: + backend.now_ms = 9 # type: ignore[attr-defined] + early = await backend.append(key, _data("widen", 0, b"early")) + backend.now_ms = 10 # type: ignore[attr-defined] + late = await backend.append(key, _data("widen", 1, b"late")) + else: + early = await backend.append(key, _data("widen", 0, b"early")) + late = await backend.append(key, _data("widen", 1, b"late")) + assert early.offset is not None and late.offset is not None + return early.offset, late.offset + + +# --- the inclusive range read ----------------------------------------------- + + +async def check_range_read_includes_both_endpoints( + backend: StreamBackend, key: StreamKey +) -> None: + """``read_range(first, last)`` is closed on both ends. + + A provider implementing it with exclusive semantics passes every live test + and drops the boundary records on the first replay, which is why this is + checked directly rather than inferred from a round trip. + """ + placed = await _append_all(backend, key, [b"a", b"b", b"c"]) + first, last = placed[0].offset, placed[-1].offset + assert first is not None and last is not None + + got = await backend.read_range(key, first, last) + + assert [r.payload for r in got] == [b"a", b"b", b"c"], ( + "read_range must be inclusive of both endpoints; it returned " + f"{[r.payload for r in got]} for the full range" + ) + + +async def check_range_read_of_one_record_returns_it( + backend: StreamBackend, key: StreamKey +) -> None: + """The degenerate range ``[x, x]`` contains exactly the record at ``x``.""" + placed = await _append_all(backend, key, [b"a", b"b"]) + only = placed[0].offset + assert only is not None + + got = await backend.read_range(key, only, only) + + assert [r.payload for r in got] == [b"a"], ( + f"read_range({only}, {only}) must return exactly that record, got " + f"{[r.payload for r in got]}" + ) + + +async def check_range_read_reports_a_missing_record_as_missing( + backend: StreamBackend, key: StreamKey +) -> None: + """A deleted record must be absent, not substituted for. + + Integrity loss must never resolve to an alternate stream result. The bug + this catches is a range read implemented as "N records starting at + ``first``" rather than as the closed interval ``[first, last]``: it returns + the right *count* by reaching past ``last``, so a trimmed stream reads back + as intact and replay delivers a record the original run never saw. + + A fourth record is appended beyond the range precisely so there is + something for such a provider to reach for. + """ + placed = await _append_all(backend, key, [b"a", b"b", b"c", b"d"]) + first, middle, last = placed[0].offset, placed[1].offset, placed[2].offset + assert first is not None and middle is not None and last is not None + + delete = getattr(backend, "delete_for_test", None) + if delete is None: + return # This provider cannot be made to lose a record on demand. + await delete(key, middle) + + got = await backend.read_range(key, first, last) + + assert [r.payload for r in got] == [b"a", b"c"], ( + "read_range must report a deleted record as absent rather than " + f"substituting another from beyond the range, got {[r.payload for r in got]}" + ) + + +# --- the exclusive watch ---------------------------------------------------- + + +async def check_watch_is_exclusive_of_its_boundary( + backend: StreamBackend, key: StreamKey +) -> None: + """``read_after(AFTER(x))`` must not re-deliver the record at ``x``. + + An inclusive watch re-delivers the last consumed record after every + resumption, which a Workflow sees as a duplicate. + """ + placed = await _append_all(backend, key, [b"a", b"b"]) + first = placed[0].offset + assert first is not None + + got = await backend.read_after(key, AFTER(first), max_records=10, block=NO_BLOCK) + + assert [r.payload for r in got] == [b"b"], ( + "read_after must return records strictly after its boundary; it " + f"returned {[r.payload for r in got]} after the first record" + ) + + +async def check_beginning_reads_from_the_start( + backend: StreamBackend, key: StreamKey +) -> None: + """``BEGINNING`` is a boundary before every record, not an offset.""" + await _append_all(backend, key, [b"a", b"b"]) + + got = await backend.read_after(key, BEGINNING, max_records=10, block=NO_BLOCK) + + assert [r.payload for r in got] == [b"a", b"b"], ( + "read_after(BEGINNING) must return every record from the start, got " + f"{[r.payload for r in got]}" + ) + + +async def check_a_consumer_parked_at_the_tail_resumes( + backend: StreamBackend, key: StreamKey +) -> None: + """The core cursor-semantics check. + + A consumer that has drained the stream holds ``AFTER(last)`` -- a boundary + past the tail. It must be able to resume from there once a record whose id + it could not have predicted arrives. A provider that requires the cursor to + name an existing record cannot express this state at all. + """ + placed = await _append_all(backend, key, [b"a"]) + tail = placed[0].offset + assert tail is not None + + # Parked: nothing beyond the boundary yet, and the boundary names a record + # that is the last one in existence. + empty = await backend.read_after(key, AFTER(tail), max_records=10, block=NO_BLOCK) + assert empty == [], ( + f"read_after(AFTER({tail})) must be empty while {tail} is the tail, got {empty}" + ) + + unpredictable = await backend.append(key, _data("later", 0, b"surprise")) + + resumed = await backend.read_after(key, AFTER(tail), max_records=10, block=NO_BLOCK) + + assert [r.payload for r in resumed] == [b"surprise"], ( + "a consumer parked at the tail must resume once a record arrives, " + f"without having to name its id in advance; got {[r.payload for r in resumed]}" + ) + assert resumed[0].offset == unpredictable.offset + + +async def check_watch_respects_max_records( + backend: StreamBackend, key: StreamKey +) -> None: + """Backpressure is the buffer bound, so the read must honour it.""" + await _append_all(backend, key, [b"a", b"b", b"c", b"d"]) + + got = await backend.read_after(key, BEGINNING, max_records=2, block=NO_BLOCK) + + assert len(got) == 2, f"read_after must return at most max_records, got {len(got)}" + + +async def check_watch_returns_empty_rather_than_hanging( + backend: StreamBackend, key: StreamKey +) -> None: + """An elapsed block returns empty; it does not raise and does not hang.""" + got = await backend.read_after( + key, BEGINNING, max_records=10, block=timedelta(milliseconds=50) + ) + + assert got == [], f"an elapsed block must return empty, got {got}" + + +# --- idempotent append ------------------------------------------------------ + + +async def check_identical_reappend_is_a_no_op( + backend: StreamBackend, key: StreamKey +) -> None: + """A retried Activity attempt must not duplicate its record.""" + first = await backend.append(key, _data("retry", 0, b"payload")) + again = await backend.append(key, _data("retry", 0, b"payload")) + + assert again.offset == first.offset, ( + "re-appending byte-identical content under the same idempotency key " + f"must return the original offset {first.offset}, got {again.offset}" + ) + + assert first.offset is not None + in_range = await backend.read_range(key, first.offset, first.offset) + assert len(in_range) == 1, ( + f"an idempotent re-append must write nothing, but {len(in_range)} " + "records occupy the original offset" + ) + + +async def check_reappend_with_different_bytes_is_rejected( + backend: StreamBackend, key: StreamKey +) -> None: + """Idempotency is on identity, not on the key alone. + + A provider that accepts this lets a retried producer attempt with changed + input overwrite what a consumer may already have delivered -- and the + divergence surfaces much later, as an unrelated nondeterminism error. + """ + await backend.append(key, _data("conflict", 0, b"original")) + + try: + await backend.append(key, _data("conflict", 0, b"different")) + except AppendConflictError: + return + raise AssertionError( + "re-appending different bytes under an existing idempotency key must " + "raise AppendConflictError; the append was accepted instead" + ) + + +async def check_different_sessions_do_not_collide( + backend: StreamBackend, key: StreamKey +) -> None: + """The key is ``(session_id, sequence)``, so sequence alone is not enough.""" + one = await backend.append(key, _data("session-a", 0, b"a")) + two = await backend.append(key, _data("session-b", 0, b"b")) + + assert one.offset != two.offset, ( + "two producers each writing their own sequence 0 must get distinct " + "offsets; the idempotency key includes the session id" + ) + + +# --- control records -------------------------------------------------------- + + +async def check_control_records_share_the_data_offset_sequence( + backend: StreamBackend, key: StreamKey +) -> None: + """A fence occupies an offset and advances the cursor like any record.""" + before = await backend.append(key, _data("s", 0, b"a")) + fence = await backend.append(key, _fence("s", 1)) + after = await backend.append(key, _data("s", 2, b"b")) + + offsets = [before.offset, fence.offset, after.offset] + assert backend.strictly_increasing(offsets), ( # type: ignore[arg-type] + f"a write fence must take its place in the offset sequence, got {offsets}" + ) + + assert before.offset is not None and after.offset is not None + got = await backend.read_range(key, before.offset, after.offset) + kinds = [r.kind for r in got] + assert kinds == [RecordKind.DATA, RecordKind.WRITE_FENCE, RecordKind.DATA], ( + f"a range read must return control records in place, got {kinds}" + ) + + +async def check_records_round_trip_unchanged( + backend: StreamBackend, key: StreamKey +) -> None: + """Every field a producer wrote must read back exactly.""" + written = StreamRecord(RecordKind.DATA, b"\x00\xff bytes", "session-x", 42) + placed = await backend.append(key, written) + + assert placed.offset is not None + (got,) = await backend.read_range(key, placed.offset, placed.offset) + + assert got.payload == written.payload, "payload bytes must round-trip exactly" + assert got.producer_session_id == written.producer_session_id + assert got.sequence == written.sequence + assert got.kind == written.kind + + +# --- registration precondition ---------------------------------------------- + + +async def check_immutability_is_declared( + backend: StreamBackend, + key: StreamKey, # pyright: ignore[reportUnusedParameter] +) -> None: + """Declared, not merely true. Undeclared is a registration failure. + + The registry is what enforces this at Worker construction; the check is + here so a provider author finds out while writing the provider rather than + when a Worker refuses to start. + """ + assert type(backend).guarantees_immutability is True, ( + f"{type(backend).__name__} must declare guarantees_immutability = True; " + f"it is {type(backend).guarantees_immutability!r}. The guarantee is a " + "precondition for registration, not a runtime mode." + ) + assert type(backend).provider_id, ( + f"{type(backend).__name__} must declare a provider_id; every annotation " + "header records it" + ) + + +# --- parking (P2b) ---------------------------------------------------------- + +SHORT_LEASE = timedelta(milliseconds=80) + + +async def check_intents_are_keyed_by_stream_and_wait_id( + backend: StreamBackend, key: StreamKey +) -> None: + """Two subscriptions to one stream must each keep their own intent. + + A stream-keyed intent looks correct until a Workflow subscribes to the same + stream twice: the second install overwrites the first, and from then on only + one of the two can ever be woken. Cursors are per subscription, so the + surviving intent also names the wrong position for the lost one. + """ + first = ParkIntent(wait_id=1, cursor=BEGINNING, park_generation=4, run_id="run-a") + second = ParkIntent( + wait_id=2, cursor=AFTER(Offset("100-0")), park_generation=4, run_id="run-a" + ) + + await backend.install_park_intent(key, first) + await backend.install_park_intent(key, second) + + got_first = await backend.park_intent(key, 1) + got_second = await backend.park_intent(key, 2) + + assert got_first == first, ( + "installing a second subscription's intent overwrote the first; park " + f"intents must be keyed (stream key, wait_id), got {got_first!r}" + ) + assert got_second == second + assert got_first.cursor != got_second.cursor # type: ignore[union-attr] + + +async def check_every_parked_subscription_is_enumerable( + backend: StreamBackend, key: StreamKey +) -> None: + """A producer must be able to find every wait it has to wake. + + ``wait_id`` is allocated by a per-Run counter inside ``subscribe()``, which no + producer can see, and the other five parking operations all take one. A + provider that enumerates only the most recent intent, or only distinct + streams, leaves the subscriptions it omits parked on records already sitting + in the stream. + """ + other = StreamKey( + key.namespace, key.workflow_id, key.first_execution_run_id, "another-stream" + ) + await backend.install_park_intent( + key, ParkIntent(1, BEGINNING, park_generation=4, run_id="run-a") + ) + await backend.install_park_intent( + key, ParkIntent(2, AFTER(Offset("100-0")), park_generation=4, run_id="run-a") + ) + await backend.install_park_intent( + other, ParkIntent(3, BEGINNING, park_generation=4, run_id="run-a") + ) + + found = await backend.parked_wait_ids(key) + + assert sorted(found) == [1, 2], ( + f"expected both parked subscriptions on this stream, got {found!r}" + ) + assert 3 not in found, ( + "enumeration must be per stream: waking a subscription parked on a " + "different stream produces a Workflow Task that finds nothing there" + ) + + # And a removed intent stops being enumerable, or a producer keeps signaling + # a wait that has already been resolved. + await backend.remove_park_intent(key, 1) + assert await backend.parked_wait_ids(key) == [2] + + +async def check_an_intent_is_removable_and_removal_is_idempotent( + backend: StreamBackend, key: StreamKey +) -> None: + """An aborted park removes every intent it installed.""" + await backend.install_park_intent( + key, ParkIntent(1, BEGINNING, park_generation=1, run_id="run-a") + ) + + await backend.remove_park_intent(key, 1) + assert await backend.park_intent(key, 1) is None + + # Idempotent: an abort may race a removal that already happened. + await backend.remove_park_intent(key, 1) + assert await backend.park_intent(key, 1) is None + + +async def check_conditional_removal_never_deletes_a_replacement( + backend: StreamBackend, key: StreamKey +) -> None: + """A predecessor's delayed cleanup must not remove a successor's park. + + And it must say which of the two "nothing was removed" cases it met. A key + that is clear is not a key someone else's park is holding: the consumer + treats the first as cleanup that happened -- possibly its own, on a call + whose reply was lost -- and the second as a park it must leave alone. + """ + successor = ParkIntent(1, BEGINNING, park_generation=2, run_id="run-b") + await backend.install_park_intent(key, successor) + + assert ( + await backend.remove_park_intent_if_matches( + key, 1, run_id="run-a", park_generation=2 + ) + is ParkIntentRemoval.MISMATCH + ) + assert await backend.park_intent(key, 1) == successor + + assert ( + await backend.remove_park_intent_if_matches( + key, 1, run_id="run-b", park_generation=1 + ) + is ParkIntentRemoval.MISMATCH + ) + assert await backend.park_intent(key, 1) == successor + + assert ( + await backend.remove_park_intent_if_matches( + key, 1, run_id="run-b", park_generation=2 + ) + is ParkIntentRemoval.REMOVED + ) + assert await backend.park_intent(key, 1) is None + + assert ( + await backend.remove_park_intent_if_matches( + key, 1, run_id="run-b", park_generation=2 + ) + is ParkIntentRemoval.ABSENT + ), "a repeat of a removal that succeeded is an absent key, not a mismatch" + assert ( + await backend.remove_park_intent_if_matches( + key, 1, run_id="run-c", park_generation=9 + ) + is ParkIntentRemoval.ABSENT + ), "an absent key is absent whoever asks about it" + + +async def check_a_new_runs_intent_replaces_its_predecessors( + backend: StreamBackend, key: StreamKey +) -> None: + """The Run ID is the intent's value, not part of its key. + + ``wait_id`` is stable across a Continue-As-New chain and the stream key + already carries the first execution Run ID, so the key is unique within the + chain. Putting the Run ID in the key instead would leave a dead intent + behind on every continuation. + """ + await backend.install_park_intent( + key, ParkIntent(1, BEGINNING, park_generation=1, run_id="run-a") + ) + await backend.install_park_intent( + key, ParkIntent(1, BEGINNING, park_generation=2, run_id="run-b") + ) + + got = await backend.park_intent(key, 1) + assert got is not None and got.run_id == "run-b", ( + f"a new Run's intent must replace its predecessor's, got {got!r}" + ) + + +async def check_recheck_sees_an_append_past_the_cursor( + backend: StreamBackend, key: StreamKey +) -> None: + """The recheck is what closes the append/park race. + + A producer appends *before* it observes the park generation, so an append is + either seen by this recheck or paired with a wake Signal. A recheck that + could not see a record appended after the intent was installed would lose + exactly the appends the ordering was designed to catch. + """ + placed = await _append_all(backend, key, [b"a"]) + tail = placed[0].offset + assert tail is not None + + await backend.install_park_intent( + key, ParkIntent(1, AFTER(tail), park_generation=1, run_id="run-a") + ) + assert not await backend.recheck(key, 1), ( + "recheck reported records with nothing past the intent's cursor" + ) + + await backend.append(key, _data("racer", 0, b"late")) + + assert await backend.recheck(key, 1), ( + "recheck missed a record appended after the intent was installed; the " + "append/park race is not closed" + ) + + +async def check_recheck_of_a_removed_intent_is_false( + backend: StreamBackend, key: StreamKey +) -> None: + await _append_all(backend, key, [b"a"]) + assert not await backend.recheck(key, 99) + + +async def check_the_current_generation_is_readable( + backend: StreamBackend, key: StreamKey +) -> None: + """A producer needs the generation to name in its wake Signal.""" + assert await backend.current_park_generation(key, 1) is None, ( + "no intent is installed, so there is no confirmed park to report; a " + "producer here must send an unparked wake rather than a stale one" + ) + + await backend.install_park_intent( + key, ParkIntent(1, BEGINNING, park_generation=7, run_id="run-a") + ) + + assert await backend.current_park_generation(key, 1) == 7 + + +async def check_a_removed_intent_reports_no_generation( + backend: StreamBackend, key: StreamKey +) -> None: + """Removing an intent is what ends a park, so it must end what readers see. + + The rule the consumer relies on is *an intent exists only while that park is + outstanding*, and the consumer can only hold up its half of it -- installing + one when a park is confirmed and removing it when the park is over. The + provider's half is that removal is visible through + ``current_park_generation``, which is the call every wake path asks and the + only thing that distinguishes a live park from none. + + A provider that answered from a remembered "last generation" beside the + intent would pass every other check here and still break both readers, in + the same direction and invisibly. A producer would name a generation Core + has already discarded, and Core ignores exactly that as stale -- so the + record is appended, the Signal is sent, and the Workflow is never woken. The + consumer's own shutdown sweep would send a parked wake where it owes the + unparked one, and a parked wake's request ID ignores sender identity by + design, so it arrives byte-identical to the wake that ended the park and the + server deduplicates it away. + """ + await backend.install_park_intent( + key, ParkIntent(1, BEGINNING, park_generation=4, run_id="run-a") + ) + assert await backend.current_park_generation(key, 1) == 4 + + await backend.remove_park_intent(key, 1) + + assert await backend.current_park_generation(key, 1) is None, ( + "a removed park intent still reports a generation; a wake naming it is " + "a claim of a park that is over, which Core discards as stale" + ) + + +async def check_a_claim_excludes_a_second_producer( + backend: StreamBackend, key: StreamKey +) -> None: + """Only meaningful for a provider that declares leased claims.""" + if not type(backend).supports_leased_claims: + return + + await backend.install_park_intent( + key, ParkIntent(1, BEGINNING, park_generation=3, run_id="run-a") + ) + + assert await backend.claim_park_generation( + key, 1, 3, claimant="producer-a", lease=SHORT_LEASE + ) + assert not await backend.claim_park_generation( + key, 1, 3, claimant="producer-b", lease=SHORT_LEASE + ), "two producers both believe they own the same park generation's wake" + + +async def check_a_claim_is_renewable_by_its_holder( + backend: StreamBackend, key: StreamKey +) -> None: + """Renewal must not look like contention. + + A producer whose append and Signal are separated by a slow call has to be + able to extend its own claim; if renewal were refused it would have to drop + the claim and race for it again. + """ + if not type(backend).supports_leased_claims: + return + + await backend.install_park_intent( + key, ParkIntent(1, BEGINNING, park_generation=3, run_id="run-a") + ) + + assert await backend.claim_park_generation( + key, 1, 3, claimant="producer-a", lease=SHORT_LEASE + ) + assert await backend.claim_park_generation( + key, 1, 3, claimant="producer-a", lease=SHORT_LEASE + ), "a claim holder could not renew its own claim" + + +async def check_an_expired_claim_is_taken_over( + backend: StreamBackend, key: StreamKey +) -> None: + """The case an unleased claim cannot recover from. + + A producer that crashes between claiming a generation and sending its wake + Signal strands the generation: every other producer sees it claimed and + concludes the wake is handled, so the parked Workflow waits forever with + data sitting in the stream. A lease turns that into a bounded delay. + """ + if not type(backend).supports_leased_claims: + return + + await backend.install_park_intent( + key, ParkIntent(1, BEGINNING, park_generation=3, run_id="run-a") + ) + assert await backend.claim_park_generation( + key, 1, 3, claimant="crashed-producer", lease=SHORT_LEASE + ) + + await asyncio.sleep(SHORT_LEASE.total_seconds() * 2.5) + + assert await backend.claim_park_generation( + key, 1, 3, claimant="producer-b", lease=SHORT_LEASE + ), ( + "an expired claim was not taken over; a producer crashing between its " + "claim and its Signal would strand this generation permanently" + ) + + +async def check_observe_only_providers_always_grant( + backend: StreamBackend, key: StreamKey +) -> None: + """A provider that cannot lease must let every producer signal. + + Duplicate Signals are harmless -- the runtime rechecks every subscription on + wakeup regardless -- and lost ones are not, so refusing a claim it cannot + police would be strictly worse than granting every one. + """ + if type(backend).supports_leased_claims: + return + + await backend.install_park_intent( + key, ParkIntent(1, BEGINNING, park_generation=3, run_id="run-a") + ) + + assert await backend.claim_park_generation( + key, 1, 3, claimant="producer-a", lease=SHORT_LEASE + ) + assert await backend.claim_park_generation( + key, 1, 3, claimant="producer-b", lease=SHORT_LEASE + ), "an observe-only provider must grant every claim so no wake is lost" + + +async def check_distinct_identities_never_share_storage( + backend: StreamBackend, key: StreamKey +) -> None: + """Two different `StreamKey` values must never address the same storage. + + Every component of the identity is a user-chosen string -- a namespace, a + Workflow ID, a Run ID, a topic name -- so a provider that renders the tuple + by joining the fields with a delimiter is not injective. These two are + different Workflows: + + ("ns", "wf", r1, f"{r2}:tokens") and ("ns", f"wf:{r1}", r2, "tokens") + + and a provider that joins on ``:`` gives them one physical location. They + then share records, idempotency state, park intents and claims: each reads + the other's data, and -- the quiet one -- each concludes the other's claim + has already taken responsibility for its wake, so neither signals and both + Runs wait forever on records that are already durable. + + The delimiter here is the one this suite happens to use. A provider using a + different separator has the same obligation against its own, which is why + this check asserts *isolation of behaviour* rather than any property of the + key strings a provider builds -- those are private to it. + """ + left = StreamKey(key.namespace, "wf", "run-a", "run-b:tokens") + right = StreamKey(key.namespace, "wf:run-a", "run-b", "tokens") + + placed_left = await backend.append(left, _data("left", 0, b"left-only")) + placed_right = await backend.append(right, _data("right", 0, b"right-only")) + + assert [ + r.payload + for r in await backend.read_after( + left, BEGINNING, max_records=10, block=NO_BLOCK + ) + ] == [b"left-only"], ( + "two distinct stream identities share one physical stream; joining the " + "identity's fields without escaping is not injective, and these two " + "unrelated Workflows now read each other's records" + ) + assert [ + r.payload + for r in await backend.read_after( + right, BEGINNING, max_records=10, block=NO_BLOCK + ) + ] == [b"right-only"] + + # Park state keyed off the same identity must be isolated for the same + # reason: an intent read across the boundary answers a generation Core has + # never heard of, and the producer's wake naming it is discarded as stale. + await backend.install_park_intent( + left, + ParkIntent(1, AFTER(placed_left.offset), park_generation=4, run_id="a"), # type: ignore[arg-type] + ) + assert await backend.park_intent(right, 1) is None, ( + "a park intent installed on one stream identity is visible on another" + ) + assert await backend.current_park_generation(right, 1) is None + assert await backend.parked_wait_ids(right) == [] + del placed_right + + +#: The parking checks, kept as their own list so a provider can adopt the core +#: contract before the parking extension. +PARKING_CONFORMANCE_CHECKS: list[Check] = [ + check_intents_are_keyed_by_stream_and_wait_id, + check_every_parked_subscription_is_enumerable, + check_an_intent_is_removable_and_removal_is_idempotent, + check_conditional_removal_never_deletes_a_replacement, + check_a_new_runs_intent_replaces_its_predecessors, + check_recheck_sees_an_append_past_the_cursor, + check_recheck_of_a_removed_intent_is_false, + check_the_current_generation_is_readable, + check_a_removed_intent_reports_no_generation, + check_distinct_identities_never_share_storage, + check_a_claim_excludes_a_second_producer, + check_a_claim_is_renewable_by_its_holder, + check_an_expired_claim_is_taken_over, + check_observe_only_providers_always_grant, +] + + +#: Every check, in the order a provider author most usefully reads them. +CONFORMANCE_CHECKS: list[Check] = [ + check_immutability_is_declared, + check_offsets_are_totally_ordered, + check_offsets_are_not_compared_lexically, + check_range_read_includes_both_endpoints, + check_range_read_of_one_record_returns_it, + check_range_read_reports_a_missing_record_as_missing, + check_watch_is_exclusive_of_its_boundary, + check_beginning_reads_from_the_start, + check_a_consumer_parked_at_the_tail_resumes, + check_watch_respects_max_records, + check_watch_returns_empty_rather_than_hanging, + check_identical_reappend_is_a_no_op, + check_reappend_with_different_bytes_is_rejected, + check_different_sessions_do_not_collide, + check_control_records_share_the_data_offset_sequence, + check_records_round_trip_unchanged, +] diff --git a/tests/contrib/external_workflow_streams/conftest.py b/tests/contrib/external_workflow_streams/conftest.py new file mode 100644 index 0000000000..d6e8483c08 --- /dev/null +++ b/tests/contrib/external_workflow_streams/conftest.py @@ -0,0 +1,96 @@ +"""Redis fixtures for External Workflow Streams tests (X2). + +One Redis server is shared by the whole suite, so every test gets its own key +prefix rather than its own server. The prefix embeds the ``pytest-xdist`` +worker id, so parallel workers cannot collide even when two of them run the +same test module. +""" + +from __future__ import annotations + +import os +import uuid +from collections.abc import AsyncGenerator +from dataclasses import dataclass + +import pytest +import pytest_asyncio + +DEFAULT_REDIS_URL = "redis://127.0.0.1:6379" + +#: Every key this suite creates starts with this, so a leaked key is +#: attributable and a global cleanup is possible without guessing. +KEY_NAMESPACE = "temporal-extstream-test" + + +def redis_url() -> str: + return os.getenv("TEMPORAL_TEST_REDIS_URL", DEFAULT_REDIS_URL) + + +@dataclass +class RedisKeyspace: + """An isolated slice of the shared Redis server. + + ``client`` is a live ``redis.asyncio.Redis``. Every key a test touches must + be built with :meth:`key`, which is what makes the teardown complete. + """ + + client: "object" # redis.asyncio.Redis, untyped to keep redis an optional import + prefix: str + + def key(self, name: str) -> str: + return f"{self.prefix}{name}" + + async def cleanup(self) -> int: + """Delete every key under this keyspace. Returns the number removed.""" + removed = 0 + batch: list[str] = [] + async for found in self.client.scan_iter(match=f"{self.prefix}*", count=500): # type: ignore[attr-defined] + batch.append(found) + if len(batch) >= 500: + removed += await self.client.delete(*batch) # type: ignore[attr-defined] + batch = [] + if batch: + removed += await self.client.delete(*batch) # type: ignore[attr-defined] + return removed + + +@pytest.fixture(scope="session") +def redis_worker_id(request: pytest.FixtureRequest) -> str: + """The ``pytest-xdist`` worker id, or ``master`` when running serially.""" + return getattr(request.config, "workerinput", {}).get("workerid", "master") + + +@pytest_asyncio.fixture(scope="session") # type: ignore[reportUntypedFunctionDecorator] +async def redis_client() -> AsyncGenerator[object, None]: + """A session-wide connection, skipping the suite if no server is reachable. + + Reachability is checked once per session rather than per test so a missing + Redis produces one clear skip reason instead of one per case. + """ + redis = pytest.importorskip("redis.asyncio", reason="redis is not installed") + + client = redis.from_url(redis_url(), decode_responses=True) + try: + await client.ping() + except Exception as err: + await client.aclose() + pytest.skip(f"Redis is not reachable at {redis_url()}: {err}") + + try: + yield client + finally: + await client.aclose() + + +@pytest_asyncio.fixture # type: ignore[reportUntypedFunctionDecorator] +async def redis_keyspace( + redis_client: object, redis_worker_id: str +) -> AsyncGenerator[RedisKeyspace, None]: + """An isolated key prefix, cleaned up whether or not the test passed.""" + prefix = f"{KEY_NAMESPACE}:{redis_worker_id}:{uuid.uuid4().hex}:" + keyspace = RedisKeyspace(client=redis_client, prefix=prefix) + try: + yield keyspace + finally: + await keyspace.cleanup() diff --git a/tests/contrib/external_workflow_streams/crash_worker.py b/tests/contrib/external_workflow_streams/crash_worker.py new file mode 100644 index 0000000000..60a164190e --- /dev/null +++ b/tests/contrib/external_workflow_streams/crash_worker.py @@ -0,0 +1,220 @@ +"""A Worker in a process of its own, so a test can kill it outright. + +A crash is not a shutdown. Every in-process way of stopping a Worker runs the +graceful path -- which, for External Workflow Streams, is exactly the path that +*finalizes the annotation and writes the marker*. Testing "the Worker went away +before its marker committed" against a Worker that politely committed its marker +on the way out would assert the opposite of the case. + +So the Worker under test runs here, in a child process that is sent ``SIGKILL`` +with a Workflow Task retained. Nothing gets a chance to run. + +Two things live in this module rather than in the test: + +- the Workflow definition, because both processes must register the *same* one; +- a provider that records every offset it reads into Redis, because a crashed + process cannot be asked afterwards what it had read. The read log is the only + evidence that survives the kill, and comparing it with the replacement + Worker's is what turns "the records arrived eventually" into "the same offsets + were read again". + +The provider import is deliberately **inside** a function: this module is +re-imported by the Workflow sandbox, and importing a provider from Workflow code +is refused outright. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +from datetime import timedelta +from typing import Any + +from temporalio import workflow +from temporalio.client import Client +from temporalio.worker import Worker + +with workflow.unsafe.imports_passed_through(): + from temporalio.contrib.external_workflow_streams._api import external_stream + from temporalio.contrib.external_workflow_streams._output_api import ( + external_output_stream, + ) + +STREAM_NAME = "tokens" + +READ_LOG = "reads" +"""Key suffix under which a recording provider logs the offsets it read.""" + +OUTPUT_STAGE_LOG = "output-stages" +"""Key suffix used to prove an output stage survived a process kill.""" + + +@workflow.defn +class CrashConsumeWorkflow: + """Consumes a fixed number of records and returns them, in order. + + Returning the values is what makes loss and duplication visible: a record + re-delivered after the crash appears twice, and one lost never appears. + """ + + @workflow.run + async def run(self, expected: int) -> list[str]: + tokens = external_stream.topic(STREAM_NAME, type=str) + seen: list[str] = [] + async for token in tokens.subscribe(): + seen.append(token) + if len(seen) >= expected: + break + return seen + + +@workflow.defn +class CrashPublishWorkflow: + """Publishes output on a Workflow Task whose Worker will be killed.""" + + @workflow.run + async def run(self) -> str: + topic = external_output_stream.topic("events", type=str) + await topic.publish("accepted-after-crash") + await topic.finish() + return "done" + + +@workflow.defn +class CrashUpdatePublishWorkflow: + """Waits for an Update that publishes while its report is withheld.""" + + def __init__(self) -> None: + self._finished = False + + @workflow.run + async def run(self) -> None: + await workflow.wait_condition(lambda: self._finished) + + @workflow.update + async def publish(self, value: str) -> str: + await external_output_stream.topic("events", type=str).publish(value) + return value + + @workflow.signal + def finish(self) -> None: + self._finished = True + + +def read_log_key(prefix: str, reader: str) -> str: + return f"{prefix}:{READ_LOG}:{reader}" + + +def output_stage_log_key(prefix: str) -> str: + return f"{prefix}:{OUTPUT_STAGE_LOG}" + + +def recording_backend(*, url: str, prefix: str, reader: str) -> Any: + """A Redis provider that appends every offset it reads to a Redis list. + + In Redis rather than in memory because the Worker doing the reading may be + killed: whatever it recorded has to outlive the process. + """ + from temporalio.contrib.external_workflow_streams._redis import RedisStreamBackend + + class ReadRecordingBackend(RedisStreamBackend): + async def read_after( # type: ignore[override] + self, + key: Any, + after: Any, + *, + max_records: int, + block: timedelta | None = None, + ) -> Any: + records = await super().read_after( + key, after, max_records=max_records, block=block + ) + await self._log(records) + return records + + async def read_range(self, key: Any, first: Any, last: Any) -> Any: # type: ignore[override] + records = await super().read_range(key, first, last) + await self._log(records) + return records + + async def _log(self, records: Any) -> None: + offsets = [str(r.offset) for r in records if r.offset is not None] + if offsets: + await self._client.rpush(read_log_key(prefix, reader), *offsets) + + return ReadRecordingBackend(url=url, key_prefix=prefix) + + +def blocking_output_backend(*, url: str, prefix: str) -> Any: + """A Redis provider that stops only after output is durably staged. + + Blocking inside ``stage_output`` places the process at the exact boundary + the crash tests need: Redis has atomically stored the pending stage, while + the Worker has not returned the activation completion to Core/server. + """ + from temporalio.contrib.external_workflow_streams._redis import RedisStreamBackend + + class DurablyStageThenBlockBackend(RedisStreamBackend): + async def stage_output(self, manifest: Any, records: Any) -> Any: + stage = await super().stage_output(manifest, records) + await self._client.rpush( + output_stage_log_key(prefix), + json.dumps( + { + "run_id": manifest.run_id, + "stage_token": manifest.stage_token, + "sub_batch_id": manifest.sub_batch_id, + }, + sort_keys=True, + ), + ) + # There is deliberately no release path. The parent sends SIGKILL, + # so neither this call nor Worker shutdown can report completion. + await asyncio.Event().wait() + return stage + + return DurablyStageThenBlockBackend(url=url, key_prefix=prefix) + + +async def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--address", required=True) + parser.add_argument("--namespace", required=True) + parser.add_argument("--task-queue", required=True) + parser.add_argument("--redis-url", required=True) + parser.add_argument("--prefix", required=True) + parser.add_argument("--reader", default="crashed") + parser.add_argument( + "--mode", + choices=("consume", "publish", "update"), + default="consume", + ) + args = parser.parse_args() + + client = await Client.connect(args.address, namespace=args.namespace) + workflows: list[type[Any]] + if args.mode == "consume": + backend = recording_backend( + url=args.redis_url, prefix=args.prefix, reader=args.reader + ) + workflows = [CrashConsumeWorkflow] + else: + backend = blocking_output_backend(url=args.redis_url, prefix=args.prefix) + workflows = ( + [CrashPublishWorkflow] + if args.mode == "publish" + else [CrashUpdatePublishWorkflow] + ) + async with Worker( + client, + task_queue=args.task_queue, + workflows=workflows, + external_stream_backend=backend, + ): + # Until killed. There is deliberately no shutdown path here. + await asyncio.Event().wait() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/contrib/external_workflow_streams/m1_gate.py b/tests/contrib/external_workflow_streams/m1_gate.py new file mode 100644 index 0000000000..0f23117b3a --- /dev/null +++ b/tests/contrib/external_workflow_streams/m1_gate.py @@ -0,0 +1,146 @@ +"""The machinery behind the required-test gates (P16a, P16b). + +The required-test lists are the milestone gates. A gate that is only a claim in +a document is not a gate at all: someone renames a test, the case it covered +quietly stops existing, and the list still says the milestone is met. So the +list is parsed from the plan itself and every case is mapped to a test that has +to actually exist. + +Three things this checks, each catching a different way the gate rots: + +- the case **count** still matches what the plan's heading declares, so a case + added to the plan cannot be silently unmapped; +- every case is either covered or recorded as open, so a new case fails the gate + rather than passing by omission -- and a partially covered one is never + counted as covered, which is how a gate stops meaning anything; +- every mapped test **exists**, so a rename or deletion fails here rather than + leaving the list pointing at nothing. + +What this deliberately does *not* do is assert that a mapped test passes. That +is the test suite's job, and duplicating it here would only mean running +everything twice. +""" + +from __future__ import annotations + +import ast +import re +from dataclasses import dataclass +from pathlib import Path + +#: The required-test lists live in the vendored Core checkout, which is the same +#: branch the Core-side work is on. Reading them from there rather than from a +#: copy is what makes the count check meaningful -- a copy would drift silently. +PLAN_DIR = ( + Path(__file__).resolve().parents[3] + / "temporalio" + / "bridge" + / "sdk-core" + / "arch_docs" + / "streaming-poc-docs" + / "required-tests" +) + +TESTS_DIR = Path(__file__).resolve().parent + +CORE_TESTS = ( + Path(__file__).resolve().parents[3] + / "temporalio" + / "bridge" + / "sdk-core" + / "crates" + / "sdk-core" + / "src" + / "core_tests" + / "external_streams.rs" +) + + +@dataclass(frozen=True) +class RequiredCase: + """One bullet of a required-test list.""" + + number: int + text: str + + +def declared_count(list_name: str) -> int: + """The count the plan's own heading declares. + + Read rather than hardcoded: the plan says to update the heading whenever a + case is added or removed, and reading it is what turns that instruction into + something enforced. + """ + heading = (PLAN_DIR / list_name).read_text().splitlines()[0] + match = re.search(r"(\d+)\s+cases", heading) + assert match, f"{list_name} heading does not declare a case count: {heading!r}" + return int(match.group(1)) + + +def required_cases(list_name: str) -> list[RequiredCase]: + """Every bullet, in document order, with continuation lines joined. + + Bullets wrap across lines in the plan, and a parser that treated each line as + a case would count the same case several times -- inflating the total to + something that happened to match by accident. + """ + cases: list[RequiredCase] = [] + current: list[str] | None = None + for line in (PLAN_DIR / list_name).read_text().splitlines(): + if line.startswith("- "): + if current is not None: + cases.append(RequiredCase(len(cases) + 1, " ".join(current))) + current = [line[2:].strip()] + elif current is not None and line.startswith(" ") and line.strip(): + current.append(line.strip()) + elif current is not None and not line.strip(): + cases.append(RequiredCase(len(cases) + 1, " ".join(current))) + current = None + if current is not None: + cases.append(RequiredCase(len(cases) + 1, " ".join(current))) + return cases + + +def _python_test_names() -> dict[str, set[str]]: + """Every test function per module, read from the source rather than pytest. + + Parsing beats collecting here: collection would import every module and run + every fixture, and this check must be able to run when something else in the + suite is broken. + """ + found: dict[str, set[str]] = {} + for path in sorted(TESTS_DIR.glob("test_*.py")): + tree = ast.parse(path.read_text()) + names = { + node.name + for node in ast.walk(tree) + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name.startswith("test_") + } + found[path.name] = names + return found + + +def _core_test_names() -> set[str]: + if not CORE_TESTS.exists(): + return set() + return set(re.findall(r"async fn (\w+)\(", CORE_TESTS.read_text())) | set( + re.findall(r"\bfn (\w+)\(\)", CORE_TESTS.read_text()) + ) + + +def unresolved(node_ids: list[str]) -> list[str]: + """The mapped tests that do not exist. Empty means the mapping is honest.""" + python = _python_test_names() + core = _core_test_names() + missing = [] + for node_id in node_ids: + module, _, name = node_id.partition("::") + if module.endswith(".rs"): + if core and name not in core: + missing.append(node_id) + elif module not in python: + missing.append(f"{node_id} (no such module)") + elif name not in python[module]: + missing.append(node_id) + return missing diff --git a/tests/contrib/external_workflow_streams/memory_backend.py b/tests/contrib/external_workflow_streams/memory_backend.py new file mode 100644 index 0000000000..7be8c3e4aa --- /dev/null +++ b/tests/contrib/external_workflow_streams/memory_backend.py @@ -0,0 +1,228 @@ +"""A conforming in-memory backend, used to validate the conformance suite. + +Its offsets are shaped like Redis stream ids -- ``-`` compared as a +numeric pair -- so the checks that catch lexical comparison and width-crossing +boundaries are meaningful here rather than only against a real Redis. The +millisecond component comes from :attr:`now_ms`, which tests set directly, so +those boundaries are reachable without sleeping. +""" + +from __future__ import annotations + +import asyncio +import time +from datetime import timedelta +from typing import ClassVar + +from temporalio.contrib.external_workflow_streams._backend import ( + DEFAULT_WATCH_BLOCK, + AppendConflictError, + ParkIntent, + ParkIntentRemoval, + StreamBackend, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._record import ( + Cursor, + IdempotencyKey, + Offset, + StreamRecord, +) + + +def parse(offset: Offset) -> tuple[int, int]: + """``"12-3"`` as ``(12, 3)``. The provider's ordering rule, in one place.""" + ms, _, seq = offset.token.partition("-") + return int(ms), int(seq or 0) + + +class MemoryStreamBackend(StreamBackend): + """A correct reference implementation.""" + + guarantees_immutability: ClassVar[bool | None] = True + provider_id = "memory" + provider_format_version = 1 + + def __init__(self) -> None: + self.now_ms = 1 + self._records: dict[StreamKey, list[StreamRecord]] = {} + self._by_key: dict[tuple[StreamKey, IdempotencyKey], StreamRecord] = {} + self._appended = asyncio.Event() + #: Every read_range call, for tests that assert replay's call count. + self.range_reads: list[tuple[Offset, Offset]] = [] + #: Park intents, keyed `(stream key, wait_id)` -- never by stream alone. + self._intents: dict[tuple[StreamKey, int], ParkIntent] = {} + #: `(stream key, wait_id) -> (claimant, generation, expires_at)`. + self._claims: dict[tuple[StreamKey, int], tuple[str, int, float]] = {} + + # --- required operations ------------------------------------------------ + + async def append(self, key: StreamKey, record: StreamRecord) -> StreamRecord: + existing = self._by_key.get((key, record.idempotency_key)) + if existing is not None: + # Idempotent on *identity*: same bytes is a no-op returning the + # original offset, different bytes is an error. + if existing.to_fields() == record.to_fields(): + return existing + raise AppendConflictError(record.idempotency_key) + + stream = self._records.setdefault(key, []) + seq = sum( + 1 + for record in stream + if record.offset is not None and parse(record.offset)[0] == self.now_ms + ) + placed = record.placed_at(Offset(f"{self.now_ms}-{seq}")) + stream.append(placed) + self._by_key[(key, record.idempotency_key)] = placed + + self._appended.set() + self._appended.clear() + return placed + + async def read_range( + self, key: StreamKey, first: Offset, last: Offset + ) -> list[StreamRecord]: + self.range_reads.append((first, last)) + lo, hi = parse(first), parse(last) + return [ + r + for r in self._records.get(key, []) + if lo <= parse(r.offset) <= hi # type: ignore[arg-type] + ] + + async def read_after( + self, + key: StreamKey, + after: Cursor, + *, + max_records: int, + block: timedelta | None = DEFAULT_WATCH_BLOCK, + ) -> list[StreamRecord]: + deadline = None if block is None else block.total_seconds() + while True: + found = self._after(key, after)[:max_records] + if found or deadline is None or deadline <= 0: + return found + waiter = asyncio.ensure_future(self._appended.wait()) + try: + await asyncio.wait_for(waiter, deadline) + except asyncio.TimeoutError: + return [] + finally: + # `wait_for` cancels the outer future but the inner `Event.wait` + # task survives it, and a leaked one keeps the event's waiter + # list growing for the life of the process. + if not waiter.done(): + waiter.cancel() + + def _after(self, key: StreamKey, after: Cursor) -> list[StreamRecord]: + stream = self._records.get(key, []) + if after.is_beginning: + return list(stream) + bound = parse(after.offset) # type: ignore[arg-type] + return [r for r in stream if parse(r.offset) > bound] # type: ignore[arg-type] + + def compare_offsets(self, left: Offset, right: Offset) -> int: + a, b = parse(left), parse(right) + return (a > b) - (a < b) + + # --- parking ------------------------------------------------------------ + + supports_leased_claims = True + + async def install_park_intent(self, key: StreamKey, intent: ParkIntent) -> None: + self._intents[(key, intent.wait_id)] = intent + + async def remove_park_intent(self, key: StreamKey, wait_id: int) -> None: + self._intents.pop((key, wait_id), None) + self._claims.pop((key, wait_id), None) + + async def remove_park_intent_if_matches( + self, + key: StreamKey, + wait_id: int, + *, + run_id: str, + park_generation: int, + ) -> ParkIntentRemoval: + intent = self._intents.get((key, wait_id)) + if intent is None: + return ParkIntentRemoval.ABSENT + if intent.run_id != run_id or intent.park_generation != park_generation: + return ParkIntentRemoval.MISMATCH + self._intents.pop((key, wait_id), None) + self._claims.pop((key, wait_id), None) + return ParkIntentRemoval.REMOVED + + async def parked_wait_ids(self, key: StreamKey) -> list[int]: + return sorted(wait_id for (stored, wait_id) in self._intents if stored == key) + + async def park_intent(self, key: StreamKey, wait_id: int) -> ParkIntent | None: + return self._intents.get((key, wait_id)) + + async def recheck(self, key: StreamKey, wait_id: int) -> bool: + intent = self._intents.get((key, wait_id)) + if intent is None: + return False + return bool(self._after(key, intent.cursor)) + + async def claim_park_generation( + self, + key: StreamKey, + wait_id: int, + park_generation: int, + *, + claimant: str, + lease: timedelta, + ) -> bool: + now = time.monotonic() + held = self._claims.get((key, wait_id)) + if held is not None: + holder, generation, expires_at = held + # A live claim held by someone else, for this generation, blocks. + # An expired one is taken over -- that is the whole point of a + # lease, and without it a producer that crashed here would strand + # the generation forever. + if ( + generation == park_generation + and holder != claimant + and expires_at > now + ): + return False + self._claims[(key, wait_id)] = ( + claimant, + park_generation, + now + lease.total_seconds(), + ) + return True + + async def current_park_generation(self, key: StreamKey, wait_id: int) -> int | None: + intent = self._intents.get((key, wait_id)) + return None if intent is None else intent.park_generation + + # --- test affordances --------------------------------------------------- + + async def expire_claims_for_test(self) -> None: + """Ages every claim past its lease. + + Stands in for a producer that crashed between claiming and signaling -- + the failure the lease exists for -- without making the test wait out a + real lease. + """ + self._claims = { + key: (holder, generation, 0.0) + for key, (holder, generation, _) in self._claims.items() + } + + async def delete_for_test(self, key: StreamKey, offset: Offset) -> None: + """Removes a record, standing in for XDEL, trimming, or retention loss. + + Deletion is permitted -- what immutability forbids is *rewriting* a + record in place, which is why replay only has to detect absence. + """ + stream = self._records.get(key, []) + self._records[key] = [r for r in stream if r.offset != offset] + + def all_records(self, key: StreamKey) -> list[StreamRecord]: + return list(self._records.get(key, [])) diff --git a/tests/contrib/external_workflow_streams/observations.py b/tests/contrib/external_workflow_streams/observations.py new file mode 100644 index 0000000000..f2755b3cb1 --- /dev/null +++ b/tests/contrib/external_workflow_streams/observations.py @@ -0,0 +1,34 @@ +"""A sink Workflow code can write its own observations to. + +A Workflow's return value is only visible for the *live* run: the ``Replayer`` +runs the same code and reports whether it failed, not what it produced. So a +test that wants to compare what a predicate saw live against what it saw on +replay needs somewhere outside the Workflow to put it. + +This module is imported through :py:func:`workflow.unsafe.imports_passed_through`, +so Workflow code inside the sandbox reaches the *same* module object the test +does. Recording is append-only per Run: the live execution appends first and the +replay appends second, which is what makes the two comparable at all. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +#: ``run_id -> [observations of each execution, in the order they ran]``. +#: +#: Deliberately untyped in its element: what is worth comparing between a live +#: run and its replay differs per test -- the states a predicate saw, the +#: records that were delivered -- and narrowing it here would only push a cast +#: into every caller. +OBSERVED: dict[str, list[list[Any]]] = {} + + +def record(run_id: str, states: Sequence[Any]) -> None: + """Appends one execution's observed sequence for a Run.""" + OBSERVED.setdefault(run_id, []).append(list(states)) + + +def executions(run_id: str) -> list[list[Any]]: + return OBSERVED.get(run_id, []) diff --git a/tests/contrib/external_workflow_streams/test_annotation.py b/tests/contrib/external_workflow_streams/test_annotation.py new file mode 100644 index 0000000000..08a7265db0 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_annotation.py @@ -0,0 +1,755 @@ +"""P5 — the replay annotation codec.""" + +from __future__ import annotations + +import pytest + +import temporalio.exceptions +from temporalio.contrib.external_workflow_streams._annotation import ( + MAX_ANNOTATION_BYTES, + ROLLOVER_HIGH_WATER, + SCHEMA_VERSION, + Annotation, + AnnotationAccumulator, + AnnotationBudgetExceeded, + AnnotationDecodeError, + AnnotationHeader, + Run, + Segment, + SegmentEndReason, + StreamBinding, + decode_annotation, + encode_annotation, + encode_bindings, + encode_segment, + encode_terminal, +) +from temporalio.contrib.external_workflow_streams._backend import StreamKey +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Cursor, + Offset, +) + +KEY = StreamKey("ns", "wf", "run-1", "tokens") +OTHER_KEY = StreamKey("ns", "wf", "run-1", "tool-events") + + +def binding(key: StreamKey, start_cursor: Cursor = BEGINNING) -> StreamBinding: + """A binding carrying the configured backend's provider identity.""" + return StreamBinding( + stream_key=key, + start_cursor=start_cursor, + provider_id="redis-streams", + provider_format_version=1, + ) + + +def header(streams: dict[int, StreamBinding] | None = None) -> AnnotationHeader: + return AnnotationHeader( + streams=streams if streams is not None else {1: binding(KEY)}, + ) + + +def offset(ms: int, seq: int = 0) -> Offset: + return Offset(f"{ms}-{seq}") + + +# --- round trips ------------------------------------------------------------ + + +def test_a_full_annotation_round_trips() -> None: + annotation = Annotation( + header=header({1: binding(KEY, AFTER(offset(50)))}), + segments=( + Segment( + runs=(Run(1, offset(100), offset(199), 100, (7, 42)),), + end_reason=SegmentEndReason.NO_DATA_AVAILABLE, + ), + Segment( + runs=(Run(1, offset(200), offset(205), 6),), + end_reason=SegmentEndReason.FENCE_REACHED, + ), + ), + terminal={1: AFTER(offset(205))}, + ) + + assert decode_annotation(encode_annotation(annotation)) == annotation + + +def test_schema_version_leads_the_encoding() -> None: + """Readable before anything version-dependent is interpreted.""" + encoded = encode_annotation(Annotation(header(), terminal={1: BEGINNING})) + + assert encoded[0] == SCHEMA_VERSION + assert decode_annotation(encoded).header.schema_version == SCHEMA_VERSION + + +def test_an_unknown_schema_version_is_refused_by_name() -> None: + encoded = bytearray(encode_annotation(Annotation(header(), terminal={}))) + encoded[0] = 99 + + with pytest.raises(AnnotationDecodeError, match="schema version 99"): + decode_annotation(bytes(encoded)) + + +def test_an_empty_segment_round_trips() -> None: + """An activation that drained and found nothing still ran one drain. + + Dropping it would make ``wait_condition`` predicates fire a different + number of times on replay than they did live. + """ + annotation = Annotation( + header(), + segments=(Segment(runs=(), end_reason=SegmentEndReason.NO_DATA_AVAILABLE),), + terminal={1: BEGINNING}, + ) + + decoded = decode_annotation(encode_annotation(annotation)) + + assert decoded == annotation + assert decoded.segments[0].runs == () + + +def test_an_annotation_with_no_segments_round_trips() -> None: + """A subscription to an empty stream: header and terminal carry everything.""" + annotation = Annotation(header(), segments=(), terminal={1: BEGINNING}) + + decoded = decode_annotation(encode_annotation(annotation)) + + assert decoded == annotation + assert decoded.segments == () + assert decoded.header.streams[1].start_cursor == BEGINNING + + +def test_cursors_and_offsets_stay_distinct_types() -> None: + """A boundary and a record identity are not interchangeable. + + ``AFTER(100-0)`` in the terminal and the record offset ``100-0`` in a run + must decode as a Cursor and an Offset respectively, not as each other. + """ + annotation = Annotation( + header(), + segments=( + Segment( + (Run(1, offset(100), offset(100), 1),), SegmentEndReason.BATCH_LIMIT + ), + ), + terminal={1: AFTER(offset(100))}, + ) + + decoded = decode_annotation(encode_annotation(annotation)) + + assert decoded.segments[0].runs[0].first_offset == offset(100) + assert decoded.terminal is not None + assert decoded.terminal[1] == AFTER(offset(100)) + assert decoded.terminal[1] != offset(100) + + +def test_beginning_is_not_the_same_as_after_a_zero_offset() -> None: + annotation = Annotation(header(), terminal={1: BEGINNING, 2: AFTER(Offset("0-0"))}) + + decoded = decode_annotation(encode_annotation(annotation)) + + assert decoded.terminal is not None + assert decoded.terminal[1].is_beginning + assert not decoded.terminal[2].is_beginning + + +def test_sparse_control_positions_round_trip() -> None: + run = Run(1, offset(100), offset(999), 500, (0, 250, 499)) + + decoded = decode_annotation( + encode_annotation( + Annotation( + header(), + segments=(Segment((run,), SegmentEndReason.FENCE_REACHED),), + terminal={1: AFTER(offset(999))}, + ) + ) + ) + + assert decoded.segments[0].runs[0].control_positions == (0, 250, 499) + + +@pytest.mark.parametrize( + ("positions", "count", "reason"), + [ + ((5,), 3, "outside a run"), + ((2, 1), 5, "strictly increasing"), + ((1, 1), 5, "strictly increasing"), + ], +) +def test_malformed_control_positions_are_rejected_at_construction( + positions: tuple[int, ...], count: int, reason: str +) -> None: + with pytest.raises(ValueError, match=reason): + Run(1, offset(1), offset(9), count, positions) + + +def test_a_run_covers_at_least_one_record() -> None: + with pytest.raises(ValueError, match="at least one record"): + Run(1, offset(1), offset(1), 0) + + +@pytest.mark.parametrize( + "corrupt", + [b"", b"\x01", b"\x01\x02", b"\x01\x01\x05abc"], + ids=["empty", "version-only", "wrong-first-frame", "truncated-header"], +) +def test_corrupt_bytes_raise_a_decode_error(corrupt: bytes) -> None: + with pytest.raises(AnnotationDecodeError): + decode_annotation(corrupt) + + +def test_two_terminals_are_refused() -> None: + doubled = encode_annotation(Annotation(header(), terminal={1: BEGINNING})) + doubled += doubled[doubled.index(b"\x03") :] + + with pytest.raises(AnnotationDecodeError, match="one terminal, not two"): + decode_annotation(doubled) + + +# --- runs are what keep markers small --------------------------------------- + + +def test_a_large_single_stream_batch_encodes_as_one_run() -> None: + """100,000 records, one run. This is the claim the whole design rests on.""" + annotation = Annotation( + header(), + segments=( + Segment( + (Run(1, offset(1), offset(100_000), 100_000),), + SegmentEndReason.NO_DATA_AVAILABLE, + ), + ), + terminal={1: AFTER(offset(100_000))}, + ) + + decoded = decode_annotation(encode_annotation(annotation)) + + assert len(decoded.segments[0].runs) == 1 + assert decoded.segments[0].runs[0].count == 100_000 + + +def test_encoded_size_stays_flat_as_a_single_stream_batch_grows() -> None: + """Byte size, not run count. + + A run-count assertion says nothing about what a future per-record field + would cost; measuring bytes is what actually catches one being added. + """ + + def encoded_size(count: int) -> int: + return len( + encode_annotation( + Annotation( + header(), + segments=( + Segment( + (Run(1, offset(1), offset(count), count),), + SegmentEndReason.NO_DATA_AVAILABLE, + ), + ), + terminal={1: AFTER(offset(count))}, + ) + ) + ) + + small, large = encoded_size(100), encoded_size(100_000) + + # Three orders of magnitude more records. The only growth permitted is in + # the varints holding the count and the offset tokens themselves. + assert large - small <= 16, ( + f"{small} bytes for 100 records but {large} for 100,000 -- something in " + "the encoding is per-record" + ) + + +def test_alternating_streams_cost_one_run_per_delivery() -> None: + """The honest worst case, stated as a test rather than assumed away.""" + runs = tuple(Run(1 + (i % 2), offset(i), offset(i), 1) for i in range(1, 21)) + annotation = Annotation( + header({1: binding(KEY), 2: binding(OTHER_KEY)}), + segments=(Segment(runs, SegmentEndReason.NO_DATA_AVAILABLE),), + terminal={1: AFTER(offset(19)), 2: AFTER(offset(20))}, + ) + + decoded = decode_annotation(encode_annotation(annotation)) + + assert len(decoded.segments[0].runs) == 20 + assert [r.wait_id for r in decoded.segments[0].runs[:4]] == [2, 1, 2, 1] + + +# --- accumulation: deltas concatenate into the annotation ------------------- + + +def test_a_delta_sequence_concatenates_into_the_annotation() -> None: + """What Core accumulates by byte-append must equal what encoding produces. + + Core never parses a delta, so if these two ever diverged nothing would + notice until replay read a marker that could not be decoded. + """ + head = header() + segments = ( + Segment((Run(1, offset(100), offset(110), 11),), SegmentEndReason.BATCH_LIMIT), + Segment((), SegmentEndReason.NO_DATA_AVAILABLE), + Segment( + (Run(1, offset(111), offset(120), 10),), SegmentEndReason.FENCE_REACHED + ), + ) + terminal = {1: AFTER(offset(120))} + + accumulator = AnnotationAccumulator(head) + deltas = [accumulator.add_segment(s) for s in segments] + deltas.append(accumulator.add_terminal(terminal)) + + core_would_hold = accumulator.accumulated() + + assert core_would_hold == encode_annotation(Annotation(head, segments, terminal)) + assert decode_annotation(core_would_hold) == Annotation(head, segments, terminal) + # And the header rode the accumulator's first emission, not a delta. + assert core_would_hold.endswith(b"".join(deltas)) + + +def test_the_accumulator_reports_the_size_the_marker_will_carry() -> None: + accumulator = AnnotationAccumulator(header()) + accumulator.add_segment( + Segment((Run(1, offset(1), offset(9), 9),), SegmentEndReason.BATCH_LIMIT) + ) + accumulator.add_terminal({1: AFTER(offset(9))}) + + assert accumulator.size == len(accumulator.accumulated()) + + +# --- a wait bound after the header ------------------------------------------ + + +def test_a_wait_bound_after_the_header_decodes_into_it() -> None: + """A subscription can be created at any activation of a retained task. + + The header frame for that task has usually already gone to Core by then, + and Core appends deltas rather than rewriting them, so the binding rides a + frame of its own. What replay reads is still one flat table: a wait that + joined late is indistinguishable, once decoded, from one that was there + from the start -- which is what lets `prepare_replay` resolve its backend + and stream key at all. + """ + accumulator = AnnotationAccumulator(header()) + accumulator.add_segment( + Segment((Run(1, offset(100), offset(100), 1),), SegmentEndReason.BATCH_LIMIT) + ) + accumulator.add_bindings({2: binding(OTHER_KEY)}) + accumulator.add_segment( + Segment( + (Run(2, offset(200), offset(200), 1),), SegmentEndReason.NO_DATA_AVAILABLE + ) + ) + accumulator.add_terminal({1: AFTER(offset(100)), 2: AFTER(offset(200))}) + + decoded = decode_annotation(accumulator.accumulated()) + + assert decoded.header.streams == {1: binding(KEY), 2: binding(OTHER_KEY)} + assert decoded.terminal == {1: AFTER(offset(100)), 2: AFTER(offset(200))} + assert [run.wait_id for segment in decoded.segments for run in segment.runs] == [ + 1, + 2, + ] + + +def test_a_late_binding_keeps_its_own_start_cursor() -> None: + """Not wherever the waits already in the header have reached. + + A wait registered part-way through a Workflow Task begins at the boundary + its own `subscribe()` was given. Recording the annotation-wide position + instead would start replay of that wait past records it in fact received. + """ + accumulator = AnnotationAccumulator(header({1: binding(KEY, AFTER(offset(100)))})) + accumulator.add_bindings({2: binding(OTHER_KEY)}) + accumulator.add_terminal({1: AFTER(offset(100)), 2: BEGINNING}) + + decoded = decode_annotation(accumulator.accumulated()) + + assert decoded.header.streams[1].start_cursor == AFTER(offset(100)) + assert decoded.header.streams[2].start_cursor == BEGINNING + + +def test_a_bindings_frame_is_the_header_frame_body_under_its_own_tag() -> None: + """Locked to exact bytes, like the golden annotation and for the reason. + + A second *header* frame would have been the smaller change and is what the + encoding deliberately does not do: a header frame also carries the schema + version, so a decoder meeting two of them has to decide whether the later + replaces the earlier. The distinct tag says one thing only. + """ + assert encode_bindings({2: binding(OTHER_KEY)}).hex() == "".join( + [ + "04", # bindings frame + "01", # one stream + "02", # wait id 2 + "02" + "ns".encode().hex(), + "02" + "wf".encode().hex(), + "05" + "run-1".encode().hex(), + "0b" + "tool-events".encode().hex(), + "00", # start cursor BEGINNING + "0d" + "redis-streams".encode().hex(), # provider id + "01", # provider format version 1 + ] + ) + + +def test_a_wait_bound_twice_is_refused() -> None: + """Two bindings for one wait leave replay choosing between stream keys. + + Whichever it chose could be the one the recorded records were not written + to, and the range read would then fail as integrity loss against a backend + nothing is wrong with. + """ + accumulator = AnnotationAccumulator(header()) + accumulator.add_bindings({1: binding(OTHER_KEY)}) + accumulator.add_terminal({1: BEGINNING}) + + with pytest.raises(AnnotationDecodeError, match="bound twice"): + decode_annotation(accumulator.accumulated()) + + +def test_bindings_after_the_terminal_are_refused() -> None: + """Both halves: the encoder cannot emit one and the decoder rejects one. + + The annotation ends at its terminal. A binding appended past it belongs to + a Workflow Task whose marker has already been decided, and Core would write + it into that marker rather than the next one. + """ + accumulator = AnnotationAccumulator(header()) + accumulator.add_terminal({1: BEGINNING}) + with pytest.raises(ValueError, match="after the terminal"): + accumulator.add_bindings({2: binding(OTHER_KEY)}) + + with pytest.raises(AnnotationDecodeError, match="follows the terminal"): + decode_annotation( + encode_annotation(Annotation(header(), terminal={1: BEGINNING})) + + encode_bindings({2: binding(OTHER_KEY)}) + ) + + +def test_a_bindings_frame_binds_at_least_one_wait() -> None: + """An empty frame costs bytes and says nothing, so it is a caller error.""" + accumulator = AnnotationAccumulator(header()) + with pytest.raises(ValueError, match="at least one wait"): + accumulator.add_bindings({}) + + +def test_a_segment_after_the_terminal_is_refused() -> None: + accumulator = AnnotationAccumulator(header()) + accumulator.add_terminal({1: BEGINNING}) + + with pytest.raises(ValueError, match="after the terminal"): + accumulator.add_segment(Segment((), SegmentEndReason.NO_DATA_AVAILABLE)) + + +def test_a_second_terminal_is_refused() -> None: + accumulator = AnnotationAccumulator(header()) + accumulator.add_terminal({1: BEGINNING}) + + with pytest.raises(ValueError, match="one terminal, not two"): + accumulator.add_terminal({1: BEGINNING}) + + +# --- the byte budget -------------------------------------------------------- + + +def test_rollover_is_requested_before_the_budget_is_reached() -> None: + accumulator = AnnotationAccumulator(header()) + assert not accumulator.request_rollover + + run = Run(1, offset(1), offset(2), 2) + while not accumulator.request_rollover: + accumulator.add_segment(Segment((run,), SegmentEndReason.BATCH_LIMIT)) + + assert accumulator.size >= MAX_ANNOTATION_BYTES * ROLLOVER_HIGH_WATER + # The margin is the point: the request has to reach Core and come back, so + # asking at 100% would arrive too late to prevent the overflow. + assert accumulator.size < MAX_ANNOTATION_BYTES + + +def test_an_alternating_two_stream_batch_asks_for_rollover_rather_than_overflowing() -> ( + None +): + """Bounded marker size is bought with additional Workflow Tasks. + + An alternating workload has a schedule transition per record, so it cannot + be range-compressed. The encoder must hit the high-water mark and ask for a + rollover, not grow past the budget. + """ + accumulator = AnnotationAccumulator( + header({1: binding(KEY), 2: binding(OTHER_KEY)}) + ) + + delivered = 0 + while not accumulator.request_rollover: + runs = tuple( + Run(1 + (i % 2), offset(delivered + i), offset(delivered + i), 1) + for i in range(50) + ) + delivered += 50 + accumulator.add_segment(Segment(runs, SegmentEndReason.BATCH_LIMIT)) + + assert accumulator.size <= MAX_ANNOTATION_BYTES + assert delivered > 0 + + +def test_exceeding_the_hard_budget_raises_rather_than_growing() -> None: + """The last resort, reached only if every prevention above it failed. + + The alternative is the server rejecting an oversized event, which surfaces + as an unexplained Workflow Task failure a long way from the cause. + """ + tiny = AnnotationAccumulator(header(), max_bytes=120, high_water=0.9) + + with pytest.raises(AnnotationBudgetExceeded, match="past the"): + for i in range(100): + tiny.add_segment( + Segment( + (Run(1, offset(i), offset(i + 1), 2),), SegmentEndReason.BATCH_LIMIT + ) + ) + + +def test_the_budget_error_fails_the_workflow_rather_than_the_task() -> None: + """A budget overflow may not become a Workflow Task that retries forever. + + ADR-007 rejects check-and-fail precisely because the encoding that overflowed + overflows again on the retry, and the server retries Workflow Task failures + regardless of cause: no marker, no terminal, no rollover, forever. A + non-retryable application failure is bounded and reported once instead. + """ + tiny = AnnotationAccumulator(header(), max_bytes=120, high_water=0.9) + + with pytest.raises(AnnotationBudgetExceeded) as caught: + for i in range(100): + tiny.add_segment( + Segment( + (Run(1, offset(i), offset(i + 1), 2),), SegmentEndReason.BATCH_LIMIT + ) + ) + + assert isinstance(caught.value, temporalio.exceptions.ApplicationError) + assert caught.value.non_retryable, ( + "a retryable budget overflow is the permanent Workflow Task retry loop " + "ADR-007 exists to rule out" + ) + + +def test_a_reserved_terminal_always_fits() -> None: + """The frame that closes the annotation may never be the one refused. + + An annotation Core writes with no terminal is durable and cannot be decoded + past the frame after it, so "the terminal did not fit" has to be arithmetically + impossible rather than merely unlikely. Reserving what it costs is what makes + it so -- segments are then refused *before* the space the terminal needs is + gone. + """ + blocked = {1: AFTER(offset(9)), 2: AFTER(offset(9))} + terminal_bytes = len(encode_terminal(blocked)) + + accumulator = AnnotationAccumulator(header(), max_bytes=4096, high_water=0.9) + accumulator.reserve(terminal_bytes) + + refused = False + for i in range(4096): + segment = Segment( + (Run(1, offset(i), offset(i + 1), 2),), SegmentEndReason.BATCH_LIMIT + ) + if not accumulator.fits(len(encode_segment(segment))): + refused = True + break + accumulator.add_segment(segment) + + assert refused, "the loop must have run into the reserve, or it proves nothing" + assert accumulator.headroom == 0 or accumulator.headroom < terminal_bytes + 32 + # The whole point: with the segments stopped, the terminal still goes in. + accumulator.add_terminal(blocked) + assert accumulator.size <= 4096 + decoded = decode_annotation(accumulator.accumulated()) + assert decoded.terminal == blocked + + +def test_the_reserve_is_only_spendable_by_the_frames_it_was_set_aside_for() -> None: + """A segment may not eat the terminal's space, however much room *looks* left.""" + accumulator = AnnotationAccumulator(header(), max_bytes=4096, high_water=0.9) + accumulator.reserve(4096 - accumulator.size) + + assert accumulator.headroom == 0 + with pytest.raises(AnnotationBudgetExceeded): + accumulator.add_segment( + Segment((Run(1, offset(0), offset(1), 2),), SegmentEndReason.BATCH_LIMIT) + ) + + +# --- golden file ------------------------------------------------------------ + +#: A fixed annotation and its exact bytes. Regenerating this to make a test +#: pass is the mistake it exists to catch: any change here is a wire-format +#: change, and markers already in History were written by the old encoder. +#: +#: The feature is private and unreleased, so this covers only the current format +#: and deliberately provides no compatibility fixture for earlier prototypes. +GOLDEN_ANNOTATION = Annotation( + header=AnnotationHeader( + streams={ + 1: StreamBinding( + stream_key=StreamKey("ns", "wf", "run-1", "tokens"), + start_cursor=BEGINNING, + provider_id="redis-streams", + provider_format_version=1, + ) + }, + ), + segments=( + Segment( + runs=(Run(1, Offset("100-0"), Offset("104-0"), 5, (2,)),), + end_reason=SegmentEndReason.FENCE_REACHED, + ), + Segment(runs=(), end_reason=SegmentEndReason.NO_DATA_AVAILABLE), + ), + terminal={1: AFTER(Offset("104-0"))}, +) + +GOLDEN_BYTES = bytes.fromhex( + "".join( + [ + "01", # schema version 1 + "01", # header frame + "01", # one stream + "01", # wait id 1 + "02" + "ns".encode().hex(), + "02" + "wf".encode().hex(), + "05" + "run-1".encode().hex(), + "06" + "tokens".encode().hex(), + "00", # start cursor BEGINNING + "0d" + "redis-streams".encode().hex(), # provider id + "01", # provider format version 1 + "02", # segment frame + "01", # one run + "01", # wait id 1 + "05" + "100-0".encode().hex(), + "05" + "104-0".encode().hex(), + "05", # count 5 + "01", # one control position + "02", # at relative index 2 + "03", # FENCE_REACHED + "02", # segment frame + "00", # no runs + "01", # NO_DATA_AVAILABLE + "03", # terminal frame + "01", # one entry + "01", # wait id 1 + "01", # AFTER + "05" + "104-0".encode().hex(), + ] + ) +) + + +def test_golden_bytes_are_unchanged() -> None: + """Silent format drift is what this catches. + + If this fails, the encoding changed. Bump `SCHEMA_VERSION` and keep the old + decoder working -- do not regenerate the constant. + """ + assert encode_annotation(GOLDEN_ANNOTATION).hex() == GOLDEN_BYTES.hex() + + +def test_golden_bytes_decode_to_the_golden_annotation() -> None: + assert decode_annotation(GOLDEN_BYTES) == GOLDEN_ANNOTATION + + +def test_encoded_size_stays_flat_with_sparse_control_records() -> None: + """The realistic shape, and the one where a per-record cost could hide. + + ``control_positions`` is the only field that carries per-record information + at all, which makes it the natural place for the encoding to start scaling + with record count. It stays sparse -- a handful of fence positions in a batch + of any size -- so the marker's cost must track the number of *control* + records, not the number of records. + + Measured in bytes rather than by inspecting the field, because what matters + is what reaches History. + """ + + def encoded_size(count: int) -> int: + # A fence every ~10,000 records: what a few long-running producers + # finishing at intervals actually looks like. + controls = tuple(range(0, count, max(1, count // 10))) + return len( + encode_annotation( + Annotation( + header(), + segments=( + Segment( + ( + Run( + 1, + offset(1), + offset(count), + count, + control_positions=controls, + ), + ), + SegmentEndReason.NO_DATA_AVAILABLE, + ), + ), + terminal={1: AFTER(offset(count))}, + ) + ) + ) + + small, large = encoded_size(100), encoded_size(100_000) + + # Both carry ten control positions; only the positions themselves grow, and + # only as varints. + assert large - small <= 32, ( + f"{small} bytes for 100 records but {large} for 100,000, both with ten " + "control records -- the control encoding is scaling with record count " + "rather than with control count" + ) + + +def test_encoded_size_grows_only_with_the_number_of_control_records() -> None: + """The other half: a marker *may* grow with fences, and only with them. + + Stated as a bound rather than left implicit, because "sparse" is a claim + about the producer's behaviour and this is what makes the encoding's side of + it checkable. + """ + + def encoded_size(controls: int) -> int: + return len( + encode_annotation( + Annotation( + header(), + segments=( + Segment( + ( + Run( + 1, + offset(1), + offset(1000), + 1000, + control_positions=tuple(range(controls)), + ), + ), + SegmentEndReason.NO_DATA_AVAILABLE, + ), + ), + terminal={1: AFTER(offset(1000))}, + ) + ) + ) + + per_control = encoded_size(101) - encoded_size(1) + + assert 100 <= per_control <= 300, ( + f"100 extra control records cost {per_control} bytes; a sparse position " + "list should cost a small constant each" + ) diff --git a/tests/contrib/external_workflow_streams/test_api.py b/tests/contrib/external_workflow_streams/test_api.py new file mode 100644 index 0000000000..08d7257618 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_api.py @@ -0,0 +1,762 @@ +"""P9 — the Workflow-facing API and its wait-id assignment.""" + +from __future__ import annotations + +import asyncio +import inspect +from datetime import timedelta +from typing import Any + +import pytest + +import temporalio.converter +import temporalio.workflow +from temporalio.contrib.external_workflow_streams import _api +from temporalio.contrib.external_workflow_streams._api import ( + DEFAULT_IDLE_TIMEOUT, + MAX_RECORDS_PER_ACTIVATION, + ExternalStreamOptions, + ExternalStreamSubscription, + ExternalStreamTopic, + _install_runtime, + external_stream, + merge, +) +from temporalio.contrib.external_workflow_streams._backend import StreamKey +from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec +from temporalio.contrib.external_workflow_streams._errors import ( + ConcurrentStreamConsumerError, + StreamDecodeError, +) +from temporalio.contrib.external_workflow_streams._manager import PreparedRecord +from temporalio.contrib.external_workflow_streams._record import ( + RecordKind, + StreamRecord, +) + + +class FakeRuntime: + """The Worker's half, standing in for the real manager handle. + + Deliberately only offers what the protocol names -- no provider instance, + no way to reach one -- because that is the point of the boundary. + """ + + def __init__(self) -> None: + self.registrations: list[tuple[int, StreamKey]] = [] + #: `wait_id -> configured idle timeout`, so a test can see that the + #: value `with_options` was given actually reached the Worker. + self.idle_timeouts: dict[int, timedelta] = {} + self.buffers: dict[int, list[StreamRecord]] = {} + self.deliveries: list[tuple[int, StreamRecord]] = [] + self.consumed: list[tuple[int, StreamRecord]] = [] + self.blocked: list[tuple[int, bool]] = [] + #: The wait ids `close()` asked the Worker to stop serving. + self.unsubscribed: list[int] = [] + self.pending: dict[int, asyncio.Future[None]] = {} + self.budget = MAX_RECORDS_PER_ACTIVATION + #: Overrides what `codec_for` hands back, so a test can control when -- + #: and whether -- decoding a record succeeds. + self.codec: Any = None + + def stream_key(self, stream_name: str) -> StreamKey: + return StreamKey("ns", "wf", "first-run", stream_name) + + def register( + self, + *, + wait_id: int, + stream_key: StreamKey, + idle_timeout: timedelta, + ) -> None: + self.registrations.append((wait_id, stream_key)) + self.idle_timeouts[wait_id] = idle_timeout + + def drain(self, wait_id: int, max_records: int | None = None) -> list[StreamRecord]: + buffered = self.buffers.get(wait_id, []) + if max_records is not None: + buffered, self.buffers[wait_id] = ( + buffered[:max_records], + buffered[max_records:], + ) + else: + self.buffers[wait_id] = [] + return buffered + + def delivery_budget_remaining(self) -> int: + return self.budget + + def record_consumption(self, wait_id: int, record: StreamRecord) -> None: + self.consumed.append((wait_id, record)) + + def codec_for( + self, + value_type: type | None, + wait_id: int, # pyright: ignore[reportUnusedParameter] + ) -> StreamPayloadCodec[Any]: + if self.codec is not None: + return self.codec + return StreamPayloadCodec( + temporalio.converter.DataConverter.default, value_type + ) + + def new_readiness_future(self) -> asyncio.Future[None]: + return asyncio.get_event_loop().create_future() + + def record_delivery(self, wait_id: int, record: StreamRecord) -> None: + self.deliveries.append((wait_id, record)) + # Charged where the real runtime charges it: the drain that moves a record + # into a ready list is the reservation, not the consumption that follows. + self.budget = max(0, self.budget - 1) + + def note_blocked(self, wait_id: int, blocked: bool) -> None: + self.blocked.append((wait_id, blocked)) + + def unsubscribe(self, wait_id: int) -> None: + self.unsubscribed.append(wait_id) + + def register_pending(self, wait_id: int, future: asyncio.Future[None]) -> None: + self.pending[wait_id] = future + + def discard_pending(self, wait_id: int) -> None: + self.pending.pop(wait_id, None) + + +class FakeInstance: + """Stands in for the user's Workflow object, which holds the per-Run state.""" + + +@pytest.fixture +def workflow_instance(monkeypatch: pytest.MonkeyPatch) -> FakeInstance: + instance = FakeInstance() + monkeypatch.setattr(temporalio.workflow, "instance", lambda: instance) + return instance + + +@pytest.fixture +def runtime(workflow_instance: FakeInstance) -> FakeRuntime: + runtime = FakeRuntime() + _install_runtime(workflow_instance, runtime) + return runtime + + +# --- options ------------------------------------------------------------------ + + +def test_the_default_idle_timeout_is_one_second() -> None: + assert external_stream.idle_timeout == DEFAULT_IDLE_TIMEOUT == timedelta(seconds=1) + + +def test_with_options_returns_a_copy() -> None: + """The module-level entry point must not be mutated by configuring it.""" + configured = external_stream.with_options(idle_timeout=timedelta(seconds=5)) + + assert configured.idle_timeout == timedelta(seconds=5) + assert external_stream.idle_timeout == timedelta(seconds=1) + assert isinstance(configured, ExternalStreamOptions) + + +@pytest.mark.parametrize("bad", [timedelta(0), timedelta(seconds=-1)]) +def test_a_non_positive_idle_timeout_is_rejected(bad: timedelta) -> None: + """A configuration error, not a request to park immediately.""" + with pytest.raises(ValueError, match="must be positive"): + external_stream.with_options(idle_timeout=bad) + + +def test_topics_inherit_their_options( + runtime: FakeRuntime, # pyright: ignore[reportUnusedParameter] +) -> None: + configured = external_stream.with_options(idle_timeout=timedelta(seconds=3)) + + subscription = configured.topic("tokens").subscribe() + + assert subscription.idle_timeout == timedelta(seconds=3) + + +# --- topics ------------------------------------------------------------------- + + +def test_a_workflow_never_sees_the_configured_backend(runtime: FakeRuntime) -> None: + topic = external_stream.topic("tokens", type=str) + + subscription = topic.subscribe() + + assert isinstance(subscription, ExternalStreamSubscription) + assert runtime.registrations == [(1, StreamKey("ns", "wf", "first-run", "tokens"))] + + +def test_a_topic_needs_a_name() -> None: + with pytest.raises(ValueError, match="non-empty name"): + external_stream.topic("") + + +def test_subscribing_without_a_configured_worker_says_so( + workflow_instance: FakeInstance, # pyright: ignore[reportUnusedParameter] +) -> None: + with pytest.raises(RuntimeError, match="external_stream_backend"): + external_stream.topic("tokens").subscribe() + + +# --- wait id assignment ------------------------------------------------------- + + +def test_wait_ids_are_assigned_from_one_in_subscribe_call_order( + runtime: FakeRuntime, +) -> None: + first = external_stream.topic("a").subscribe() + second = external_stream.topic("b").subscribe() + third = external_stream.topic("c").subscribe() + + assert [first.wait_id, second.wait_id, third.wait_id] == [1, 2, 3] + assert [wait_id for wait_id, _ in runtime.registrations] == [1, 2, 3] + + +def test_wait_ids_reproduce_across_two_runs_of_the_same_code( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Replay depends on this exactly as it depends on timer sequence numbers. + + A renumbered wait produces an annotation mismatch rather than a silently + different stream result -- but only because the numbering is reproducible + in the first place. + """ + + def run_once() -> list[int]: + instance = FakeInstance() + monkeypatch.setattr(temporalio.workflow, "instance", lambda: instance) + _install_runtime(instance, FakeRuntime()) + return [ + external_stream.topic(name).subscribe().wait_id + for name in ("tokens", "tool-events", "tokens") + ] + + assert run_once() == run_once() == [1, 2, 3] + + +def test_two_subscriptions_to_one_stream_get_distinct_wait_ids( + runtime: FakeRuntime, +) -> None: + """They are two independent waits, each with its own cursor and park intent. + + Delivery is broadcast, so each sees every record from its own position -- + which is only expressible if they are numbered apart. + """ + topic = external_stream.topic("tokens") + + first = topic.subscribe() + second = topic.subscribe() + + assert first.wait_id != second.wait_id + assert first.stream_key == second.stream_key + assert [wait_id for wait_id, _ in runtime.registrations] == [1, 2] + + +def test_a_second_run_restarts_the_counter(monkeypatch: pytest.MonkeyPatch) -> None: + """Per-Run state lives on the instance, so an evicted Run takes it with it. + + A module global would outlive the Run and hand its wait ids to the next one. + """ + first_instance = FakeInstance() + monkeypatch.setattr(temporalio.workflow, "instance", lambda: first_instance) + _install_runtime(first_instance, FakeRuntime()) + external_stream.topic("tokens").subscribe() + + second_instance = FakeInstance() + monkeypatch.setattr(temporalio.workflow, "instance", lambda: second_instance) + _install_runtime(second_instance, FakeRuntime()) + + assert external_stream.topic("tokens").subscribe().wait_id == 1 + + +# --- iteration ---------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_iteration_yields_decoded_values_from_the_buffer( + runtime: FakeRuntime, +) -> None: + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + subscription = external_stream.topic("tokens", type=str).subscribe() + runtime.buffers[subscription.wait_id] = [ + StreamRecord(RecordKind.DATA, await codec.encode(v), "s", i) + for i, v in enumerate(["a", "b"]) + ] + + seen = [] + async for value in subscription: + seen.append(value) + if len(seen) == 2: + break + + assert seen == ["a", "b"] + + +@pytest.mark.asyncio +async def test_control_records_are_never_yielded(runtime: FakeRuntime) -> None: + """A fence advances the cursor but is the runtime's, not the Workflow's.""" + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + subscription = external_stream.topic("tokens", type=str).subscribe() + runtime.buffers[subscription.wait_id] = [ + StreamRecord(RecordKind.DATA, await codec.encode("a"), "s", 0), + StreamRecord(RecordKind.WRITE_FENCE, b"", "s", 1), + StreamRecord(RecordKind.DATA, await codec.encode("b"), "s", 2), + ] + + seen = [] + async for value in subscription: + seen.append(value) + if len(seen) == 2: + break + + assert seen == ["a", "b"] + + +@pytest.mark.asyncio +async def test_an_empty_buffer_blocks_on_a_readiness_future( + runtime: FakeRuntime, # pyright: ignore[reportUnusedParameter] +) -> None: + """Iteration never polls the backend; only Core can say when to look again.""" + subscription = external_stream.topic("tokens", type=str).subscribe() + + iterator = subscription.__aiter__() + pending = asyncio.ensure_future(iterator.__anext__()) + await asyncio.sleep(0.05) + + assert not pending.done(), "iteration must block rather than spin on the backend" + pending.cancel() + try: + await pending + except asyncio.CancelledError: + pass + + +def test_closing_a_subscription_unsubscribes_it(runtime: FakeRuntime) -> None: + """The Workflow side going quiet is only half of closing. + + Everything the Workflow itself can observe is already right without this + call -- iteration ends, the wait leaves the blocked set, undelivered records + stay unconsumed -- so nothing on this side of the boundary would notice it + missing. What is left behind is on the Worker: a watcher still prefetching + into a buffer nobody will drain, and a park intent still in the backend for + every producer that asks. + """ + subscription = external_stream.topic("tokens", type=str).subscribe() + + subscription.close() + + assert runtime.unsubscribed == [subscription.wait_id], ( + "closing left the Worker serving a wait no Workflow code is reading" + ) + # Idempotent, like the rest of `close`: the ordinary shape is a `finally` + # that cannot know whether the iterator already ended, and by a second call + # the Worker has already dropped this wait. + subscription.close() + assert runtime.unsubscribed == [subscription.wait_id] + + +# --- one subscription, one consumer ------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_second_coroutine_waiting_on_one_subscription_is_refused( + runtime: FakeRuntime, +) -> None: + """Sharing a single-slot wait between two waiters strands one of them forever. + + ``__aiter__`` returns a new generator every time, but everything a blocked + wait is found through is the subscription's: one ``_pending_future`` and one + entry in the runtime's pending map, keyed by wait id. A second waiter + therefore *replaced* the first in both places -- readiness resolved only the + newer one, its ``finally`` removed the map entry, and the older future became + unreachable by the readiness activation and by ``close()`` alike. Observed as + a coroutine still pending after a record had been buffered, readiness + resolved, and the subscription closed: permanently stuck, with the shared + blocked flag saying the wait was not even blocked. + """ + subscription = external_stream.topic("tokens", type=str).subscribe() + + first = asyncio.ensure_future(subscription.__aiter__().__anext__()) + await asyncio.sleep(0.05) + assert subscription.wait_id in runtime.pending, ( + "the first consumer has to be registered and blocked before a second one " + "can overwrite it" + ) + + with pytest.raises(ConcurrentStreamConsumerError, match="single consumer"): + await subscription.__aiter__().__anext__() + + # The refusal changed nothing about the consumer that was already there. + assert runtime.pending[subscription.wait_id] is subscription._pending_future + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + runtime.buffers[subscription.wait_id] = [ + StreamRecord(RecordKind.DATA, await codec.encode("a"), "s", 0) + ] + runtime.pending[subscription.wait_id].set_result(None) + + assert await asyncio.wait_for(first, 1) == "a" + assert not runtime.pending, "the resolved wait was left registered" + assert subscription._pending_future is None + + +@pytest.mark.asyncio +async def test_iterating_again_after_the_first_consumer_stopped_is_allowed( + runtime: FakeRuntime, +) -> None: + """The shape the refusal must not catch, which is why it is not on the iterator. + + Taking a few records, doing something else, and coming back to the same + subscription is ordinary code, and a ``break`` leaves the generator *suspended* + rather than closed -- so a guard that claimed the subscription for an iterator + could not tell that shape from two live consumers, and would refuse it. + """ + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + subscription = external_stream.topic("tokens", type=str).subscribe() + runtime.buffers[subscription.wait_id] = [ + StreamRecord(RecordKind.DATA, await codec.encode(value), "s", i) + for i, value in enumerate(["a", "b"]) + ] + + async for value in subscription: + assert value == "a" + break + + # The abandoned generator is still suspended at its `yield`; nothing closed + # it. A second pass must still work, and must resume rather than restart. + async for value in subscription: + assert value == "b" + break + + +@pytest.mark.asyncio +async def test_a_merge_cannot_take_a_wait_another_consumer_is_blocked_on( + runtime: FakeRuntime, +) -> None: + """``merge()`` registers the same single slot, so it is the same defect. + + And it fails *before* registering anything, because a merge that refused + half-way would leave its earlier waits blocked with no coroutine behind them + -- which is the state that asks Core to retain a Workflow Task for nobody. + """ + first = external_stream.topic("a", type=str).subscribe() + second = external_stream.topic("b", type=str).subscribe() + + solo = asyncio.ensure_future(second.__aiter__().__anext__()) + await asyncio.sleep(0.05) + assert set(runtime.pending) == {second.wait_id} + + async def consume_merged() -> None: + async for _ in merge(first, second): + pass + + with pytest.raises(ConcurrentStreamConsumerError): + await consume_merged() + + assert set(runtime.pending) == {second.wait_id}, ( + "the refused merge left a wait registered, or took the one that was " + "already blocked" + ) + assert first._pending_future is None + solo.cancel() + try: + await solo + except asyncio.CancelledError: + pass + + +# --- names -------------------------------------------------------------------- + + +def test_no_exported_name_collides_with_the_shipped_contrib_feature() -> None: + """The two coexist and are mirror images; nothing here may shadow that one.""" + import temporalio.contrib.external_workflow_streams as package + + for module in (package, _api): + for name in dir(module): + assert not name.startswith("__temporal_workflow_stream"), ( + f"{module.__name__}.{name} collides with the reserved namespace of " + "temporalio.contrib.workflow_streams" + ) + + +def test_the_reserved_instance_attribute_is_in_this_features_namespace() -> None: + assert _api._RUN_STATE_ATTR.startswith("__temporal_external_stream") + assert not _api._RUN_STATE_ATTR.startswith("__temporal_workflow_stream") + + +def test_the_handle_types_are_named_as_the_design_says() -> None: + assert ExternalStreamTopic.__name__ == "ExternalStreamTopic" + assert not hasattr(ExternalStreamTopic, "publish") + assert hasattr(ExternalStreamTopic, "subscribe") + + +def test_the_package_exports_both_directions_without_merging_their_roles() -> None: + import temporalio.contrib.external_workflow_streams as package + + assert { + "ExternalOutputStreamClient", + "ExternalOutputStreamProducer", + "ExternalOutputStreamTopic", + "OutputStreamBackend", + "StreamDirection", + "external_output_stream", + } <= set(package.__all__) + assert not hasattr(package.ExternalOutputStreamTopic, "subscribe") + assert hasattr(package.ExternalOutputStreamTopic, "publish") + + +@pytest.mark.asyncio +async def test_a_record_buffered_while_not_iterating_is_still_delivered( + runtime: FakeRuntime, +) -> None: + """The readiness for it was reported to nobody, and none is coming. + + A record that arrives while Workflow code is doing something else -- a + timer, an activity, another stream -- is buffered and its readiness is + reported and consumed immediately. By the time the Workflow comes back and + asks for the next value, no further notification will ever be sent: the + watcher has moved its prefetch cursor past that record. Blocking without + looking at the buffer first strands the Workflow forever on a record that is + already sitting in front of it. + """ + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + subscription = external_stream.topic("tokens", type=str).subscribe() + runtime.buffers[subscription.wait_id] = [ + StreamRecord(RecordKind.DATA, await codec.encode("first"), "s", 0) + ] + + iterator = subscription.__aiter__() + assert await iterator.__anext__() == "first" + + # Arrives while the Workflow is off doing something else. Nothing resolves a + # readiness future for it, because nothing was waiting when it landed. + runtime.buffers[subscription.wait_id] = [ + StreamRecord(RecordKind.DATA, await codec.encode("second"), "s", 1) + ] + + assert await asyncio.wait_for(iterator.__anext__(), 1) == "second", ( + "the iterator blocked on a readiness notification that had already been " + "spent, with the record buffered in front of it" + ) + + +@pytest.mark.asyncio +async def test_a_record_buffered_after_blocking_begins_still_resolves( + runtime: FakeRuntime, +) -> None: + """The other side of the window, which must keep working. + + Looking at the buffer before blocking must not replace the readiness path: + a record that genuinely arrives later has no buffered copy to find, and the + future is what delivers it. + """ + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + subscription = external_stream.topic("tokens", type=str).subscribe() + + iterator = subscription.__aiter__() + pending = asyncio.ensure_future(iterator.__anext__()) + await asyncio.sleep(0.05) + assert not pending.done() + + runtime.buffers[subscription.wait_id] = [ + StreamRecord(RecordKind.DATA, await codec.encode("late"), "s", 0) + ] + runtime.pending[subscription.wait_id].set_result(None) + + assert await asyncio.wait_for(pending, 1) == "late" + + +# --- consumption is committed only once a value exists ------------------------ + + +class FailingCodec: + """A converter mismatch: the stream is fine, the configuration is not. + + Fails in :meth:`convert`, the half that runs on the Workflow thread, since + that is where a type hint meets a payload the producer did not write. + """ + + def __init__(self, value_type: type | None = str) -> None: + self._inner: StreamPayloadCodec[Any] = StreamPayloadCodec( + temporalio.converter.DataConverter.default, value_type + ) + self.calls = 0 + + def parse_unprepared(self, payload: bytes) -> Any: + return self._inner.parse_unprepared(payload) + + def convert( + self, + prepared: Any, # pyright: ignore[reportUnusedParameter] + ) -> Any: + self.calls += 1 + raise RuntimeError("this converter cannot read this payload") + + +@pytest.mark.asyncio +async def test_decoding_on_the_workflow_thread_cannot_suspend( + runtime: FakeRuntime, +) -> None: + """The Workflow thread converts; it does not await a converter. + + ``DataConverter.decode`` is retrieval, then a user ``PayloadCodec``, then + ``from_payloads``. The first two are arbitrary asynchronous work -- a + network fetch, a KMS round trip -- and the Worker awaits them on its own + loop before a record is buffered, exactly as it does for every other payload + an activation carries. What is left here is the third, which needs the + topic's type and performs no I/O. + + Asserted structurally as well as behaviourally, because the property is + structural: a coroutine on this path is a suspension point inside + ``activate()``, and a suspension point is where I/O, a synthesized Workflow + command, or a deadlock timeout can appear. + """ + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + subscription = external_stream.topic("tokens", type=str).subscribe() + runtime.buffers[subscription.wait_id] = [ + StreamRecord(RecordKind.DATA, await codec.encode("a"), "s", 0) + ] + + assert not inspect.iscoroutinefunction(ExternalStreamSubscription._decode), ( + "decoding a record awaits on the Workflow thread, so a payload codec " + "or an external payload fetch runs inside activate()" + ) + + record = runtime.buffers[subscription.wait_id][0] + subscription._fill() + assert subscription._decode(record) == "a" + + +@pytest.mark.asyncio +async def test_a_cancelled_delivery_leaves_the_record_unconsumed( + runtime: FakeRuntime, +) -> None: + """Consumption is a claim about what Workflow code received. + + Decoding no longer suspends, so the one place a cancellation can still land + between a record arriving and Workflow code receiving it is the readiness + wait. A cancellation there -- a stream raced against a timer, most often -- + must leave the record exactly where a later pass finds it: recording + consumption for a value nothing yielded makes the claim false in the case + that matters, because the consumption cursor is what a Continue-As-New + successor resumes from. + """ + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + subscription = external_stream.topic("tokens", type=str).subscribe() + + iterator = subscription.__aiter__() + pending = asyncio.ensure_future(iterator.__anext__()) + # Blocked: nothing is buffered yet, and only Core can say when to look. + await asyncio.sleep(0.05) + assert not pending.done() + + # Buffered while the wait is outstanding, then cancelled before the + # readiness that would have delivered it. + runtime.buffers[subscription.wait_id] = [ + StreamRecord(RecordKind.DATA, await codec.encode("a"), "s", 0) + ] + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + + assert runtime.consumed == [], ( + "a cancelled wait recorded a consumption, so the cursor claims the " + "Workflow received a record it never saw" + ) + + # And it is still there to be taken. + assert await asyncio.wait_for(subscription.__aiter__().__anext__(), 1) == "a" + assert [wait_id for wait_id, _ in runtime.consumed] == [subscription.wait_id] + + +@pytest.mark.asyncio +async def test_a_failed_decode_leaves_the_record_unconsumed( + runtime: FakeRuntime, +) -> None: + """The same ordering, reached by the failure a converter mismatch produces. + + The budget is spent on the record either way; the question is whether the + Run also records having handed it over. + """ + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + runtime.codec = FailingCodec() + subscription = external_stream.topic("tokens", type=str).subscribe() + runtime.buffers[subscription.wait_id] = [ + StreamRecord(RecordKind.DATA, await codec.encode("a"), "s", 0) + ] + + with pytest.raises(StreamDecodeError, match="could not be decoded"): + await subscription.__aiter__().__anext__() + + assert runtime.consumed == [], ( + "a record whose decode raised was marked consumed, so the value is lost " + "and the cursor claims the Workflow received it" + ) + + +@pytest.mark.asyncio +async def test_a_preparation_failure_is_raised_where_the_record_would_arrive( + runtime: FakeRuntime, +) -> None: + """A codec that fails on the Worker's loop still fails *this* Workflow. + + Preparation happens in the watcher, which has no Workflow Task to fail and + no Workflow to tell -- and which may be preparing a record the Workflow + never asks for. So the failure travels with the record and is raised by the + delivery that would have yielded its value, as the taxonomy's third row: + the bytes are intact and the consumer's converter cannot read them. + + Raised, and *not* consumed: nothing was received, so the cursor must not say + otherwise. + """ + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + subscription = external_stream.topic("tokens", type=str).subscribe() + runtime.buffers[subscription.wait_id] = [ + PreparedRecord.of( + StreamRecord(RecordKind.DATA, await codec.encode("a"), "s", 0), + None, + RuntimeError("the codec could not decrypt this payload"), + ) + ] + + with pytest.raises(StreamDecodeError, match="could not be decoded") as caught: + await subscription.__aiter__().__anext__() + assert isinstance(caught.value.__cause__, RuntimeError) + + assert runtime.consumed == [], ( + "a record whose preparation failed was marked consumed, so the cursor " + "claims the Workflow received a value that never existed" + ) + + +@pytest.mark.asyncio +async def test_an_unprepared_record_is_refused_rather_than_decoded_late( + runtime: FakeRuntime, +) -> None: + """A codec-bearing converter has an asynchronous half that must have run. + + If a record reaches the Workflow thread without it, running it here is the + defect this split exists to remove, and converting the raw bytes anyway + would hand Workflow code whatever the codec's output happens to look like. + Refused instead -- and refused as a decode failure, because that is what the + Workflow can be told. + """ + + class NeverDecodes(temporalio.converter.PayloadCodec): + async def encode(self, payloads: Any) -> Any: + return list(payloads) + + async def decode(self, payloads: Any) -> Any: + raise AssertionError("the Workflow thread ran a payload codec") + + plain = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + runtime.codec = StreamPayloadCodec( + temporalio.converter.DataConverter(payload_codec=NeverDecodes()), str + ) + subscription = external_stream.topic("tokens", type=str).subscribe() + runtime.buffers[subscription.wait_id] = [ + StreamRecord(RecordKind.DATA, await plain.encode("a"), "s", 0) + ] + + with pytest.raises(StreamDecodeError, match="could not be decoded"): + await subscription.__aiter__().__anext__() + assert runtime.consumed == [] diff --git a/tests/contrib/external_workflow_streams/test_backend_conformance.py b/tests/contrib/external_workflow_streams/test_backend_conformance.py new file mode 100644 index 0000000000..9adf60d3aa --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_backend_conformance.py @@ -0,0 +1,318 @@ +"""P2 — the conformance suite, and proof that it fails a broken backend. + +The first half runs every check against the conforming reference backend. The +second half runs a *single* check against a stub broken in exactly one way and +asserts it fails with the message naming that obligation. A suite that only +ever passes is not evidence of anything. +""" + +from __future__ import annotations + +import uuid +from collections.abc import AsyncGenerator +from datetime import timedelta +from typing import cast + +import pytest + +from temporalio.contrib.external_workflow_streams._backend import ( + StreamBackend, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + Cursor, + Offset, + RecordKind, + StreamRecord, +) +from tests.contrib.external_workflow_streams import conformance +from tests.contrib.external_workflow_streams.conformance import ( + CONFORMANCE_CHECKS, + Check, +) +from tests.contrib.external_workflow_streams.memory_backend import ( + MemoryStreamBackend, + parse, +) + + +@pytest.fixture +def stream_key() -> StreamKey: + return StreamKey("ns", "wf", uuid.uuid4().hex, "tokens") + + +@pytest.fixture +def backend() -> MemoryStreamBackend: + return MemoryStreamBackend() + + +# --- the reference backend conforms ----------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("check", CONFORMANCE_CHECKS, ids=lambda c: c.__name__) +async def test_reference_backend_conforms( + check: Check, backend: MemoryStreamBackend, stream_key: StreamKey +) -> None: + await check(backend, stream_key) + + +def test_the_suite_is_not_silently_empty() -> None: + assert len(CONFORMANCE_CHECKS) >= 16 + assert len({c.__name__ for c in CONFORMANCE_CHECKS}) == len(CONFORMANCE_CHECKS) + + +# --- deliberately broken backends ------------------------------------------- + + +class ExclusiveRangeBackend(MemoryStreamBackend): + """Implements the inclusive range read with exclusive-of-first semantics. + + The realistic mistake: reaching for the same primitive the watch uses. + """ + + async def read_range( + self, key: StreamKey, first: Offset, last: Offset + ) -> list[StreamRecord]: + got = await super().read_range(key, first, last) + return [r for r in got if r.offset != first] + + +class LexicalOffsetBackend(MemoryStreamBackend): + """Compares offsets as strings.""" + + def compare_offsets(self, left: Offset, right: Offset) -> int: + return (left.token > right.token) - (left.token < right.token) + + +class NameableCursorBackend(MemoryStreamBackend): + """Validates the resume token against a record that must follow it. + + Stands for a provider whose cursor is a record identity rather than a + boundary: a position with nothing after it is not a position it can + express, so a consumer that has drained to the tail has nothing legal to + hold. This is the state parking puts every subscription in, which is why + the contract forbids it. + """ + + async def read_after( + self, + key: StreamKey, + after: Cursor, + *, + max_records: int, + block: timedelta | None = None, + ) -> list[StreamRecord]: + if not after.is_beginning: + bound = parse(after.offset) # type: ignore[arg-type] + if not any(parse(r.offset) > bound for r in self.all_records(key)): # type: ignore[arg-type] + raise ValueError( + f"cursor {after.offset} names no position: nothing follows it" + ) + return await super().read_after( + key, after, max_records=max_records, block=block + ) + + +class KeyOnlyIdempotencyBackend(MemoryStreamBackend): + """Idempotent on the key alone, ignoring whether the bytes match.""" + + async def append(self, key: StreamKey, record: StreamRecord) -> StreamRecord: + existing = self._by_key.get((key, record.idempotency_key)) + if existing is not None: + return existing + return await super().append(key, record) + + +class InclusiveWatchBackend(MemoryStreamBackend): + """Re-delivers the record its boundary names.""" + + async def read_after( + self, + key: StreamKey, + after: Cursor, + *, + max_records: int, + block: timedelta | None = None, + ) -> list[StreamRecord]: + if after.is_beginning: + return await super().read_after( + key, after, max_records=max_records, block=block + ) + bound = parse(after.offset) # type: ignore[arg-type] + return [r for r in self.all_records(key) if parse(r.offset) >= bound][ # type: ignore[arg-type] + :max_records + ] + + +class CountingRangeBackend(MemoryStreamBackend): + """Reads "however many records fit between the endpoints, starting at first". + + The realistic mistake: treating ``last`` as a length hint rather than as a + closing endpoint, so a gap is silently filled from beyond the range and a + trimmed stream reads back with the right count. + """ + + async def read_range( + self, key: StreamKey, first: Offset, last: Offset + ) -> list[StreamRecord]: + want = len( + [ + r + for r in self.all_records(key) + if parse(first) <= parse(r.offset) <= parse(last) # type: ignore[arg-type] + ] + ) + # Whatever was asked for, deliver `want` records starting at `first` -- + # reaching past `last` when the interval has holes in it. + from_first = [ + r + for r in self.all_records(key) + if parse(r.offset) >= parse(first) # type: ignore[arg-type] + ] + return from_first[: max(want, 3)] + + +class UndeclaredImmutabilityBackend(MemoryStreamBackend): + """Never considered the question.""" + + guarantees_immutability = None + + +class DeniedImmutabilityBackend(MemoryStreamBackend): + """Considered it and cannot make the guarantee.""" + + guarantees_immutability = False + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("broken", "check", "reason"), + [ + pytest.param( + ExclusiveRangeBackend, + conformance.check_range_read_includes_both_endpoints, + "inclusive of both endpoints", + id="exclusive-range-read", + ), + pytest.param( + LexicalOffsetBackend, + conformance.check_offsets_are_not_compared_lexically, + "lexically", + id="lexical-offset-comparison", + ), + pytest.param( + InclusiveWatchBackend, + conformance.check_watch_is_exclusive_of_its_boundary, + "strictly after", + id="inclusive-watch", + ), + pytest.param( + KeyOnlyIdempotencyBackend, + conformance.check_reappend_with_different_bytes_is_rejected, + "must raise AppendConflictError", + id="idempotent-on-key-alone", + ), + pytest.param( + CountingRangeBackend, + conformance.check_range_read_reports_a_missing_record_as_missing, + "rather than substituting another", + id="substitutes-for-a-deleted-record", + ), + pytest.param( + UndeclaredImmutabilityBackend, + conformance.check_immutability_is_declared, + "must declare guarantees_immutability", + id="immutability-undeclared", + ), + pytest.param( + DeniedImmutabilityBackend, + conformance.check_immutability_is_declared, + "must declare guarantees_immutability", + id="immutability-denied", + ), + ], +) +async def test_a_broken_backend_fails_for_the_right_reason( + broken: type[StreamBackend], check: Check, reason: str, stream_key: StreamKey +) -> None: + with pytest.raises(AssertionError, match=reason): + await check(broken(), stream_key) + + +@pytest.mark.asyncio +async def test_a_backend_needing_a_nameable_cursor_fails_the_tail_check( + stream_key: StreamKey, +) -> None: + """Fails, though not with an assertion -- it cannot express the state at all. + + Worth its own case: the failure mode is "the provider refuses the cursor", + not "the provider returned the wrong records", and a suite that only caught + AssertionError would let it through. + """ + with pytest.raises((AssertionError, ValueError)): + await conformance.check_a_consumer_parked_at_the_tail_resumes( + NameableCursorBackend(), stream_key + ) + + +# --- the derived helpers on the ABC ----------------------------------------- + + +@pytest.mark.asyncio +async def test_watch_iterates_across_batches( + backend: MemoryStreamBackend, stream_key: StreamKey +) -> None: + """The cursor advances past what was delivered, not to the current tail.""" + for i, payload in enumerate([b"a", b"b", b"c"]): + await backend.append(stream_key, StreamRecord(RecordKind.DATA, payload, "s", i)) + + seen = [] + watcher = cast( + AsyncGenerator[StreamRecord, None], + backend.watch( + stream_key, Cursor(), max_records=1, block=timedelta(milliseconds=10) + ), + ) + async for record in watcher: + seen.append(record.payload) + if len(seen) == 3: + break + await watcher.aclose() + + assert seen == [b"a", b"b", b"c"] + + +@pytest.mark.asyncio +async def test_strictly_increasing_uses_the_provider_comparator( + backend: MemoryStreamBackend, +) -> None: + assert backend.strictly_increasing([Offset("9-0"), Offset("10-0")]) + assert not backend.strictly_increasing([Offset("10-0"), Offset("9-0")]) + assert not backend.strictly_increasing([Offset("9-0"), Offset("9-0")]) + + +@pytest.mark.asyncio +async def test_read_after_blocks_until_a_record_arrives( + backend: MemoryStreamBackend, stream_key: StreamKey +) -> None: + """Blocking is the provider's, so latency never reaches the Workflow thread.""" + import asyncio + + placed = await backend.append( + stream_key, StreamRecord(RecordKind.DATA, b"a", "s", 0) + ) + assert placed.offset is not None + + async def append_soon() -> None: + await asyncio.sleep(0.05) + await backend.append(stream_key, StreamRecord(RecordKind.DATA, b"b", "s", 1)) + + task = asyncio.ensure_future(append_soon()) + got = await backend.read_after( + stream_key, AFTER(placed.offset), max_records=10, block=timedelta(seconds=2) + ) + await task + + assert [r.payload for r in got] == [b"b"] diff --git a/tests/contrib/external_workflow_streams/test_bridge.py b/tests/contrib/external_workflow_streams/test_bridge.py new file mode 100644 index 0000000000..a383741d5e --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_bridge.py @@ -0,0 +1,221 @@ +"""P7 — the readiness call and the run-status probe, reachable from Python. + +PYTEST_DONT_REWRITE: sandboxed fixture Workflows re-import this module, so pytest's +injected imports would make sandbox validation depend on pytest's import locks. +""" + +from __future__ import annotations + +import asyncio +import threading +import uuid + +import pytest + +from temporalio import workflow +from temporalio.bridge.worker import ( + ExternalStreamReadyResult, + ExternalStreamRunStatus, +) +from temporalio.client import Client +from temporalio.worker import Worker + + +@workflow.defn +class IdleWorkflow: + @workflow.run + async def run(self) -> None: + await workflow.wait_condition(lambda: False) + + +def bridge_worker(worker: Worker): # type: ignore[no-untyped-def] + """The bridge worker underneath a Python Worker.""" + return worker._workflow_worker._bridge_worker() # type: ignore[union-attr] + + +async def _wait_for_first_task_completed(handle, timeout: float = 30) -> None: # type: ignore[no-untyped-def] + """Waits until this Worker has run the Workflow's first task and cached the Run. + + Polled rather than slept: how long a Worker takes to poll, execute, and + complete one Workflow Task is a property of the machine, and a fixed wait is a + guess about it that a loaded machine invalidates. + """ + deadline = asyncio.get_running_loop().time() + timeout + while asyncio.get_running_loop().time() < deadline: + events = [event async for event in handle.fetch_history_events()] + if any(e.HasField("workflow_task_completed_event_attributes") for e in events): + return + await asyncio.sleep(0.05) + raise AssertionError( + "the Workflow's first task never completed, so the Run was never cached " + "and the probe has no settled state to be repeatable about" + ) + + +@pytest.fixture +def unknown_run_id() -> str: + return f"no-such-run-{uuid.uuid4()}" + + +async def test_readiness_for_an_unknown_run_returns_rather_than_raising( + client: Client, unknown_run_id: str +) -> None: + """A watcher must be able to *handle* a missing Run, not catch an exception. + + ``RunNotFound`` is a normal, expected answer -- it is what a watcher sees + every time its Run is evicted -- so surfacing it as an exception would put + ordinary lifecycle in the error path. + """ + async with Worker( + client, task_queue=f"tq-{uuid.uuid4()}", workflows=[IdleWorkflow] + ) as worker: + bridge = bridge_worker(worker) + + result = await bridge.notify_external_stream_ready(unknown_run_id, 1, 0) + + assert result is ExternalStreamReadyResult.RUN_NOT_FOUND + assert result.needs_wake_signal + + +async def test_status_for_an_unknown_run_returns_rather_than_raising( + client: Client, unknown_run_id: str +) -> None: + async with Worker( + client, task_queue=f"tq-{uuid.uuid4()}", workflows=[IdleWorkflow] + ) as worker: + status = await bridge_worker(worker).external_stream_run_status(unknown_run_id) + + assert status is ExternalStreamRunStatus.RUN_NOT_FOUND + + +async def test_both_calls_surface_their_full_result_enums( + client: Client, # pyright: ignore[reportUnusedParameter] +) -> None: + """The five and four values reach Python as values, not strings. + + A watcher branches on all five, and a stringly-typed result would let a + typo silently take the "do nothing" branch for a Run that needed a Signal. + """ + assert {r.value for r in ExternalStreamReadyResult} == { + "Accepted", + "Stale", + "Parked", + "NoOpenWorkflowTask", + "RunNotFound", + } + assert {s.value for s in ExternalStreamRunStatus} == { + "WftOpen", + "Parked", + "NoOpenWorkflowTask", + "RunNotFound", + } + + +async def test_the_readiness_call_is_safe_from_several_threads( + client: Client, unknown_run_id: str +) -> None: + """One watcher per subscription means concurrent callers by construction. + + Each thread drives its own event loop, so this exercises the call from + genuinely separate threads rather than from concurrent tasks on one. The + joins run off the main loop for the reason + :py:meth:`notify_external_stream_ready_sync` documents: blocking the loop + that owns the Worker would stop it polling, and the answer comes from a + lane that Worker drives. + """ + async with Worker( + client, task_queue=f"tq-{uuid.uuid4()}", workflows=[IdleWorkflow] + ) as worker: + bridge = bridge_worker(worker) + # One call first, so the Worker is provably polling before anything + # blocks -- otherwise this test could pass or hang on timing alone. + assert ( + await bridge.notify_external_stream_ready(unknown_run_id, 1, 0) + ) is ExternalStreamReadyResult.RUN_NOT_FOUND + + results: list[object] = [] + errors: list[BaseException] = [] + lock = threading.Lock() + + def call_from_own_loop(wait_id: int) -> None: + async def go() -> None: + for _ in range(5): + got = await bridge.notify_external_stream_ready( + unknown_run_id, wait_id, 0 + ) + with lock: + results.append(got) + + try: + asyncio.run(go()) + except BaseException as err: # noqa: BLE001 -- recorded and asserted on + with lock: + errors.append(err) + + def run_all_threads() -> None: + threads = [ + threading.Thread(target=call_from_own_loop, args=(wait_id,)) + for wait_id in range(1, 9) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=20) + with lock: + errors.extend( + RuntimeError(f"thread {i} did not finish") + for i, thread in enumerate(threads) + if thread.is_alive() + ) + + await asyncio.to_thread(run_all_threads) + + assert not errors, f"concurrent readiness calls failed: {errors}" + assert len(results) == 8 * 5 + assert all(r is ExternalStreamReadyResult.RUN_NOT_FOUND for r in results) + + +async def test_the_status_probe_is_repeatable(client: Client) -> None: + """Asking must change nothing, however many times it is asked. + + The Run has to be **settled** before the first question, or the test is about + something else. `start_workflow` returns when the server accepts the Workflow, + which is before this Worker has polled it, run its first Workflow Task, and + cached the Run -- and until that has happened the probe honestly answers + `RunNotFound`. Answers of `RunNotFound` then `NoOpenWorkflowTask` are Core + telling the truth twice about two different states, not a probe that changed + something. + + So the first Workflow Task's completion is awaited first. `IdleWorkflow` blocks + forever, so after it the Run is cached with no task open, and that state does + not move again on its own -- which is what makes "the same answer every time" a + property of the probe rather than of the timing. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, task_queue=task_queue, workflows=[IdleWorkflow] + ) as worker: + handle = await client.start_workflow( + IdleWorkflow.run, id=f"wf-{uuid.uuid4()}", task_queue=task_queue + ) + try: + description = await handle.describe() + bridge = bridge_worker(worker) + + await _wait_for_first_task_completed(handle) + settled = await bridge.external_stream_run_status(description.run_id) + assert settled is ExternalStreamRunStatus.NO_OPEN_WORKFLOW_TASK, ( + "the Run is cached and blocked with no Workflow Task open, so this " + f"is the only honest answer; got {settled}" + ) + + answers = [ + await bridge.external_stream_run_status(description.run_id) + for _ in range(5) + ] + + # No subscriptions exist, so the honest answer is the same one every + # time -- and in particular the probe never manufactured a task. + assert answers == [settled] * 5 + finally: + await handle.terminate() diff --git a/tests/contrib/external_workflow_streams/test_codec.py b/tests/contrib/external_workflow_streams/test_codec.py new file mode 100644 index 0000000000..4b25393673 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_codec.py @@ -0,0 +1,119 @@ +"""P4 — payload encoding through the Workflow's DataConverter.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import pytest + +import temporalio.api.common.v1 +import temporalio.converter +from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec +from temporalio.contrib.external_workflow_streams._record import ( + Offset, + RecordKind, + StreamRecord, +) + + +@dataclass +class Token: + """A custom type, to prove the topic's declared type reaches the converter.""" + + text: str + index: int + + +DEFAULT = temporalio.converter.DataConverter.default + + +@pytest.mark.asyncio +async def test_str_round_trips() -> None: + codec = StreamPayloadCodec(DEFAULT, str) + + assert await codec.decode(await codec.encode("hello")) == "hello" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", ["", "unicode: é中", "a" * 10_000]) +async def test_awkward_strings_round_trip(value: str) -> None: + codec = StreamPayloadCodec(DEFAULT, str) + + assert await codec.decode(await codec.encode(value)) == value + + +@pytest.mark.asyncio +async def test_a_custom_type_round_trips_as_that_type() -> None: + codec = StreamPayloadCodec(DEFAULT, Token) + + decoded = await codec.decode(await codec.encode(Token("hi", 3))) + + # Not merely equal-looking: the declared type is what makes this a `Token` + # rather than the `dict` the default converter would otherwise infer. + assert isinstance(decoded, Token) + assert decoded == Token("hi", 3) + + +@pytest.mark.asyncio +async def test_without_a_declared_type_the_converter_infers() -> None: + untyped: StreamPayloadCodec[Any] = StreamPayloadCodec(DEFAULT, None) + + decoded = await untyped.decode(await untyped.encode(Token("hi", 3))) + + assert decoded == {"text": "hi", "index": 3} + + +@pytest.mark.asyncio +async def test_encoded_bytes_are_a_payload_envelope() -> None: + """The envelope is what lets a consumer decode without out-of-band metadata.""" + codec = StreamPayloadCodec(DEFAULT, str) + + payload = temporalio.api.common.v1.Payload() + payload.ParseFromString(await codec.encode("hello")) + + assert payload.metadata["encoding"] == b"json/plain" + + +@pytest.mark.asyncio +async def test_with_type_rebinds_without_touching_the_converter() -> None: + codec = StreamPayloadCodec(DEFAULT, str) + rebound = codec.with_type(Token) + + assert rebound.data_converter is codec.data_converter + assert rebound.value_type is Token + assert codec.value_type is str + + +@pytest.mark.asyncio +async def test_a_payload_encoded_for_one_type_decodes_under_another() -> None: + """A converter mismatch surfaces here, on the consumer, not as damage. + + The bytes are intact; only the reader's expectation is wrong. Classifying + that as a decode failure rather than integrity loss is what keeps an + operator from being sent to repair an undamaged backend. + """ + produced = await StreamPayloadCodec(DEFAULT, str).encode("not a token") + + with pytest.raises(Exception): + await StreamPayloadCodec(DEFAULT, Token).decode(produced) + + +@pytest.mark.asyncio +async def test_undecodable_bytes_raise_rather_than_returning_a_value() -> None: + with pytest.raises(Exception): + await StreamPayloadCodec(DEFAULT, str).decode(b"\xff\xfe not a payload") + + +@pytest.mark.asyncio +async def test_a_record_carries_the_encoded_payload_unchanged() -> None: + """The codec's output survives the record's own field round trip.""" + codec = StreamPayloadCodec(DEFAULT, Token) + encoded = await codec.encode(Token("round", 1)) + + record = StreamRecord(RecordKind.DATA, encoded, "session-a", 0).placed_at( + Offset("1-0") + ) + rebuilt = StreamRecord.from_fields(Offset("1-0"), record.to_fields()) + + assert await codec.decode(rebuilt.payload) == Token("round", 1) diff --git a/tests/contrib/external_workflow_streams/test_continuation.py b/tests/contrib/external_workflow_streams/test_continuation.py new file mode 100644 index 0000000000..f3dba0b607 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_continuation.py @@ -0,0 +1,969 @@ +"""P15 — the Continue-As-New cursor. + +A stream spans a whole chain, so a new Run must resume where its predecessor +stopped. The position travels in a reserved internal header and is restored from +History, because a cursor derived from mutable backend state would give replay +whatever the stream holds *now* rather than what the Run started from -- and two +replays of one history could then diverge (ADR-022). + +PYTEST_DONT_REWRITE: sandboxed fixture Workflows re-import this module, so pytest's +injected imports would make sandbox validation depend on pytest's import locks. +""" + +# pyright: reportMissingParameterType=false, reportUnusedFunction=false +from __future__ import annotations + +import asyncio +import dataclasses +import uuid +from datetime import timedelta +from typing import Any + +import pytest + +import temporalio.converter +from temporalio import workflow +from temporalio.bridge.proto.workflow_activation import ( + InitializeWorkflow, + WorkflowActivation, +) +from temporalio.client import Client, WorkflowHandle +from temporalio.contrib.external_workflow_streams._annotation import ( + AnnotationDecodeError, +) +from temporalio.contrib.external_workflow_streams._backend import StreamKey +from temporalio.contrib.external_workflow_streams._continuation import ( + CONTINUATION_HEADER, + Continuation, + decode_continuation, + encode_continuation, + read_continuation_header, + write_continuation_header, +) +from temporalio.contrib.external_workflow_streams._errors import StreamStorageError +from temporalio.contrib.external_workflow_streams._manager import ( + ReadinessResult, +) +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Offset, + RecordKind, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._runtime import WorkflowStreamRuntime +from temporalio.exceptions import ApplicationError +from temporalio.worker import Replayer, Worker +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend + +with workflow.unsafe.imports_passed_through(): + from temporalio.contrib.external_workflow_streams._api import external_stream + + +async def _notify( + run_id: str, # pyright: ignore[reportUnusedParameter] + wait_id: int, # pyright: ignore[reportUnusedParameter] + generation: int, # pyright: ignore[reportUnusedParameter] +) -> str: + return ReadinessResult.ACCEPTED + + +def continuation( + cursors, + stream_names, + provider_ids=None, + provider_format_versions=None, +): # type: ignore[no-untyped-def] + provider_ids = provider_ids or { + wait_id: MemoryStreamBackend.provider_id for wait_id in cursors + } + provider_format_versions = provider_format_versions or { + wait_id: MemoryStreamBackend.provider_format_version for wait_id in cursors + } + return Continuation( + cursors=cursors, + stream_names=stream_names, + provider_ids=provider_ids, + provider_format_versions=provider_format_versions, + ) + + +def make_runtime(manager, backend, continuation=None): # type: ignore[no-untyped-def] + return WorkflowStreamRuntime( + manager=manager, + backend=backend, + run_id="run-2", + namespace="ns", + workflow_id="wf", + first_execution_run_id="first-run", + data_converter=temporalio.converter.DataConverter.default, + default_idle_timeout=timedelta(seconds=1), + continuation=continuation, + ) + + +@pytest.fixture +def backend() -> MemoryStreamBackend: + return MemoryStreamBackend() + + +class OtherProviderBackend(MemoryStreamBackend): + """The same name, mapped to a different implementation. + + A deployment can do this without touching Workflow code, which is exactly + why the continuation has to carry what the cursor was produced by rather + than only what it was produced for. + """ + + provider_id = "other-memory" + + +class NewerFormatBackend(MemoryStreamBackend): + """The recorded provider, at a format version that reads offsets differently.""" + + provider_format_version = 2 + + +class StubManager: + def cancel_from_workflow_thread(self, run_id, wait_id): # type: ignore[no-untyped-def] + pass + + """Records registrations and starts no watcher. + + What these tests are about is the cursor a subscription is registered *with*, + which is decided before any watching happens. A real manager here would spawn + prefetch loops against a backend nothing is writing to, and its failures + would be noise rather than signal. + """ + + def __init__(self) -> None: + self.registered: list[tuple[int, object]] = [] + + def register(self, *, run_id, wait_id, stream_key, start_cursor): # type: ignore[no-untyped-def] + self.registered.append((wait_id, start_cursor)) + + def note_wait_generation(self, run_id, wait_id, generation) -> None: # type: ignore[no-untyped-def] + pass + + +@pytest.fixture +def manager() -> StubManager: + return StubManager() + + +# --- the header's encoding ---------------------------------------------------- + + +def test_a_continuation_round_trips() -> None: + """Binding included: a cursor without one cannot be checked on restoration.""" + original = continuation( + cursors={1: AFTER(Offset("100-3")), 2: BEGINNING}, + stream_names={1: "tokens", 2: "tool-events"}, + provider_ids={1: "memory", 2: "redis"}, + provider_format_versions={1: 1, 2: 3}, + ) + + assert decode_continuation(encode_continuation(original)) == original + + +def test_the_encoding_is_stable_for_one_state() -> None: + """The header is re-derived from live state every time the command is built. + + Encoding the same state two different ways -- by iterating a dict in + insertion order, say -- would hand the successor a different cursor on a + Workflow Task the server retried. + """ + forwards = continuation({1: BEGINNING, 2: BEGINNING}, {1: "a", 2: "b"}) + backwards = continuation({2: BEGINNING, 1: BEGINNING}, {2: "b", 1: "a"}) + + assert encode_continuation(forwards) == encode_continuation(backwards) + + +def test_a_newer_schema_version_is_reported_not_guessed_at() -> None: + """Silently starting at BEGINNING would redeliver the whole stream.""" + future = bytearray(encode_continuation(continuation({1: BEGINNING}, {1: "a"}))) + future[0] = 99 + + with pytest.raises(AnnotationDecodeError, match="schema version 99"): + decode_continuation(bytes(future)) + + +def test_the_header_bypasses_the_user_data_converter() -> None: + """One Run writes it and the next reads it, possibly on different config. + + A chain whose Runs were deployed with different converter configuration + would otherwise restart at an unreadable cursor -- silently, since an + unreadable cursor looks exactly like no cursor at all. + """ + payload = write_continuation_header(continuation({1: BEGINNING}, {1: "a"})) + + assert payload.metadata["encoding"] == b"binary/plain" + assert decode_continuation(payload.data).cursors == {1: BEGINNING} + + +def test_the_header_name_is_in_this_features_namespace() -> None: + assert CONTINUATION_HEADER.startswith("__temporal_external_stream") + assert not CONTINUATION_HEADER.startswith("__temporal_workflow_stream") + + +def test_no_header_is_a_first_execution_not_a_failure() -> None: + assert read_continuation_header(None) is None + assert read_continuation_header({}) is None + assert ( + read_continuation_header( + {"other": write_continuation_header(continuation({1: BEGINNING}, {1: "a"}))} + ) + is None + ) + + +# --- restoration -------------------------------------------------------------- + + +def test_a_first_execution_starts_at_the_beginning( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + """The same field, filled the same way -- replay reads a boundary either way. + + Including when the stream was empty for the subscription's entire life, + which is the case an implicit "start wherever" would leave unrecorded. + """ + runtime = make_runtime(manager, backend) + + runtime.register(wait_id=1, stream_key=runtime.stream_key("tokens")) + + assert runtime._subscriptions[1].start_cursor == BEGINNING + + +def test_a_successor_run_restores_its_predecessors_cursor( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + runtime = make_runtime( + manager, + backend, + continuation({1: AFTER(Offset("100-3"))}, {1: "tokens"}), + ) + + runtime.register(wait_id=1, stream_key=runtime.stream_key("tokens")) + + assert runtime._subscriptions[1].start_cursor == AFTER(Offset("100-3")) + + +def test_restoration_reads_nothing_from_the_backend( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + """The criterion, and the reason the header exists at all. + + A cursor read from the backend is whatever the stream holds now, so two + replays of one history could restore different positions and diverge. + """ + runtime = make_runtime( + manager, backend, continuation({1: AFTER(Offset("100-3"))}, {1: "tokens"}) + ) + + runtime.register(wait_id=1, stream_key=runtime.stream_key("tokens")) + + assert backend.range_reads == [], "restoration must not read the stream" + + +def test_two_same_stream_subscriptions_restore_independently( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + """They are separate waits with their own cursors. + + A stream-keyed continuation would restart one of them at the other's + position, replaying records it had already consumed or skipping records it + had not. + """ + runtime = make_runtime( + manager, + backend, + continuation( + {1: AFTER(Offset("100-0")), 2: AFTER(Offset("500-0"))}, + {1: "tokens", 2: "tokens"}, + ), + ) + key = runtime.stream_key("tokens") + + runtime.register(wait_id=1, stream_key=key) + runtime.register(wait_id=2, stream_key=key) + + assert runtime._subscriptions[1].start_cursor == AFTER(Offset("100-0")) + assert runtime._subscriptions[2].start_cursor == AFTER(Offset("500-0")) + + +def test_a_subscription_the_predecessor_lacked_starts_at_the_beginning( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + """Adding one on a path the chain has not reached yet is a supported change.""" + runtime = make_runtime( + manager, backend, continuation({1: AFTER(Offset("100-0"))}, {1: "tokens"}) + ) + + runtime.register(wait_id=2, stream_key=runtime.stream_key("later")) + + assert runtime._subscriptions[2].start_cursor == BEGINNING + + +def test_a_renumbered_subscription_is_caught_rather_than_resumed_wrong( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + """Restoring a cursor onto a differently-numbered wait resumes the wrong stream. + + The backend would accept the offset and no later check would catch it, so + this is reported here -- as nondeterminism, with the same remedy an inserted + timer needs. Not as integrity loss: the cursor is exactly what the + predecessor committed, and it is the Workflow code that moved. + """ + runtime = make_runtime( + manager, backend, continuation({1: AFTER(Offset("100-0"))}, {1: "tokens"}) + ) + + with pytest.raises(workflow.NondeterminismError, match="workflow.patched"): + runtime.register( + wait_id=1, + stream_key=runtime.stream_key("tool-events"), + ) + + +def test_a_different_configured_provider_is_a_storage_failure( + manager: StubManager, +) -> None: + """The Workflow is unchanged and neither store is damaged. + + Marker replay reports this against a recorded range; here it is reported + against a restored cursor, which is the same question one Workflow Task + earlier -- and before any read, so the wrong implementation never gets to + interpret the boundary at all. + """ + backend = OtherProviderBackend() + runtime = make_runtime( + manager, + backend, + continuation( + {1: AFTER(Offset("200-0"))}, + {1: "tokens"}, + {1: "memory"}, + {1: 1}, + ), + ) + + with pytest.raises(StreamStorageError, match="'other-memory'"): + runtime.register(wait_id=1, stream_key=runtime.stream_key("tokens")) + + assert manager.registered == [] + assert backend.range_reads == [], "raised before any backend read" + + +def test_a_newer_provider_format_version_is_a_storage_failure( + manager: StubManager, +) -> None: + """The implementation is the recorded one but reads offsets differently. + + Resuming would interpret a boundary under a format it was not written in, + which is not something the offset itself can reveal. + """ + backend = NewerFormatBackend() + runtime = make_runtime( + manager, + backend, + continuation( + {1: AFTER(Offset("200-0"))}, + {1: "tokens"}, + {1: "memory"}, + {1: 1}, + ), + ) + + with pytest.raises(StreamStorageError, match="format version"): + runtime.register(wait_id=1, stream_key=runtime.stream_key("tokens")) + + assert manager.registered == [] + assert backend.range_reads == [], "raised before any backend read" + + +def test_a_recorded_format_version_of_zero_is_still_compared( + manager: StubManager, +) -> None: + """Zero is a format version, not a "nothing was recorded" sentinel. + + The backend contract types ``provider_format_version`` as a plain integer + and reserves no value, and the header represents zero exactly. Deciding + whether a version was recorded by truthiness therefore skipped the + comparison for the one value that looks falsey, which made Continue-As-New + less safe than marker replay for the same binding: replay compares exactly. + """ + backend = MemoryStreamBackend() + assert type(backend).provider_format_version == 1, ( + "the fixture must declare a version the recorded zero disagrees with" + ) + runtime = make_runtime( + manager, + backend, + continuation( + {1: AFTER(Offset("200-0"))}, + {1: "tokens"}, + {1: "memory"}, + {1: 0}, + ), + ) + + with pytest.raises(StreamStorageError, match="format version"): + runtime.register(wait_id=1, stream_key=runtime.stream_key("tokens")) + + assert manager.registered == [], "the cursor must not reach the manager" + assert backend.range_reads == [], "raised before any backend read" + + +async def test_a_worker_without_a_backend_still_decodes_the_continuation( + client: Client, +) -> None: + async with Worker( + client, + task_queue=f"tq-{uuid.uuid4()}", + workflows=[ChainedConsumerWorkflow], + ) as worker: + assert worker._workflow_worker is not None + init = InitializeWorkflow( + workflow_id="workflow", first_execution_run_id="first" + ) + init.headers[CONTINUATION_HEADER].data = b"\x03" + + with pytest.raises(AnnotationDecodeError, match="schema version 3"): + worker._workflow_worker._create_external_stream_runtime( + WorkflowActivation(run_id="run"), init + ) + + init.headers[CONTINUATION_HEADER].CopyFrom( + write_continuation_header(continuation({1: BEGINNING}, {1: "tokens"})) + ) + with pytest.raises(RuntimeError, match="external_stream_backend"): + worker._workflow_worker._create_external_stream_runtime( + WorkflowActivation(run_id="run"), init + ) + + +# --- what gets committed ------------------------------------------------------ + + +def test_the_continuation_reports_what_was_consumed_not_what_was_delivered( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + """The successor must not start past records the Workflow never took. + + Delivery advances by whole drained batches, because that is what the + annotation records and what replay must reproduce. A Workflow that stops + iterating part-way through a batch has consumed only its prefix, and the + rest dies with the Run's buffer -- so a continuation taken from the delivery + cursor would step over records nothing had ever shown to Workflow code. + """ + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=runtime.stream_key("tokens")) + batch = [ + StreamRecord(RecordKind.DATA, b"a", "s", i).placed_at(Offset(f"10-{i}")) + for i in range(3) + ] + for record in batch: + runtime.record_delivery(1, record) + # The Workflow took only the first of the three, then stopped iterating. + runtime.record_consumption(1, batch[0]) + + assert runtime.continuation().cursors[1] == AFTER(Offset("10-0")), ( + "the successor must resume at the first record this Run did not consume" + ) + + +def test_a_control_record_still_advances_consumption( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + """It is never yielded, but it is finished with. + + Leaving it behind would restart the successor at a fence it cannot act on, + and it would be redelivered on every Run of the chain. + """ + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=runtime.stream_key("tokens")) + fence = StreamRecord(RecordKind.WRITE_FENCE, b"", "s", 0).placed_at(Offset("10-0")) + + runtime.record_consumption(1, fence) + + assert runtime.continuation().cursors[1] == AFTER(Offset("10-0")) + + +def test_a_run_that_delivered_nothing_still_reports_its_start( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + """Otherwise the successor would restart at BEGINNING and redeliver.""" + runtime = make_runtime( + manager, backend, continuation({1: AFTER(Offset("100-0"))}, {1: "tokens"}) + ) + runtime.register(wait_id=1, stream_key=runtime.stream_key("tokens")) + + assert runtime.continuation().cursors[1] == AFTER(Offset("100-0")) + + +# --- through a real chain ----------------------------------------------------- + + +@workflow.defn +class ChainedConsumerWorkflow: + """Consumes two records, continues as new, consumes two more. + + The whole point is the second Run must not see the first Run's records. A + Workflow that returned everything it saw would hide that: only counting what + the *successor* observed makes a redelivery visible. + """ + + @workflow.run + async def run(self, remaining: int) -> list[str]: + tokens = external_stream.with_options(idle_timeout=timedelta(seconds=30)).topic( + "tokens", type=str + ) + seen: list[str] = [] + async for token in tokens.subscribe(): + seen.append(token) + if len(seen) >= 2: + break + if remaining > 1: + workflow.continue_as_new(remaining - 1) + return seen + + +async def publish(backend: MemoryStreamBackend, key: StreamKey, values: list[str]): + from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec + + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + for i, value in enumerate(values): + await backend.append( + key, + StreamRecord(RecordKind.DATA, await codec.encode(value), "producer", i), + ) + + +async def test_a_chain_resumes_where_its_predecessor_stopped( + client: Client, backend: MemoryStreamBackend +) -> None: + """The successor consumes the records the first Run did not, and no others.""" + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[ChainedConsumerWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + ChainedConsumerWorkflow.run, + 2, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + description = await handle.describe() + key = StreamKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + "tokens", + ) + await asyncio.sleep(1) + await publish(backend, key, ["a", "b", "c", "d"]) + + assert await asyncio.wait_for(handle.result(), 60) == ["c", "d"], ( + "the successor Run redelivered records its predecessor had already " + "consumed, so the continuation cursor did not survive Continue-As-New" + ) + + +async def test_a_chain_carries_the_provider_binding_through_history( + client: Client, backend: MemoryStreamBackend +) -> None: + """The binding reaches the successor as History, and is checked on the way in. + + This shows that a live Worker puts the binding on the command, the server + persists it into the successor's ``WorkflowExecutionStarted``, and the + successor validates it rather than merely tolerating it. + + The last of those is asserted against the *recorded* bytes with only the + provider identity moved, so the run that must fail differs from the run that + must succeed in exactly the thing the binding exists to detect. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[ChainedConsumerWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + ChainedConsumerWorkflow.run, + 2, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + description = await handle.describe() + first_run_id = description.raw_description.workflow_execution_info.first_run_id + key = StreamKey(client.namespace, handle.id, first_run_id, "tokens") + await asyncio.sleep(1) + # One call, carrying its own sequence: a second call would restart the + # producer sequence at zero and the backend would reject the appends as + # non-identical repeats, which presents as a hung Workflow. + await publish(backend, key, ["a", "b", "c", "d"]) + + assert await asyncio.wait_for(handle.result(), 60) == ["c", "d"] + predecessor = await client.get_workflow_handle( + handle.id, run_id=first_run_id + ).fetch_history() + + continued = [ + e + for e in predecessor.events + if e.HasField("workflow_execution_continued_as_new_event_attributes") + ] + assert continued, "the predecessor did not continue as new" + attributes = continued[0].workflow_execution_continued_as_new_event_attributes + recorded = decode_continuation(attributes.header.fields[CONTINUATION_HEADER].data) + + assert recorded.provider_ids == {1: MemoryStreamBackend.provider_id} + assert recorded.provider_format_versions == { + 1: MemoryStreamBackend.provider_format_version + } + + # The bytes the successor actually reads. `read_continuation_header` is + # handed this event's headers, not the command's, so a header the server + # dropped or rewrote between the two would leave the successor starting at + # BEGINNING with nothing to say so. + successor = await client.get_workflow_handle( + handle.id, run_id=attributes.new_execution_run_id + ).fetch_history() + started = successor.events[0].workflow_execution_started_event_attributes + assert ( + decode_continuation(started.header.fields[CONTINUATION_HEADER].data) == recorded + ) + + # Replayed against the provider that produced the cursor, the recorded + # binding has to agree; this is also what makes the mismatch below evidence + # about the binding rather than about the history. + matched = await Replayer( + workflows=[ChainedConsumerWorkflow], + external_stream_backend=backend, + ).replay_workflow(successor, raise_on_replay_failure=False) + + assert matched.replay_failure is None, ( + f"the successor's own history no longer replays: {matched.replay_failure}" + ) + + # Same history, same records, and same stream -- only the configured + # implementation has moved. Accepting the cursor would skip records in + # silence because the boundary means nothing in the replacement store. + # + # Reported by *marker* replay, not by the continuation: the successor's + # `ReplayExternalStreams` job is handled before its Workflow code runs, so on + # a history that already carries a marker the marker binding is checked + # first. This is therefore the wrong layer to ask the continuation's own + # check about -- disabling `_verify_restored_provider` does not change this + # outcome, which is why that check is asked for separately below instead of + # being inferred from this failure. + foreign = OtherProviderBackend() + foreign._records = {k: list(v) for k, v in backend._records.items()} + moved = await Replayer( + workflows=[ChainedConsumerWorkflow], + external_stream_backend=foreign, + ).replay_workflow(successor, raise_on_replay_failure=False) + + assert moved.replay_failure is not None, ( + "the successor read records out of a store that did not produce them" + ) + assert "other-memory" in str(moved.replay_failure), ( + f"the failure does not name the provider that was found: {moved.replay_failure}" + ) + + # The continuation's own check, on the one Workflow Task that has to make it: + # the successor's *first*, which restores the cursor before any marker exists + # for marker replay to check in its place. Asked with the bytes History + # actually holds rather than with a hand-built `Continuation`, so what is + # validated is the recorded binding and not a shape a test chose. + first_task = make_runtime( + StubManager(), + foreign, + recorded, + ) + + with pytest.raises(StreamStorageError, match="'other-memory'"): + first_task.register( + wait_id=1, + stream_key=first_task.stream_key("tokens"), + ) + + assert foreign.range_reads == [], ( + "the mismatch was reported after a read, so the wrong implementation " + "already interpreted the recorded boundary" + ) + + +# --- the boundary the header is taken at -------------------------------------- +# +# Creating the Continue-As-New command does not end the activation. The event +# loop keeps draining what is already ready, and a stream consumer that runs in +# that tail consumes records the predecessor's final marker *does* record -- so a +# header taken when the command was created describes an earlier boundary than +# History does, and the successor is handed those records a second time. + + +@workflow.defn +class LateConsumerContinueAsNewWorkflow: + """Schedules a consumer, then continues as new without yielding to it. + + The scheduled task takes a record that is already buffered, so it finishes + without blocking -- it just does so after the terminal command exists. + + Only the successor's records are returned. A Run that reported everything it + saw would hide the redelivery inside the predecessor's own list. + """ + + @workflow.run + async def run(self, remaining: int) -> list[str]: + tokens = external_stream.with_options(idle_timeout=timedelta(seconds=30)).topic( + "tokens", type=str + ) + subscription = tokens.subscribe() + iterator = subscription.__aiter__() + first = await iterator.__anext__() + if remaining <= 1: + return [first, await iterator.__anext__()] + _require_a_buffered_record(subscription) + + async def take_one() -> None: + await iterator.__anext__() + + # On the ready queue, and able to finish without blocking. `_run_once` + # gives it its turn after the command below is created. + asyncio.create_task(take_one()) + workflow.continue_as_new(remaining - 1) + + +@workflow.defn +class SignalContinueAsNewWorkflow: + """A signal handler continues as new while a consumer is being unblocked. + + The consumer is parked on a condition the handler sets, and conditions are + re-checked *after* the ready queue drains -- which is after the terminal + command was created. So this reaches the same tail by the route ordinary + application structure reaches it by, rather than by scheduling a task next + to the Continue-As-New call. + """ + + def __init__(self) -> None: + self._release = False + self._staged = False + + @workflow.run + async def run(self, remaining: int) -> list[str]: + tokens = external_stream.with_options(idle_timeout=timedelta(seconds=30)).topic( + "tokens", type=str + ) + subscription = tokens.subscribe() + iterator = subscription.__aiter__() + first = await iterator.__anext__() + if remaining <= 1: + return [first, await iterator.__anext__()] + _require_a_buffered_record(subscription) + + async def take_one() -> None: + await workflow.wait_condition(lambda: self._release) + await iterator.__anext__() + + asyncio.create_task(take_one()) + self._staged = True + # The handler ends this Run; nothing else does. + await workflow.wait_condition(lambda: False) + raise AssertionError("unreachable") + + @workflow.query + def staged(self) -> bool: + """Whether the consumer is parked yet. + + Signalling before it is would continue as new with nothing else ready, + and the test would pass without exercising anything. + """ + return self._staged + + @workflow.signal + async def wrap_up(self, remaining: int) -> None: + self._release = True + workflow.continue_as_new(remaining) + + +def _require_a_buffered_record(subscription: object) -> None: + """Fails the Run rather than letting a thin test pass. + + The whole shape depends on a record still sitting in the subscription's + ready list when the terminal command is created. If the batch arrived one + record at a time there is nothing for the tail to consume, and the + assertions below would hold for the wrong reason. + """ + if not subscription._ready: # type: ignore[attr-defined] + raise ApplicationError( + "no record was buffered when the terminal command was created, so " + "nothing was consumed after it and this Run proves nothing", + non_retryable=True, + ) + + +async def stage_stream( + client: Client, + backend: MemoryStreamBackend, + handle: WorkflowHandle[Any, Any], + values: list[str], +) -> StreamKey: + """Puts the whole batch in the stream before any Worker can read it. + + A record that arrives an activation later is delivered to the successor by + the ordinary path and says nothing about this boundary. Publishing before + the Worker exists is what makes the Run's first drain take the batch whole. + """ + description = await handle.describe() + key = StreamKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + "tokens", + ) + await publish(backend, key, values) + return key + + +async def _until_staged(handle: WorkflowHandle[Any, Any]) -> None: + for _ in range(300): + if await handle.query(SignalContinueAsNewWorkflow.staged): + return + await asyncio.sleep(0.1) + raise AssertionError("the consumer never parked on its condition") + + +async def test_a_consumer_that_runs_after_the_terminal_command_is_still_consumed( + client: Client, backend: MemoryStreamBackend +) -> None: + """The predecessor's marker and the successor's header are one boundary. + + Both records the predecessor took are recorded in its final marker. A + continuation snapshot taken when the command was created holds only the + first, and the successor then starts one record too early. + """ + task_queue = f"tq-{uuid.uuid4()}" + handle = await client.start_workflow( + LateConsumerContinueAsNewWorkflow.run, + 2, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + await stage_stream(client, backend, handle, ["a", "b", "c", "d"]) + + async with Worker( + client, + task_queue=task_queue, + workflows=[LateConsumerContinueAsNewWorkflow], + external_stream_backend=backend, + ): + assert await asyncio.wait_for(handle.result(), 30) == ["c", "d"], ( + "the successor was handed a record its predecessor consumed after " + "the Continue-As-New command was created, so the continuation " + "header was taken before the activation had finished" + ) + + +async def test_a_signal_handler_continuing_as_new_waits_for_the_consumer( + client: Client, backend: MemoryStreamBackend +) -> None: + """The same boundary, reached through a condition rather than a new task.""" + task_queue = f"tq-{uuid.uuid4()}" + handle = await client.start_workflow( + SignalContinueAsNewWorkflow.run, + 2, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + await stage_stream(client, backend, handle, ["a", "b", "c", "d"]) + + async with Worker( + client, + task_queue=task_queue, + workflows=[SignalContinueAsNewWorkflow], + external_stream_backend=backend, + ): + await _until_staged(handle) + await handle.signal(SignalContinueAsNewWorkflow.wrap_up, 1) + + assert await asyncio.wait_for(handle.result(), 30) == ["c", "d"], ( + "the consumer the signal handler unblocked consumed a record after " + "the Continue-As-New command was created, and the successor " + "received it again" + ) + + +async def test_a_history_written_at_the_earlier_boundary_still_replays( + client: Client, backend: MemoryStreamBackend +) -> None: + """A chain that continued as new before this Worker was deployed. + + Replay regenerates the header at the *later* boundary, so it no longer + matches the recorded one. That is allowed and needs no compatibility flag: + Core matches a Continue-As-New command to its + ``WorkflowExecutionContinuedAsNew`` event by command type alone and never + compares headers. The successor of such a chain is unaffected either way -- + it reads its cursor from its own ``WorkflowExecutionStarted``, which + replaying the predecessor does not rewrite. + """ + task_queue = f"tq-{uuid.uuid4()}" + handle = await client.start_workflow( + LateConsumerContinueAsNewWorkflow.run, + 2, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + key = await stage_stream(client, backend, handle, ["a", "b", "c", "d"]) + first_run_id = handle.first_execution_run_id + assert first_run_id is not None + + async with Worker( + client, + task_queue=task_queue, + workflows=[LateConsumerContinueAsNewWorkflow], + external_stream_backend=backend, + ): + await asyncio.wait_for(handle.result(), 30) + history = await client.get_workflow_handle( + handle.id, run_id=first_run_id + ).fetch_history() + + continued = [ + e + for e in history.events + if e.HasField("workflow_execution_continued_as_new_event_attributes") + ] + assert continued, "the predecessor did not continue as new" + attributes = continued[0].workflow_execution_continued_as_new_event_attributes + recorded = decode_continuation(attributes.header.fields[CONTINUATION_HEADER].data) + + # What the earlier snapshot held: the first record, taken before the tail of + # the activation consumed the second. + first_offset = backend._records[key][0].offset + assert first_offset is not None + early = dataclasses.replace( + recorded, cursors={wait_id: AFTER(first_offset) for wait_id in recorded.cursors} + ) + assert early != recorded, ( + "the recorded header already holds the earlier boundary, so this " + "history is not the pre-fix one it is meant to stand in for" + ) + attributes.header.fields[CONTINUATION_HEADER].CopyFrom( + write_continuation_header(early) + ) + + result = await Replayer( + workflows=[LateConsumerContinueAsNewWorkflow], + external_stream_backend=backend, + ).replay_workflow(history, raise_on_replay_failure=False) + + assert result.replay_failure is None, ( + "a history whose Continue-As-New header holds the earlier boundary no " + f"longer replays: {result.replay_failure}" + ) diff --git a/tests/contrib/external_workflow_streams/test_delivery_budget.py b/tests/contrib/external_workflow_streams/test_delivery_budget.py new file mode 100644 index 0000000000..79386d6fee --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_delivery_budget.py @@ -0,0 +1,1397 @@ +"""The per-activation delivery budget. + +A producer that never lets a subscription's buffer run dry makes one +``activate()`` call never return: the iterator re-fills before every record and +always finds one. Activations run on a thread-pool executor under a **2-second +deadlock timeout**, so the Workflow Task fails -- and every retry meets the same +producer, so the Workflow is stuck permanently rather than merely slow. Driving +the real iterator against a runtime whose ``drain`` always returned a full batch +consumed 316,086 records in two seconds and blocked zero times. + +Everything here asserts the *mechanism*: that the iterator gives control back +after a fixed number of records, that the count spans a merged set rather than +each member of it, that the records the budget left behind get their readiness +re-reported, and that replay -- whose boundaries are already recorded -- is not +subject to any of it. +""" + +from __future__ import annotations + +import asyncio +import time +import uuid +from collections import Counter +from datetime import timedelta +from typing import Any, cast + +import pytest + +import temporalio.converter +import temporalio.workflow +from temporalio.contrib.external_workflow_streams import _runtime as _runtime_module +from temporalio.contrib.external_workflow_streams._annotation import ( + MAX_ANNOTATION_BYTES, + SegmentEndReason, + decode_annotation, +) +from temporalio.contrib.external_workflow_streams._api import ( + MAX_RECORDS_PER_ACTIVATION, + _install_runtime, + external_stream, + merge, +) +from temporalio.contrib.external_workflow_streams._backend import StreamKey +from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec +from temporalio.contrib.external_workflow_streams._manager import ( + ReadinessResult, + StreamSubscriptionManager, +) +from temporalio.contrib.external_workflow_streams._record import ( + Offset, + RecordKind, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._runtime import WorkflowStreamRuntime +from temporalio.worker._workflow_instance import _WorkflowInstanceImpl +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend + +RUN_ID = "run-1" + +#: How far past the budget an unbounded fake is allowed to run before the test +#: gives up. Large enough that a missing cap is unmistakable, small enough that +#: the failure arrives quickly. +_UNBOUNDED_BATCH = 1000 + + +# --- fixtures ----------------------------------------------------------------- + + +class FakeInstance: + """Stands in for the user's Workflow object, which holds the per-Run state.""" + + +@pytest.fixture +def workflow_instance(monkeypatch: pytest.MonkeyPatch) -> FakeInstance: + instance = FakeInstance() + monkeypatch.setattr(temporalio.workflow, "instance", lambda: instance) + return instance + + +@pytest.fixture +def backend() -> MemoryStreamBackend: + return MemoryStreamBackend() + + +def make_runtime(manager: Any, backend: MemoryStreamBackend) -> WorkflowStreamRuntime: + return WorkflowStreamRuntime( + manager=manager, + backend=backend, + run_id=RUN_ID, + namespace="ns", + workflow_id="wf", + first_execution_run_id="first-run", + data_converter=temporalio.converter.DataConverter.default, + default_idle_timeout=timedelta(seconds=1), + ) + + +async def encoded(value: str) -> bytes: + codec: StreamPayloadCodec[str] = StreamPayloadCodec( + temporalio.converter.DataConverter.default, str + ) + return await codec.encode(value) + + +class NeverEmptyManager: + """A manager whose buffers never run dry. + + The exact shape that makes an activation never return. It honours + ``max_records`` because a real manager does; when asked for an unbounded + drain it hands back a large batch, so a budget that failed to reach the drain + shows up as an iterator that runs past the cap rather than as a fake that + quietly stops. + """ + + def __init__(self, payload: bytes) -> None: + self.payload = payload + self.max_records_seen: list[int | None] = [] + #: Every wait the runtime actually asked about. A wait that never + #: appears was not merely served nothing -- the Worker was never + #: enquired of on its behalf at all. + self.drained_waits: list[int] = [] + self.rearmed: list[str] = [] + self._next: dict[int, int] = {} + + def register(self, *, run_id, wait_id, stream_key, start_cursor): # type: ignore[no-untyped-def] + self._next.setdefault(wait_id, 0) + + def note_wait_generation(self, run_id, wait_id, generation) -> None: # type: ignore[no-untyped-def] + pass + + def drain( + self, + run_id: str, # pyright: ignore[reportUnusedParameter] + wait_id: int, + max_records: int | None = None, + ) -> list[StreamRecord]: + self.max_records_seen.append(max_records) + self.drained_waits.append(wait_id) + count = _UNBOUNDED_BATCH if max_records is None else max_records + start = self._next.get(wait_id, 0) + self._next[wait_id] = start + count + return [ + StreamRecord(RecordKind.DATA, self.payload, f"p{wait_id}", i).placed_at( + Offset(f"{wait_id:02d}-{i:08d}") + ) + for i in range(start, start + count) + ] + + def rearm_ready(self, run_id: str) -> None: + self.rearmed.append(run_id) + + +async def consume(iterator: Any, seen: list[Any]) -> None: + async for value in iterator: + seen.append(value) + + +async def settle(task: asyncio.Task[None], seen: list[Any], expected: int) -> None: + """Runs the consumer until it stops making progress, then a little longer. + + The extra window is what turns "it has not finished yet" into "it blocked": + an iterator that ignored the budget would keep going right through it. + """ + deadline = time.monotonic() + 5 + while len(seen) < expected and not task.done() and time.monotonic() < deadline: + await asyncio.sleep(0.005) + await asyncio.sleep(0.1) + + +async def cancel(task: asyncio.Task[None]) -> None: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + +# --- one subscription --------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_never_empty_subscription_yields_control_after_the_budget( + workflow_instance: FakeInstance, backend: MemoryStreamBackend +) -> None: + """The measured failure, and the cap that ends it. + + Without a bound this loop never returns and the Workflow Task dies on the + deadlock timeout -- on every retry, because the producer is still there. + """ + manager = NeverEmptyManager(await encoded("x")) + runtime = make_runtime(manager, backend) + _install_runtime(workflow_instance, runtime) + runtime.begin_activation() + subscription = external_stream.topic("tokens", type=str).subscribe() + + seen: list[str] = [] + task = asyncio.ensure_future(consume(subscription.__aiter__(), seen)) + await settle(task, seen, MAX_RECORDS_PER_ACTIVATION) + + assert not task.done(), ( + "the iterator ran past the budget against a buffer that never empties; " + "this activation would never return and the Workflow Task would fail on " + "the 2-second deadlock timeout" + ) + assert len(seen) == MAX_RECORDS_PER_ACTIVATION + assert runtime.delivery_budget_exhausted() + assert manager.max_records_seen and all( + limit is not None for limit in manager.max_records_seen + ), "every drain must carry the remaining budget" + assert 0 not in manager.max_records_seen, ( + "a spent budget must stop the drain from being attempted at all, rather " + "than reaching the manager and asking it for nothing" + ) + await cancel(task) + + +@pytest.mark.asyncio +async def test_the_budget_blocks_even_though_records_are_buffered( + workflow_instance: FakeInstance, backend: MemoryStreamBackend +) -> None: + """Blocking is the whole point: it is what ends the activation. + + The double-check refill inside ``_await_readiness`` is the hazard here. It + exists to catch a record buffered while the wait was being registered, and if + it ignored the budget it would hand this iterator a record immediately and + the activation would go right back to never returning. + """ + manager = NeverEmptyManager(await encoded("x")) + runtime = make_runtime(manager, backend) + _install_runtime(workflow_instance, runtime) + runtime.begin_activation() + subscription = external_stream.topic("tokens", type=str).subscribe() + + seen: list[str] = [] + task = asyncio.ensure_future(consume(subscription.__aiter__(), seen)) + await settle(task, seen, MAX_RECORDS_PER_ACTIVATION) + + # Blocked on a readiness future, with the manager still holding records. + assert runtime.resolve_all_pending() == 1, ( + "the iterator must be parked on a readiness future, which is what lets " + "the activation return" + ) + await cancel(task) + + +@pytest.mark.asyncio +async def test_the_next_activation_gets_a_fresh_budget( + workflow_instance: FakeInstance, backend: MemoryStreamBackend +) -> None: + """The cap bounds one activation, not the Run. + + A budget that were never reset would deliver 256 records and then never + another one, which is a different way of hanging the same Workflow. + """ + manager = NeverEmptyManager(await encoded("x")) + runtime = make_runtime(manager, backend) + _install_runtime(workflow_instance, runtime) + runtime.begin_activation() + subscription = external_stream.topic("tokens", type=str).subscribe() + + seen: list[str] = [] + task = asyncio.ensure_future(consume(subscription.__aiter__(), seen)) + await settle(task, seen, MAX_RECORDS_PER_ACTIVATION) + assert len(seen) == MAX_RECORDS_PER_ACTIVATION + + # The next activation begins: budget reset, readiness resolved. + runtime.begin_activation() + assert runtime.delivery_budget_remaining() == MAX_RECORDS_PER_ACTIVATION + runtime.resolve_all_pending() + await settle(task, seen, 2 * MAX_RECORDS_PER_ACTIVATION) + + assert len(seen) == 2 * MAX_RECORDS_PER_ACTIVATION + assert not task.done() + await cancel(task) + + +# --- the budget spans subscriptions ------------------------------------------- + + +@pytest.mark.asyncio +async def test_merge_spends_one_budget_across_every_subscription( + workflow_instance: FakeInstance, backend: MemoryStreamBackend +) -> None: + """A per-subscription budget would let n streams run n times as long. + + That is the same deadlock with a larger constant in front of it, which is + why the counter lives on the runtime rather than on a subscription. + """ + manager = NeverEmptyManager(await encoded("x")) + runtime = make_runtime(manager, backend) + _install_runtime(workflow_instance, runtime) + runtime.begin_activation() + first = external_stream.topic("a", type=str).subscribe() + second = external_stream.topic("b", type=str).subscribe() + + seen: list[tuple[int, str]] = [] + + async def consume_merged() -> None: + async for subscription, value in merge(first, second): + seen.append((subscription.wait_id, value)) + + task = asyncio.ensure_future(consume_merged()) + await settle(task, seen, MAX_RECORDS_PER_ACTIVATION) # type: ignore[arg-type] + + assert not task.done() + assert len(seen) == MAX_RECORDS_PER_ACTIVATION, ( + "the budget covers the whole merged set; two never-empty streams must " + "not buy twice the records" + ) + await cancel(task) + + +@pytest.mark.asyncio +async def test_merge_blocks_on_every_wait_once_the_budget_is_spent( + workflow_instance: FakeInstance, backend: MemoryStreamBackend +) -> None: + """``_await_any_readiness`` takes the same last look, under the same cap. + + A refill there that ignored the budget would resume the merge immediately and + the activation would never end. + """ + manager = NeverEmptyManager(await encoded("x")) + runtime = make_runtime(manager, backend) + _install_runtime(workflow_instance, runtime) + runtime.begin_activation() + first = external_stream.topic("a", type=str).subscribe() + second = external_stream.topic("b", type=str).subscribe() + + seen: list[tuple[int, str]] = [] + + async def consume_merged() -> None: + async for subscription, value in merge(first, second): + seen.append((subscription.wait_id, value)) + + task = asyncio.ensure_future(consume_merged()) + await settle(task, seen, MAX_RECORDS_PER_ACTIVATION) # type: ignore[arg-type] + + assert runtime.resolve_all_pending() == 2, ( + "the complete set must be blocked, not just the stream that happened to " + "spend the last of the budget" + ) + await cancel(task) + + +async def settle_all( + tasks: list[asyncio.Task[None]], counted: list[list[Any]], expected: int +) -> None: + """`settle`, over several consumers that share one budget between them.""" + deadline = time.monotonic() + 5 + while ( + sum(len(seen) for seen in counted) < expected + and not any(task.done() for task in tasks) + and time.monotonic() < deadline + ): + await asyncio.sleep(0.005) + await asyncio.sleep(0.1) + + +@pytest.mark.asyncio +async def test_two_independent_consumers_share_one_budget( + workflow_instance: FakeInstance, backend: MemoryStreamBackend +) -> None: + """The same property as the merge case, without `merge()` to enforce it. + + `merge()` fills one record at a time, so it spends the budget in the same + places it checks it. Two ordinary subscriptions consumed by two ordinary + coroutines do not: each one drains a whole batch into its own ready list, and + a budget that is *checked* at the drain but *charged* at consumption is no + budget at all. The first subscription takes all 256 slots into its private + list and yields one, so the second sees 255 still available and takes those, + and both then drain lists nothing will check again. Measured at 511 records + before either coroutine blocked, and it is `n` times the cap for `n` + consumers -- the same deadlock the merge case is about, reached without + merging anything. + """ + manager = NeverEmptyManager(await encoded("x")) + runtime = make_runtime(manager, backend) + _install_runtime(workflow_instance, runtime) + runtime.begin_activation() + first = external_stream.topic("a", type=str).subscribe() + second = external_stream.topic("b", type=str).subscribe() + + seen_first: list[str] = [] + seen_second: list[str] = [] + + async def consume_yielding(subscription: Any, seen: list[str]) -> None: + # Awaiting after every value is what lets the other coroutine reach its + # own drain, which is the whole shape: two private ready lists, filled + # from one budget, consumed independently. + async for value in subscription: + seen.append(value) + await asyncio.sleep(0) + + tasks = [ + asyncio.ensure_future(consume_yielding(first, seen_first)), + asyncio.ensure_future(consume_yielding(second, seen_second)), + ] + await settle_all(tasks, [seen_first, seen_second], MAX_RECORDS_PER_ACTIVATION) + + try: + assert len(seen_first) + len(seen_second) == MAX_RECORDS_PER_ACTIVATION, ( + "two independently consumed subscriptions spent more than one " + "activation's budget between them; with enough streams this is the " + "activation that never returns" + ) + assert all(not task.done() for task in tasks) + assert runtime.delivery_budget_remaining() == 0 + assert runtime.delivery_budget_exhausted() + # Both are blocked, which is what ends the activation. A consumer still + # holding a ready list would not be, and the assertion above would have + # been reached with records left to hand over. + assert runtime.resolve_all_pending() == 2, ( + "both consumers must be parked on a readiness future once the budget " + "is spent, not sitting on a private ready list the cap cannot see" + ) + # And what the activation *recorded* matches what it delivered: the + # segment is the schedule replay reproduces, so a cap that bounded + # consumption while the annotation recorded delivery bounded the wrong + # quantity. + delta = runtime.take_observation_delta() + assert delta is not None + encoded_annotation = delta + runtime.add_terminal() + annotation = decode_annotation(encoded_annotation) + recorded = sum( + run.count for segment in annotation.segments for run in segment.runs + ) + assert recorded == MAX_RECORDS_PER_ACTIVATION + # And it still fits. Delivering two speculative batches put both in the + # current annotation before the counter caught up, so the same schedule + # could also overrun the affordability calculation and turn the intended + # rollover into a Workflow Task failure that repeats on every attempt. + assert len(encoded_annotation) <= MAX_ANNOTATION_BYTES + finally: + for task in tasks: + await cancel(task) + + +@pytest.mark.asyncio +async def test_a_carried_over_ready_list_is_charged_to_the_next_activation( + workflow_instance: FakeInstance, backend: MemoryStreamBackend +) -> None: + """The same hole from the other side, and the reason the reset is not to zero. + + A batch is delivered whole and consumed one record at a time, so a consumer + that stops part-way through leaves the rest in its ready list -- where the + next activation takes it with no drain, and so with no budget check. Charging + only at the drain would make that remainder free, and it accumulates: one + subscription per activation can leave a nearly full list behind, so `n` of + them arrive holding `n` times the cap and hand all of it over in one + `activate()` call. + + Nothing is redelivered here and no drain is needed, so what the budget must + bound is what Workflow code can *take*: the carry-over plus whatever is newly + delivered, together, is one cap. + """ + manager = NeverEmptyManager(await encoded("x")) + runtime = make_runtime(manager, backend) + _install_runtime(workflow_instance, runtime) + runtime.begin_activation() + subscription = external_stream.topic("tokens", type=str).subscribe() + + # One activation that delivers a whole batch and consumes a single record, + # which is the ordinary shape of a consumer that does real work per record: + # the loop body blocks on something else and the activation ends. + seen: list[str] = [] + iterator = subscription.__aiter__() + seen.append(await iterator.__anext__()) + assert runtime.delivery_budget_remaining() == 0 + carried = MAX_RECORDS_PER_ACTIVATION - 1 + + runtime.begin_activation() + + assert ( + runtime.delivery_budget_remaining() == MAX_RECORDS_PER_ACTIVATION - carried + ), ( + f"the next activation was handed a whole fresh budget on top of {carried} " + "records already in Workflow code's hands; that remainder is consumed " + "with no drain and so with no check, and it accumulates across streams" + ) + # Everything carried over is still handed over -- the records are already here + # and refusing them would strand them, since nothing re-announces a record + # already popped out of the manager's buffer -- and the whole activation still + # totals one cap rather than the carry-over plus another one. + task = asyncio.ensure_future(consume(iterator, seen)) + await settle(task, seen, MAX_RECORDS_PER_ACTIVATION + 1) + assert len(seen) - 1 == MAX_RECORDS_PER_ACTIVATION, ( + "the second activation handed over the carry-over and a fresh cap on top " + "of it; what it may hand over is one cap in total" + ) + assert not task.done() + await cancel(task) + + +class OneBusyOneQuietManager: + """Wait 1 never runs dry; wait 2 holds a single record and then nothing. + + The shape a merge has to survive: one stream a producer keeps saturated, and + one that has something to say exactly once. A merge that drains the busy + stream to the bottom of the budget before it looks at the quiet one never + reaches the quiet one at all, and never will -- every fresh budget starts at + the same lowest wait id. + """ + + def __init__(self, payload: bytes) -> None: + self.payload = payload + self.rearmed: list[str] = [] + self._next = 0 + self._quiet_pending = True + + def register(self, *, run_id, wait_id, stream_key, start_cursor): # type: ignore[no-untyped-def] + pass + + def note_wait_generation(self, run_id, wait_id, generation) -> None: # type: ignore[no-untyped-def] + pass + + def drain( + self, + run_id: str, # pyright: ignore[reportUnusedParameter] + wait_id: int, + max_records: int | None = None, + ) -> list[StreamRecord]: + if max_records is not None and max_records <= 0: + return [] + if wait_id == 2: + if not self._quiet_pending: + return [] + self._quiet_pending = False + return [ + StreamRecord(RecordKind.DATA, self.payload, "p2", 0).placed_at( + Offset("02-00000000") + ) + ] + count = _UNBOUNDED_BATCH if max_records is None else max_records + start = self._next + self._next = start + count + return [ + StreamRecord(RecordKind.DATA, self.payload, "p1", i).placed_at( + Offset(f"01-{i:08d}") + ) + for i in range(start, start + count) + ] + + def rearm_ready(self, run_id: str) -> None: + self.rearmed.append(run_id) + + +@pytest.mark.asyncio +async def test_merge_does_not_starve_a_stream_behind_a_saturated_one( + workflow_instance: FakeInstance, backend: MemoryStreamBackend +) -> None: + """A merge that is not fair is not a merge. + + Draining one subscription's whole ready list before looking at the next lets + the lowest wait id spend the entire activation budget by itself, and the next + activation begins the same pass in the same order. A record sitting ready on + a later wait is then never yielded -- deterministically, but that only makes + it reproducible rather than acceptable. + + Three budget cycles is well past any plausible "it was just slow": one pass + over the set is enough for a fair merge. + """ + manager = OneBusyOneQuietManager(await encoded("x")) + runtime = make_runtime(manager, backend) + _install_runtime(workflow_instance, runtime) + runtime.begin_activation() + busy = external_stream.topic("busy", type=str).subscribe() + quiet = external_stream.topic("quiet", type=str).subscribe() + assert (busy.wait_id, quiet.wait_id) == (1, 2) + + seen: list[tuple[int, str]] = [] + + async def consume_merged() -> None: + async for subscription, value in merge(busy, quiet): + seen.append((subscription.wait_id, value)) + + task = asyncio.ensure_future(consume_merged()) + for cycle in range(1, 4): + await settle(task, seen, cycle * MAX_RECORDS_PER_ACTIVATION) # type: ignore[arg-type] + # What the Worker's completion path does: a fresh budget, and readiness + # re-armed for the records the budget left buffered. + runtime.begin_activation() + runtime.resolve_all_pending() + + await cancel(task) + assert quiet.wait_id in {wait_id for wait_id, _ in seen}, ( + "the quiet stream's one ready record was never yielded across three " + "whole activation budgets; the busy stream spent every one of them" + ) + assert len(seen) <= 3 * MAX_RECORDS_PER_ACTIVATION + 1, ( + "the budget must still bound the merged set as a whole" + ) + + +# --- fairness across activations ---------------------------------------------- + +#: The budget the two fairness tests below run against, in place of the real +#: 256. What they prove is a relationship between the *subscription count* and +#: the budget -- more ready waits than the budget can pay for, or a count that +#: does not divide it -- and the real constant only sets the scale at which that +#: relationship bites. Standing up 257 subscriptions and handing out 768 records +#: to watch it costs more than the property is worth, and destabilised the +#: end-to-end tests that share the run. +_SMALL_BUDGET = 8 + + +@pytest.fixture +def small_budget(monkeypatch: pytest.MonkeyPatch) -> int: + """Shrinks the delivery budget for one test, and returns what it is now. + + Patched on ``_runtime`` rather than on ``_api``, which is where the constant + is defined: ``_runtime`` binds the name at import, so it reads its own + module global, and a patch on the defining module would leave every check + that actually enforces the budget looking at 256. + """ + monkeypatch.setattr(_runtime_module, "MAX_RECORDS_PER_ACTIVATION", _SMALL_BUDGET) + return _SMALL_BUDGET + + +#: A subscription count that does not divide the budget, at the shrunk scale: 8 +#: records over 3 always-ready waits is two whole passes and two waits of a +#: third, so the pass is always cut in the middle -- which is the arrangement a +#: start position fixed at the lowest wait id turns into permanent, growing +#: unfairness. The real-scale case is 100 waits against 256 records: two whole +#: passes and 56 waits of a third, the same cut in the same place. +_SKEWED_STREAMS = 3 + +#: How many activations the skew is watched over. Whatever the scale, a pass +#: cut in the middle hands the favoured waits exactly one record more per +#: activation than the rest, so the gap is the activation count and nothing +#: bounds it: five is five times the tolerance rather than a coin toss. +_SKEW_CYCLES = 5 + + +async def run_budget_cycles( + runtime: WorkflowStreamRuntime, + task: asyncio.Task[None], + seen: list[Any], + cycles: int, + budget: int, +) -> None: + """Runs `cycles` whole activations, doing what the Worker's completion does. + + A fresh budget and readiness re-armed for the records the budget left + buffered -- which is what makes the next pass a *new activation* rather + than a continuation of this one, and so is the only place a start position + that never moves can show itself. + + `budget` is passed rather than read from the constant because these callers + run against a patched one; a never-empty set delivers exactly that many + records per activation, which is how each cycle knows it has finished. + """ + for cycle in range(1, cycles + 1): + await settle(task, seen, cycle * budget) + runtime.begin_activation() + runtime.resolve_all_pending() + + +@pytest.mark.asyncio +async def test_merge_reaches_a_wait_the_budget_cannot_pay_for_in_one_pass( + workflow_instance: FakeInstance, + backend: MemoryStreamBackend, + small_budget: int, +) -> None: + """More ready subscriptions than one activation has records to hand out. + + The budget is charged per record handed to Workflow code, once per + subscription per pass, so a pass over *budget + 1* ready waits runs out + exactly one short. A pass that began at the lowest wait id on every + activation ran out in the same place on every one of them: the last wait was + never yielded, and -- because ``_fill`` returns on the spent budget before it + reaches the manager -- never even drained. The Worker was not asked about + that stream at all, for the life of the Run. + + Run against a budget of ``_SMALL_BUDGET`` rather than the real 256, because + "one more ready wait than the budget can pay for" is the whole property and + the constant only says how many that is: at full scale the starved case is + 257 subscriptions, and it starves for exactly the reason nine do here. + """ + manager = NeverEmptyManager(await encoded("x")) + runtime = make_runtime(manager, backend) + _install_runtime(workflow_instance, runtime) + runtime.begin_activation() + # The shrunk budget has to have reached the binding the runtime reads, or + # the counts below stand in a relationship to 256 that they were not chosen + # for and the test passes without exercising anything. + assert runtime.delivery_budget_remaining() == small_budget + subscriptions = [ + external_stream.topic(f"s{i}", type=str).subscribe() + for i in range(small_budget + 1) + ] + last = subscriptions[-1] + + seen: list[int] = [] + + async def consume_merged() -> None: + async for subscription, _ in merge(*subscriptions): + seen.append(subscription.wait_id) + + task = asyncio.ensure_future(consume_merged()) + await run_budget_cycles(runtime, task, seen, 3, small_budget) + await cancel(task) + + assert last.wait_id in set(seen), ( + "the wait one past the budget was never yielded across three whole " + "activations; every one of them spent itself on the same prefix" + ) + assert last.wait_id in manager.drained_waits, ( + "the starved wait was never even drained -- the spent budget stopped " + "`_fill` before it reached the manager, so the Worker was never asked " + "whether that stream had anything at all" + ) + assert len(seen) <= 3 * small_budget, ( + "the budget must still bound the merged set as a whole" + ) + + +@pytest.mark.asyncio +async def test_merge_keeps_the_skew_bounded_when_the_count_splits_the_budget( + workflow_instance: FakeInstance, + backend: MemoryStreamBackend, + small_budget: int, +) -> None: + """The general defect, of which total starvation is only the extreme. + + Nothing has to be starved outright for the merge to be unfair: a count that + does not divide the budget is enough. At full scale, 256 records over 100 + always-ready waits leaves the pass cut at the 56th, and a pass that restarted + at the lowest wait id was cut there on every activation -- the first 56 + taking three records per activation and the other 44 taking two, forever. The + gap between the best- and worst-served stream then grows by one per + activation and nothing bounds it: measured at 8 after eight activations, and + it would be 800 after eight hundred. + + Run here at ``_SMALL_BUDGET`` over ``_SKEWED_STREAMS`` waits, which is the + same arrangement in miniature -- two whole passes and a cut in the middle of + a third -- because what makes the skew unbounded is the remainder, not its + size. The gap still grows by exactly one per activation, so + ``_SKEW_CYCLES`` activations separate a merge that resumes where the last + pass stopped from one that does not. + + Taking at most one record per turn bounds the skew *within* a pass. Only + resuming after the wait that last took one bounds it *across* passes, which + is what the single-record bound was always claimed to be. + """ + manager = NeverEmptyManager(await encoded("x")) + runtime = make_runtime(manager, backend) + _install_runtime(workflow_instance, runtime) + runtime.begin_activation() + # The shrunk budget has to have reached the binding the runtime reads, or + # the counts below stand in a relationship to 256 that they were not chosen + # for and the test passes without exercising anything. + assert runtime.delivery_budget_remaining() == small_budget + subscriptions = [ + external_stream.topic(f"s{i}", type=str).subscribe() + for i in range(_SKEWED_STREAMS) + ] + + seen: list[int] = [] + + async def consume_merged() -> None: + async for subscription, _ in merge(*subscriptions): + seen.append(subscription.wait_id) + + task = asyncio.ensure_future(consume_merged()) + await run_budget_cycles(runtime, task, seen, _SKEW_CYCLES, small_budget) + await cancel(task) + + per_wait = Counter(seen) + counts = [per_wait[s.wait_id] for s in subscriptions] + assert max(counts) - min(counts) <= 1, ( + f"after {_SKEW_CYCLES} activations the best-served stream had " + f"{max(counts)} records and the worst-served {min(counts)}; the skew " + "grows by one per activation and is bounded by nothing, so a long-" + "running merge reads its later streams arbitrarily far behind its " + "earlier ones" + ) + assert min(counts) > 0 + assert len(seen) <= _SKEW_CYCLES * small_budget, ( + "the budget must still bound the merged set as a whole" + ) + + +# --- re-arming readiness ------------------------------------------------------ + + +class RecordingNotifier: + """Stands in for Core's readiness call, recording every notification.""" + + def __init__(self) -> None: + self.calls: list[tuple[str, int, int]] = [] + self.notified = asyncio.Event() + + async def __call__(self, run_id: str, wait_id: int, wait_generation: int) -> str: + self.calls.append((run_id, wait_id, wait_generation)) + self.notified.set() + return ReadinessResult.ACCEPTED + + +@pytest.mark.asyncio +async def test_rearming_re_reports_readiness_for_a_non_empty_buffer() -> None: + """Without this the Workflow waits forever on records already in front of it. + + Readiness is reported once, when the watcher buffers a record, and that + watcher has long since moved its prefetch cursor past whatever the budget + left behind. Nothing else will ever announce those records. + """ + stream_key = StreamKey("ns", "wf", uuid.uuid4().hex, "tokens") + backend = MemoryStreamBackend() + notifier = RecordingNotifier() + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=notifier, + watch_block=timedelta(milliseconds=20), + ) + try: + subscription = manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await backend.append( + stream_key, StreamRecord(RecordKind.DATA, b"a", "session", 0) + ) + await asyncio.wait_for(notifier.notified.wait(), 2) + # The activation drained nothing: the budget was already spent elsewhere. + assert subscription.buffered == 1 + first_round = len(notifier.calls) + + manager.rearm_ready(RUN_ID) + await asyncio.sleep(0.1) + + assert len(notifier.calls) > first_round, ( + "the records the budget left buffered got no second notification, so " + "the Workflow would block on data that is already at the Worker" + ) + assert notifier.calls[-1] == (RUN_ID, 1, 0) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_rearming_says_nothing_about_an_empty_buffer() -> None: + """Readiness means *buffered*, and re-arming may not weaken that. + + A notification for an empty buffer produces an activation whose drain finds + nothing, which is the spurious Workflow Task the "only after buffered" rule + exists to prevent. + """ + stream_key = StreamKey("ns", "wf", uuid.uuid4().hex, "tokens") + backend = MemoryStreamBackend() + notifier = RecordingNotifier() + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=notifier, + watch_block=timedelta(milliseconds=20), + ) + try: + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await asyncio.sleep(0.05) + assert notifier.calls == [] + + manager.rearm_ready(RUN_ID) + await asyncio.sleep(0.1) + + assert notifier.calls == [] + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_a_record_arriving_after_the_last_drain_is_re_announced() -> None: + """The stall this closes, driven through the real manager and completion path. + + A record that arrives after an activation's final drain is announced to Core + at the wait generation that was current when the watcher buffered it. The + Workflow then re-blocks on the way out of that activation, which bumps the + generation -- and Core keeps readiness across a quiescent snapshot only where + the generations match, because a bumped one ordinarily means the Workflow + drained and re-blocked and the readiness refers to a block already resolved. + Here it means the opposite, and Core cannot tell the difference: the record was + never seen. + + Nothing else would announce it. The watcher reports only after a *new* + non-empty read and its prefetch cursor is already past this record; the idle + timer would park a Workflow Task whose data is sitting at the Worker. So the + completion has to re-announce, and it has to do so whether or not a delivery + budget was involved -- which is what this asserts. + """ + stream_key = StreamKey("ns", "wf", uuid.uuid4().hex, "tokens") + backend = MemoryStreamBackend() + notifier = RecordingNotifier() + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=notifier, + watch_block=timedelta(milliseconds=20), + ) + try: + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=stream_key) + # A wait is blocked from creation, so reaching generation 1 means taking a + # record and blocking again -- which is the ordinary shape and the one the + # generation counter exists for. + runtime.note_blocked(1, False) + runtime.note_blocked(1, True) + + await backend.append( + stream_key, StreamRecord(RecordKind.DATA, await encoded("late"), "s", 0) + ) + await asyncio.wait_for(notifier.notified.wait(), 5) + announced_at = list(notifier.calls) + assert announced_at and announced_at[-1][2] == 1, ( + f"the watcher did not announce under generation 1: {announced_at}" + ) + subscription = manager.subscription(RUN_ID, 1) + assert subscription is not None + assert subscription.buffered == 1, "the record is not buffered at the Worker" + + # This activation delivered nothing -- the record landed after its last + # drain -- and the Workflow re-blocks on the way out. That is the + # generation bump Core reads as "already resolved". + runtime.note_blocked(1, False) + runtime.note_blocked(1, True) + assert subscription.current_wait_generation() == 2 + + before = len(notifier.calls) + _WorkflowInstanceImpl._emit_external_stream_commands( + cast(Any, _CompletionStub(runtime)) + ) + await asyncio.sleep(0.2) + + assert len(notifier.calls) > before, ( + "the completion announced nothing, so the record Core dropped the " + "readiness for is announced by nobody and the Workflow waits on data " + "it is already holding" + ) + assert notifier.calls[-1] == (RUN_ID, 1, 2), ( + "the re-announcement has to name the *current* generation, or Core " + f"discards it for the same reason: {notifier.calls[-1]}" + ) + finally: + await manager.shutdown() + + +class _CompletionStub: + """Just enough of the Workflow instance to drive the completion path. + + The real method is reached unbound, so what is under test is the code that + ships rather than a copy of it. + """ + + def __init__(self, runtime: Any, *, replaying: bool = False) -> None: + self._external_stream_runtime = runtime + self._deleting = False + self._is_replaying = replaying + import temporalio.bridge.proto.workflow_completion + + self._current_completion = ( + temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion() + ) + self._current_completion.successful.SetInParent() + + def _add_command(self) -> Any: + return self._current_completion.successful.commands.add() + + # Reached unbound for the same reason `_emit_external_stream_commands` is: + # the completion path calls these, and a stub that reimplemented them would + # let the shipped versions drift. Neither reads the runtime unless the + # completion carries a Continue-As-New command, which none of these do. + def _refresh_external_stream_continuations(self) -> None: + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + + _WorkflowInstanceImpl._refresh_external_stream_continuations(self) # type: ignore[arg-type] + + def _attach_external_stream_continuation(self, command: Any) -> None: + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + + _WorkflowInstanceImpl._attach_external_stream_continuation(self, command) # type: ignore[arg-type] + + +class _RecordingRuntime: + """Answers the two questions the completion path asks, and records the call.""" + + def __init__(self, *, exhausted: bool) -> None: + self._exhausted = exhausted + self.rearms = 0 + + def delivery_budget_exhausted(self) -> bool: + return self._exhausted + + def rearm_readiness(self) -> None: + self.rearms += 1 + + def take_observation_delta(self) -> bytes | None: + return None + + def quiescent_snapshot(self) -> None: + return None + + def start_new_annotation(self) -> None: + pass + + @property + def annotation_started(self) -> bool: + # Nothing was ever accumulated, so there is no annotation to terminate. + return False + + @property + def request_rollover(self) -> bool: + return False + + +@pytest.mark.parametrize("exhausted", [True, False]) +def test_every_completion_rearms_readiness_not_only_a_budget_stop( + exhausted: bool, +) -> None: + """The window must be closed on **every** completion, budget or not. + + It has to happen on the completion path: the waits involved are marked + blocked, so they enter the quiescent snapshot, and a snapshot is what lets + Core start the idle timer and eventually park. A Workflow Task parked with + records in the local buffer would wait out an idle timeout for data that had + already arrived. + + Gating it on the budget was a real stall, not a tidiness point, because the + budget is not the only way an activation ends with a full buffer. Core keeps + readiness accepted during an open activation **only if the quiescent snapshot + that follows reports the same wait generation** -- same generation means the + Workflow never saw the record, a bumped one means it drained and re-blocked. + So a record arriving after this activation's last drain is accepted at + generation G, the wait re-blocks to G+1 on the way out, and the snapshot + legitimately drops a readiness that now refers to a resolved block. The + watcher will not report it again either: it only reports after a *new* + non-empty read, and its prefetch cursor is already past that record. Nothing + else announces it, so the Workflow waits on data it is already holding. + + The re-arm is cheap where it is not needed -- the manager skips every + subscription whose buffer is empty, which is the ordinary case. + """ + runtime = _RecordingRuntime(exhausted=exhausted) + stub = _CompletionStub(runtime) + + _WorkflowInstanceImpl._emit_external_stream_commands(stub) # type: ignore[arg-type] + + assert runtime.rearms == 1, ( + "a completion that did not hit the delivery budget skipped the re-arm, " + "which is the only thing that re-announces a record Core dropped the " + "readiness for" + ) + + +def test_the_completion_rearms_even_when_it_emits_nothing_else() -> None: + """The early returns below it must not swallow the re-arm. + + A completion carrying the Workflow's own commands returns before asking for + retention. If the re-arm sat behind that return, a Workflow with records still + buffered while also starting a timer would never hear about them again. + """ + runtime = _RecordingRuntime(exhausted=True) + stub = _CompletionStub(runtime) + stub._add_command() # something server-bound rides along + + _WorkflowInstanceImpl._emit_external_stream_commands(stub) # type: ignore[arg-type] + + assert runtime.rearms == 1 + + +def test_a_terminal_completion_does_not_rearm_buffered_readiness() -> None: + """A Workflow that is ending needs no wake for records it leaves behind. + + Re-arming after the terminal command has already been produced reports the + remaining buffer after Core closes the Workflow Task. ``NoOpenWorkflowTask`` + then sends a fresh Signal into the execution's closing window, where it can + contend with the terminal report and force another replay of the same + completion. Complete, fail, cancel, and Continue-As-New all end this Run; + none can consume a later activation, so none has a liveness reason to wake. + """ + runtime = _RecordingRuntime(exhausted=False) + stub = _CompletionStub(runtime) + stub._add_command().complete_workflow_execution.SetInParent() + + _WorkflowInstanceImpl._emit_external_stream_commands(stub) # type: ignore[arg-type] + + assert runtime.rearms == 0 + + +def test_a_wait_the_budget_stopped_is_not_immediately_parkable( + backend: MemoryStreamBackend, +) -> None: + """Blocked is not quiescent when the budget is what stopped it. + + ``immediately_parkable`` invites Core to park without waiting out the idle + timer. A subscription the budget stopped still has records at the Worker, so + parking it would strand them. + """ + + class StubManager: + def cancel_from_workflow_thread(self, run_id, wait_id): # type: ignore[no-untyped-def] + pass + + def register( + self, + **kwargs: Any, # pyright: ignore[reportUnusedParameter] + ) -> None: + pass + + def note_wait_generation(self, run_id, wait_id, generation) -> None: # type: ignore[no-untyped-def] + pass + + runtime = make_runtime(StubManager(), backend) + runtime.begin_activation() + runtime.register(wait_id=1, stream_key=runtime.stream_key("tokens")) + record = StreamRecord(RecordKind.WRITE_FENCE, b"", "s", 0).placed_at(Offset("1")) + runtime.record_delivery(1, record) + runtime.note_blocked(1, True) + + fenced = runtime.quiescent_snapshot() + assert fenced is not None and fenced[0].immediately_parkable + + # Spent through `record_delivery`, which is where the budget is charged: a + # record moving into a subscription's ready list is the reservation, and the + # consumption that follows it costs nothing. + for _ in range(MAX_RECORDS_PER_ACTIVATION): + runtime.record_delivery(1, record) + + stopped = runtime.quiescent_snapshot() + assert stopped is not None and not stopped[0].immediately_parkable + + +# --- the double-check refill still works -------------------------------------- + + +class LateArrivalManager: + """Empty on the first look, holding a record on the second. + + Exactly the window the double-check inside ``_await_readiness`` exists for: + the record landed after ``_iterate`` looked and before the wait was + registered, so its readiness was reported to nobody and none is coming. + """ + + def __init__(self, record: StreamRecord) -> None: + self.record: StreamRecord | None = record + self.calls = 0 + + def register( + self, + **kwargs: Any, # pyright: ignore[reportUnusedParameter] + ) -> None: + pass + + def note_wait_generation(self, run_id, wait_id, generation) -> None: # type: ignore[no-untyped-def] + pass + + def drain( + self, + run_id: str, # pyright: ignore[reportUnusedParameter] + wait_id: int, # pyright: ignore[reportUnusedParameter] + max_records: int | None = None, # pyright: ignore[reportUnusedParameter] + ) -> list[StreamRecord]: + self.calls += 1 + if self.calls == 1 or self.record is None: + return [] + record, self.record = self.record, None + return [record] + + +@pytest.mark.asyncio +async def test_the_double_check_refill_still_delivers_a_late_record( + workflow_instance: FakeInstance, backend: MemoryStreamBackend +) -> None: + """Budgeting the double-check must not disable it. + + It fixes a separate stall: a record buffered while Workflow code was + elsewhere has already had its readiness reported and consumed, so blocking + without looking again strands the Workflow on a record sitting in front of + it. + """ + record = StreamRecord(RecordKind.DATA, await encoded("late"), "s", 0).placed_at( + Offset("1") + ) + manager = LateArrivalManager(record) + runtime = make_runtime(manager, backend) + _install_runtime(workflow_instance, runtime) + runtime.begin_activation() + subscription = external_stream.topic("tokens", type=str).subscribe() + + iterator = subscription.__aiter__() + + assert await asyncio.wait_for(iterator.__anext__(), 1) == "late" + assert manager.calls >= 2, ( + "the record must have been found by the look taken *after* the wait was " + "registered, which is the window that double-check closes" + ) + + +# --- replay is unaffected ----------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_replay_segment_larger_than_the_budget_is_delivered_in_full( + workflow_instance: FakeInstance, backend: MemoryStreamBackend +) -> None: + """The recorded boundaries already say how many records each activation got. + + Re-cutting them at the budget would deliver a different schedule than the + live run recorded, which is divergence rather than protection: replay reads + from a finite recorded segment, not from a producer that can outrun it. + """ + payload = await encoded("r") + count = MAX_RECORDS_PER_ACTIVATION + 44 + + class StubManager: + def cancel_from_workflow_thread(self, run_id, wait_id): # type: ignore[no-untyped-def] + pass + + def register( + self, + **kwargs: Any, # pyright: ignore[reportUnusedParameter] + ) -> None: + pass + + def note_wait_generation(self, run_id, wait_id, generation) -> None: # type: ignore[no-untyped-def] + pass + + def drain(self, run_id: str, wait_id: int, max_records: int | None = None): # type: ignore[no-untyped-def] + return [] + + runtime = make_runtime(StubManager(), backend) + _install_runtime(workflow_instance, runtime) + runtime.begin_activation() + subscription = external_stream.topic("tokens", type=str).subscribe() + runtime.begin_replay_segment( + [ + ( + subscription.wait_id, + StreamRecord(RecordKind.DATA, payload, "s", i).placed_at( + Offset(f"{i:08d}") + ), + ) + for i in range(count) + ] + ) + + seen: list[str] = [] + task = asyncio.ensure_future(consume(subscription.__aiter__(), seen)) + await settle(task, seen, count) + + assert len(seen) == count, ( + "the recorded segment was cut short by the live delivery budget; replay " + "would then diverge from the run it is replaying" + ) + + # Back on the live buffer, which is where the budget applies again. + runtime.end_replay() + assert runtime.delivery_budget_remaining() == MAX_RECORDS_PER_ACTIVATION, ( + "replayed records must not spend the live budget either -- a later live " + "delivery in the same activation would start with none left, and the " + "completion would re-arm readiness for a purely replayed Workflow Task" + ) + assert not runtime.delivery_budget_exhausted() + await cancel(task) + + +@pytest.mark.asyncio +async def test_replay_delivers_in_full_even_when_the_live_budget_is_spent( + workflow_instance: FakeInstance, backend: MemoryStreamBackend +) -> None: + """The bypass has to be the replay state, not an empty counter. + + A budget that merely happened to be untouched would work by accident. What + must hold is that a recorded segment is delivered whole no matter what the + live counter says, because its boundaries are already in History and cutting + them somewhere else is divergence. + """ + payload = await encoded("r") + count = MAX_RECORDS_PER_ACTIVATION + 44 + + class StubManager: + def cancel_from_workflow_thread(self, run_id, wait_id): # type: ignore[no-untyped-def] + pass + + def register( + self, + **kwargs: Any, # pyright: ignore[reportUnusedParameter] + ) -> None: + pass + + def note_wait_generation(self, run_id, wait_id, generation) -> None: # type: ignore[no-untyped-def] + pass + + def drain(self, run_id: str, wait_id: int, max_records: int | None = None): # type: ignore[no-untyped-def] + return [] + + runtime = make_runtime(StubManager(), backend) + _install_runtime(workflow_instance, runtime) + runtime.begin_activation() + subscription = external_stream.topic("tokens", type=str).subscribe() + + # Spend the whole live budget first, through the delivery that charges it. + spent = StreamRecord(RecordKind.DATA, payload, "s", 0).placed_at(Offset("live")) + for _ in range(MAX_RECORDS_PER_ACTIVATION): + runtime.record_delivery(subscription.wait_id, spent) + runtime.record_consumption(subscription.wait_id, spent) + assert runtime.delivery_budget_remaining() == 0 + + runtime.begin_replay_segment( + [ + ( + subscription.wait_id, + StreamRecord(RecordKind.DATA, payload, "s", i).placed_at( + Offset(f"{i:08d}") + ), + ) + for i in range(count) + ] + ) + + seen: list[str] = [] + task = asyncio.ensure_future(consume(subscription.__aiter__(), seen)) + await settle(task, seen, count) + + assert len(seen) == count, ( + "a spent live budget stopped a recorded segment, so replay delivered " + "fewer records than the marker says the live run received" + ) + await cancel(task) + + +# --- the recorded reason ------------------------------------------------------ + + +def _annotation_reasons(runtime: WorkflowStreamRuntime) -> list[SegmentEndReason]: + delta = runtime.take_observation_delta() + assert delta is not None + annotation = decode_annotation(delta + runtime.add_terminal()) + return [segment.end_reason for segment in annotation.segments] + + +def test_a_budget_stop_records_batch_limit(backend: MemoryStreamBackend) -> None: + """``NO_DATA_AVAILABLE`` would be a false statement, and a durable one. + + The end reason is the only thing in the annotation that says why an + activation stopped. Recording "the stream ran dry" for an activation that + stopped with records still buffered puts a claim in History that was never + true. + """ + + class StubManager: + def cancel_from_workflow_thread(self, run_id, wait_id): # type: ignore[no-untyped-def] + pass + + def register( + self, + **kwargs: Any, # pyright: ignore[reportUnusedParameter] + ) -> None: + pass + + def note_wait_generation(self, run_id, wait_id, generation) -> None: # type: ignore[no-untyped-def] + pass + + runtime = make_runtime(StubManager(), backend) + runtime.begin_activation() + runtime.register(wait_id=1, stream_key=runtime.stream_key("tokens")) + for i in range(MAX_RECORDS_PER_ACTIVATION): + record = StreamRecord(RecordKind.DATA, b"x", "s", i).placed_at( + Offset(f"{i:08d}") + ) + runtime.record_delivery(1, record) + runtime.record_consumption(1, record) + + assert _annotation_reasons(runtime)[0] is SegmentEndReason.BATCH_LIMIT + + +def test_an_activation_that_ran_out_of_records_still_records_no_data( + backend: MemoryStreamBackend, +) -> None: + """The other direction, without which "always BATCH_LIMIT" would pass.""" + + class StubManager: + def cancel_from_workflow_thread(self, run_id, wait_id): # type: ignore[no-untyped-def] + pass + + def register( + self, + **kwargs: Any, # pyright: ignore[reportUnusedParameter] + ) -> None: + pass + + def note_wait_generation(self, run_id, wait_id, generation) -> None: # type: ignore[no-untyped-def] + pass + + runtime = make_runtime(StubManager(), backend) + runtime.begin_activation() + runtime.register(wait_id=1, stream_key=runtime.stream_key("tokens")) + for i in range(5): + record = StreamRecord(RecordKind.DATA, b"x", "s", i).placed_at( + Offset(f"{i:08d}") + ) + runtime.record_delivery(1, record) + runtime.record_consumption(1, record) + + assert _annotation_reasons(runtime)[0] is SegmentEndReason.NO_DATA_AVAILABLE diff --git a/tests/contrib/external_workflow_streams/test_errors.py b/tests/contrib/external_workflow_streams/test_errors.py new file mode 100644 index 0000000000..6368b0fd22 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_errors.py @@ -0,0 +1,197 @@ +"""P18 — the failure taxonomy, its metrics, and the classification rule.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import pytest + +import temporalio.common +from temporalio.contrib.external_workflow_streams._errors import ( + METRIC_DECODE, + METRIC_INTEGRITY, + METRIC_SHUTDOWN_WAKE_FAILED, + METRIC_STORAGE, + StreamDecodeError, + StreamError, + StreamIntegrityError, + StreamMetrics, + StreamStorageError, + classify_read_failure, +) + + +@dataclass +class RecordingCounter(temporalio.common.MetricCounter): + counter_name: str + added: list[int] = field(default_factory=list) + + @property + def name(self) -> str: + return self.counter_name + + @property + def description(self) -> str | None: + return None + + @property + def unit(self) -> str | None: + return None + + def add( + self, + value: int, + additional_attributes: temporalio.common.MetricAttributes | None = None, + ) -> None: + self.added.append(value) + + def with_additional_attributes( + self, additional_attributes: temporalio.common.MetricAttributes + ) -> RecordingCounter: + return self + + +class RecordingMeter(temporalio.common.MetricMeter): + def __init__(self) -> None: + self.counters: dict[str, RecordingCounter] = {} + + def create_counter( + self, name: str, description: str | None = None, unit: str | None = None + ) -> temporalio.common.MetricCounter: + counter = RecordingCounter(name) + self.counters[name] = counter + return counter + + def create_histogram(self, name, description=None, unit=None): # type: ignore[no-untyped-def] + raise NotImplementedError + + def create_histogram_float(self, name, description=None, unit=None): # type: ignore[no-untyped-def] + raise NotImplementedError + + def create_histogram_timedelta(self, name, description=None, unit=None): # type: ignore[no-untyped-def] + raise NotImplementedError + + def create_gauge(self, name, description=None, unit=None): # type: ignore[no-untyped-def] + raise NotImplementedError + + def create_gauge_float(self, name, description=None, unit=None): # type: ignore[no-untyped-def] + raise NotImplementedError + + def with_additional_attributes(self, additional_attributes): # type: ignore[no-untyped-def] + return self + + +# --- the types -------------------------------------------------------------- + + +def test_the_two_types_are_distinct_and_neither_subclasses_the_other() -> None: + """Collapsing them would send an operator to repair an undamaged backend.""" + assert issubclass(StreamIntegrityError, StreamError) + assert issubclass(StreamDecodeError, StreamError) + assert not issubclass(StreamIntegrityError, StreamDecodeError) + assert not issubclass(StreamDecodeError, StreamIntegrityError) + + +# --- the classification rule, in both directions ---------------------------- + + +def test_a_failure_after_the_range_validated_is_a_decode_failure() -> None: + cause = ValueError("bad JSON") + + err = classify_read_failure(range_validated=True, cause=cause) + + assert isinstance(err, StreamDecodeError) + assert err.__cause__ is cause + + +def test_a_failure_before_the_range_validated_is_integrity_loss() -> None: + cause = LookupError("offset 100-0 is missing") + + err = classify_read_failure(range_validated=False, cause=cause) + + assert isinstance(err, StreamIntegrityError) + assert err.__cause__ is cause + + +def test_the_same_cause_classifies_differently_by_validation_alone() -> None: + """The rule is mechanical: only `range_validated` decides which row it is.""" + cause = ValueError("undecodable") + + assert isinstance( + classify_read_failure(range_validated=True, cause=cause), StreamDecodeError + ) + assert isinstance( + classify_read_failure(range_validated=False, cause=cause), StreamIntegrityError + ) + + +def test_a_storage_failure_is_never_reclassified() -> None: + """There was no range to read, so validation has nothing to say about it.""" + cause = StreamStorageError("connection refused") + + for validated in (True, False): + assert classify_read_failure(range_validated=validated, cause=cause) is cause + + +# --- metrics ---------------------------------------------------------------- + + +def test_each_class_gets_its_own_counter() -> None: + meter = RecordingMeter() + metrics = StreamMetrics.create(meter) + + assert set(meter.counters) == { + METRIC_INTEGRITY, + METRIC_DECODE, + METRIC_STORAGE, + METRIC_SHUTDOWN_WAKE_FAILED, + } + # Separate series, not one counter with a `reason` attribute -- otherwise an + # integrity alert is diluted by a backend outage sharing the series. + assert len({id(c) for c in meter.counters.values()}) == 4 + assert metrics.integrity is not metrics.decode + + +def test_recording_increments_only_the_matching_counter() -> None: + meter = RecordingMeter() + metrics = StreamMetrics.create(meter) + + metrics.record(StreamIntegrityError("gone")) + metrics.record(StreamDecodeError("garbled")) + metrics.record(StreamDecodeError("garbled again")) + metrics.record(StreamStorageError("timeout")) + + assert meter.counters[METRIC_INTEGRITY].added == [1] + assert meter.counters[METRIC_DECODE].added == [1, 1] + assert meter.counters[METRIC_STORAGE].added == [1] + assert meter.counters[METRIC_SHUTDOWN_WAKE_FAILED].added == [] + + +def test_an_unrelated_error_increments_nothing() -> None: + meter = RecordingMeter() + metrics = StreamMetrics.create(meter) + + metrics.record(RuntimeError("something else entirely")) + + assert all(not c.added for c in meter.counters.values()) + + +# --- no terminal-failure opt-in --------------------------------------------- + + +@pytest.mark.parametrize( + "err_type", [StreamIntegrityError, StreamDecodeError, StreamStorageError] +) +def test_no_stream_error_is_a_temporal_failure(err_type: type[StreamError]) -> None: + """Integrity loss blocks a Workflow; it does not terminate one (ADR-014). + + A blocked Workflow can be resumed once the backend is repaired, and a + failed one cannot. These stay plain exceptions rather than + ``FailureError``s so that the SDK's failure conversion cannot turn one into + a Workflow *execution* failure -- which is what a + ``workflow_failure_exception_types`` registration would do. + """ + import temporalio.exceptions + + assert not issubclass(err_type, temporalio.exceptions.FailureError) + assert not issubclass(err_type, temporalio.exceptions.ApplicationError) diff --git a/tests/contrib/external_workflow_streams/test_m1_gate.py b/tests/contrib/external_workflow_streams/test_m1_gate.py new file mode 100644 index 0000000000..78ee8854a1 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_m1_gate.py @@ -0,0 +1,409 @@ +"""P16a/P16b — the milestone gates. + +The required-test lists are what "the milestone is met" means. This module makes +them enforceable rather than aspirational: each case in the plan is mapped to the +test that covers it, and the mapping is checked against both the plan and the +test suite. + +A case that no test fully covers yet maps to its concrete missing boundary in +the corresponding gaps table. That is deliberately louder than leaving it out: +an unmapped case would pass by omission, while an open one is counted and +reported every time the gate runs. +""" + +from __future__ import annotations + +import pytest + +from tests.contrib.external_workflow_streams.m1_gate import ( + declared_count, + required_cases, + unresolved, +) + +#: Case number -> the test(s) that fully cover it. A case may need +#: more than one test: several of them name two independent properties, and +#: mapping such a case to whichever half happened to be written first would +#: claim coverage the suite does not have. +#: Numbers follow `tests-m1.md` in document order. +M1_COVERAGE: dict[int, str | tuple[str, ...]] = { + 1: "external_streams.rs::a_normal_completion_commits_its_marker", + 2: ( + "external_streams.rs::an_activity_command_writes_its_marker_ordered_before_it", + "external_streams.rs::replaying_a_marker_before_an_activity_delivers_its_record_exactly_once", + ), + 3: ( + "external_streams.rs::a_failed_workflow_writes_its_marker_ordered_before_the_failure", + "external_streams.rs::a_continue_as_new_writes_its_marker_ordered_before_it", + ), + 4: "external_streams.rs::several_progress_reports_collapse_into_one_marker", + 5: "external_streams.rs::every_completion_path_writes_exactly_one_marker_ending_in_a_terminal", + 6: "test_runtime.py::test_a_first_subscription_to_an_empty_stream_still_emits", + 7: "test_runtime.py::test_an_activation_that_drained_nothing_emits_an_empty_segment", + 8: ( + "external_streams.rs::a_budget_driven_split_writes_two_markers_rather_than_one_oversized_one", + "test_replay.py::test_two_markers_reassemble_in_workflow_task_order", + ), + 9: "test_annotation.py::test_a_large_single_stream_batch_encodes_as_one_run", + 10: ( + "test_annotation.py::test_encoded_size_stays_flat_with_sparse_control_records", + "test_annotation.py::test_encoded_size_grows_only_with_the_number_of_control_records", + ), + 11: "external_streams.rs::readiness_before_the_idle_timer_expires_cancels_it", + 12: "external_streams.rs::a_confirmed_idle_park_writes_one_marker_and_completes_the_task", + 13: "test_runtime_only_jobs.py::test_a_recheck_that_finds_records_abandons_the_whole_park", + 14: ( + "test_wake.py::test_a_producer_that_finds_a_wakeable_generation_sends_exactly_one_signal", + "test_runtime.py::test_readiness_after_a_re_block_names_the_current_generation", + ), + 15: "external_streams.rs::a_stale_generation_produces_no_activation", + 16: "external_streams.rs::an_all_fenced_snapshot_parks_without_waiting_out_the_idle_timeout", + 17: "external_streams.rs::a_core_decided_boundary_asks_for_a_terminal_before_writing_anything", + 18: "external_streams.rs::a_core_decided_boundary_asks_for_a_terminal_before_writing_anything", + 19: "test_rollover_integration.py::test_a_continuously_fed_stream_survives_a_rollover", + 20: "external_streams.rs::the_rollover_deadline_fires_on_a_workflow_only_worker", + 21: "test_rollover_integration.py::test_a_signal_into_a_retained_task_lands_by_the_rollover_deadline", + 22: ( + "test_worker_integration.py::test_an_append_with_no_open_task_wakes_the_workflow", + "test_api.py::test_a_record_buffered_while_not_iterating_is_still_delivered", + ), + 23: "test_rollover_integration.py::test_an_append_after_a_rollover_completion_wakes_the_subscription", + 24: "test_manager.py::test_undeliverable_readiness_owes_a_wake_and_keeps_the_right_watchers", + 25: "external_streams.rs::an_unknown_envelope_version_is_ignored_harmlessly", + 26: ( + "test_wake.py::test_two_producers_retrying_one_wake_derive_the_identical_request_id", + "test_wake.py::test_two_producers_retrying_one_wake_are_deduplicated_by_the_server", + ), + 27: ( + "test_wake.py::test_two_unparked_wakes_from_different_senders_differ", + "test_wake.py::test_two_workers_unparked_wakes_are_both_delivered", + "test_wake.py::test_one_senders_retry_stays_a_single_wake", + ), + 28: "external_streams.rs::a_second_worker_reconstructs_the_subscription_from_the_shutdown_marker", + 29: "test_worker_handoff.py::test_shutdown_with_no_open_task_hands_the_run_to_another_worker", + 30: "test_worker_handoff.py::test_a_finalization_that_cannot_be_answered_writes_no_marker", + 31: "external_streams.rs::an_unwritten_annotation_exists_only_while_a_workflow_task_is_open", + 32: ( + "test_shutdown_sweep.py::test_a_momentarily_failing_wake_is_retried_and_succeeds", + "test_shutdown_sweep.py::test_the_retry_is_bounded_and_then_reported", + "test_shutdown_sweep.py::test_the_retry_is_the_same_wake_not_a_second_one", + "test_shutdown_sweep.py::test_an_unacknowledged_wake_surfaces_on_the_metric", + ), + 33: "test_runtime_only_jobs.py::test_finalization_touches_no_provider_at_all", + 34: "test_replay.py::test_replay_performs_no_live_waiting", + 35: "test_replay_end_to_end.py::test_replaying_a_stream_history_reproduces_the_same_observations", + 36: ( + "test_replay_end_to_end.py::test_an_empty_stream_parked_and_evicted_replays_from_the_recorded_cursor", + "test_replay_end_to_end.py::test_an_empty_stream_replays_from_its_recorded_boundary", + ), + 37: "test_replay.py::test_an_unreachable_backend_fails_as_transient_storage", + 38: "test_replay.py::test_each_deletion_position_fails_a_different_check", + 39: "test_replay.py::test_an_intact_but_undecodable_record_is_a_decode_error", + 40: "test_replay.py::test_a_marker_naming_an_unknown_wait_is_nondeterminism_not_integrity_loss", + 41: "test_worker_crash.py::test_a_crash_before_the_marker_makes_the_next_worker_re_read", + 42: "test_continuation.py::test_a_first_execution_starts_at_the_beginning", + 43: "external_streams.rs::a_pending_timer_suppresses_retention_but_its_subscriptions_survive_it", + 44: "test_backend_conformance.py::test_a_broken_backend_fails_for_the_right_reason", + 45: "test_backend_conformance.py::test_a_backend_needing_a_nameable_cursor_fails_the_tail_check", + 46: "test_redis_backend.py::test_offsets_compare_numerically_not_lexically", + 47: "test_producer.py::test_republishing_different_content_under_one_key_is_an_error", + 48: "test_manager.py::test_a_backend_slower_than_the_deadlock_timeout_delays_readiness", + 49: "test_runtime_only_jobs.py::test_a_park_slower_than_the_deadlock_timeout_is_still_answered", + 50: "test_manager.py::test_the_provider_is_never_called_from_the_workflow_thread", + 51: "test_manager.py::test_a_full_buffer_stops_prefetch_without_dropping_or_blocking", + 52: "test_manager.py::test_an_evicted_run_re_delivers_the_records_it_had_already_seen", + 53: "test_shutdown_sweep.py::test_teardown_removes_every_run", + 54: "test_registry.py::test_a_backend_without_the_guarantee_fails_worker_construction", + 55: ( + "test_redis_backend.py::test_a_trimmed_range_is_integrity_loss_not_a_storage_failure", + "test_redis_backend.py::test_a_deleted_write_fence_is_integrity_loss", + ), + # Cases 56-62 -- "Boundaries an implementation review found unprotected". + # These are mapped here before the vendored copy of the list carries them: + # the gate parses the plan out of `temporalio/bridge/sdk-core`, so it counts + # 55 cases until the submodule pointer moves to the commit that adds them. + # Mapping ahead of that is deliberate -- an unmapped case is what the gate is + # designed to fail on, so the map has to arrive no later than the list does. + 56: "test_replay.py::test_the_first_live_drain_after_a_replay_needs_no_loop_turn", + 57: "test_replay.py::test_a_marker_drains_once_per_segment_including_the_activations_own", + 58: "test_runtime.py::test_the_segment_that_crosses_the_mark_asks_for_rollover_itself", + 59: ( + "test_runtime.py::test_a_frame_larger_than_the_slack_rolls_over_instead_of_raising", + "test_runtime.py::test_a_subscription_set_too_large_to_record_is_refused_at_subscribe", + ), + 60: ( + "test_wake.py::test_a_coordination_failure_after_the_append_is_still_unacknowledged", + "test_wake.py::test_a_coordination_failure_after_a_fence_is_still_unacknowledged", + ), + 61: ( + "test_shutdown_sweep.py::test_a_grace_period_expiry_counts_every_wake_it_abandons", + "test_shutdown_sweep.py::test_a_probe_that_cannot_answer_is_not_reported_as_nothing_owed", + "test_shutdown_sweep.py::test_a_hanging_probe_counts_the_runs_it_never_answered_for", + "test_shutdown_sweep.py::test_a_run_with_nothing_owed_is_not_counted_as_a_failure", + ), + 62: "test_replay.py::test_an_unreachable_payload_store_is_a_storage_failure_not_a_decode_one", + # 63 and 64 came out of reviewing the fixes for 56-62 rather than out of the + # review itself: the first fix for the byte budget stopped delivery without + # obliging the same completion to ask for the rollover, and the first fix for + # the sweep counted subscriptions on a manager that had no sweep to perform. + 63: "test_runtime.py::test_an_activation_the_annotation_budget_stopped_is_not_wedged_by_it", + 64: "test_shutdown_sweep.py::test_a_manager_with_no_probe_owes_nothing_and_reports_nothing", + 65: "test_replay.py::test_a_marker_with_no_segments_closes_before_the_activations_own_drain", + # 66-69 came out of an independent review of the fixes for 56-62, which found + # four more defects in this round's own code. The shape is the one worth + # remembering: each was correct for the case it was aimed at. + 66: "test_runtime.py::test_the_first_record_of_an_activation_is_priced_from_a_measurement", + 67: "test_runtime.py::test_a_run_too_large_for_any_annotation_says_so_and_fails_the_workflow", + 68: "test_runtime.py::test_the_capacity_floor_covers_everything_an_empty_annotation_carries", + 69: ( + "test_replay.py::test_a_failing_activation_reports_its_own_error_not_the_replays", + "test_shutdown_sweep.py::test_a_wake_the_live_path_delivered_is_not_counted_as_abandoned", + ), + # 70-71 pin both halves of replay-to-live repositioning: a read already in + # flight cannot refill a retracted buffer, and replay delivers no record the + # marker omitted. + 70: "test_replay.py::test_a_read_in_flight_across_a_reposition_cannot_be_appended", + 71: "test_replay.py::test_a_replay_activation_never_delivers_a_record_the_marker_omits", + # 72-76 are the fourth review's five findings. Four of them are the same shape + # as the third round's: a guarantee written for the path it was aimed at and + # not for the adjacent one -- a budget spent where it was not reserved, a key + # drawn after an await that reorders it, a handler for `Exception` on a path + # cancellation also takes, and a retry whose answer was discarded. + 72: ( + "test_delivery_budget.py::test_two_independent_consumers_share_one_budget", + "test_delivery_budget.py::test_a_carried_over_ready_list_is_charged_to_the_next_activation", + ), + 73: ( + "test_producer.py::test_concurrent_publishes_take_their_sequence_in_invocation_order", + "test_producer.py::test_reordered_encodes_cannot_duplicate_across_two_topics", + ), + 74: ( + "test_wake.py::test_cancellation_after_the_append_is_an_unacknowledged_wake", + "test_wake.py::test_cancellation_after_a_fence_is_an_unacknowledged_wake", + "test_wake.py::test_cancellation_before_the_append_stays_a_cancellation", + ), + 75: ( + "test_api.py::test_a_second_coroutine_waiting_on_one_subscription_is_refused", + "test_api.py::test_iterating_again_after_the_first_consumer_stopped_is_allowed", + "test_api.py::test_a_merge_cannot_take_a_wait_another_consumer_is_blocked_on", + ), + 76: "test_manager.py::test_a_stale_retry_that_finds_the_run_gone_tears_the_watcher_down", + # 77 came out of reviewing case 74's fix, and is the same shape once more: a + # recovery written for the boundary it was aimed at -- cancellation after + # `append()` returned -- and not for the one just before it, where a remote + # backend has already committed and the answer is what was lost. + 77: ( + "test_wake.py::test_cancellation_after_backend_commit_before_append_ack_is_recoverable", + "test_wake.py::test_a_lost_append_response_is_recoverable_the_same_way", + "test_wake.py::test_an_unresolved_append_refuses_the_publish_that_would_duplicate_it", + "test_wake.py::test_settling_an_unacknowledged_append_wakes_exactly_once", + "test_wake.py::test_an_unacknowledged_fence_settles_to_exactly_one_fence", + "test_wake.py::test_settling_an_append_that_never_landed_appends_it_once", + "test_wake.py::test_a_conflicting_append_is_a_refusal_not_an_unknown_outcome", + ), + # 78 is case 77's own review. The outcome was right and its recovery was + # under-bound: it carried the record but not the operation, and validated + # neither the stream the record belongs to nor the producer instance that + # still holds the session's counters. + 78: ( + "test_wake.py::test_a_refused_append_preserves_the_unresolved_operations_recovery", + "test_wake.py::test_an_unknown_append_can_only_be_resolved_on_its_originating_topic", + "test_wake.py::test_recovery_is_bound_to_the_producer_instance_that_made_the_append", + ), + # 79 follows the same operation through a second unknown outcome. The + # recovery attempt's effective wake and lease supersede the original + # publish's, while cancellation delivered on either attempt remains owed. + 79: "test_wake.py::test_a_reinterrupted_append_recovery_preserves_its_latest_wake", + # 80-101 promote Workflow-originated output. Only cases with every named + # boundary asserted belong here; partial unit coverage stays in M1_GAPS. + 80: "test_output_worker_integration.py::test_reader_waits_at_pending_barrier_until_marker_commit", + 81: "test_output_worker_integration.py::test_rejected_post_stage_completion_never_exposes_phantom_output", + 82: "external_streams.rs::speculative_output_redelivery_accepts_a_fresh_token_once", + 83: "test_output_client.py::test_first_task_boundary_without_token_aborts_pending_stage", + 84: "external_streams.rs::exact_output_floor_excludes_the_previous_workflow_task_close", + 85: ( + "test_worker_crash.py::test_disconnected_update_leaves_live_staged_output_undecided_until_timeout", + "test_output_client.py::test_disconnected_speculative_update_stays_pending_until_history_loss", + ), + 86: "test_worker_crash.py::test_real_crash_after_output_stage_aborts_old_token_and_exposes_one_retry", + 87: "test_output_worker_integration.py::test_cold_client_repairs_post_report_commit_failure", + 88: ( + "test_redis_output.py::test_exact_stage_retry_returns_actual_status_after_terminal", + "test_output_client.py::test_repeating_the_same_history_reconciliation_is_idempotent", + ), + 89: "test_redis_output.py::test_reusing_stage_identity_with_another_manifest_conflicts_atomically", + 90: "test_output_runtime.py::test_output_replay_performs_no_io_token_mint_or_live_policy_split", + 91: "test_output_runtime.py::test_replay_keeps_recorded_segments_when_codec_wire_size_changes", + 92: "test_output_worker_integration.py::test_output_latency_flushes_a_retained_workflow_task", + 93: "test_output_worker_integration.py::test_three_retained_output_windows_write_three_markers_and_wfts", + 94: "test_redis_output.py::test_pending_stage_blocks_later_direct_output_until_commit", + 95: "test_output_producer.py::test_ambiguous_append_retries_exact_record_without_duplicate", + 96: "test_output_runtime.py::test_changed_codec_bytes_keep_logical_retry_identity_and_first_bytes", + 97: "test_output_codec.py::test_metadata_insertion_order_does_not_change_frame_or_fingerprint", + 98: "test_output_worker_integration.py::test_concurrent_updates_preserve_their_own_turn_ids", + 99: "test_output_worker_integration.py::test_stage_outage_blocks_wft_and_reports_external_storage_cause", + 100: ( + "test_redis_output.py::test_missing_staged_record_is_an_integrity_failure", + "test_output_client.py::test_history_loss_is_integrity_but_an_outage_is_storage", + ), + 101: ( + "test_output_runtime.py::test_oversized_marker_manifest_is_rejected_before_external_io", + "test_output_runtime.py::test_one_oversized_logical_record_is_refused_without_poisoning_batch", + ), + # 102 restores the public default to the required contract after the design + # documents were split into guide, specification, and rationale sets. + 102: "test_api.py::test_the_default_idle_timeout_is_one_second", +} + +#: Case number -> what is still missing, for cases the suite does not fully +#: cover. A case is in exactly one of the two maps: claiming a partially covered +#: case as covered is how a gate stops meaning anything. +M1_GAPS: dict[int, str] = {} + +#: The same, for `tests-m2.md`. +M2_COVERAGE: dict[int, str | tuple[str, ...]] = { + 1: "test_m2_required.py::test_readiness_on_one_of_several_streams_resets_global_quiescence", + 2: "test_multi_stream.py::test_one_idle_stream_cannot_park_while_another_is_active", + 3: "test_multi_stream.py::test_a_fence_on_one_stream_alone_does_not_make_the_set_parkable", + 4: "test_multi_stream.py::test_all_fenced_streams_make_the_whole_set_parkable", + 5: ( + "test_multi_stream.py::test_an_alternating_two_stream_batch_encodes_one_run_per_delivery", + "test_m2_required.py::test_an_alternating_two_stream_batch_rolls_over_within_budget", + ), + 6: "test_m2_required.py::test_simultaneously_ready_streams_are_drained_in_one_pass", + 7: "test_runtime.py::test_alternating_streams_produce_one_run_per_delivery", + 8: ( + "test_multi_stream.py::test_each_same_stream_subscription_receives_every_record", + "test_continuation.py::test_two_same_stream_subscriptions_restore_independently", + ), + 9: ( + "test_multi_stream.py::test_two_same_stream_subscriptions_install_distinct_park_intents", + "test_m2_required.py::test_two_same_stream_subscriptions_never_overwrite_each_other", + ), + 10: ( + "test_multi_stream.py::test_differing_idle_timeouts_reduce_to_the_minimum", + "test_m2_required.py::test_the_idle_timeout_reduction_reproduces_exactly", + ), + 11: ( + "test_m2_required.py::test_a_wake_for_one_stream_resolves_every_blocked_wait", + "test_m2_required.py::test_resolving_twice_is_harmless", + ), + 12: "test_continuation.py::test_a_chain_resumes_where_its_predecessor_stopped", + 13: "test_replay_end_to_end.py::test_live_and_replay_share_one_input_output_drain_schedule", + 14: "test_output_worker_integration.py::test_cursor_resume_crosses_rollover_and_continue_as_new", + 15: "test_redis_output.py::test_input_and_output_with_same_user_identity_are_physically_isolated", + 16: ( + "test_output_worker_integration.py::test_output_deadline_wins_park_race_and_removes_every_intent", + "external_streams.rs::confirmed_park_flushes_output_in_its_marker_without_forcing_a_task", + "external_streams.rs::output_deadline_invalidates_an_issued_park_and_flushes_once", + ), + 17: "test_output_worker_integration.py::test_finished_output_survives_continue_as_new_without_backend_reads", +} + +#: The same, for `tests-m2.md`. +M2_GAPS: dict[int, str] = {} + + +def _coverage(list_name: str) -> dict[int, str | tuple[str, ...]]: + return M1_COVERAGE if list_name == "tests-m1.md" else M2_COVERAGE + + +def _gaps(list_name: str) -> dict[int, str]: + return M1_GAPS if list_name == "tests-m1.md" else M2_GAPS + + +def _node_ids(coverage: dict[int, str | tuple[str, ...]]) -> list[str]: + """Flattens the map, since a case may name more than one test.""" + flat: list[str] = [] + for value in coverage.values(): + flat.extend([value] if isinstance(value, str) else list(value)) + return flat + + +@pytest.mark.parametrize("list_name", ["tests-m1.md", "tests-m2.md"]) +def test_the_parsed_case_count_matches_the_plans_own_heading(list_name: str) -> None: + """Catches a case added to the plan without the heading being updated. + + The two numbers are maintained by hand in the same document, so they are + exactly the kind of pair that drifts. If they disagree, every count below is + measuring the wrong thing. + """ + cases = required_cases(list_name) + + assert len(cases) == declared_count(list_name), ( + f"{list_name} declares {declared_count(list_name)} cases but contains " + f"{len(cases)}" + ) + + +@pytest.mark.parametrize("list_name", ["tests-m1.md", "tests-m2.md"]) +def test_every_required_case_is_accounted_for(list_name: str) -> None: + """A new case must fail the gate rather than pass by not being mentioned.""" + cases = required_cases(list_name) + coverage = _coverage(list_name) + + gaps = _gaps(list_name) + unmapped = [c for c in cases if c.number not in coverage and c.number not in gaps] + + assert not unmapped, ( + "these required cases appear in neither map -- add the test that covers " + "each, or record what is still missing in the gaps map:\n" + + "\n".join(f" {c.number}. {c.text}" for c in unmapped) + ) + + +@pytest.mark.parametrize("list_name", ["tests-m1.md", "tests-m2.md"]) +def test_every_mapped_test_exists(list_name: str) -> None: + """Catches a rename or deletion that would leave the gate pointing at nothing. + + This is the failure the whole mechanism exists for: without it the map keeps + claiming coverage that was deleted, and the milestone stays "met" on paper. + """ + coverage = _coverage(list_name) + + missing = unresolved(_node_ids(coverage)) + + assert not missing, "the coverage map names tests that do not exist:\n" + "\n".join( + f" {node_id}" for node_id in missing + ) + + +@pytest.mark.parametrize("list_name", ["tests-m1.md", "tests-m2.md"]) +def test_no_case_is_both_covered_and_open(list_name: str) -> None: + """A case in both maps would be counted twice and read as covered. + + Claiming a partially covered case as covered is how a gate stops meaning + anything, so the two maps must partition the list rather than overlap. + """ + overlap = sorted(set(_coverage(list_name)) & set(_gaps(list_name))) + + assert not overlap, f"cases claimed as both covered and open: {overlap}" + + +@pytest.mark.parametrize("list_name", ["tests-m1.md", "tests-m2.md"]) +def test_every_open_case_says_what_is_missing(list_name: str) -> None: + """An empty reason is indistinguishable from "nobody looked".""" + empty = sorted(n for n, reason in _gaps(list_name).items() if not reason.strip()) + + assert not empty, f"open cases with no stated gap: {empty}" + + +def test_the_gate_reports_how_far_each_milestone_is_from_met() -> None: + """The gate's actual output, and the thing a human reads. + + Coverage and open cases are reported independently of the consistency + assertions above. If a future requirement is added before its test, its gap + remains visible without turning this reporting test into the failure itself. + """ + for list_name in ("tests-m1.md", "tests-m2.md"): + cases = {c.number: c.text for c in required_cases(list_name)} + gaps = _gaps(list_name) + covered = len(cases) - len(gaps) + blocked = {n for n, why in gaps.items() if why.startswith("BLOCKED")} + print( + f"\n{list_name}: {covered}/{len(cases)} covered, {len(gaps)} open " + f"({len(blocked)} of them blocked on a deliverable)" + ) + for number in sorted(gaps): + print(f" {number}. {cases[number][:70]}") + print(f" {gaps[number]}") diff --git a/tests/contrib/external_workflow_streams/test_m2_required.py b/tests/contrib/external_workflow_streams/test_m2_required.py new file mode 100644 index 0000000000..f5cf37c4a8 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_m2_required.py @@ -0,0 +1,341 @@ +"""P16b — the Milestone 2 cases not already covered elsewhere. + +The rest of the twelve are covered by the deliverables that produced them (P21, +P15, P10b); these are the ones that fall between deliverables, which is exactly +why a milestone gate is a separate list rather than a sum of "done when"s. +""" + +# pyright: reportMissingParameterType=false +from __future__ import annotations + +import asyncio +import uuid +from datetime import timedelta + +import pytest +import pytest_asyncio + +import temporalio.converter +from temporalio.contrib.external_workflow_streams._annotation import ( + MAX_ANNOTATION_BYTES, +) +from temporalio.contrib.external_workflow_streams._backend import ParkIntent, StreamKey +from temporalio.contrib.external_workflow_streams._continuation import Continuation +from temporalio.contrib.external_workflow_streams._manager import ( + ReadinessResult, + StreamSubscriptionManager, +) +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Offset, + RecordKind, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._runtime import WorkflowStreamRuntime +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend + +RUN_ID = "run-1" + + +async def _notify( + run_id: str, # pyright: ignore[reportUnusedParameter] + wait_id: int, # pyright: ignore[reportUnusedParameter] + generation: int, # pyright: ignore[reportUnusedParameter] +) -> str: + return ReadinessResult.ACCEPTED + + +@pytest.fixture +def backend() -> MemoryStreamBackend: + return MemoryStreamBackend() + + +@pytest_asyncio.fixture # type: ignore[reportUntypedFunctionDecorator] +async def manager(backend: MemoryStreamBackend): + mgr = StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=10), + ) + yield mgr + await mgr.shutdown() + + +def make_runtime(manager, backend, continuation=None): # type: ignore[no-untyped-def] + return WorkflowStreamRuntime( + manager=manager, + backend=backend, + run_id=RUN_ID, + namespace="ns", + workflow_id="wf", + first_execution_run_id=uuid.uuid4().hex, + data_converter=temporalio.converter.DataConverter.default, + default_idle_timeout=timedelta(seconds=1), + continuation=continuation, + ) + + +def data(offset: str) -> StreamRecord: + return StreamRecord(RecordKind.DATA, b"x", "s", 0).placed_at(Offset(offset)) + + +# --- case 1: readiness on one stream resets global quiescence ----------------- + + +@pytest.mark.asyncio +async def test_readiness_on_one_of_several_streams_resets_global_quiescence( + manager, backend: MemoryStreamBackend +) -> None: + """The idle timer covers the set, so any member's record restarts it. + + Core starts a fresh idle timer when the generations it is retaining for + change. If a record on one stream left every generation untouched, the timer + that was already running would keep counting down toward parking a Workflow + that had just been given work. + """ + runtime = make_runtime(manager, backend) + for wait_id, name in ((1, "a"), (2, "b")): + runtime.register( + wait_id=wait_id, + stream_key=runtime.stream_key(name), + ) + runtime.note_blocked(wait_id, True) + before = {w.wait_id: w.generation for w in runtime.quiescent_snapshot()} # type: ignore[union-attr] + + # A record arrives on the first stream: it stops being blocked, Workflow code + # consumes it, and then blocks again. + runtime.note_blocked(1, False) + runtime.record_delivery(1, data("10-0")) + runtime.note_blocked(1, True) + + after = {w.wait_id: w.generation for w in runtime.quiescent_snapshot()} # type: ignore[union-attr] + assert after[1] > before[1], ( + "the wait that received a record must re-enter the blocked state with a " + "new generation, which is what makes Core restart the idle timer" + ) + assert after[2] == before[2], ( + "the other stream did not move and must keep its generation, or Core " + "could not tell a stale readiness for it from a live one" + ) + + +# --- case 5: an alternating batch rolls over rather than exceeding the budget -- + + +@pytest.mark.asyncio +async def test_an_alternating_two_stream_batch_rolls_over_within_budget( + manager, backend: MemoryStreamBackend +) -> None: + """Alternation is the worst case for annotation size, and must still be bounded. + + Every delivery starts a new run, so an alternating batch grows the annotation + per *record* rather than per batch -- the one shape where the byte budget is + reachable in normal use. It must ask for rollover rather than grow past the + limit, because an oversized marker cannot be written at all. + """ + runtime = make_runtime(manager, backend) + for wait_id, name in ((1, "a"), (2, "b")): + runtime.register( + wait_id=wait_id, + stream_key=runtime.stream_key(name), + ) + + delivered = 0 + while not runtime.request_rollover and delivered < 20_000: + wait_id = 1 + (delivered % 2) + runtime.record_delivery(wait_id, data(f"{delivered}-0")) + runtime.take_observation_delta() + delivered += 1 + + assert runtime.request_rollover, ( + f"an alternating batch of {delivered} records never asked for rollover" + ) + annotation = b"".join( + [runtime.take_observation_delta() or b"", runtime.add_terminal()] + ) + assert len(annotation) <= MAX_ANNOTATION_BYTES, ( + f"the annotation reached {len(annotation)} bytes, past the " + f"{MAX_ANNOTATION_BYTES} budget: rollover was requested too late to be " + "of any use" + ) + + +# --- case 6: simultaneously ready streams are observable together ------------- + + +@pytest.mark.asyncio +async def test_simultaneously_ready_streams_are_drained_in_one_pass( + manager, backend: MemoryStreamBackend +) -> None: + """Two records that arrived together cost one activation, not two. + + Waiting for readiness between them would turn every simultaneous arrival + into a second Workflow Task -- and the whole point of retention is that a + batch costs one. + """ + from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec + + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + first = StreamKey("ns", "wf", "chain", "a") + second = StreamKey("ns", "wf", "chain", "b") + for wait_id, key in ((1, first), (2, second)): + manager.register( + run_id=RUN_ID, + wait_id=wait_id, + stream_key=key, + start_cursor=BEGINNING, + ) + await backend.append( + first, StreamRecord(RecordKind.DATA, await codec.encode("a1"), "p", 0) + ) + await backend.append( + second, StreamRecord(RecordKind.DATA, await codec.encode("b1"), "p", 0) + ) + await asyncio.sleep(0.2) + + # Both are buffered at the same moment, so one drain pass sees both. + assert manager.drain(RUN_ID, 1) and manager.drain(RUN_ID, 2), ( + "both streams must be buffered together; if one had to wait for a " + "separate readiness round trip the batch would cost two activations" + ) + + +# --- case 9: two same-stream subscriptions stay independent throughout -------- + + +@pytest.mark.asyncio +async def test_two_same_stream_subscriptions_never_overwrite_each_other( + manager, backend: MemoryStreamBackend +) -> None: + """Park, wake, cancel, and Continue-As-New, verified against the backend. + + Every step is a chance for one subscription's state to land under the + other's key. The park intents are read back rather than assumed, because a + collapsed intent looks like success from the caller's side -- the second + install simply returns. + """ + key = StreamKey("ns", "wf", "chain", "tokens") + for wait_id, cursor in ((1, BEGINNING), (2, AFTER(Offset("50-0")))): + await backend.install_park_intent( + key, + ParkIntent( + wait_id=wait_id, cursor=cursor, park_generation=3, run_id="run-1" + ), + ) + + # Parked: two intents, two cursors. + first = await backend.park_intent(key, 1) + second = await backend.park_intent(key, 2) + assert first is not None and second is not None + assert first.cursor != second.cursor + assert sorted(await backend.parked_wait_ids(key)) == [1, 2] + + # Cancelled: removing one leaves the other exactly as it was. + await backend.remove_park_intent(key, 1) + assert await backend.park_intent(key, 1) is None + assert await backend.park_intent(key, 2) == second + assert await backend.parked_wait_ids(key) == [2] + + # Continued as new: each restores its own cursor, not the other's. + runtime = make_runtime( + manager, + backend, + Continuation( + {1: AFTER(Offset("10-0")), 2: AFTER(Offset("90-0"))}, + {1: "tokens", 2: "tokens"}, + {1: "memory", 2: "memory"}, + {1: 1, 2: 1}, + ), + ) + runtime.register(wait_id=1, stream_key=key) + runtime.register(wait_id=2, stream_key=key) + assert runtime._subscriptions[1].start_cursor == AFTER(Offset("10-0")) + assert runtime._subscriptions[2].start_cursor == AFTER(Offset("90-0")) + + +# --- case 10: the reduction reproduces on replay ------------------------------ + + +@pytest.mark.asyncio +async def test_the_idle_timeout_reduction_reproduces_exactly( + manager, backend: MemoryStreamBackend +) -> None: + """It rides on a command replay compares against History. + + The inputs are the quiescent set's configured values and nothing else, which + is what makes the result a pure function of state replay reconstructs. A + reduction that consulted wall-clock time, arrival order, or anything else + live would produce a different `WorkflowStreamQuiescent` on replay and fail + as nondeterminism. + """ + + def reduce_over(order: list[int]) -> timedelta: + runtime = make_runtime(manager, backend) + for wait_id in order: + runtime.register( + wait_id=wait_id, + stream_key=runtime.stream_key(f"s{wait_id}"), + idle_timeout=timedelta(seconds=wait_id), + ) + runtime.note_blocked(wait_id, True) + return runtime.effective_idle_timeout() + + # Registration order differs; the set does not. Replay reconstructs the set, + # so only the set may matter. + assert reduce_over([1, 2, 3]) == reduce_over([3, 2, 1]) == timedelta(seconds=1) + + +# --- case 11: a wake for one stream rechecks all ------------------------------ + + +@pytest.mark.asyncio +async def test_a_wake_for_one_stream_resolves_every_blocked_wait( + manager, backend: MemoryStreamBackend +) -> None: + """The Signal names one stream; the runtime rechecks all of them. + + A wake can only name the stream whose producer sent it, but by the time the + Workflow Task exists other streams may have records too. Resolving only the + named wait would leave those buffered until something else happened to wake + the Workflow -- and the wake it needed had already been spent. + """ + runtime = make_runtime(manager, backend) + loop = asyncio.get_running_loop() + futures = {} + for wait_id, name in ((1, "named"), (2, "other")): + runtime.register( + wait_id=wait_id, + stream_key=runtime.stream_key(name), + ) + runtime.note_blocked(wait_id, True) + future: asyncio.Future[None] = loop.create_future() + runtime.register_pending(wait_id, future) + futures[wait_id] = future + + # What the resolve activation does, regardless of which wait it names. + runtime.resolve_all_pending() + + assert all(f.done() for f in futures.values()), ( + "every blocked wait must be resolved, not only the one the Signal named" + ) + + +@pytest.mark.asyncio +async def test_resolving_twice_is_harmless( + manager, backend: MemoryStreamBackend +) -> None: + """Duplicate wake Signals are expected, and must not raise. + + Two producers racing one generation both signal; the server deduplicates the + request but a stale Signal can still arrive alongside a live one. + """ + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=runtime.stream_key("a")) + future: asyncio.Future[None] = asyncio.get_running_loop().create_future() + runtime.register_pending(1, future) + + runtime.resolve_all_pending() + runtime.resolve_all_pending() + + assert future.done() diff --git a/tests/contrib/external_workflow_streams/test_manager.py b/tests/contrib/external_workflow_streams/test_manager.py new file mode 100644 index 0000000000..e3f4b35948 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_manager.py @@ -0,0 +1,2241 @@ +"""P8 — the subscription manager, driven directly. + +No Workflow API and no Core activations: neither is in this deliverable's +closure, and both would obscure what is actually under test -- that backend I/O +never reaches the thread ``_apply`` runs on, and that readiness means *buffered*. +""" + +from __future__ import annotations + +import asyncio +import inspect +import threading +import time +import uuid +from collections.abc import Awaitable, Callable +from datetime import timedelta + +import pytest + +from temporalio.contrib.external_workflow_streams._backend import ( + ParkIntent, + ParkIntentRemoval, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._errors import ( + StreamStorageError, +) +from temporalio.contrib.external_workflow_streams._manager import ( + PARK_REMOVAL_ATTEMPTS, + ReadinessResult, + StreamSubscriptionManager, + Subscription, +) +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Cursor, + Offset, + RecordKind, + StreamRecord, +) +from temporalio.worker.workflow_sandbox._restrictions import SandboxRestrictions +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend + +RUN_ID = "run-1" + + +@pytest.fixture +def stream_key() -> StreamKey: + return StreamKey("ns", "wf", uuid.uuid4().hex, "tokens") + + +class RecordingNotifier: + """Stands in for Core's readiness call, recording every notification.""" + + def __init__(self, answer: str = ReadinessResult.ACCEPTED) -> None: + self.answer = answer + self.calls: list[tuple[str, int, int]] = [] + self.notified = asyncio.Event() + + async def __call__(self, run_id: str, wait_id: int, wait_generation: int) -> str: + self.calls.append((run_id, wait_id, wait_generation)) + self.notified.set() + return self.answer + + +def make_manager( + backend: MemoryStreamBackend, + notifier: RecordingNotifier, + *, + buffer_size: int = 256, + watch_block: timedelta = timedelta(milliseconds=20), + send_wake: object | None = None, +) -> StreamSubscriptionManager: + return StreamSubscriptionManager( + backend=backend, + notify_ready=notifier, + send_wake=send_wake, # type: ignore[arg-type] + buffer_size=buffer_size, + watch_block=watch_block, + ) + + +async def append( + backend: MemoryStreamBackend, key: StreamKey, *payloads: bytes +) -> None: + for i, payload in enumerate(payloads): + await backend.append( + key, StreamRecord(RecordKind.DATA, payload, f"s{payload.decode()}", i) + ) + + +async def until( + condition: Callable[[], object | Awaitable[object]], + message: str, + timeout: float = 2.0, +) -> None: + """Waits for something a background task is on its way to doing. + + Retries, reconciliations and watchers all run on the manager's own loop, so + a single ``sleep`` long enough to be reliable is also long enough to make + every one of these tests slow. Accepts an awaitable condition because most + of what is being waited for is a question for the backend. + """ + deadline = time.monotonic() + timeout + while True: + outcome = condition() + if inspect.isawaitable(outcome): + outcome = await outcome + if outcome: + return + if time.monotonic() >= deadline: + pytest.fail(message) + await asyncio.sleep(0.01) + + +# --- readiness means buffered ------------------------------------------------ + + +@pytest.mark.asyncio +async def test_readiness_is_reported_only_after_a_record_is_buffered( + stream_key: StreamKey, +) -> None: + """The rule the whole structure rests on. + + Readiness for an unbuffered record would produce an activation whose drain + must block -- exactly the deadlock hazard the out-of-thread buffer removes. + """ + backend = MemoryStreamBackend() + notifier = RecordingNotifier() + manager = make_manager(backend, notifier) + subscription = manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + try: + await asyncio.sleep(0.05) + assert notifier.calls == [], "nothing is buffered, so nothing is ready" + + await append(backend, stream_key, b"a") + await asyncio.wait_for(notifier.notified.wait(), 2) + + assert notifier.calls == [(RUN_ID, 1, 0)] + assert subscription.buffered == 1, ( + "the record must already be in the buffer when readiness is reported" + ) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_a_drain_returns_buffered_records_without_touching_the_backend( + stream_key: StreamKey, +) -> None: + backend = MemoryStreamBackend() + notifier = RecordingNotifier() + manager = make_manager(backend, notifier) + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + try: + await append(backend, stream_key, b"a", b"b") + await asyncio.wait_for(notifier.notified.wait(), 2) + reads_before = len(backend.range_reads) + + drained = manager.drain(RUN_ID, 1) + + assert [r.payload for r in drained] == [b"a", b"b"] + assert len(backend.range_reads) == reads_before + assert manager.drain(RUN_ID, 1) == [] + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_draining_advances_only_the_delivery_cursor( + stream_key: StreamKey, +) -> None: + """Consuming is not committing. Only a marker moves ``committed``.""" + backend = MemoryStreamBackend() + notifier = RecordingNotifier() + manager = make_manager(backend, notifier) + subscription = manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + try: + await append(backend, stream_key, b"a") + await asyncio.wait_for(notifier.notified.wait(), 2) + manager.drain(RUN_ID, 1) + + assert subscription.committed_cursor == BEGINNING + assert subscription.delivery_cursor != BEGINNING + assert subscription.prefetch_cursor != BEGINNING + finally: + await manager.shutdown() + + +# --- backend latency never reaches the Workflow thread ----------------------- + + +class SlowBackend(MemoryStreamBackend): + """Every read takes longer than the Workflow deadlock timeout.""" + + delay_seconds = 2.5 + + async def read_after(self, key, after, *, max_records, block=None): # type: ignore[no-untyped-def] + await asyncio.sleep(self.delay_seconds) + return await super().read_after( + key, after, max_records=max_records, block=block + ) + + +@pytest.mark.asyncio +async def test_a_backend_slower_than_the_deadlock_timeout_delays_readiness( + stream_key: StreamKey, +) -> None: + """It delays the *report*, not the caller. + + The Workflow thread runs under a 2-second deadlock timeout. A provider + slower than that must therefore be invisible to it -- and it is, because + the drain the Workflow thread performs only ever touches the buffer. + """ + backend = SlowBackend() + notifier = RecordingNotifier() + manager = make_manager(backend, notifier, watch_block=timedelta(milliseconds=20)) + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + try: + await append(backend, stream_key, b"a") + + # Well past the deadlock timeout: no readiness yet, because nothing is + # buffered yet. + await asyncio.sleep(2.2) + assert notifier.calls == [] + + # And the Workflow thread's own call returns immediately regardless. + start = asyncio.get_running_loop().time() + assert manager.drain(RUN_ID, 1) == [] + assert asyncio.get_running_loop().time() - start < 0.5, ( + "the drain must not wait on the provider" + ) + + await asyncio.wait_for(notifier.notified.wait(), 5) + assert manager.drain(RUN_ID, 1) != [] + finally: + await manager.shutdown() + + +# --- backpressure ------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_a_full_buffer_stops_prefetch_without_dropping_or_blocking( + stream_key: StreamKey, +) -> None: + """Backpressure *is* the buffer bound. Nothing is dropped, nobody blocks.""" + backend = MemoryStreamBackend() + notifier = RecordingNotifier() + manager = make_manager(backend, notifier, buffer_size=3) + subscription = manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + try: + await append(backend, stream_key, *[bytes([i]) for i in range(10)]) + await asyncio.sleep(0.2) + + assert subscription.buffered == 3, ( + f"prefetch must stop at the bound, buffered {subscription.buffered}" + ) + + # Nothing was dropped: draining frees room and the rest arrives, in order. + seen: list[bytes] = [] + for _ in range(10): + seen.extend(r.payload for r in manager.drain(RUN_ID, 1)) + if len(seen) == 10: + break + await asyncio.sleep(0.1) + + assert seen == [bytes([i]) for i in range(10)] + finally: + await manager.shutdown() + + +# --- the manager's loop owns every provider call ----------------------------- + + +class ThreadCheckingBackend(MemoryStreamBackend): + """Raises if called from any thread other than the one it was made on.""" + + def __init__(self) -> None: + super().__init__() + self.owning_thread = threading.get_ident() + self.foreign_calls: list[str] = [] + + def _check(self, name: str) -> None: + if threading.get_ident() != self.owning_thread: + self.foreign_calls.append(name) + raise AssertionError( + f"{name} was called from thread {threading.get_ident()}, not the " + f"manager's loop thread {self.owning_thread}" + ) + + async def read_after(self, key, after, *, max_records, block=None): # type: ignore[no-untyped-def] + self._check("read_after") + return await super().read_after( + key, after, max_records=max_records, block=block + ) + + async def read_range(self, key, first, last): # type: ignore[no-untyped-def] + self._check("read_range") + return await super().read_range(key, first, last) + + +@pytest.mark.asyncio +async def test_the_provider_is_never_called_from_the_workflow_thread( + stream_key: StreamKey, +) -> None: + """Driven from a real second thread, standing in for the executor. + + ``activate()`` runs on a thread-pool executor, so "the Workflow thread" is + genuinely a different OS thread -- not merely a different task. + """ + backend = ThreadCheckingBackend() + notifier = RecordingNotifier() + manager = make_manager(backend, notifier) + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + try: + await append(backend, stream_key, b"a", b"b") + await asyncio.wait_for(notifier.notified.wait(), 2) + + drained: list[bytes] = [] + + def workflow_thread() -> None: + for record in manager.drain(RUN_ID, 1): + drained.append(record.payload) + + thread = threading.Thread(target=workflow_thread) + thread.start() + thread.join(timeout=5) + + assert drained == [b"a", b"b"] + assert backend.foreign_calls == [], ( + f"the provider was called off the manager's loop: {backend.foreign_calls}" + ) + finally: + await manager.shutdown() + + +# --- eviction ---------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_eviction_discards_prefetch_state_and_restarts_from_committed( + stream_key: StreamKey, +) -> None: + """Reading is not consuming, and consuming is not committing. + + Nothing prefetched or delivered was ever a claim, which is exactly why "no + cursor advances unless the marker commits" is safe to state. + """ + backend = MemoryStreamBackend() + notifier = RecordingNotifier() + manager = make_manager(backend, notifier) + subscription = manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await append(backend, stream_key, b"a", b"b", b"c") + await asyncio.wait_for(notifier.notified.wait(), 2) + manager.drain(RUN_ID, 1, max_records=1) + + assert subscription.delivery_cursor != BEGINNING + assert subscription.prefetch_cursor != BEGINNING + + await manager.evict_run(RUN_ID) + + assert subscription.buffered == 0 + assert subscription.delivery_cursor == subscription.committed_cursor == BEGINNING + assert subscription.prefetch_cursor == BEGINNING + assert manager.subscriptions(RUN_ID) == [] + + +@pytest.mark.asyncio +async def test_a_committed_cursor_is_where_a_restart_resumes( + stream_key: StreamKey, +) -> None: + backend = MemoryStreamBackend() + notifier = RecordingNotifier() + manager = make_manager(backend, notifier) + subscription = manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + try: + await append(backend, stream_key, b"a", b"b") + await asyncio.wait_for(notifier.notified.wait(), 2) + drained = manager.drain(RUN_ID, 1) + assert drained[0].offset is not None + + # A marker commits the first record only. + subscription.commit(AFTER(drained[0].offset)) + subscription.reset_to_committed() + + assert subscription.delivery_cursor == AFTER(drained[0].offset) + assert subscription.prefetch_cursor == AFTER(drained[0].offset) + finally: + await manager.shutdown() + + +# --- what the watcher does with each readiness answer ------------------------ + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("answer", "keeps_watcher"), + [ + (ReadinessResult.PARKED, True), + (ReadinessResult.NO_OPEN_WORKFLOW_TASK, True), + (ReadinessResult.RUN_NOT_FOUND, False), + ], +) +async def test_undeliverable_readiness_owes_a_wake_and_keeps_the_right_watchers( + stream_key: StreamKey, answer: str, keeps_watcher: bool +) -> None: + """All three send a Signal; they differ in what happens afterwards. + + Tearing a watcher down for ``NoOpenWorkflowTask`` would remove the only + thing watching during the normal window between Workflow Tasks. + """ + backend = MemoryStreamBackend() + notifier = RecordingNotifier(answer) + woken: list[int] = [] + + async def send_wake(subscription) -> None: # type: ignore[no-untyped-def] + woken.append(subscription.wait_id) + + manager = make_manager(backend, notifier, send_wake=send_wake) + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + try: + await append(backend, stream_key, b"a") + await asyncio.wait_for(notifier.notified.wait(), 2) + await asyncio.sleep(0.1) + + assert woken == [1] + assert bool(manager.subscription(RUN_ID, 1)) is keeps_watcher + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("answer", [ReadinessResult.ACCEPTED, ReadinessResult.STALE]) +async def test_deliverable_readiness_owes_no_wake( + stream_key: StreamKey, answer: str +) -> None: + """A Signal sent here would be a wake nobody needed. + + ``Stale`` in particular must not signal: the wait moved on, so the right + response is to re-probe, not to manufacture a Workflow Task. + """ + backend = MemoryStreamBackend() + notifier = RecordingNotifier(answer) + woken: list[int] = [] + + async def send_wake(subscription) -> None: # type: ignore[no-untyped-def] + woken.append(subscription.wait_id) + + manager = make_manager(backend, notifier, send_wake=send_wake) + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + try: + await append(backend, stream_key, b"a") + await asyncio.wait_for(notifier.notified.wait(), 2) + await asyncio.sleep(0.1) + + assert woken == [] + assert manager.subscription(RUN_ID, 1) is not None + finally: + await manager.shutdown() + + +# --- the blocked snapshot ---------------------------------------------------- + + +@pytest.mark.asyncio +async def test_the_blocked_snapshot_is_read_from_manager_state_alone( + stream_key: StreamKey, +) -> None: + """Finalization performs no backend I/O (ADR-010). + + The boundary is not "wherever the stream is now"; it is where this Workflow + Task's deliveries stopped, which is already fixed. Refreshing it against the + backend could name a position replay must not reproduce. + """ + backend = ThreadCheckingBackend() + notifier = RecordingNotifier() + manager = make_manager(backend, notifier) + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + manager.register(run_id=RUN_ID, wait_id=2, stream_key=stream_key) + try: + await append(backend, stream_key, b"a") + await asyncio.wait_for(notifier.notified.wait(), 2) + drained = manager.drain(RUN_ID, 1) + reads_before = len(backend.range_reads) + + snapshot = manager.blocked_snapshot(RUN_ID) + + assert set(snapshot) == {1, 2} + assert drained[0].offset is not None + assert snapshot[1] == AFTER(drained[0].offset) + assert snapshot[2] == BEGINNING, ( + "a subscription that delivered nothing is blocked at its start cursor" + ) + assert len(backend.range_reads) == reads_before + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_a_start_cursor_is_where_a_continued_chain_resumes( + stream_key: StreamKey, +) -> None: + """All three cursors begin at the restored boundary, not at BEGINNING.""" + backend = MemoryStreamBackend() + notifier = RecordingNotifier() + manager = make_manager(backend, notifier) + restored: Cursor = AFTER(Offset("500-0")) + subscription = manager.register( + run_id=RUN_ID, + wait_id=1, + stream_key=stream_key, + start_cursor=restored, + ) + try: + assert subscription.committed_cursor == restored + assert subscription.delivery_cursor == restored + assert subscription.prefetch_cursor == restored + finally: + await manager.shutdown() + + +# --- teardown ---------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_cancelling_one_subscription_leaves_the_others( + stream_key: StreamKey, +) -> None: + backend = MemoryStreamBackend() + notifier = RecordingNotifier() + manager = make_manager(backend, notifier) + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + manager.register(run_id=RUN_ID, wait_id=2, stream_key=stream_key) + try: + await manager.cancel(RUN_ID, 1) + + assert manager.subscription(RUN_ID, 1) is None + assert manager.subscription(RUN_ID, 2) is not None + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_cancelling_a_wait_takes_back_its_park_intent( + stream_key: StreamKey, +) -> None: + """The intent has to go here, because nothing later can be asked to retry. + + The resolve path can afford to log a failure and try again, since the + subscription it is working on stays registered. This one drops it, so an + intent left behind is left behind for good -- and a stale intent is what + `current_park_generation` answers to every producer that asks. A producer + naming a generation Core has discarded sends a wake Core ignores as stale: + the record is appended, the Signal is sent, and the Workflow is never woken. + """ + backend = MemoryStreamBackend() + manager = make_manager(backend, RecordingNotifier()) + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + try: + assert not await manager.prepare_park(RUN_ID, 4, {1: BEGINNING}) + assert await backend.parked_wait_ids(stream_key) == [1] + + await manager.cancel(RUN_ID, 1) + + assert await backend.parked_wait_ids(stream_key) == [], ( + "the closed wait's intent is still in the backend, advertising a " + "park generation Core has discarded" + ) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_cancelling_a_wait_stops_its_watcher(stream_key: StreamKey) -> None: + """Dropping the subscription is what makes an orphaned watcher invisible. + + Once it is out of the Run's map nothing can reach it to notice, and it goes + on prefetching into a buffer nobody can drain -- holding a backend + connection for the rest of the Worker's life. + """ + backend = MemoryStreamBackend() + manager = make_manager(backend, RecordingNotifier()) + subscription = manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + try: + await asyncio.sleep(0.05) + assert subscription._watcher is not None and not subscription._watcher.done() + + await manager.cancel(RUN_ID, 1) + + assert subscription._cancelled, "the cancelled subscription must be marked dead" + assert subscription._watcher.done(), ( + "its watcher outlived it, so it is still reading a stream nothing " + "is reading back" + ) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_a_cancel_from_the_workflow_thread_reaches_a_loop_with_nothing_to_do( + stream_key: StreamKey, +) -> None: + """Driven from a real second thread, standing in for the executor. + + ``close()`` runs inside the synchronous ``activate()``, and ``create_task`` + from there does not raise: the task is appended to the loop's ready queue + and the loop is never woken, so it runs only if something else happens to + wake the loop. Between activations nothing does -- the watchers are sitting + in blocking reads -- which is why the watch block below is long and why the + other thread, not this coroutine, is what waits. Awaiting anything with a + timeout here would arm a timer, and the wake that timer produces would run + the cancel whether or not the loop was ever told about it. + """ + backend = MemoryStreamBackend() + manager = make_manager( + backend, RecordingNotifier(), watch_block=timedelta(seconds=30) + ) + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + try: + # Long enough for the watcher to settle into its blocking read. + await asyncio.sleep(0.1) + + def workflow_thread() -> bool: + # Waited out here rather than on the loop, and waited out at all: + # a task appended while the loop is still on its way into the + # select is picked up on the way in, which hides the missing wake. + time.sleep(0.2) + manager.cancel_from_workflow_thread(RUN_ID, 1) + deadline = time.monotonic() + 2 + while time.monotonic() < deadline: + if manager.subscription(RUN_ID, 1) is None: + return True + time.sleep(0.01) + return False + + cancelled = await asyncio.get_running_loop().run_in_executor( + None, workflow_thread + ) + + assert cancelled, ( + "the cancel was queued onto a loop nothing woke, so a watcher that " + "was supposed to stop keeps running with nothing to say so" + ) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_shutdown_tears_down_every_run(stream_key: StreamKey) -> None: + backend = MemoryStreamBackend() + notifier = RecordingNotifier() + manager = make_manager(backend, notifier) + for run_id in ("run-a", "run-b"): + manager.register(run_id=run_id, wait_id=1, stream_key=stream_key) + + await manager.shutdown() + + assert manager.runs_with_subscriptions() == [] + + +@pytest.mark.asyncio +async def test_manager_state_is_keyed_by_run(stream_key: StreamKey) -> None: + """So a stale Run cannot leak connections into a live one.""" + backend = MemoryStreamBackend() + notifier = RecordingNotifier() + manager = make_manager(backend, notifier) + manager.register(run_id="run-a", wait_id=1, stream_key=stream_key) + manager.register(run_id="run-b", wait_id=1, stream_key=stream_key) + try: + await manager.evict_run("run-a") + + assert manager.subscriptions("run-a") == [] + assert len(manager.subscriptions("run-b")) == 1 + finally: + await manager.shutdown() + + +# --- the sandbox ------------------------------------------------------------- + + +def test_the_manager_module_passes_through_the_sandbox() -> None: + """Re-importing it inside would give the Workflow a manager watching nothing. + + The real one owns the Worker's connections, watcher tasks, and buffers; only + an opaque handle to it may cross the boundary. + """ + assert ( + "temporalio.contrib.external_workflow_streams._manager" + in SandboxRestrictions.passthrough_modules_minimum + ) + assert ( + "temporalio.contrib.external_workflow_streams._manager" + in SandboxRestrictions.default.passthrough_modules + ) + + +@pytest.mark.asyncio +async def test_re_registering_a_wait_stops_the_watcher_it_replaced( + stream_key: StreamKey, +) -> None: + """The replaced subscription is unreachable and must not keep polling. + + Nothing else would notice: the new subscription works, records are + delivered, and the only symptom is a backend connection that never closes + and a task that outlives its Run -- visible, if at all, as a Worker whose + memory grows. + """ + backend = MemoryStreamBackend() + notifier = RecordingNotifier() + manager = make_manager(backend, notifier) + try: + first = manager.register( + run_id="run-1", + wait_id=1, + stream_key=stream_key, + ) + await asyncio.sleep(0.05) + assert first._watcher is not None and not first._watcher.done() + + second = manager.register( + run_id="run-1", + wait_id=1, + stream_key=stream_key, + ) + await asyncio.sleep(0.05) + + assert first._cancelled, "the replaced subscription must be marked dead" + assert first._watcher.done(), "its watcher must be stopped, not orphaned" + assert second._watcher is not None and not second._watcher.done(), ( + "the replacement's own watcher must still be running" + ) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_an_evicted_run_re_delivers_the_records_it_had_already_seen( + stream_key: StreamKey, +) -> None: + """Discarding prefetch state is only half of it; the records must come back. + + Cursors reset to the committed boundary is the mechanism, but what makes it + correct is the outcome: a Run that consumed records without committing a + marker must see those same records again, or eviction would silently drop + everything between the last marker and the eviction. + """ + backend = MemoryStreamBackend() + notifier = RecordingNotifier() + manager = make_manager(backend, notifier) + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await append(backend, stream_key, b"a", b"b", b"c") + await asyncio.wait_for(notifier.notified.wait(), 2) + first_pass = [r.offset for r in manager.drain(RUN_ID, 1)] + assert len(first_pass) == 3 + + # Evicted with no marker committed: nothing it saw was ever a claim. + await manager.evict_run(RUN_ID) + + notifier.notified.clear() + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await asyncio.wait_for(notifier.notified.wait(), 2) + second_pass = [r.offset for r in manager.drain(RUN_ID, 1)] + + try: + assert second_pass == first_pass, ( + "the restarted subscription must re-read the same offsets; resuming " + "past them would drop every record between the last marker and the " + "eviction" + ) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_a_stale_answer_is_re_reported_rather_than_dropped( + stream_key: StreamKey, +) -> None: + """`Stale` means Core held a newer generation, not that the record landed. + + The watcher calls back here only after a *new* non-empty read, and its + prefetch cursor is already past the buffered record, so a `Stale` treated as + delivered announces that record to nobody. The Workflow then blocks forever + on data it is already holding. + """ + backend = MemoryStreamBackend() + + class StaleThenAccepted(RecordingNotifier): + async def __call__(self, run_id: str, wait_id: int, generation: int) -> str: + self.calls.append((run_id, wait_id, generation)) + self.notified.set() + return ( + ReadinessResult.STALE + if len(self.calls) == 1 + else ReadinessResult.ACCEPTED + ) + + notifier = StaleThenAccepted() + manager = make_manager(backend, notifier) + try: + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await append(backend, stream_key, b"a") + await asyncio.wait_for(notifier.notified.wait(), 2) + await asyncio.sleep(0.3) + + assert len(notifier.calls) >= 2, ( + "a stale answer must be re-reported; the watcher will not announce " + f"this record again, got {len(notifier.calls)} report(s)" + ) + assert manager.subscriptions(RUN_ID)[0].buffered == 1, ( + "the record must still be there to deliver" + ) + finally: + await manager.shutdown() + + +@pytest.mark.parametrize("position", [1, 2, 3]) +@pytest.mark.asyncio +async def test_a_stale_retry_that_finds_the_run_gone_tears_the_watcher_down( + stream_key: StreamKey, position: int +) -> None: + """The retries' answer is the answer, and one of them requires a teardown. + + A first report can race a wait-generation change and be answered `Stale`; the + Run can then be evicted before the delayed retry that follows. The retry loop + kept only whether *some* attempt was accepted and threw the rest away, so + control returned with the original `Stale` still in hand: the owed wake went + out (right) and the `RunNotFound` teardown never ran (wrong). What is left + behind is a watcher, a buffer, a backend read loop and a `_runs` entry for a + Run this Worker no longer owns -- and every later readiness report and wake + attempt is made on its behalf. + + Parameterized over which retry discovers the eviction, because nothing makes + the first one special. + """ + backend = MemoryStreamBackend() + + class StaleThenGone(RecordingNotifier): + async def __call__(self, run_id: str, wait_id: int, generation: int) -> str: + self.calls.append((run_id, wait_id, generation)) + self.notified.set() + if len(self.calls) < position + 1: + return ReadinessResult.STALE + return ReadinessResult.RUN_NOT_FOUND + + notifier = StaleThenGone() + wake = CountingWake(failures=0) + manager = make_manager(backend, notifier, send_wake=wake) + try: + subscription = manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await append(backend, stream_key, b"a") + + await until( + lambda: manager.subscription(RUN_ID, 1) is None, + "a `RunNotFound` reached on a stale retry left the subscription in " + "`_runs`, so this Worker keeps a watcher, a buffer and a read loop " + "for a Run it no longer holds", + ) + assert subscription._cancelled, "the dropped subscription is still live" + await until( + lambda: subscription._watcher is not None and subscription._watcher.done(), + "the watcher for a Run that is gone is still reading from the backend", + ) + assert len(wake.counters) == 1, ( + "the wake `RunNotFound` requires must still be sent exactly once, " + f"got {len(wake.counters)}" + ) + assert len(notifier.calls) == position + 1, ( + "the retries did not stop at `RunNotFound`; a Run that is gone cannot " + "come back, and each further attempt delays the wake the record needs" + ) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_a_readiness_transport_failure_does_not_kill_the_watcher( + stream_key: StreamKey, +) -> None: + """An exception is a transport problem, not an answer. + + Letting it escape ends the watcher for good: the subscription stays + registered, its buffer keeps its records, and nothing announces them again. + Swallowing it would instead claim the record was announced. + """ + backend = MemoryStreamBackend() + + class FailsOnce(RecordingNotifier): + async def __call__(self, run_id: str, wait_id: int, generation: int) -> str: + self.calls.append((run_id, wait_id, generation)) + self.notified.set() + if len(self.calls) == 1: + raise ConnectionError("core unreachable") + return ReadinessResult.ACCEPTED + + notifier = FailsOnce() + manager = make_manager(backend, notifier) + try: + subscription = manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await append(backend, stream_key, b"a") + await asyncio.wait_for(notifier.notified.wait(), 2) + await asyncio.sleep(0.5) + + assert len(notifier.calls) >= 2, "the failing report must be retried" + watcher = subscription._watcher + assert watcher is not None and not watcher.done(), ( + "the watcher must survive a readiness failure; if it ends, this " + "subscription is never served again" + ) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_a_wake_failure_does_not_kill_the_watcher( + stream_key: StreamKey, +) -> None: + """The wake sender reports failure by raising, and that must stop here. + + An owed wake that could not be sent stays owed for the shutdown sweep. An + exception escaping instead takes the watcher with it, which loses every + later record too. + """ + backend = MemoryStreamBackend() + notifier = RecordingNotifier(answer=ReadinessResult.NO_OPEN_WORKFLOW_TASK) + + async def failing_wake(subscription) -> None: # type: ignore[no-untyped-def] + raise ConnectionError("service unavailable") + + manager = make_manager(backend, notifier, send_wake=failing_wake) + try: + subscription = manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await append(backend, stream_key, b"a") + await asyncio.wait_for(notifier.notified.wait(), 2) + await asyncio.sleep(0.3) + + watcher = subscription._watcher + assert watcher is not None and not watcher.done(), ( + "a failed wake must not end the watcher" + ) + assert subscription.wakes_owed >= 1, "the wake stays owed for the sweep" + finally: + await manager.shutdown() + + +# --- the park handshake's wait set (P19) -------------------------------------- +# +# One Core activation is constructed here, and only here. The set that gets +# parked is decided in two places -- the Worker's job handler turns Core's +# `PrepareExternalStreamPark.waits` into the map it hands over, and the manager +# installs and rechecks against it -- so a test on either half alone cannot show +# that the two agree on what Core asked for. + + +def _park_job(run_id: str, quiescence_generation: int, *wait_ids: int): # type: ignore[no-untyped-def] + """The activation Core sends to open a park handshake.""" + from temporalio.bridge.proto.workflow_activation import WorkflowActivation + + activation = WorkflowActivation() + activation.run_id = run_id + job = activation.jobs.add().prepare_external_stream_park + job.quiescence_generation = quiescence_generation + for wait_id in wait_ids: + job.waits.add().wait_id = wait_id + return activation + + +class _StubRuntime: + """The runtime's half of the park path: cursors in, terminal out.""" + + def __init__(self, cursors: dict[int, Cursor]) -> None: + self._cursors = cursors + self.terminals = 0 + + def blocked_snapshot(self) -> dict[int, Cursor]: + return dict(self._cursors) + + def add_terminal(self) -> bytes: + self.terminals += 1 + return b"terminal" + + +class _StubWorker: + """Only the two things ``_handle_external_stream_jobs`` reaches for.""" + + def __init__(self, run_id: str, runtime: _StubRuntime, manager: object) -> None: + self._external_stream_runtimes = {run_id: runtime} + self._manager = manager + + def _stream_manager(self) -> object: + return self._manager + + +@pytest.mark.asyncio +async def test_the_park_set_is_cores_wait_set_and_not_every_registration() -> None: + """A registered subscription is not necessarily one Core is parking. + + ``PrepareExternalStreamPark.waits`` is Core's complete *blocked* snapshot. + The runtime's own registration list is a superset of it: a subscription that + delivered a record and was not awaited again is registered and not blocked, + so it is absent from the quiescent snapshot Core parks. + + Parking the superset is wrong in both directions. The recheck for a wait + Core is not parking finds that wait's records -- which is *not* news, since + nothing is waiting on them -- and aborts a park that was entirely + legitimate, so the Workflow Task never parks and the handshake runs again on + the next idle timeout, and again. And an intent installed for a wait outside + the park set is an intent with no park behind it, which is precisely the + thing `backend-contract.md` forbids leaving in a backend. + """ + from temporalio.worker._workflow import _WorkflowWorker + + backend = MemoryStreamBackend() + manager = make_manager(backend, RecordingNotifier()) + parked = StreamKey("ns", "wf", "first-run", "tokens") + driving = StreamKey("ns", "wf", "first-run", "tool-events") + try: + manager.register(run_id=RUN_ID, wait_id=1, stream_key=parked) + manager.register(run_id=RUN_ID, wait_id=2, stream_key=driving) + # Wait 2 has a record sitting in it. Workflow code is not blocked on it, + # so it is not in the set Core asks to park. + await append(backend, driving, b"b") + + runtime = _StubRuntime({1: BEGINNING, 2: BEGINNING}) + worker = _StubWorker(RUN_ID, runtime, manager) + completion = await _WorkflowWorker._handle_external_stream_jobs( + worker, # type: ignore[arg-type] + _park_job(RUN_ID, 4, 1), + None, # type: ignore[arg-type] + ) + + assert completion is not None + result = completion.successful.commands[0].external_stream_park_result + assert result.WhichOneof("outcome") == "confirmed", ( + "a wait Core is not parking aborted the park of the wait it is" + ) + assert await backend.park_intent(parked, 1) is not None + assert await backend.park_intent(driving, 2) is None, ( + "an intent was installed for a wait Core never asked to park, so a " + "producer on that stream reads a park generation nothing is sitting in" + ) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failing", ["install", "recheck"]) +async def test_a_failed_park_leaves_no_externally_visible_half(failing: str) -> None: + """Parking is all-or-nothing, and a raised activation is not an exception to it. + + Intents are installed one at a time, so any failure part-way through -- the + second install, or a recheck once every install has landed -- leaves the + completed ones visible to producers while the activation itself fails. Core + parks nothing, so those intents describe a park that does not exist; an + eviction then takes the local bookkeeping away and nothing can remove them + at all. + """ + + class Failing(MemoryStreamBackend): + async def install_park_intent(self, key, intent): # type: ignore[no-untyped-def] + if failing == "install" and intent.wait_id == 2: + raise ConnectionError("backend unavailable") + return await super().install_park_intent(key, intent) + + async def recheck(self, key, wait_id): # type: ignore[no-untyped-def] + if failing == "recheck": + raise ConnectionError("backend unavailable") + return await super().recheck(key, wait_id) + + backend = Failing() + manager = make_manager(backend, RecordingNotifier()) + first = StreamKey("ns", "wf", "first-run", "tokens") + second = StreamKey("ns", "wf", "first-run", "tool-events") + try: + manager.register(run_id=RUN_ID, wait_id=1, stream_key=first) + manager.register(run_id=RUN_ID, wait_id=2, stream_key=second) + + # Reported as the taxonomy's transient row -- a backend that is + # unreachable is nothing for an operator to do anything about -- with + # the provider's own error kept as the cause, so the log still says + # which call failed and why. + with pytest.raises(StreamStorageError) as failure: + await manager.prepare_park(RUN_ID, 4, {1: BEGINNING, 2: BEGINNING}) + assert isinstance(failure.value.__cause__, ConnectionError) + + assert await backend.parked_wait_ids(first) == [], ( + "the park failed, and wait 1's intent is still in the backend " + "advertising a park Core never confirmed" + ) + assert await backend.parked_wait_ids(second) == [] + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_a_cancelled_park_takes_back_the_intents_it_had_installed() -> None: + """Cancellation is the *most* likely way a half-installed park is abandoned. + + Core withdrawing the activation and the Worker shutting down both arrive as + ``CancelledError``, which derives from ``BaseException`` -- so a rollback + reached only by ``except Exception`` is skipped for exactly the failures + most likely to leave a park half-installed. The intents that stay describe a + park Core never confirmed, and the eviction that follows takes the local + mirror away while they remain. + """ + + class Blocking(MemoryStreamBackend): + """Holds the second install open, so the first one is already visible.""" + + def __init__(self) -> None: + super().__init__() + self.reached = asyncio.Event() + self.release = asyncio.Event() + + async def install_park_intent(self, key, intent): # type: ignore[no-untyped-def] + if intent.wait_id == 2: + self.reached.set() + await self.release.wait() + return await super().install_park_intent(key, intent) + + backend = Blocking() + manager = make_manager(backend, RecordingNotifier()) + key = StreamKey("ns", "wf", "first-run", "tokens") + try: + manager.register(run_id=RUN_ID, wait_id=1, stream_key=key) + manager.register(run_id=RUN_ID, wait_id=2, stream_key=key) + parking = asyncio.get_running_loop().create_task( + manager.prepare_park(RUN_ID, 4, {1: BEGINNING, 2: BEGINNING}) + ) + await asyncio.wait_for(backend.reached.wait(), 2) + assert await backend.parked_wait_ids(key) == [1], ( + "this case is only meaningful once one intent is installed" + ) + + parking.cancel() + with pytest.raises(asyncio.CancelledError): + await parking + + assert await backend.parked_wait_ids(key) == [], ( + "the park was abandoned mid-install and wait 1's intent is still in " + "the backend, advertising a park Core never confirmed" + ) + finally: + backend.release.set() + await manager.shutdown() + + +# --- the removal a failure leaves owed ---------------------------------------- +# +# A park intent is durable backend state; the manager's knowledge of it is not. +# `installed_park_generation` is a mirror that only the Worker which installed +# the park holds, and every removal path reaches its intent *through* the +# Subscription that carries it -- which the close, the eviction and the hand-off +# all throw away. So a failed removal recorded on the Subscription is a failed +# removal recorded on the object the next step discards, and the manager keeps a +# per-Run ledger of owed removals instead. + + +class FailingRemovals(MemoryStreamBackend): + """A backend whose park-intent removals fail on demand. + + ``failures`` counts down, so a test can ask for a single blip or -- with a + number no retry can reach -- for a window it closes itself by setting it + back to zero. + """ + + def __init__(self, failures: int = 0) -> None: + super().__init__() + self.failures = failures + self.removal_attempts = 0 + + async def remove_park_intent_if_matches( + self, + key: StreamKey, + wait_id: int, + *, + run_id: str, + park_generation: int, + ) -> ParkIntentRemoval: + self.removal_attempts += 1 + if self.failures > 0: + self.failures -= 1 + raise ConnectionError("backend unavailable") + return await super().remove_park_intent_if_matches( + key, + wait_id, + run_id=run_id, + park_generation=park_generation, + ) + + +async def _nothing_parked(backend: MemoryStreamBackend, key: StreamKey) -> bool: + return await backend.parked_wait_ids(key) == [] + + +async def _only_wait_two_is_parked( + backend: MemoryStreamBackend, key: StreamKey +) -> bool: + return await backend.parked_wait_ids(key) == [2] + + +async def _inherit(backend: MemoryStreamBackend, key: StreamKey, run_id: str) -> None: + """Leaves behind the intent of a park whose Worker is gone. + + Installed straight into the backend rather than through a manager, because + that is exactly what makes it inherited: no mirror of it exists anywhere on + the Worker that finds it. + """ + await backend.install_park_intent( + key, + ParkIntent(wait_id=1, cursor=BEGINNING, park_generation=7, run_id=run_id), + ) + + +@pytest.mark.asyncio +async def test_one_blip_does_not_end_the_inherited_park_reconciliation( + stream_key: StreamKey, +) -> None: + """Registration is the only moment an inherited intent is looked for. + + It runs at most once per Subscription, so a single transient backend error + used to leave the intent installed with nothing left to try it again. What + that costs is the invariant's whole point: ``current_park_generation`` goes + on answering a generation Core has discarded, and because a parked wake's + request ID ignores sender identity the wake naming it is byte-identical to + the one that already ended that generation -- so the server deduplicates it + and the Workflow is never told about a record that is durably present. + """ + backend = FailingRemovals(failures=1) + await _inherit(backend, stream_key, RUN_ID) + manager = make_manager(backend, RecordingNotifier()) + try: + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + + await until( + lambda: _nothing_parked(backend, stream_key), + "the inherited intent is still installed after one failed removal, " + "and only another registration of this same wait would try again", + ) + assert backend.removal_attempts >= 2, "the failed removal must be retried" + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_an_inherited_intent_is_retried_autonomously_after_recovery( + stream_key: StreamKey, +) -> None: + """Retries are the cheap first line; the ledger is what makes them optional. + + Bounding the retries is only survivable because giving up records the + removal rather than forgetting it. Waiting instead for "the next time this + wait is registered" is a coincidence of eviction, not a mechanism -- and the + Run that most needs the removal is the one cached and blocked on something + other than this stream, which never registers this wait again at all. + """ + backend = FailingRemovals(failures=99) + await _inherit(backend, stream_key, RUN_ID) + manager = make_manager(backend, RecordingNotifier()) + try: + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await until( + lambda: ( + not manager._reconciliations + and backend.removal_attempts >= PARK_REMOVAL_ATTEMPTS + ), + "the reconciliation must exhaust its attempts before the ledger is " + "the only thing left holding the removal", + ) + backend.failures = 0 + + await until( + lambda: _nothing_parked(backend, stream_key), + "the inherited intent stayed installed after backend recovery even " + "though no park, resolve, registration, or eviction occurred", + ) + + assert await backend.parked_wait_ids(stream_key) == [], ( + "the reconciliation exhausted its inline attempts and its ledger " + "had no autonomous retry" + ) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_a_close_whose_removal_fails_leaves_the_removal_owed( + stream_key: StreamKey, +) -> None: + """A close drops the subscription, and used to drop the removal with it. + + Nothing could retry afterwards: the resolve path iterates *registered* + subscriptions, eviction and the shutdown sweep remove no intents, and + another wait's park cannot touch a per-wait key. + + And the damage is not confined to the closed wait. A stale intent keeps + ``parked_wait_ids`` non-empty, which suppresses the unparked-wake fallback + for the whole stream: with no live wait parked, the producer sends only the + dead generation, Core discards it as stale, and dedup silences every later + publish -- so wait 2 here loses its wakes to wait 1's leftover. + """ + backend = FailingRemovals(failures=99) + manager = make_manager(backend, RecordingNotifier()) + try: + for wait_id in (1, 2): + manager.register( + run_id=RUN_ID, + wait_id=wait_id, + stream_key=stream_key, + ) + assert not await manager.prepare_park(RUN_ID, 4, {1: BEGINNING, 2: BEGINNING}) + + await manager.cancel(RUN_ID, 1) + + assert manager.subscription(RUN_ID, 1) is None, ( + "the close must still drop the subscription -- one kept registered " + "while a removal is retried is an orphaned watcher" + ) + assert await backend.parked_wait_ids(stream_key) == [1, 2], ( + "this case is only meaningful while the removal is genuinely failing" + ) + backend.failures = 0 + + await manager.resolve_park(RUN_ID) + + assert await backend.parked_wait_ids(stream_key) == [], ( + "the closed wait's intent survived the only path that could still " + "have removed it, and it suppresses the unparked wake for wait 2 too" + ) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_a_failed_resolve_removal_retries_without_another_event( + stream_key: StreamKey, +) -> None: + """Backend recovery alone retires a resolved park's stale intent.""" + backend = FailingRemovals(failures=99) + manager = make_manager(backend, RecordingNotifier()) + try: + manager.register( + run_id=RUN_ID, + wait_id=1, + stream_key=stream_key, + ) + assert not await manager.prepare_park(RUN_ID, 4, {1: BEGINNING}) + + await manager.resolve_park(RUN_ID) + assert await backend.current_park_generation(stream_key, 1) == 4 + backend.failures = 0 + + await until( + lambda: _nothing_parked(backend, stream_key), + "the resolved intent stayed installed after backend recovery even " + "though no park, resolve, registration, or eviction occurred", + ) + assert manager._owed_removals == {} + finally: + await manager.shutdown() + + +class BlockingAutonomousRemoval(MemoryStreamBackend): + """Fails the eager attempt, then blocks the manager-owned retry.""" + + def __init__(self) -> None: + super().__init__() + self.attempts = 0 + self.retry_started = asyncio.Event() + self.retry_cancelled = asyncio.Event() + + async def remove_park_intent_if_matches( + self, + key: StreamKey, + wait_id: int, + *, + run_id: str, + park_generation: int, + ) -> ParkIntentRemoval: + self.attempts += 1 + if self.attempts == 1: + raise ConnectionError("backend unavailable") + self.retry_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.retry_cancelled.set() + raise + return ParkIntentRemoval.ABSENT + + +@pytest.mark.asyncio +async def test_shutdown_cancels_and_awaits_an_owed_removal_retry( + stream_key: StreamKey, +) -> None: + backend = BlockingAutonomousRemoval() + manager = make_manager(backend, RecordingNotifier()) + manager.register( + run_id=RUN_ID, + wait_id=1, + stream_key=stream_key, + ) + assert not await manager.prepare_park(RUN_ID, 4, {1: BEGINNING}) + await manager.resolve_park(RUN_ID) + await asyncio.wait_for(backend.retry_started.wait(), 2) + + await asyncio.wait_for(manager.shutdown(grace=timedelta(milliseconds=200)), 1) + + assert backend.retry_cancelled.is_set() + assert manager._owed_removal_retries == {} + + +class HangsOnceThenRecovers(MemoryStreamBackend): + """One removal call that never returns, and healthy ones after it.""" + + def __init__(self) -> None: + super().__init__() + self.attempts = 0 + self.hanging = asyncio.Event() + + async def remove_park_intent_if_matches( # type: ignore[override] + self, + key: StreamKey, + wait_id: int, + *, + run_id: str, + park_generation: int, + ) -> ParkIntentRemoval: + self.attempts += 1 + if self.attempts == 1: + raise ConnectionError("backend unavailable") + if self.attempts == 2: + # A call that hangs rather than raises: no exception to catch, no + # return to act on, and nothing counting attempts can bound it. + self.hanging.set() + await asyncio.Event().wait() + return await super().remove_park_intent_if_matches( + key, wait_id, run_id=run_id, park_generation=park_generation + ) + + +@pytest.mark.asyncio +async def test_a_removal_call_that_hangs_does_not_end_the_autonomous_retry( + stream_key: StreamKey, +) -> None: + """The loop's own liveness, which counting attempts cannot give it. + + The retry loop is the only thing left holding a removal after the inline + attempts are spent, so a single call that never returns is not one lost + attempt -- it is the mechanism stopping. Nothing restarts the loop either: + `_owe_removal` starts one only when the task it finds is absent or done, and + a task awaiting a hung call is neither. The intent would stay installed on a + backend that had already recovered. + + It also holds the Run's park lock while it waits, which every park, resolve + and eviction of this Run needs. + """ + backend = HangsOnceThenRecovers() + manager = make_manager(backend, RecordingNotifier()) + try: + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + assert not await manager.prepare_park(RUN_ID, 4, {1: BEGINNING}) + await manager.resolve_park(RUN_ID) + await asyncio.wait_for(backend.hanging.wait(), 2) + assert await backend.parked_wait_ids(stream_key) == [1] + + await until( + lambda: _nothing_parked(backend, stream_key), + "the retry gave up its park lock and came back around after the " + "call that hung, or it did not", + timeout=4.0, + ) + # And the lock an activation needs is reachable again: nothing is owed + # by now, so this resolve wants the lock and nothing else. + await asyncio.wait_for(manager.resolve_park(RUN_ID), 2) + finally: + await manager.shutdown(grace=timedelta(milliseconds=200)) + + +@pytest.mark.asyncio +async def test_shutdown_makes_the_last_attempt_the_retry_loop_cannot( + stream_key: StreamKey, +) -> None: + """A backend that recovered during a backoff, on a Worker that is exiting. + + The retry loop is cancelled by shutdown, and cancelling it during its + backoff throws away an attempt that would have succeeded -- the failure that + started the backoff is not the state of the backend now. Eviction used to be + that last attempt and deliberately stands aside while shutting down, so the + pass after the loops are cancelled is what has to make it. + """ + backend = FailingRemovals(failures=99) + manager = make_manager(backend, RecordingNotifier()) + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + assert not await manager.prepare_park(RUN_ID, 4, {1: BEGINNING}) + await manager.resolve_park(RUN_ID) + # Long enough for the backoff to have grown past the grace period below, so + # a sleeping loop cannot be what removes it. + await until( + lambda: backend.removal_attempts >= 4, + "the retry loop must have failed enough times to be backing off", + timeout=4.0, + ) + attempts = backend.removal_attempts + backend.failures = 0 + + await asyncio.wait_for(manager.shutdown(grace=timedelta(seconds=1)), 3) + + assert backend.removal_attempts > attempts, ( + "shutdown cancelled the retry and exited without trying once itself" + ) + assert await backend.parked_wait_ids(stream_key) == [], ( + "the intent outlived a Worker whose backend was healthy again" + ) + + +@pytest.mark.asyncio +async def test_cleanup_re_announces_the_record_the_stale_intent_silenced( + stream_key: StreamKey, +) -> None: + """Removing the intent is not the same as delivering what it suppressed. + + Every wake sent while the stale intent was installed named the generation + behind it, and Core discards a non-zero generation that is not the park it + holds. The record that arrived in that window is buffered here, its wake is + counted as sent, and the watcher will not report it again without a *new* + append -- so retiring the intent stops the suppression and announces + nothing. The Run stays blocked on a record this Worker is already holding. + """ + backend = FailingRemovals(failures=99) + sent: list[int] = [] + + async def send_wake(subscription: Subscription) -> None: + # What the Worker's sender does: read the generation at send time. + sent.append( + await subscription.backend.current_park_generation( + subscription.stream_key, subscription.wait_id + ) + or 0 + ) + + # No open Workflow Task, so local readiness cannot be delivered and each + # announcement becomes a wake Signal. + manager = make_manager( + backend, + RecordingNotifier(ReadinessResult.NO_OPEN_WORKFLOW_TASK), + send_wake=send_wake, + ) + try: + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + assert not await manager.prepare_park(RUN_ID, 4, {1: BEGINNING}) + await manager.resolve_park(RUN_ID) + assert await backend.current_park_generation(stream_key, 1) == 4 + + await append(backend, stream_key, b"1") + await until(lambda: sent, "the record's wake must be sent") + assert sent == [4], "the wake named the generation Core has discarded" + + backend.failures = 0 + await until( + lambda: _nothing_parked(backend, stream_key), + "the autonomous retry must retire the stale intent", + ) + await until( + lambda: len(sent) > 1, + "nothing announced the record again once the suppression ended", + ) + assert sent[1] == 0, ( + "the wake that follows the cleanup must be the unparked one Core " + "accepts unconditionally" + ) + finally: + await manager.shutdown(grace=timedelta(milliseconds=200)) + + +class DelaysTheIntentRead(MemoryStreamBackend): + """Holds the reconciliation's read until the watcher has sent its wake.""" + + def __init__(self) -> None: + super().__init__() + self.may_read = asyncio.Event() + + async def park_intent(self, key: StreamKey, wait_id: int) -> ParkIntent | None: + await self.may_read.wait() + return await super().park_intent(key, wait_id) + + +@pytest.mark.asyncio +async def test_reconciling_an_inherited_intent_re_announces_what_it_silenced( + stream_key: StreamKey, +) -> None: + """The ordinary reconciliation races the watcher, and can lose it a record. + + `_start_watcher` starts the reconciliation and the watcher as two tasks, so + the watcher can read a record, find no open Workflow Task and send its wake + -- naming the inherited generation, which Core discards -- while the + reconciliation is still in flight. The removal then succeeds *first time*, + which is the path that never goes near the owed-removal ledger, so the + reannouncement the retry path does would never happen here. The record sits + buffered with its prefetch cursor already past it. + """ + backend = DelaysTheIntentRead() + sent: list[int] = [] + + async def send_wake(subscription: Subscription) -> None: + sent.append( + await subscription.backend.current_park_generation( + subscription.stream_key, subscription.wait_id + ) + or 0 + ) + # Only now may the reconciliation read: the wake naming the inherited + # generation is already gone. + backend.may_read.set() + + manager = make_manager( + backend, + RecordingNotifier(ReadinessResult.NO_OPEN_WORKFLOW_TASK), + send_wake=send_wake, + ) + try: + await _inherit(backend, stream_key, RUN_ID) + await append(backend, stream_key, b"1") + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + + await until(lambda: sent, "the watcher must announce the buffered record") + assert sent[0] == 7, "the wake named the inherited generation" + + await until( + lambda: _nothing_parked(backend, stream_key), + "the reconciliation must remove the intent it inherited", + ) + await until( + lambda: len(sent) > 1, + "the ordinary reconciliation removed the intent and left the record " + "it had silenced unannounced", + ) + assert sent[1] == 0, "the wake after cleanup must be the unparked one" + finally: + await manager.shutdown(grace=timedelta(milliseconds=200)) + + +class LosesTheReplyAfterCommitting(MemoryStreamBackend): + """Deletes the intent, then fails before the reply reaches the Worker.""" + + def __init__(self) -> None: + super().__init__() + self.calls = 0 + self.removing = asyncio.Event() + self.wake_sent = asyncio.Event() + + async def remove_park_intent_if_matches( # type: ignore[override] + self, + key: StreamKey, + wait_id: int, + *, + run_id: str, + park_generation: int, + ) -> ParkIntentRemoval: + self.calls += 1 + if self.calls == 1: + self.removing.set() + await self.wake_sent.wait() + await super().remove_park_intent_if_matches( + key, wait_id, run_id=run_id, park_generation=park_generation + ) + raise ConnectionError("the reply never arrived") + return await super().remove_park_intent_if_matches( + key, wait_id, run_id=run_id, park_generation=park_generation + ) + + +@pytest.mark.asyncio +async def test_a_removal_whose_reply_was_lost_still_re_announces( + stream_key: StreamKey, +) -> None: + """An absent key is cleanup that happened, not a park to leave alone. + + A provider can commit the compare-and-delete and lose the connection before + its reply arrives. The removal raises, so it stays owed; the retry finds the + key clear. Reading that as "someone else's intent is there" would leave the + record whose wake named the dead generation buffered and unannounced -- + which is why the outcome is three-valued rather than a Boolean. + """ + backend = LosesTheReplyAfterCommitting() + sent: list[int] = [] + + async def send_wake(subscription: Subscription) -> None: + sent.append( + await subscription.backend.current_park_generation( + subscription.stream_key, subscription.wait_id + ) + or 0 + ) + + manager = make_manager( + backend, + RecordingNotifier(ReadinessResult.NO_OPEN_WORKFLOW_TASK), + send_wake=send_wake, + ) + try: + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + assert not await manager.prepare_park(RUN_ID, 4, {1: BEGINNING}) + + # Core has ended the park, so generation 4 is dead the moment this + # resolve starts -- and its removal is in flight when the record lands. + resolve = asyncio.ensure_future(manager.resolve_park(RUN_ID)) + await asyncio.wait_for(backend.removing.wait(), 2) + await append(backend, stream_key, b"1") + await until(lambda: sent, "the record's wake must be sent") + assert sent == [4], "the wake named the generation Core has discarded" + + backend.wake_sent.set() + await asyncio.wait_for(resolve, 2) + + await until( + lambda: len(sent) > 1, + "the retry read the committed delete as a mismatch and left the " + "record it had silenced unannounced", + ) + assert sent[1] == 0, "the wake after cleanup must be the unparked one" + assert manager._owed_removals == {} + finally: + await manager.shutdown(grace=timedelta(milliseconds=200)) + + +class FailingIntentReads(MemoryStreamBackend): + """Cannot say whether an intent exists until it is told to recover.""" + + def __init__(self) -> None: + super().__init__() + self.healthy = False + self.reads = 0 + + async def park_intent(self, key: StreamKey, wait_id: int) -> ParkIntent | None: + self.reads += 1 + if not self.healthy: + raise ConnectionError("backend unavailable") + return await super().park_intent(key, wait_id) + + +@pytest.mark.asyncio +async def test_discovering_an_inherited_intent_is_retried_after_recovery( + stream_key: StreamKey, +) -> None: + """The ledger cannot hold what was never read, so the read needs the retry. + + An owed removal is recorded from an intent's *identity*, and a read that + never succeeds produces no identity -- so a reconciliation that gave up + after its cheap attempts left no ledger entry, no retry task and no owner, + while logging that an autonomous retry was continuing. Recovery then needed + exactly the unrelated Workflow or Core event the ledger exists to stop + needing. + """ + backend = FailingIntentReads() + manager = make_manager(backend, RecordingNotifier()) + try: + await _inherit(backend, stream_key, RUN_ID) + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await until( + lambda: backend.reads >= PARK_REMOVAL_ATTEMPTS, + "the inline attempts must be spent before recovery", + ) + assert manager._owed_removals == {}, "nothing was read, so nothing can be owed" + backend.healthy = True + + await until( + lambda: _nothing_parked(backend, stream_key), + "the inherited intent outlived a backend that recovered, with no " + "park, resolve, registration or eviction to prompt another read", + timeout=4.0, + ) + finally: + await manager.shutdown(grace=timedelta(milliseconds=200)) + + +@pytest.mark.asyncio +async def test_closing_a_wait_removes_an_intent_it_only_inherited( + stream_key: StreamKey, +) -> None: + """The second door, which needs no failure of its own to open. + + ``installed_park_generation`` is set only by a park *this* Worker installed, + so for an inherited intent it is ``None`` and the removal short-circuits + silently -- the close attempts nothing at all. That is why the close is not + a backstop for a reconciliation that has already failed: it is disabled by + exactly the condition that made the reconciliation necessary. + """ + backend = FailingRemovals(failures=99) + await _inherit(backend, stream_key, RUN_ID) + manager = make_manager(backend, RecordingNotifier()) + try: + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await until( + lambda: backend.removal_attempts >= PARK_REMOVAL_ATTEMPTS, + "the reconciliation must have given up before the close is asked to " + "cover for it", + ) + backend.failures = 0 + + await manager.cancel(RUN_ID, 1) + + assert await backend.parked_wait_ids(stream_key) == [], ( + "the close attempted no removal, because the intent it inherited is " + "mirrored in no generation of its own" + ) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_an_owed_removal_outlives_eviction_and_keeps_retrying( + stream_key: StreamKey, +) -> None: + """Cache ownership ending does not end backend-cleanup ownership. + + Installed intents are deliberately untouched by the same path: an eviction + is not the end of a park, and the intents of one that really is outstanding + have to survive the Worker losing the Run. Only a removal already decided on + is retried after it. + """ + backend = FailingRemovals(failures=99) + manager = make_manager(backend, RecordingNotifier()) + try: + for wait_id in (1, 2): + manager.register( + run_id=RUN_ID, + wait_id=wait_id, + stream_key=stream_key, + ) + assert not await manager.prepare_park(RUN_ID, 4, {1: BEGINNING, 2: BEGINNING}) + await manager.cancel(RUN_ID, 1) + + await manager.evict_run(RUN_ID) + assert await backend.parked_wait_ids(stream_key) == [1, 2] + backend.failures = 0 + + await until( + lambda: _only_wait_two_is_parked(backend, stream_key), + "wait 1's owed removal stopped retrying when eviction dropped the " + "Run, while wait 2's outstanding park still had to survive", + ) + + assert await backend.parked_wait_ids(stream_key) == [2], ( + "the owed removal did not outlive the cached Run" + ) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_a_cleanup_after_eviction_still_announces_what_it_silenced( + stream_key: StreamKey, +) -> None: + """The removal outlives the Subscription; the wake it owes has to as well. + + Eviction is what makes the ledger worth having -- it drops the Subscription + and deliberately leaves the entry and its retry alive -- so by the time a + recovered backend lets the removal through, the object every announcement + path goes through is gone. Announcing only when the wait happens to still be + registered therefore kept the leak fixed and put the *record loss* straight + back: while the stale intent was installed a producer saw a non-empty parked + set and sent only generation 4, Core discarded it, and with no cached Run + there was no local watcher to send an unparked wake of its own. + + The wake is unconditional on purpose. From out here the buffer is gone and a + remote producer's append was never visible at all, so "a record arrived" and + "none did" are the same observation -- and only one of them is silent if the + guess goes the wrong way. The cost is one empty Workflow Task. + """ + backend = FailingRemovals(failures=99) + sent: list[tuple[int, int]] = [] + + async def send_wake(target: object) -> None: + # What the Worker's real sender composes: the manager decides which park + # a wake names, because only it can tell an intent Core is still parked + # on from one already decided against. + sent.append( + ( + target.wait_id, # type: ignore[attr-defined] + await manager.wake_park_generation(target) or 0, # type: ignore[arg-type] + ) + ) + + manager = make_manager(backend, RecordingNotifier(), send_wake=send_wake) + try: + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + assert not await manager.prepare_park(RUN_ID, 4, {1: BEGINNING}) + await manager.resolve_park(RUN_ID) + assert await backend.current_park_generation(stream_key, 1) == 4, ( + "the failed removal must leave the stale intent installed" + ) + + await manager.evict_run(RUN_ID) + assert manager.subscription(RUN_ID, 1) is None, ( + "eviction must drop the Subscription -- that is the premise" + ) + assert await backend.parked_wait_ids(stream_key) == [1], ( + "and must leave the intent and its ledger entry behind" + ) + + # The record that gets lost. Appended by a producer this Worker cannot + # see, into a stream whose parked set is non-empty because of the stale + # intent -- so the only wake it sends names generation 4. + await append(backend, stream_key, b"1") + + backend.failures = 0 + await until( + lambda: _nothing_parked(backend, stream_key), + "the autonomous retry must retire the stale intent once the backend " + "recovers", + ) + await until( + lambda: sent, + "the removal ended the suppression and announced nothing, so the " + "record stays undelivered until some unrelated event happens to " + "create a Workflow Task", + ) + assert sent == [(1, 0)], ( + "the cleanup handoff must be the unparked wake Core accepts " + "unconditionally, since the generation it would otherwise name is " + "the one it just removed" + ) + finally: + await manager.shutdown(grace=timedelta(milliseconds=200)) + + +@pytest.mark.asyncio +async def test_a_drain_never_removes_the_intent_of_a_park_that_replaced_it( + stream_key: StreamKey, +) -> None: + """An entry records a removal, not the claim that some intent exists. + + A Continue-As-New successor re-uses this stream key with wait ids that start + again at 1, so an entry a predecessor Run left behind can name a live park's + key. Retiring it on the key alone would take the intent out from under a + park that really is outstanding -- manufacturing exactly the unwakeable Run + the ledger exists to prevent. + """ + backend = FailingRemovals(failures=99) + manager = make_manager(backend, RecordingNotifier()) + try: + manager.register(run_id="run-a", wait_id=1, stream_key=stream_key) + assert not await manager.prepare_park("run-a", 4, {1: BEGINNING}) + await manager.cancel("run-a", 1) + backend.failures = 0 + + # The successor's park, installed directly: what matters is that it is a + # different Run's intent sitting at the key run-a still owes a removal + # for, not how it got there. + await backend.install_park_intent( + stream_key, + ParkIntent(wait_id=1, cursor=BEGINNING, park_generation=1, run_id="run-b"), + ) + + await manager.resolve_park("run-a") + + installed = await backend.park_intent(stream_key, 1) + assert installed is not None and installed.run_id == "run-b", ( + "run-a's owed removal was retired against whatever it found, taking " + "out the intent of a park run-b is still sitting in" + ) + finally: + await manager.shutdown() + + +# --- a wake owed on the live path --------------------------------------------- + + +class CountingWake: + """Records the counter each attempt would derive its request ID from. + + A retry that re-counts is not a retry: the request ID moves with the + counter, so the server deduplicates nothing and answers with a second, empty + Workflow Task instead of resolving the wake that may in fact have arrived. + """ + + def __init__(self, failures: int) -> None: + self.failures = failures + self.counters: list[int] = [] + #: Asked at each attempt, for tests about what has happened to the + #: subscription by the time an attempt is made. Assigned after the + #: manager exists, since that is what it usually asks about. + self.observe: Callable[[], object] | None = None + self.observed: list[object] = [] + + async def __call__(self, subscription) -> None: # type: ignore[no-untyped-def] + self.counters.append(subscription.wake_counter) + if self.observe is not None: + self.observed.append(self.observe()) + if self.failures > 0: + self.failures -= 1 + raise ConnectionError("service unavailable") + + +@pytest.mark.asyncio +async def test_a_failed_live_wake_is_retried_as_the_same_wake( + stream_key: StreamKey, +) -> None: + """One attempt on the live path is one attempt with nothing behind it. + + The watcher has already moved ``prefetch_cursor`` past the buffered record + and returns on an empty read, so it never comes back here without a *new* + append. ``rearm_ready`` needs the activation the lost wake was supposed to + cause. And the idle timer only runs while a Workflow Task is retained, which + ``NoOpenWorkflowTask`` says there is not -- so no park handshake happens + either. The shutdown sweep already does this correctly, which is why the + retry belongs to both and not just to it. + """ + backend = MemoryStreamBackend() + notifier = RecordingNotifier(answer=ReadinessResult.NO_OPEN_WORKFLOW_TASK) + wake = CountingWake(failures=1) + manager = make_manager(backend, notifier, send_wake=wake) + try: + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await append(backend, stream_key, b"a") + + await until( + lambda: len(wake.counters) >= 2, + "the failed wake was logged and forgotten; nothing on this path " + "ever attempts it again, so the buffered record is never announced", + ) + assert wake.counters == [1, 1], ( + "the retry derived a different request ID, so it asks for a second " + "empty Workflow Task rather than re-sending the wake that was owed" + ) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_replayed_readiness_is_coalesced_until_the_task_completes( + stream_key: StreamKey, +) -> None: + """A failed Workflow Task must not recursively ask for another task. + + Eviction discards the subscription and its speculative buffer. Replaying + the failed task rebuilds both, but the wake that caused it remains + outstanding because Core never accepted a successful completion. Once the + replayed task does complete, any buffer it leaves behind needs a fresh wake. + """ + backend = MemoryStreamBackend() + notifier = RecordingNotifier(answer=ReadinessResult.NO_OPEN_WORKFLOW_TASK) + wake = CountingWake(failures=0) + manager = make_manager(backend, notifier, send_wake=wake) + try: + # The readiness races the completion of a task that was already open. + # That completion cannot have been caused by the wake and must not + # release its gate. + manager.note_workflow_task_started(RUN_ID) + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await append(backend, stream_key, b"a") + await until( + lambda: len(wake.counters) >= 1, + "the original buffered readiness never owed a wake", + ) + manager.note_workflow_task_completed(RUN_ID, terminal=False) + + await manager.evict_run(RUN_ID) + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await until( + lambda: len(notifier.calls) >= 2, + "replay never re-reported the reconstructed buffered readiness", + ) + + assert wake.counters == [1], ( + "replay turned one outstanding wake into another, so each failed " + f"task can recursively create the next: {wake.counters}" + ) + + manager.note_workflow_task_started(RUN_ID) + manager.note_workflow_task_completed(RUN_ID, terminal=False) + await until( + lambda: len(wake.counters) >= 2, + "a buffer left after successful replay was never announced again", + ) + assert wake.counters == [1, 2] + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_a_failed_wake_attempt_cannot_claim_a_completion_for_its_retry( + stream_key: StreamKey, +) -> None: + """Each retry correlates only with activations that begin during that attempt. + + The first Signal attempt can lose a closing-window race while a different + Workflow Task completes. If the later successful retry inherited that + completion, the gate would open before the retry's Signal had produced its + own task and another buffered generation could send a competing wake. + """ + backend = MemoryStreamBackend() + notifier = RecordingNotifier(answer=ReadinessResult.NO_OPEN_WORKFLOW_TASK) + started = [asyncio.Event() for _ in range(3)] + release = [asyncio.Event() for _ in range(3)] + finished = [asyncio.Event() for _ in range(3)] + counters: list[int] = [] + + async def raced_wake(subscription: Subscription) -> None: + attempt = len(counters) + counters.append(subscription.wake_counter) + started[attempt].set() + try: + await release[attempt].wait() + if attempt == 0: + raise ConnectionError("closing window") + finally: + finished[attempt].set() + + manager = make_manager(backend, notifier, send_wake=raced_wake) + try: + manager.note_workflow_task_started(RUN_ID) + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await append(backend, stream_key, b"a") + await asyncio.wait_for(started[0].wait(), 2) + + # A different task starts and completes while attempt 1 is in flight. + manager.note_workflow_task_started(RUN_ID) + manager.note_workflow_task_completed(RUN_ID, terminal=False) + release[0].set() + + await asyncio.wait_for(started[1].wait(), 2) + release[1].set() + await asyncio.wait_for(finished[1].wait(), 2) + + # The successful retry has not caused an activation yet. Re-reporting + # the buffer must remain coalesced behind it. + manager.rearm_ready(RUN_ID) + await asyncio.sleep(0.05) + assert counters == [1, 1] + + manager.note_workflow_task_started(RUN_ID) + manager.note_workflow_task_completed(RUN_ID, terminal=False) + await asyncio.wait_for(started[2].wait(), 2) + assert counters == [1, 1, 2] + release[2].set() + finally: + for event in release: + event.set() + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_a_new_buffered_range_after_eviction_gets_a_new_wake_counter( + stream_key: StreamKey, +) -> None: + """A later range is a new ask even when the Run was evicted between them. + + The unparked request ID is derived from the sender's identity and its + counter, and the identity is fixed for the Worker's lifetime -- so the + counter is the only thing keeping two of this Worker's wakes apart. A + counter held on the `Subscription` restarts at one every time an evicted Run + is rebuilt, which re-derives the request ID the previous incarnation already + used: the server deduplicates it, no Workflow Task is created, and a Run + holding buffered records waits for a wake that was thrown away as a + duplicate. The sequence therefore belongs to the manager, which outlives the + Run. + """ + backend = MemoryStreamBackend() + notifier = RecordingNotifier(answer=ReadinessResult.NO_OPEN_WORKFLOW_TASK) + wake = CountingWake(failures=0) + manager = make_manager(backend, notifier, send_wake=wake) + try: + first = await backend.append( + stream_key, StreamRecord(RecordKind.DATA, b"a", "sa", 0) + ) + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await until( + lambda: len(wake.counters) >= 1, + "the first incarnation never owed its wake, so there is nothing for " + "the second one to collide with", + ) + manager.drain(RUN_ID, 1) + manager.note_workflow_task_started(RUN_ID) + manager.note_workflow_task_completed(RUN_ID, terminal=False) + + # The Run is evicted and comes back after the first range's boundary. + # This is not replay of the first readiness: the committed start cursor + # and the range waiting behind it are both new. + await manager.evict_run(RUN_ID) + assert first.offset is not None + manager.register( + run_id=RUN_ID, + wait_id=1, + stream_key=stream_key, + start_cursor=AFTER(first.offset), + ) + await append(backend, stream_key, b"b") + await until( + lambda: len(wake.counters) >= 2, + "the rebuilt subscription never owed a wake for the record waiting for it", + ) + + assert wake.counters[0] != wake.counters[1], ( + "the rebuilt Run drew the counter its predecessor had already used, " + "so both wakes derive one request ID and the server deduplicates " + f"the second away: {wake.counters}" + ) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_a_wake_owed_by_a_vanished_run_is_retried_before_it_is_dropped( + stream_key: StreamKey, +) -> None: + """The sharpest sub-case: after the pop there is no backstop at all. + + ``RunNotFound`` takes the subscription out of ``_runs``, and the shutdown + sweep is driven by ``_runs`` -- so a wake given up on before that pop is not + merely delayed, it is unreachable by every remaining mechanism. + """ + backend = MemoryStreamBackend() + notifier = RecordingNotifier(answer=ReadinessResult.RUN_NOT_FOUND) + wake = CountingWake(failures=1) + manager = make_manager(backend, notifier, send_wake=wake) + wake.observe = lambda: manager.subscription(RUN_ID, 1) is not None + try: + manager.register(run_id=RUN_ID, wait_id=1, stream_key=stream_key) + await append(backend, stream_key, b"a") + + await until( + lambda: len(wake.counters) >= 2, + "the wake failed once and the subscription was dropped on top of " + "it, putting it out of reach of the sweep that would have retried it", + ) + assert wake.counters == [1, 1] + assert wake.observed == [True, True], ( + "an attempt was made after the subscription had already left " + "`_runs`, where the sweep that is supposed to back it up cannot " + "see it either" + ) + await until( + lambda: manager.subscription(RUN_ID, 1) is None, + "a Run this Worker no longer holds must still be dropped once the " + "wake it owed has been attempted", + ) + finally: + await manager.shutdown() diff --git a/tests/contrib/external_workflow_streams/test_multi_stream.py b/tests/contrib/external_workflow_streams/test_multi_stream.py new file mode 100644 index 0000000000..c90acb21c6 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_multi_stream.py @@ -0,0 +1,1043 @@ +"""P21 — several streams, ``merge``, and two subscriptions to one stream. + +The idle timeout is a **Workflow-Task policy**, not a per-subscription one: it +applies to the complete set the Workflow is blocked on, so one idle stream must +not park the task while another is still delivering. Everything below is about +that set behaving as a set. + +PYTEST_DONT_REWRITE: sandboxed fixture Workflows re-import this module, so pytest's +injected imports would make sandbox validation depend on pytest's import locks. +""" + +# pyright: reportMissingParameterType=false, reportUnusedVariable=false +from __future__ import annotations + +import asyncio +import uuid +from datetime import timedelta + +import pytest +import pytest_asyncio + +import temporalio.converter +import temporalio.workflow +from temporalio import workflow +from temporalio.client import Client +from temporalio.contrib.external_workflow_streams._annotation import decode_annotation +from temporalio.contrib.external_workflow_streams._api import ( + MAX_RECORDS_PER_ACTIVATION, + _install_runtime, + external_stream, + merge, +) +from temporalio.contrib.external_workflow_streams._backend import StreamKey +from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec +from temporalio.contrib.external_workflow_streams._manager import ( + ReadinessResult, + StreamSubscriptionManager, +) +from temporalio.contrib.external_workflow_streams._record import ( + Offset, + RecordKind, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._runtime import WorkflowStreamRuntime +from temporalio.worker import Worker +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend + +with workflow.unsafe.imports_passed_through(): + pass + +RUN_ID = "run-1" + + +async def _notify( + run_id: str, # pyright: ignore[reportUnusedParameter] + wait_id: int, # pyright: ignore[reportUnusedParameter] + generation: int, # pyright: ignore[reportUnusedParameter] +) -> str: + return ReadinessResult.ACCEPTED + + +@pytest.fixture +def backend() -> MemoryStreamBackend: + return MemoryStreamBackend() + + +class StubManager: + """Records registrations and cancellations, and starts no watcher. + + The snapshot and annotation tests are about what the runtime decides, which + happens before any watching. A real manager would spawn prefetch loops + against a backend nothing is writing to, and their failures would be noise. + """ + + def __init__(self) -> None: + #: `(run_id, wait_id)` for every wait the runtime asked to stop serving. + self.cancelled: list[tuple[str, int]] = [] + + def register(self, *, run_id, wait_id, stream_key, start_cursor): # type: ignore[no-untyped-def] + pass + + def cancel_from_workflow_thread(self, run_id, wait_id): # type: ignore[no-untyped-def] + self.cancelled.append((run_id, wait_id)) + + def note_wait_generation(self, run_id, wait_id, generation) -> None: # type: ignore[no-untyped-def] + pass + + def drain(self, run_id, wait_id, max_records=None): # type: ignore[no-untyped-def] + return [] + + +@pytest.fixture +def manager() -> StubManager: + return StubManager() + + +@pytest_asyncio.fixture # type: ignore[reportUntypedFunctionDecorator] +async def live_manager(backend: MemoryStreamBackend): + """A real manager, for the tests that need records to actually move.""" + mgr = StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=10), + ) + yield mgr + await mgr.shutdown() + + +def make_runtime(manager, backend, idle=timedelta(seconds=1)): # type: ignore[no-untyped-def] + return WorkflowStreamRuntime( + manager=manager, + backend=backend, + run_id=RUN_ID, + namespace="ns", + workflow_id="wf", + first_execution_run_id="first-run", + data_converter=temporalio.converter.DataConverter.default, + default_idle_timeout=idle, + ) + + +def data(seq: int, offset: str) -> StreamRecord: + return StreamRecord(RecordKind.DATA, b"x", "s", seq).placed_at(Offset(offset)) + + +def fence(seq: int, offset: str) -> StreamRecord: + return StreamRecord(RecordKind.WRITE_FENCE, b"", "s", seq).placed_at(Offset(offset)) + + +# --- the idle timeout is a property of the set -------------------------------- + + +def test_one_idle_stream_cannot_park_while_another_is_active( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + """The criterion the whole "complete set" rule exists for. + + Parking on the idle one alone would strand records already sitting in the + active one, and the Workflow would wait out a timeout for a stream that was + never quiet. + """ + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=runtime.stream_key("idle")) + runtime.register(wait_id=2, stream_key=runtime.stream_key("busy")) + runtime.note_blocked(1, True) + # The second is *not* blocked -- Workflow code is still working through it. + runtime.note_blocked(2, False) + + snapshot = runtime.quiescent_snapshot() + + assert snapshot is not None + assert [w.wait_id for w in snapshot] == [1], ( + "the quiescent set must name only the waits Workflow code is actually " + "blocked on; naming an active one would ask Core to park it" + ) + + +def test_the_quiescent_set_is_complete_when_everything_is_blocked( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + """A set missing a member would let Core park a wait that was still live.""" + runtime = make_runtime(manager, backend) + for wait_id, name in ((1, "a"), (2, "b"), (3, "c")): + runtime.register( + wait_id=wait_id, + stream_key=runtime.stream_key(name), + ) + runtime.note_blocked(wait_id, True) + + snapshot = runtime.quiescent_snapshot() + + assert snapshot is not None + assert [w.wait_id for w in snapshot] == [1, 2, 3] + + +def test_differing_idle_timeouts_reduce_to_the_minimum( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + """``min``, in ``wait_id`` order, over the quiescent set and nothing else. + + Taking the maximum would let one lax subscription hold the Workflow Task + open past what another was configured to tolerate; taking the *configured + values of the whole set* rather than of the blocked subset would make the + result depend on state replay does not reproduce. + """ + runtime = make_runtime(manager, backend) + for wait_id, seconds in ((1, 5), (2, 2), (3, 9)): + runtime.register( + wait_id=wait_id, + stream_key=runtime.stream_key(f"s{wait_id}"), + idle_timeout=timedelta(seconds=seconds), + ) + runtime.note_blocked(wait_id, True) + + assert runtime.effective_idle_timeout() == timedelta(seconds=2) + + +def test_with_options_timeouts_reach_the_reduction_through_the_public_api( + manager: StubManager, + backend: MemoryStreamBackend, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The reduction is unreachable unless ``subscribe()`` passes the value on. + + The test above registers with an explicit ``idle_timeout``, which proves + ``min`` reduces but says nothing about whether a configured value can ever + get there. ``with_options`` is the only route a user has, and a + ``subscribe()`` that dropped the option would leave every quiescent set on + the one-second default while that test, and + ``ExternalStreamSubscription.idle_timeout``, both kept passing -- the option + would be silently decorative. + """ + + class Instance: + """Stands in for the Workflow object the per-Run state hangs off.""" + + instance = Instance() + monkeypatch.setattr(temporalio.workflow, "instance", lambda: instance) + runtime = make_runtime(manager, backend, idle=timedelta(seconds=1)) + _install_runtime(instance, runtime) # type: ignore[arg-type] + + for name, seconds in (("slow", 30), ("quick", 4), ("slower", 60)): + subscription = ( + external_stream.with_options(idle_timeout=timedelta(seconds=seconds)) + .topic(name, type=str) + .subscribe() + ) + runtime.note_blocked(subscription.wait_id, True) + + assert runtime.effective_idle_timeout() == timedelta(seconds=4), ( + "no configured timeout reached the runtime, so the whole quiescent set " + "fell back to the default and every with_options() call was ignored" + ) + + +def test_the_reduction_ignores_waits_that_are_not_blocked( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + """The inputs are the quiescent set's configured values and nothing else.""" + runtime = make_runtime(manager, backend) + runtime.register( + wait_id=1, + stream_key=runtime.stream_key("a"), + idle_timeout=timedelta(seconds=7), + ) + runtime.register( + wait_id=2, + stream_key=runtime.stream_key("b"), + idle_timeout=timedelta(milliseconds=50), + ) + runtime.note_blocked(1, True) + runtime.note_blocked(2, False) + + assert runtime.effective_idle_timeout() == timedelta(seconds=7) + + +# --- fences are per stream, parking is per set -------------------------------- + + +def test_a_fence_on_one_stream_alone_does_not_make_the_set_parkable( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + """It marks *that* stream parkable and says nothing about the others. + + A fence means only that one producer session has finished writing. Treating + it as a signal for the whole set would park a Workflow Task while another + stream was still being written to. + """ + runtime = make_runtime(manager, backend) + for wait_id, name in ((1, "fenced"), (2, "open")): + runtime.register( + wait_id=wait_id, + stream_key=runtime.stream_key(name), + ) + runtime.note_blocked(wait_id, True) + runtime.record_delivery(1, fence(0, "10-0")) + + snapshot = runtime.quiescent_snapshot() + + assert snapshot is not None + parkable = {w.wait_id: w.immediately_parkable for w in snapshot} + assert parkable == {1: True, 2: False}, ( + "the fenced stream is immediately parkable and the open one is not; " + "Core parks the task early only when every wait says True" + ) + + +def test_all_fenced_streams_make_the_whole_set_parkable( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + runtime = make_runtime(manager, backend) + for wait_id, name in ((1, "a"), (2, "b")): + runtime.register( + wait_id=wait_id, + stream_key=runtime.stream_key(name), + ) + runtime.note_blocked(wait_id, True) + runtime.record_delivery(wait_id, fence(0, f"1{wait_id}-0")) + + snapshot = runtime.quiescent_snapshot() + + assert snapshot is not None + assert all(w.immediately_parkable for w in snapshot) + + +def test_a_record_after_a_fence_reopens_the_stream( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + """A later record does not violate a fence; it simply resumes consumption. + + The fence spoke only for one producer session. Treating the stream as closed + would strand every record another producer went on to write. + """ + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=runtime.stream_key("a")) + runtime.note_blocked(1, True) + runtime.record_delivery(1, fence(0, "10-0")) + assert runtime.quiescent_snapshot()[0].immediately_parkable # type: ignore[index] + + runtime.record_delivery(1, data(1, "10-1")) + + assert not runtime.quiescent_snapshot()[0].immediately_parkable # type: ignore[index] + + +# --- two subscriptions to one stream ------------------------------------------ + + +@pytest.mark.asyncio +async def test_two_same_stream_subscriptions_install_distinct_park_intents( + live_manager, backend: MemoryStreamBackend +) -> None: + """Verified by reading both back, not by trusting the call. + + Intents keyed by stream alone would have the second install overwrite the + first, after which only one of the two could ever be woken -- and the + surviving intent would name the wrong cursor for the lost one. + """ + runtime = make_runtime(live_manager, backend) + key = runtime.stream_key("tokens") + runtime.register(wait_id=1, stream_key=key) + runtime.register(wait_id=2, stream_key=key) + # Two different positions, so a collapsed intent is visible rather than + # merely suspected. + runtime.record_delivery(1, data(0, "10-0")) + runtime.note_blocked(1, True) + runtime.note_blocked(2, True) + + await live_manager.prepare_park(RUN_ID, 1, runtime.blocked_snapshot()) + + first = await backend.park_intent(key, 1) + second = await backend.park_intent(key, 2) + assert first is not None and second is not None + assert first.cursor != second.cursor, ( + "each subscription parks at its own position; a shared intent would " + "resume one of them at the other's cursor" + ) + + +@pytest.mark.asyncio +async def test_each_same_stream_subscription_receives_every_record( + live_manager, backend: MemoryStreamBackend +) -> None: + """Delivery is broadcast, not competing-consumer. + + Two subscriptions to one stream are two independent readers. Splitting the + records between them would make either one's view of the stream depend on + the other's existence. + """ + runtime = make_runtime(live_manager, backend) + key = runtime.stream_key("tokens") + runtime.register(wait_id=1, stream_key=key) + runtime.register(wait_id=2, stream_key=key) + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + for i, value in enumerate(["a", "b"]): + await backend.append( + key, StreamRecord(RecordKind.DATA, await codec.encode(value), "p", i) + ) + + await asyncio.sleep(0.2) + + first = live_manager.drain(RUN_ID, 1) + second = live_manager.drain(RUN_ID, 2) + assert len(first) == 2 and len(second) == 2, ( + f"each subscription must see both records, got {len(first)} and {len(second)}" + ) + + +# --- the recorded delivery schedule ------------------------------------------- + + +def test_an_alternating_two_stream_batch_encodes_one_run_per_delivery( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + """A run is a maximal consecutive stretch from a *single* wait. + + Collapsing the alternation into two runs -- one per stream -- would record a + delivery order that never happened, and replay would hand Workflow code all + of one stream before any of the other. + """ + runtime = make_runtime(manager, backend) + for wait_id, name in ((1, "a"), (2, "b")): + runtime.register( + wait_id=wait_id, + stream_key=runtime.stream_key(name), + ) + for i in range(2): + runtime.record_delivery(1, data(i, f"1{i}-0")) + runtime.record_delivery(2, data(i, f"2{i}-0")) + + delta = runtime.take_observation_delta() + assert delta is not None + annotation = decode_annotation(delta + runtime.add_terminal()) + + runs = [run for segment in annotation.segments for run in segment.runs] + assert [run.wait_id for run in runs] == [1, 2, 1, 2], ( + f"expected one run per delivery in the observed order, got " + f"{[(r.wait_id, r.count) for r in runs]}" + ) + assert all(run.count == 1 for run in runs) + + +def test_consecutive_records_from_one_stream_are_one_run( + manager: StubManager, backend: MemoryStreamBackend +) -> None: + """The other half of the same rule, and what keeps the annotation compact. + + If every record became its own run the annotation would grow per item, which + is the cost model this feature exists to avoid. + """ + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=runtime.stream_key("a")) + for i in range(5): + runtime.record_delivery(1, data(i, f"10-{i}")) + + delta = runtime.take_observation_delta() + assert delta is not None + annotation = decode_annotation(delta + runtime.add_terminal()) + + runs = [run for segment in annotation.segments for run in segment.runs] + assert len(runs) == 1 + assert runs[0].count == 5 + + +# --- merge -------------------------------------------------------------------- + + +class FakeRuntime: + """Enough of the runtime protocol for the merge iterator, and no more.""" + + def __init__(self) -> None: + self.buffers: dict[int, list[StreamRecord]] = {} + self.blocked: list[tuple[int, bool]] = [] + self.pending: dict[int, asyncio.Future[None]] = {} + self.deliveries: list[tuple[int, StreamRecord]] = [] + self.consumed: list[tuple[int, StreamRecord]] = [] + self.registrations: list[tuple[int, StreamKey]] = [] + self.budget = MAX_RECORDS_PER_ACTIVATION + + def stream_key(self, stream_name: str) -> StreamKey: + return StreamKey("ns", "wf", "first-run", stream_name) + + def register( + self, + *, + wait_id: int, + stream_key: StreamKey, + idle_timeout: timedelta, # pyright: ignore[reportUnusedParameter] + ): + self.registrations.append((wait_id, stream_key)) + + def drain(self, wait_id: int, max_records: int | None = None): + buffered = self.buffers.get(wait_id, []) + if max_records is not None: + buffered, self.buffers[wait_id] = ( + buffered[:max_records], + buffered[max_records:], + ) + else: + self.buffers[wait_id] = [] + return buffered + + def delivery_budget_remaining(self) -> int: + return self.budget + + def codec_for( + self, + value_type: type | None, + wait_id: int, # pyright: ignore[reportUnusedParameter] + ): + return StreamPayloadCodec( + temporalio.converter.DataConverter.default, value_type + ) + + def new_readiness_future(self) -> asyncio.Future[None]: + return asyncio.get_event_loop().create_future() + + def record_delivery(self, wait_id: int, record: StreamRecord) -> None: + self.deliveries.append((wait_id, record)) + + def record_consumption(self, wait_id: int, record: StreamRecord) -> None: + self.consumed.append((wait_id, record)) + self.budget = max(0, self.budget - 1) + + def note_blocked(self, wait_id: int, blocked: bool) -> None: + self.blocked.append((wait_id, blocked)) + + def register_pending(self, wait_id: int, future: asyncio.Future[None]) -> None: + self.pending[wait_id] = future + + def discard_pending(self, wait_id: int) -> None: + self.pending.pop(wait_id, None) + + +class FakeInstance: + pass + + +@pytest.fixture +def fake_runtime(monkeypatch: pytest.MonkeyPatch) -> FakeRuntime: + instance = FakeInstance() + monkeypatch.setattr(temporalio.workflow, "instance", lambda: instance) + runtime = FakeRuntime() + _install_runtime(instance, runtime) # type: ignore[arg-type] + return runtime + + +async def encoded(*values: str) -> list[StreamRecord]: + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + return [ + StreamRecord(RecordKind.DATA, await codec.encode(v), "p", i) + for i, v in enumerate(values) + ] + + +@pytest.mark.asyncio +async def test_merge_yields_from_every_subscription(fake_runtime: FakeRuntime) -> None: + """One record per subscription per pass, in ``wait_id`` order. + + The second stream's record comes out between the first stream's two, not + behind them: a pass that emptied one subscription before looking at the next + would let a backlogged first stream spend a whole activation budget on + itself and reach the second one never. + """ + first = external_stream.topic("a", type=str).subscribe() + second = external_stream.topic("b", type=str).subscribe() + fake_runtime.buffers[first.wait_id] = await encoded("a1", "a2") + fake_runtime.buffers[second.wait_id] = await encoded("b1") + + seen = [] + async for subscription, value in merge(first, second): + seen.append((subscription.wait_id, value)) + if len(seen) == 3: + break + + assert seen == [(1, "a1"), (2, "b1"), (1, "a2")] + + +@pytest.mark.asyncio +async def test_merge_drains_in_wait_id_order_regardless_of_argument_order( + fake_runtime: FakeRuntime, +) -> None: + """Records that arrived in one batch have no inherent order between them. + + An order that depended on argument order -- or on which watcher happened to + run first -- would replay differently than it ran. + """ + first = external_stream.topic("a", type=str).subscribe() + second = external_stream.topic("b", type=str).subscribe() + fake_runtime.buffers[first.wait_id] = await encoded("a1") + fake_runtime.buffers[second.wait_id] = await encoded("b1") + + seen = [] + async for subscription, value in merge(second, first): + seen.append(value) + if len(seen) == 2: + break + + assert seen == ["a1", "b1"] + + +@pytest.mark.asyncio +async def test_merge_marks_every_wait_blocked_when_nothing_is_ready( + fake_runtime: FakeRuntime, +) -> None: + """The quiescent set must name the complete set. + + Blocking on only the first would leave the second out of the snapshot, and + Core would park a wait the Workflow was still waiting on. + """ + first = external_stream.topic("a", type=str).subscribe() + second = external_stream.topic("b", type=str).subscribe() + + iterator = merge(first, second) + pending = asyncio.ensure_future(iterator.__anext__()) + await asyncio.sleep(0.05) + + assert set(fake_runtime.pending) == {1, 2} + assert (1, True) in fake_runtime.blocked and (2, True) in fake_runtime.blocked + pending.cancel() + try: + await pending + except asyncio.CancelledError: + pass + assert dict(fake_runtime.blocked) == {1: False, 2: False}, ( + "cancelling the merged wait left members in the blocked set" + ) + assert not fake_runtime.pending, "cancelling left readiness futures registered" + + +@pytest.mark.asyncio +async def test_merge_resumes_when_any_one_wait_is_resolved( + fake_runtime: FakeRuntime, +) -> None: + """One record on one stream is enough; it need not wait for the others.""" + first = external_stream.topic("a", type=str).subscribe() + second = external_stream.topic("b", type=str).subscribe() + + iterator = merge(first, second) + pending = asyncio.ensure_future(iterator.__anext__()) + await asyncio.sleep(0.05) + + fake_runtime.buffers[second.wait_id] = await encoded("b1") + fake_runtime.pending[second.wait_id].set_result(None) + + assert await asyncio.wait_for(pending, 1) == (second, "b1") + + +@pytest.mark.asyncio +async def test_merging_a_subscription_with_itself_is_refused( + fake_runtime: FakeRuntime, # pyright: ignore[reportUnusedParameter] +) -> None: + """It would deliver every record to that wait twice.""" + only = external_stream.topic("a", type=str).subscribe() + + with pytest.raises(ValueError, match="same subscription twice"): + await merge(only, only).__anext__() + + +@pytest.mark.asyncio +async def test_merging_nothing_is_refused( + fake_runtime: FakeRuntime, # pyright: ignore[reportUnusedParameter] +) -> None: + with pytest.raises(ValueError, match="at least one"): + await merge().__anext__() + + +@pytest.mark.asyncio +async def test_merge_never_yields_control_records(fake_runtime: FakeRuntime) -> None: + """A fence advances the cursor but belongs to the runtime, not the Workflow.""" + first = external_stream.topic("a", type=str).subscribe() + second = external_stream.topic("b", type=str).subscribe() + records = await encoded("a1", "a2") + fake_runtime.buffers[first.wait_id] = [ + records[0], + StreamRecord(RecordKind.WRITE_FENCE, b"", "p", 2), + records[1], + ] + fake_runtime.buffers[second.wait_id] = await encoded("b1") + + seen = [] + async for subscription, value in merge(first, second): + seen.append(value) + if len(seen) == 3: + break + + # The fence spends the first stream's turn on the second pass, so its "a2" + # arrives on the third -- consumed, and never yielded. + assert seen == ["a1", "b1", "a2"] + # The fence was still consumed -- it occupies an offset inside a run. + assert any(r.is_control for _, r in fake_runtime.consumed) + + +@pytest.mark.asyncio +async def test_a_replayed_segment_is_yielded_in_its_recorded_order( + public_api_runtime: WorkflowStreamRuntime, +) -> None: + """Rotation changes which wait is *asked* first, never what is *served*. + + ``merge`` resumes each pass after the wait that last took a record, and + under replay the budget is unbounded, so the passes are not cut where they + were cut live and the rotation reaches positions the live run never started + from. That is safe only because a replay drain serves from the **front** of + the recorded segment and only while the front belongs to the asking wait: a + wait asked out of turn is told nothing and the record stays for whoever + asks next. Every active wait is still asked exactly once per pass, so the + front's owner is always reached and the recorded global order comes out + whole. + + A recorded order that interleaves three waits and repeats one of them is + what would expose an ask order leaking into a yield order -- and a yield + order that differed from History is nondeterminism, which is the one thing + a fairness fix may not buy its fairness with. + """ + recorded = [2, 1, 3, 1, 2, 3, 3] + subscriptions = [ + external_stream.topic(name, type=str).subscribe() for name in ("a", "b", "c") + ] + assert [s.wait_id for s in subscriptions] == [1, 2, 3] + records = await encoded(*(f"r{i}" for i in range(len(recorded)))) + public_api_runtime.begin_activation() + public_api_runtime.begin_replay_segment( + [ + (wait_id, record.placed_at(Offset(f"{index:08d}"))) + for index, (wait_id, record) in enumerate(zip(recorded, records)) + ] + ) + + iterator = merge(*subscriptions) + seen: list[tuple[int, str]] = [] + for _ in recorded: + subscription, value = await asyncio.wait_for(iterator.__anext__(), 1) + seen.append((subscription.wait_id, value)) + await iterator.aclose() + + assert [wait_id for wait_id, _ in seen] == recorded, ( + "the replayed segment came out in an order the live run never " + "delivered; the merge is asking in a different order than it recorded " + "in and the drain is obliging" + ) + assert [value for _, value in seen] == [f"r{i}" for i in range(len(recorded))] + + +# --- through a real Worker ---------------------------------------------------- + + +@workflow.defn +class TwoStreamWorkflow: + """Consumes from two streams as one wait set, reporting the order it saw.""" + + @workflow.run + async def run(self, expected: int) -> list[str]: + from temporalio.contrib.external_workflow_streams._api import ( + external_stream as es, + ) + from temporalio.contrib.external_workflow_streams._api import merge as m + + options = es.with_options(idle_timeout=timedelta(seconds=30)) + left = options.topic("left", type=str).subscribe() + right = options.topic("right", type=str).subscribe() + + seen: list[str] = [] + async for _, value in m(left, right): + seen.append(value) + if len(seen) >= expected: + break + return seen + + +async def test_a_workflow_consumes_two_streams_as_one_wait_set( + client: Client, backend: MemoryStreamBackend +) -> None: + """One idle stream must not stall the other. + + Only the ``right`` stream is ever written to. If the two were waited on + independently the Workflow would block on ``left`` forever and never see a + single record. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[TwoStreamWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + TwoStreamWorkflow.run, + 2, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + description = await handle.describe() + first_run = description.raw_description.workflow_execution_info.first_run_id + right = StreamKey(client.namespace, handle.id, first_run, "right") + await asyncio.sleep(1) + + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + for i, value in enumerate(["r1", "r2"]): + await backend.append( + right, + StreamRecord(RecordKind.DATA, await codec.encode(value), "p", i), + ) + + assert await asyncio.wait_for(handle.result(), 60) == ["r1", "r2"] + + +# --- a wait nothing is awaiting ------------------------------------------------ + + +@pytest.fixture +def public_api_runtime( + manager: StubManager, + backend: MemoryStreamBackend, + monkeypatch: pytest.MonkeyPatch, +) -> WorkflowStreamRuntime: + """The real runtime, reachable through ``subscribe()``. + + The quiescent snapshot is the thing under test here, and only the real + runtime has one. + """ + + class Instance: + pass + + instance = Instance() + monkeypatch.setattr(temporalio.workflow, "instance", lambda: instance) + runtime = make_runtime(manager, backend) + _install_runtime(instance, runtime) # type: ignore[arg-type] + return runtime + + +def test_a_subscription_nobody_has_iterated_is_not_blocked( + public_api_runtime: WorkflowStreamRuntime, +) -> None: + """Blocked means *Workflow code is waiting*, not *a subscription exists*. + + A snapshot is a request to Core to retain the Workflow Task and, once the + idle timer expires, to park it. Constructing a subscription and then doing + something else entirely -- a timer, an activity, a signal handler -- must not + on its own produce a wait for Core to hold the task open for. + """ + external_stream.topic("tokens", type=str).subscribe() + + assert public_api_runtime.quiescent_snapshot() is None, ( + "a subscription that has never been iterated reported itself as a " + "quiescent wait, so Core would retain and eventually park the Workflow " + "Task for a wait no coroutine is awaiting" + ) + + +@pytest.mark.asyncio +async def test_a_cancelled_wait_leaves_the_quiescent_set( + public_api_runtime: WorkflowStreamRuntime, +) -> None: + """Racing a stream against a timer is ordinary code, and the loser is cancelled. + + ``asyncio.wait(..., FIRST_COMPLETED)`` followed by cancelling the pending + half is the standard shape. If the cancelled half stays in the blocked set, + every such race leaves a ghost wait behind, and the ghosts accumulate. + """ + subscription = external_stream.topic("tokens", type=str).subscribe() + + pending = asyncio.ensure_future(subscription.__aiter__().__anext__()) + await asyncio.sleep(0.05) + assert public_api_runtime.quiescent_snapshot() is not None, ( + "a wait that is genuinely being awaited must be in the snapshot" + ) + + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + + assert public_api_runtime.quiescent_snapshot() is None, ( + "the cancelled wait stayed in the quiescent set, so Core is asked to " + "hold the Workflow Task open for a coroutine that no longer exists" + ) + + +@pytest.mark.asyncio +async def test_closing_a_subscription_ends_its_wait_and_its_iteration( + public_api_runtime: WorkflowStreamRuntime, +) -> None: + """``_finished`` has to be reachable, or iteration has no end at all.""" + subscription = external_stream.topic("tokens", type=str).subscribe() + + iterator = subscription.__aiter__() + pending = asyncio.ensure_future(iterator.__anext__()) + await asyncio.sleep(0.05) + assert not pending.done() + + subscription.close() + + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(pending, 1) + assert public_api_runtime.quiescent_snapshot() is None + # Idempotent: closing twice is not an error, because the ordinary shape is a + # `finally` that cannot know whether the iterator already ended. + subscription.close() + + +@pytest.mark.asyncio +async def test_closing_a_subscription_tells_the_worker_to_stop_serving_it( + public_api_runtime: WorkflowStreamRuntime, # pyright: ignore[reportUnusedParameter] + manager: StubManager, +) -> None: + """Ending the wait is what the Workflow sees; the Worker has to be told too. + + Nothing above this line would notice the message never being sent: iteration + ends, the wait leaves the quiescent set, and the Workflow runs on. What is + left behind is a watcher still prefetching into a buffer nobody can drain + and a park intent still in the backend, and only the manager can take either + of them back. + """ + subscription = external_stream.topic("tokens", type=str).subscribe() + + subscription.close() + + assert manager.cancelled == [(RUN_ID, subscription.wait_id)], ( + "closing left the Worker serving a wait no Workflow code is reading" + ) + # Cancelled once, however often it is closed: by the second call the Worker + # has already dropped this wait, and its state here is kept only so that + # replay and a Continue-As-New successor can still read it. + subscription.close() + assert manager.cancelled == [(RUN_ID, subscription.wait_id)] + + +@pytest.mark.asyncio +async def test_a_resolved_merge_leaves_no_nonwinning_wait_blocked( + public_api_runtime: WorkflowStreamRuntime, + manager: StubManager, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A successful group wait ends for every member, not only its winner. + + The readiness activation resolves the group's futures together, while a + record may exist for only one wait. The others have no coroutine behind + them once the group returns, so leaving one in the quiescent snapshot asks + Core to retain and eventually park a Workflow Task for nobody. + """ + first = external_stream.topic("a", type=str).subscribe() + second = external_stream.topic("b", type=str).subscribe() + iterator = merge(first, second) + pending = asyncio.ensure_future(iterator.__anext__()) + await asyncio.sleep(0.05) + + snapshot = public_api_runtime.quiescent_snapshot() + assert snapshot is not None and [wait.wait_id for wait in snapshot] == [1, 2] + + record = (await encoded("b1"))[0].placed_at(Offset("00000001")) + delivered = False + + def drain(run_id, wait_id, max_records=None): # type: ignore[no-untyped-def] + nonlocal delivered + if wait_id == second.wait_id and not delivered: + delivered = True + return [record] + return [] + + monkeypatch.setattr(manager, "drain", drain) + public_api_runtime.resolve_all_pending() + + assert await asyncio.wait_for(pending, 1) == (second, "b1") + snapshot_after_yield = public_api_runtime.quiescent_snapshot() + await iterator.aclose() + + assert snapshot_after_yield is None, ( + "the merge yielded one member but left its non-winning member blocked, " + "so Core would retain the Workflow Task for a discarded future" + ) + + +@pytest.mark.asyncio +async def test_a_merge_double_check_leaves_no_nonwinning_wait_blocked( + public_api_runtime: WorkflowStreamRuntime, + manager: StubManager, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The post-registration early return has the same group cleanup contract. + + A record can arrive after the initial pass found every stream empty but + before the pending futures are awaited. `_await_any_readiness()` finds it in + its double-check and returns early; that successful exit must still unblock + every member. + """ + first = external_stream.topic("a", type=str).subscribe() + second = external_stream.topic("b", type=str).subscribe() + record = (await encoded("b1"))[0].placed_at(Offset("00000001")) + drain_calls: dict[int, int] = {} + + def drain(run_id, wait_id, max_records=None): # type: ignore[no-untyped-def] + drain_calls[wait_id] = drain_calls.get(wait_id, 0) + 1 + # First call: the merge's initial pass. Second call: the registered + # wait's double-check immediately before `asyncio.wait()`. + if wait_id == second.wait_id and drain_calls[wait_id] == 2: + return [record] + return [] + + monkeypatch.setattr(manager, "drain", drain) + iterator = merge(first, second) + + assert await asyncio.wait_for(iterator.__anext__(), 1) == (second, "b1") + snapshot_after_yield = public_api_runtime.quiescent_snapshot() + await iterator.aclose() + + assert snapshot_after_yield is None, ( + "the double-check returned a record but left another group member in " + "the quiescent snapshot" + ) + + +@pytest.mark.asyncio +async def test_a_half_registered_merge_leaves_no_wait_blocked( + public_api_runtime: WorkflowStreamRuntime, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Registering the group is covered by the cleanup that ends it. + + The refusal pass is careful to check every member before claiming any, so + that a refused merge claims nothing. A raise part-way through the claim + itself reaches the same state by the other route: the members already + marked blocked have futures nothing can resolve, and the snapshot asks Core + to retain a Workflow Task on their behalf. + """ + first = external_stream.topic("a", type=str).subscribe() + second = external_stream.topic("b", type=str).subscribe() + + real = public_api_runtime.new_readiness_future + made = 0 + + def fail_on_the_second() -> asyncio.Future[None]: + nonlocal made + made += 1 + if made == 2: + raise RuntimeError("no future for you") + return real() + + monkeypatch.setattr(public_api_runtime, "new_readiness_future", fail_on_the_second) + + with pytest.raises(RuntimeError, match="no future for you"): + await merge(first, second).__anext__() + + assert public_api_runtime.quiescent_snapshot() is None, ( + "a merge that failed while claiming its group left a wait blocked with " + "no future behind it, so Core would retain the Workflow Task for it" + ) + assert public_api_runtime.resolve_all_pending() == 0, ( + "a merge that failed while claiming its group left a readiness future " + "registered, which the next activation would resolve for nobody" + ) + assert first._pending_future is None and second._pending_future is None + + # And the group is claimable again rather than permanently refused: the + # single-slot claim each member holds is what the refusal reads. + monkeypatch.setattr(public_api_runtime, "new_readiness_future", real) + pending = asyncio.ensure_future(merge(first, second).__anext__()) + await asyncio.sleep(0.05) + snapshot = public_api_runtime.quiescent_snapshot() + assert snapshot is not None and [wait.wait_id for wait in snapshot] == [1, 2] + pending.cancel() + try: + await pending + except asyncio.CancelledError: + pass diff --git a/tests/contrib/external_workflow_streams/test_output_backend.py b/tests/contrib/external_workflow_streams/test_output_backend.py new file mode 100644 index 0000000000..7243bc0758 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_output_backend.py @@ -0,0 +1,185 @@ +"""Provider-neutral output identity, records, and staging contract.""" + +from __future__ import annotations + +import dataclasses + +import pytest + +from temporalio.contrib.external_workflow_streams._backend import ( + StreamBackend, + StreamDirection, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._output_backend import ( + OutputReadResult, + OutputStage, + OutputStageManifest, + OutputStageStatus, + OutputStreamBackend, + OutputStreamRecord, + PendingOutputBarrier, + StagedOutputRecord, +) +from temporalio.contrib.external_workflow_streams._producer import WorkflowChainKey +from temporalio.contrib.external_workflow_streams._record import ( + Offset, + RecordKind, + StreamRecord, +) + + +def output_key(name: str = "events") -> StreamKey: + return StreamKey( + "namespace", + "workflow", + "first-run", + name, + direction=StreamDirection.OUTPUT, + ) + + +def manifest(**changes: object) -> OutputStageManifest: + values: dict[str, object] = { + "stream_key": output_key(), + "provider_id": "memory", + "provider_format_version": 1, + "stage_token": "stage-token", + "run_id": "current-run", + "history_floor_event_id": 7, + "sub_batch_id": 0, + "fingerprint_version": 1, + "fingerprint": b"f" * 32, + "record_count": 2, + "logical_byte_count": 41, + } + values.update(changes) + return OutputStageManifest(**values) # type: ignore[arg-type] + + +def test_stream_direction_defaults_to_the_existing_input_identity() -> None: + existing = StreamKey("namespace", "workflow", "first-run", "events") + + assert existing.direction is StreamDirection.INPUT + assert existing == StreamKey( + "namespace", + "workflow", + "first-run", + "events", + direction=StreamDirection.INPUT, + ) + assert str(existing) == "namespace/workflow/first-run/input/events" + + +def test_input_and_output_topics_with_the_same_name_are_distinct() -> None: + input_key = StreamKey("namespace", "workflow", "first-run", "events") + output = output_key() + + assert input_key != output + assert len({input_key, output}) == 2 + assert str(output) == "namespace/workflow/first-run/output/events" + + +def test_workflow_chain_key_builds_either_direction() -> None: + chain = WorkflowChainKey("namespace", "workflow", "first-run") + + assert chain.stream_key("events") == StreamKey( + "namespace", "workflow", "first-run", "events" + ) + assert chain.stream_key("events", direction=StreamDirection.OUTPUT) == output_key() + + +def test_stream_key_rejects_a_raw_direction_string() -> None: + with pytest.raises(TypeError, match="must be a StreamDirection"): + StreamKey( + "namespace", + "workflow", + "first-run", + "events", + direction="output", # type: ignore[arg-type] + ) + + +def test_output_manifest_is_bound_to_an_output_key() -> None: + with pytest.raises(ValueError, match="requires an OUTPUT stream key"): + manifest(stream_key=StreamKey("namespace", "workflow", "first-run", "events")) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("provider_id", "", "non-empty provider_id"), + ("stage_token", "", "non-empty stage_token"), + ("run_id", "", "non-empty run_id"), + ("provider_format_version", 0, "positive provider_format_version"), + ("fingerprint_version", 0, "positive fingerprint_version"), + ("record_count", 0, "positive record_count"), + ("history_floor_event_id", 0, "positive history_floor_event_id"), + ("sub_batch_id", -1, "non-negative sub_batch_id"), + ("logical_byte_count", -1, "non-negative logical_byte_count"), + ("fingerprint", b"short", "must be 32 bytes"), + ], +) +def test_output_manifest_rejects_incomplete_identity( + field: str, value: object, message: str +) -> None: + with pytest.raises(ValueError, match=message): + manifest(**{field: value}) + + +def test_staged_and_placed_control_records_carry_no_payload() -> None: + staged = StagedOutputRecord(0, RecordKind.FINISH, b"") + placed = OutputStreamRecord(RecordKind.FINISH, b"", Offset("1-0")) + + assert staged.kind is RecordKind.FINISH + assert placed.kind.is_control + with pytest.raises(ValueError, match="carries no payload"): + StagedOutputRecord(0, RecordKind.FINISH, b"not-empty") + with pytest.raises(ValueError, match="carries no payload"): + OutputStreamRecord(RecordKind.WRITE_FENCE, b"not-empty", Offset("1-1")) + + +def test_finish_round_trips_through_provider_neutral_record_fields() -> None: + finish = StreamRecord(RecordKind.FINISH, b"", "session", 3).placed_at(Offset("2-0")) + + assert finish.is_control + assert StreamRecord.from_fields(Offset("2-0"), finish.to_fields()) == finish + + +def test_output_stage_record_count_must_match_its_manifest() -> None: + records = ( + OutputStreamRecord(RecordKind.DATA, b"one", Offset("1-0")), + OutputStreamRecord(RecordKind.DATA, b"two", Offset("1-1")), + ) + + stage = OutputStage(manifest(), records, OutputStageStatus.PENDING) + assert stage.records == records + with pytest.raises(ValueError, match="record count does not match"): + OutputStage(manifest(record_count=1), records, OutputStageStatus.PENDING) + + +def test_output_models_are_frozen_and_use_immutable_record_collections() -> None: + stage_manifest = manifest(record_count=1) + record = OutputStreamRecord(RecordKind.DATA, b"data", Offset("1-0")) + stage = OutputStage(stage_manifest, (record,), OutputStageStatus.PENDING) + barrier = PendingOutputBarrier(stage_manifest, Offset("1-0")) + result = OutputReadResult((record,), barrier) + + assert isinstance(stage.records, tuple) + assert isinstance(result.records, tuple) + with pytest.raises(dataclasses.FrozenInstanceError): + result.pending = None # type: ignore[misc] + + +def test_output_backend_is_an_additive_capability() -> None: + assert not issubclass(OutputStreamBackend, StreamBackend) + assert OutputStreamBackend.__abstractmethods__ == { + "abort_output", + "append_output", + "commit_output", + "compare_offsets", + "output_stage", + "output_tail", + "read_output_after", + "stage_output", + } diff --git a/tests/contrib/external_workflow_streams/test_output_client.py b/tests/contrib/external_workflow_streams/test_output_client.py new file mode 100644 index 0000000000..ab8a7ea51c --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_output_client.py @@ -0,0 +1,543 @@ +"""Focused client and pending-stage reconciliation tests for output streams.""" + +from __future__ import annotations + +from collections import deque +from collections.abc import AsyncIterator, Sequence +from types import SimpleNamespace +from typing import Any + +import pytest + +import temporalio.api.history.v1 +import temporalio.converter +import temporalio.service +from temporalio.contrib.external_workflow_streams import _output_client +from temporalio.contrib.external_workflow_streams._backend import ( + StreamDirection, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec +from temporalio.contrib.external_workflow_streams._errors import ( + StreamDecodeError, + StreamIntegrityError, + StreamStorageError, +) +from temporalio.contrib.external_workflow_streams._output_backend import ( + OutputReadResult, + OutputStage, + OutputStageManifest, + OutputStageStatus, + OutputStreamBackend, + OutputStreamRecord, + PendingOutputBarrier, + StagedOutputRecord, +) +from temporalio.contrib.external_workflow_streams._output_client import ( + ExternalOutputStreamClient, + ExternalOutputStreamClientTopic, +) +from temporalio.contrib.external_workflow_streams._producer import WorkflowChainKey +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Cursor, + Offset, + RecordKind, + StreamRecord, +) + + +class FakeOutputBackend(OutputStreamBackend): + guarantees_immutability = True + provider_id = "fake-output" + provider_format_version = 1 + + def __init__(self) -> None: + self.read_results: deque[OutputReadResult | BaseException] = deque() + self.read_cursors: list[Cursor] = [] + self.tail = BEGINNING + self.tail_calls = 0 + self.stages: dict[OutputStageManifest, OutputStage] = {} + self.committed: list[OutputStageManifest] = [] + self.aborted: list[OutputStageManifest] = [] + + async def stage_output( + self, + manifest: OutputStageManifest, + records: Sequence[StagedOutputRecord], + ) -> OutputStage: + raise NotImplementedError + + async def commit_output(self, manifest: OutputStageManifest) -> OutputStage: + self.committed.append(manifest) + stage = self.stages[manifest] + resolved = OutputStage(manifest, stage.records, OutputStageStatus.COMMITTED) + self.stages[manifest] = resolved + return resolved + + async def abort_output(self, manifest: OutputStageManifest) -> OutputStage: + self.aborted.append(manifest) + stage = self.stages[manifest] + resolved = OutputStage(manifest, stage.records, OutputStageStatus.ABORTED) + self.stages[manifest] = resolved + return resolved + + async def output_stage(self, manifest: OutputStageManifest) -> OutputStage | None: + return self.stages.get(manifest) + + async def append_output( + self, key: StreamKey, record: StreamRecord + ) -> OutputStreamRecord: + raise NotImplementedError + + async def read_output_after( + self, + key: StreamKey, + after: Cursor, + *, + max_records: int, + block: Any = None, + ) -> OutputReadResult: + assert key.direction is StreamDirection.OUTPUT + assert max_records > 0 + self.read_cursors.append(after) + result = self.read_results.popleft() + if isinstance(result, BaseException): + raise result + return result + + async def output_tail(self, key: StreamKey) -> Cursor: + assert key.direction is StreamDirection.OUTPUT + self.tail_calls += 1 + return self.tail + + def compare_offsets(self, left: Offset, right: Offset) -> int: + return int(left.token) - int(right.token) + + +class FakeHandle: + def __init__( + self, + *, + first_run_id: str, + events: Sequence[temporalio.api.history.v1.HistoryEvent] = (), + history_error: BaseException | None = None, + ) -> None: + self._first_run_id = first_run_id + self._events = events + self._history_error = history_error + + async def describe(self) -> Any: + return SimpleNamespace( + raw_description=SimpleNamespace( + workflow_execution_info=SimpleNamespace(first_run_id=self._first_run_id) + ) + ) + + def fetch_history_events(self) -> AsyncIterator[Any]: + async def iterate() -> AsyncIterator[Any]: + if self._history_error is not None: + raise self._history_error + for event in self._events: + yield event + + return iterate() + + +class FakeClient: + namespace = "ns" + data_converter = temporalio.converter.DataConverter.default + + def __init__(self) -> None: + self.events: Sequence[temporalio.api.history.v1.HistoryEvent] = () + self.history_error: BaseException | None = None + self.history_run_ids: list[str] = [] + + def get_workflow_handle( + self, workflow_id: str, *, run_id: str | None = None + ) -> FakeHandle: + assert workflow_id == "wf" + if run_id is not None: + self.history_run_ids.append(run_id) + return FakeHandle( + first_run_id="first-run", + events=self.events if run_id is not None else (), + history_error=self.history_error if run_id is not None else None, + ) + + +CHAIN = WorkflowChainKey("ns", "wf", "first-run") +OUTPUT_KEY = CHAIN.stream_key("events", direction=StreamDirection.OUTPUT) + + +def history_event(event_id: int, attributes: str) -> Any: + event = temporalio.api.history.v1.HistoryEvent(event_id=event_id) + getattr(event, attributes).SetInParent() + return event + + +def manifest(*, floor: int = 1, token: str = "stage-1") -> OutputStageManifest: + return OutputStageManifest( + stream_key=OUTPUT_KEY, + provider_id="fake-output", + provider_format_version=1, + stage_token=token, + run_id="current-run", + history_floor_event_id=floor, + sub_batch_id=0, + fingerprint_version=1, + fingerprint=b"f" * 32, + record_count=1, + logical_byte_count=10, + ) + + +def output_record( + offset: int, payload: bytes = b"", *, kind: RecordKind = RecordKind.DATA +) -> OutputStreamRecord: + return OutputStreamRecord(kind=kind, payload=payload, offset=Offset(str(offset))) + + +async def connected( + backend: FakeOutputBackend, client: FakeClient +) -> ExternalOutputStreamClient: + return await ExternalOutputStreamClient.connect( + backend=backend, + workflow=CHAIN, + client=client, # type: ignore[arg-type] + ) + + +async def test_committed_read_is_typed_resumable_and_backend_only() -> None: + backend = FakeOutputBackend() + client = FakeClient() + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + first = output_record(4, await codec.encode("first")) + second = output_record(5, await codec.encode("second")) + finish = output_record(6, kind=RecordKind.FINISH) + backend.read_results.extend( + [ + OutputReadResult((first,)), + OutputReadResult((second, finish)), + ] + ) + + topic = (await connected(backend, client)).topic("events", type=str) + items = [item async for item in topic.subscribe(after=AFTER(Offset("3")))] + + assert [(item.data, item.offset) for item in items] == [ + ("first", Offset("4")), + ("second", Offset("5")), + ] + assert backend.read_cursors == [AFTER(Offset("3")), AFTER(Offset("4"))] + assert client.history_run_ids == [] + + +async def test_an_item_offset_can_be_passed_directly_as_the_resume_cursor() -> None: + backend = FakeOutputBackend() + client = FakeClient() + backend.read_results.append( + OutputReadResult((output_record(4, kind=RecordKind.FINISH),)) + ) + + topic: ExternalOutputStreamClientTopic[Any] = ( + await connected(backend, client) + ).topic("events") + assert [item async for item in topic.subscribe(after=Offset("3"))] == [] + + assert backend.read_cursors == [AFTER(Offset("3"))] + + +async def test_tail_is_the_backend_committed_boundary_without_history() -> None: + backend = FakeOutputBackend() + backend.tail = AFTER(Offset("12")) + client = FakeClient() + + topic: ExternalOutputStreamClientTopic[Any] = ( + await connected(backend, client) + ).topic("events") + + assert await topic.tail() == AFTER(Offset("12")) + assert backend.tail_calls == 1 + assert client.history_run_ids == [] + + +async def test_pending_token_in_exact_run_history_commits_before_read_continues( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend = FakeOutputBackend() + client = FakeClient() + staged = manifest() + placed = output_record(10, b"payload") + backend.stages[staged] = OutputStage(staged, (placed,), OutputStageStatus.PENDING) + barrier = PendingOutputBarrier(staged, Offset("10")) + client.events = ( + history_event(1, "workflow_execution_started_event_attributes"), + history_event(2, "workflow_task_completed_event_attributes"), + history_event(3, "marker_recorded_event_attributes"), + ) + monkeypatch.setattr( + _output_client, + "_event_has_output_stage_token", + lambda event, token: event.event_id == 3 and token == "stage-1", + ) + + topic: ExternalOutputStreamClientTopic[Any] = ( + await connected(backend, client) + ).topic("events") + + assert await topic._resolve_pending(barrier) + assert backend.committed == [staged] + assert backend.aborted == [] + assert client.history_run_ids == ["current-run"] + + +async def test_repeating_the_same_history_reconciliation_is_idempotent( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend = FakeOutputBackend() + client = FakeClient() + staged = manifest() + placed = output_record(10, b"payload") + backend.stages[staged] = OutputStage(staged, (placed,), OutputStageStatus.PENDING) + barrier = PendingOutputBarrier(staged, placed.offset) + client.events = ( + history_event(1, "workflow_execution_started_event_attributes"), + history_event(2, "workflow_task_completed_event_attributes"), + history_event(3, "marker_recorded_event_attributes"), + ) + monkeypatch.setattr( + _output_client, + "_event_has_output_stage_token", + lambda event, token: event.event_id == 3 and token == "stage-1", + ) + topic: ExternalOutputStreamClientTopic[Any] = ( + await connected(backend, client) + ).topic("events") + + assert await topic._resolve_pending(barrier) + assert await topic._resolve_pending(barrier) + assert backend.committed == [staged, staged] + assert backend.stages[staged].status is OutputStageStatus.COMMITTED + + +def test_real_shared_marker_envelope_exposes_the_output_stage_token() -> None: + from temporalio.bridge.proto.external_data import ExternalStreamMarkerData + + marker = ExternalStreamMarkerData() + marker.output.stage_token = "stage-1" + event = history_event(3, "marker_recorded_event_attributes") + attributes = event.marker_recorded_event_attributes + attributes.marker_name = "core_external_stream" + attributes.details[ + "external_stream" + ].payloads.add().data = marker.SerializeToString() + + assert _output_client._event_has_output_stage_token(event, "stage-1") + assert not _output_client._event_has_output_stage_token(event, "another-stage") + + +async def test_subscribe_stops_at_barrier_and_resumes_from_committed_prefix( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend = FakeOutputBackend() + client = FakeClient() + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + prefix = output_record(5, await codec.encode("before")) + staged_record = output_record(10, await codec.encode("staged")) + finish = output_record(11, kind=RecordKind.FINISH) + staged = manifest() + backend.stages[staged] = OutputStage( + staged, (staged_record,), OutputStageStatus.PENDING + ) + backend.read_results.extend( + [ + OutputReadResult( + (prefix,), PendingOutputBarrier(staged, staged_record.offset) + ), + OutputReadResult((staged_record, finish)), + ] + ) + client.events = ( + history_event(1, "workflow_execution_started_event_attributes"), + history_event(2, "workflow_task_completed_event_attributes"), + history_event(3, "marker_recorded_event_attributes"), + ) + monkeypatch.setattr( + _output_client, + "_event_has_output_stage_token", + lambda event, token: event.event_id == 3 and token == "stage-1", + ) + topic = (await connected(backend, client)).topic("events", type=str) + + items = [item async for item in topic.subscribe()] + + assert [item.data for item in items] == ["before", "staged"] + assert backend.read_cursors == [BEGINNING, AFTER(prefix.offset)] + assert backend.committed == [staged] + + +async def test_first_task_boundary_without_token_aborts_pending_stage( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend = FakeOutputBackend() + client = FakeClient() + staged = manifest() + backend.stages[staged] = OutputStage( + staged, (output_record(10),), OutputStageStatus.PENDING + ) + client.events = ( + history_event(1, "workflow_execution_started_event_attributes"), + history_event(2, "workflow_task_failed_event_attributes"), + ) + monkeypatch.setattr( + _output_client, + "_event_has_output_stage_token", + lambda event, token: False, + ) + + topic: ExternalOutputStreamClientTopic[Any] = ( + await connected(backend, client) + ).topic("events") + + assert await topic._resolve_pending(PendingOutputBarrier(staged, Offset("10"))) + assert backend.aborted == [staged] + assert backend.committed == [] + + +async def test_no_durable_boundary_leaves_the_stage_pending( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend = FakeOutputBackend() + client = FakeClient() + staged = manifest() + backend.stages[staged] = OutputStage( + staged, (output_record(10),), OutputStageStatus.PENDING + ) + client.events = ( + history_event(1, "workflow_execution_started_event_attributes"), + history_event(2, "workflow_task_scheduled_event_attributes"), + history_event(3, "workflow_task_started_event_attributes"), + ) + monkeypatch.setattr( + _output_client, + "_event_has_output_stage_token", + lambda event, token: False, + ) + + topic: ExternalOutputStreamClientTopic[Any] = ( + await connected(backend, client) + ).topic("events") + + assert not await topic._resolve_pending(PendingOutputBarrier(staged, Offset("10"))) + assert backend.committed == [] and backend.aborted == [] + + +async def test_disconnected_speculative_update_stays_pending_until_history_loss( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Elapsed retries cannot guess the outcome of an unreported Update task.""" + backend = FakeOutputBackend() + client = FakeClient() + staged = manifest() + placed = output_record(10, b"speculative-update-output") + backend.stages[staged] = OutputStage(staged, (placed,), OutputStageStatus.PENDING) + barrier = PendingOutputBarrier(staged, placed.offset) + + # The admitted Update exists durably, but there is no deciding Workflow + # Task terminal event above the exact floor. This is the shape left by a + # client disconnect while the Update's speculative task is still in flight. + client.events = ( + history_event(1, "workflow_execution_started_event_attributes"), + history_event(2, "workflow_execution_update_admitted_event_attributes"), + history_event(3, "workflow_task_scheduled_event_attributes"), + history_event(4, "workflow_task_started_event_attributes"), + ) + monkeypatch.setattr( + _output_client, + "_event_has_output_stage_token", + lambda event, token: False, + ) + topic: ExternalOutputStreamClientTopic[Any] = ( + await connected(backend, client) + ).topic("events") + + assert not await topic._resolve_pending(barrier) + assert not await topic._resolve_pending(barrier) + assert backend.stages[staged].status is OutputStageStatus.PENDING + assert backend.committed == [] and backend.aborted == [] + assert client.history_run_ids == ["current-run", "current-run"] + + # Once retention removes that exact Run, lack of evidence becomes an + # integrity failure. It still must not be guessed into an abort. + client.history_error = temporalio.service.RPCError( + "history expired", + temporalio.service.RPCStatusCode.NOT_FOUND, + b"", + ) + with pytest.raises(StreamIntegrityError, match="no longer available"): + await topic._resolve_pending(barrier) + assert backend.stages[staged].status is OutputStageStatus.PENDING + assert backend.committed == [] and backend.aborted == [] + + +async def test_history_loss_is_integrity_but_an_outage_is_storage() -> None: + backend = FakeOutputBackend() + client = FakeClient() + topic: ExternalOutputStreamClientTopic[Any] = ( + await connected(backend, client) + ).topic("events") + + client.events = (history_event(2, "workflow_task_completed_event_attributes"),) + with pytest.raises(StreamIntegrityError, match="exact floor event"): + await topic._history_decision(manifest()) + + client.history_error = temporalio.service.RPCError( + "server unavailable", + temporalio.service.RPCStatusCode.UNAVAILABLE, + b"", + ) + with pytest.raises(StreamStorageError, match="Temporal History"): + await topic._history_decision(manifest()) + + client.history_error = temporalio.service.RPCError( + "history expired", + temporalio.service.RPCStatusCode.NOT_FOUND, + b"", + ) + with pytest.raises(StreamIntegrityError, match="no longer available"): + await topic._history_decision(manifest()) + + +async def test_provider_cannot_return_a_record_beyond_pending_barrier() -> None: + backend = FakeOutputBackend() + client = FakeClient() + staged = manifest() + backend.read_results.append( + OutputReadResult( + (output_record(11),), + PendingOutputBarrier(staged, Offset("10")), + ) + ) + topic: ExternalOutputStreamClientTopic[Any] = ( + await connected(backend, client) + ).topic("events") + + with pytest.raises(StreamIntegrityError, match="pending barrier"): + await anext(topic.subscribe()) + assert client.history_run_ids == [] + + +async def test_backend_and_decode_failures_keep_the_existing_taxonomy() -> None: + backend = FakeOutputBackend() + client = FakeClient() + topic = (await connected(backend, client)).topic("events", type=str) + backend.read_results.append(ConnectionError("backend down")) + + with pytest.raises(StreamStorageError, match="backend down"): + await anext(topic.subscribe()) + + backend.read_results.append(OutputReadResult((output_record(1, b"not-proto"),))) + with pytest.raises(StreamDecodeError): + await anext(topic.subscribe()) diff --git a/tests/contrib/external_workflow_streams/test_output_codec.py b/tests/contrib/external_workflow_streams/test_output_codec.py new file mode 100644 index 0000000000..f6385930b0 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_output_codec.py @@ -0,0 +1,150 @@ +"""Canonical pre-codec output frames and fingerprints.""" + +from __future__ import annotations + +import hashlib +import struct + +import pytest + +import temporalio.api.common.v1 +from temporalio.contrib.external_workflow_streams._output_codec import ( + LOGICAL_FINGERPRINT_VERSION, + canonical_logical_record_frame, + fingerprint_logical_frames, + fingerprint_logical_records, +) +from temporalio.contrib.external_workflow_streams._record import RecordKind + + +def lp(value: bytes) -> bytes: + return struct.pack(">Q", len(value)) + value + + +def test_canonical_frame_has_an_exact_version_one_encoding() -> None: + payload = temporalio.api.common.v1.Payload( + metadata={"z": b"last", "a": b"first"}, + data=b"payload", + ) + + frame = canonical_logical_record_frame("events", RecordKind.DATA, payload) + + assert frame == b"".join( + ( + lp(b"events"), + lp(b"\x01"), + struct.pack(">Q", 2), + lp(b"a"), + lp(b"first"), + lp(b"z"), + lp(b"last"), + lp(b"payload"), + ) + ) + + +def test_metadata_insertion_order_does_not_change_frame_or_fingerprint() -> None: + first = temporalio.api.common.v1.Payload(data=b"data") + first.metadata["z"] = b"last" + first.metadata["a"] = b"first" + second = temporalio.api.common.v1.Payload(data=b"data") + second.metadata["a"] = b"first" + second.metadata["z"] = b"last" + + first_frame = canonical_logical_record_frame("events", RecordKind.DATA, first) + second_frame = canonical_logical_record_frame("events", RecordKind.DATA, second) + + assert first_frame == second_frame + assert fingerprint_logical_frames([first_frame]) == fingerprint_logical_frames( + [second_frame] + ) + + +def test_batch_hashes_ordered_length_prefixed_frames() -> None: + frames = [b"a", b"bc"] + + fingerprint = fingerprint_logical_frames(frames) + + assert fingerprint.version == LOGICAL_FINGERPRINT_VERSION == 1 + assert fingerprint.digest == hashlib.sha256(lp(b"a") + lp(b"bc")).digest() + assert fingerprint.logical_byte_count == 3 + assert fingerprint.record_count == 2 + + +def test_fingerprint_changes_with_order_topic_kind_metadata_or_data() -> None: + base = temporalio.api.common.v1.Payload( + metadata={"encoding": b"json/plain"}, data=b"one" + ) + changed_metadata = temporalio.api.common.v1.Payload( + metadata={"encoding": b"binary/plain"}, data=b"one" + ) + changed_data = temporalio.api.common.v1.Payload( + metadata={"encoding": b"json/plain"}, data=b"two" + ) + variants: list[ + list[ + tuple[ + str, + RecordKind, + temporalio.api.common.v1.Payload | None, + ] + ] + ] = [ + [("other", RecordKind.DATA, base)], + [("events", RecordKind.DATA, changed_metadata)], + [("events", RecordKind.DATA, changed_data)], + [ + ("events", RecordKind.DATA, base), + ("events", RecordKind.FINISH, None), + ], + [ + ("events", RecordKind.FINISH, None), + ("events", RecordKind.DATA, base), + ], + ] + baseline = fingerprint_logical_records([("events", RecordKind.DATA, base)]) + + assert all(fingerprint_logical_records(variant) != baseline for variant in variants) + assert fingerprint_logical_records(variants[-1]) != fingerprint_logical_records( + variants[-2] + ) + + +def test_length_prefixes_prevent_field_and_record_repartitioning() -> None: + topic_a = canonical_logical_record_frame( + "a", RecordKind.DATA, temporalio.api.common.v1.Payload(data=b"bc") + ) + topic_ab = canonical_logical_record_frame( + "ab", RecordKind.DATA, temporalio.api.common.v1.Payload(data=b"c") + ) + + assert topic_a != topic_ab + assert fingerprint_logical_frames([b"a", b"bc"]) != fingerprint_logical_frames( + [b"ab", b"c"] + ) + + +def test_finish_has_a_canonical_payload_free_frame() -> None: + frame = canonical_logical_record_frame("events", RecordKind.FINISH, None) + + assert frame == b"".join( + ( + lp(b"events"), + lp(b"\x03"), + struct.pack(">Q", 0), + lp(b""), + ) + ) + + +def test_logical_frame_rejects_invalid_payload_shapes() -> None: + with pytest.raises(ValueError, match="DATA.*needs a Payload"): + canonical_logical_record_frame("events", RecordKind.DATA, None) + with pytest.raises(ValueError, match="FINISH.*carries no Payload"): + canonical_logical_record_frame( + "events", RecordKind.FINISH, temporalio.api.common.v1.Payload() + ) + with pytest.raises(ValueError, match="non-empty topic"): + canonical_logical_record_frame( + "", RecordKind.DATA, temporalio.api.common.v1.Payload() + ) diff --git a/tests/contrib/external_workflow_streams/test_output_continuation.py b/tests/contrib/external_workflow_streams/test_output_continuation.py new file mode 100644 index 0000000000..6b87dd2686 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_output_continuation.py @@ -0,0 +1,50 @@ +"""Continue-As-New state for explicitly finished output topics.""" + +from __future__ import annotations + +import pytest + +from temporalio.contrib.external_workflow_streams._annotation import ( + AnnotationDecodeError, +) +from temporalio.contrib.external_workflow_streams._output_continuation import ( + OUTPUT_CONTINUATION_HEADER, + OutputContinuation, + decode_output_continuation, + encode_output_continuation, + read_output_continuation_header, + write_output_continuation_header, +) + + +def state(*topics: str) -> OutputContinuation: + return OutputContinuation(frozenset(topics), "redis-streams", 1) + + +def test_output_continuation_round_trips_stably() -> None: + first = state("status", "events") + second = state("events", "status") + + assert encode_output_continuation(first) == encode_output_continuation(second) + assert decode_output_continuation(encode_output_continuation(first)) == first + + +def test_output_continuation_uses_a_separate_must_understand_header() -> None: + payload = write_output_continuation_header(state("events")) + + assert read_output_continuation_header( + {OUTPUT_CONTINUATION_HEADER: payload} + ) == state("events") + assert read_output_continuation_header({}) is None + + +def test_output_continuation_rejects_unknown_schema_or_encoding() -> None: + payload = write_output_continuation_header(state("events")) + payload.data = b"\x02" + payload.data[1:] + with pytest.raises(AnnotationDecodeError, match="schema version 2"): + read_output_continuation_header({OUTPUT_CONTINUATION_HEADER: payload}) + + payload = write_output_continuation_header(state("events")) + payload.metadata["encoding"] = b"json/plain" + with pytest.raises(AnnotationDecodeError, match="unsupported encoding"): + read_output_continuation_header({OUTPUT_CONTINUATION_HEADER: payload}) diff --git a/tests/contrib/external_workflow_streams/test_output_producer.py b/tests/contrib/external_workflow_streams/test_output_producer.py new file mode 100644 index 0000000000..257f711736 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_output_producer.py @@ -0,0 +1,258 @@ +"""Direct Activity/external producer tests for output streams.""" + +from __future__ import annotations + +import asyncio +import dataclasses +from collections.abc import Sequence +from typing import Any + +import pytest + +import temporalio.api.common.v1 +import temporalio.converter +from temporalio.contrib.external_workflow_streams._backend import ( + AppendConflictError, + StreamDirection, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._output_backend import ( + OutputReadResult, + OutputStage, + OutputStageManifest, + OutputStreamBackend, + OutputStreamRecord, + StagedOutputRecord, +) +from temporalio.contrib.external_workflow_streams._output_producer import ( + ExternalOutputStreamProducer, + OutputAppendNotAcknowledgedError, +) +from temporalio.contrib.external_workflow_streams._producer import WorkflowChainKey +from temporalio.contrib.external_workflow_streams._record import ( + BEGINNING, + Cursor, + Offset, + RecordKind, + StreamRecord, +) + + +class FakeOutputBackend(OutputStreamBackend): + guarantees_immutability = True + provider_id = "fake-output" + provider_format_version = 1 + + def __init__(self) -> None: + self.append_calls: list[tuple[StreamKey, StreamRecord]] = [] + self.stored: dict[ + tuple[StreamKey, str, int], tuple[StreamRecord, OutputStreamRecord] + ] = {} + self.fail_before_store_once = False + self.fail_after_store_once = False + self.block_data = False + self.data_append_started = asyncio.Event() + self.release_data_append = asyncio.Event() + self._next_offset = 1 + + async def stage_output( + self, + manifest: OutputStageManifest, + records: Sequence[StagedOutputRecord], + ) -> OutputStage: + raise NotImplementedError + + async def commit_output(self, manifest: OutputStageManifest) -> OutputStage: + raise NotImplementedError + + async def abort_output(self, manifest: OutputStageManifest) -> OutputStage: + raise NotImplementedError + + async def output_stage(self, manifest: OutputStageManifest) -> OutputStage | None: + raise NotImplementedError + + async def append_output( + self, key: StreamKey, record: StreamRecord + ) -> OutputStreamRecord: + assert key.direction is StreamDirection.OUTPUT + self.append_calls.append((key, record)) + if self.fail_before_store_once: + self.fail_before_store_once = False + raise ConnectionError("append outcome unknown before reply") + + identity = (key, record.producer_session_id, record.sequence) + existing = self.stored.get(identity) + if existing is not None: + original, placed = existing + if original != record: + raise AppendConflictError(record.idempotency_key) + return placed + + if self.block_data and record.kind is RecordKind.DATA: + self.data_append_started.set() + await self.release_data_append.wait() + placed = OutputStreamRecord( + kind=record.kind, + payload=record.payload, + offset=Offset(str(self._next_offset)), + ) + self._next_offset += 1 + self.stored[identity] = (record, placed) + if self.fail_after_store_once: + self.fail_after_store_once = False + raise ConnectionError("append committed but acknowledgement was lost") + return placed + + async def read_output_after( + self, + key: StreamKey, + after: Cursor, + *, + max_records: int, + block: Any = None, + ) -> OutputReadResult: + raise NotImplementedError + + async def output_tail(self, key: StreamKey) -> Cursor: + return BEGINNING + + def compare_offsets(self, left: Offset, right: Offset) -> int: + return int(left.token) - int(right.token) + + +CHAIN = WorkflowChainKey("ns", "wf", "first-run") + + +def producer( + backend: FakeOutputBackend, + *, + data_converter: temporalio.converter.DataConverter = temporalio.converter.DataConverter.default, +) -> ExternalOutputStreamProducer: + return ExternalOutputStreamProducer( + backend=backend, + workflow=CHAIN, + data_converter=data_converter, + session_id="activity-session", + ) + + +async def test_publish_uses_output_identity_and_workflow_serialization_context() -> ( + None +): + observed: list[temporalio.converter.SerializationContext | None] = [] + + class RecordingCodec( + temporalio.converter.PayloadCodec, + temporalio.converter.WithSerializationContext, + ): + def __init__( + self, + context: temporalio.converter.SerializationContext | None = None, + ) -> None: + self.context = context + + def with_context( + self, context: temporalio.converter.SerializationContext + ) -> RecordingCodec: + return RecordingCodec(context) + + async def encode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + observed.append(self.context) + return list(payloads) + + async def decode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + return list(payloads) + + backend = FakeOutputBackend() + output = producer( + backend, + data_converter=temporalio.converter.DataConverter( + payload_codec=RecordingCodec() + ), + ) + + placed = await output.topic("events", type=str).publish("hello") + + key, record = backend.append_calls[0] + assert key.direction is StreamDirection.OUTPUT + assert key.stream_name == "events" + assert record.kind is RecordKind.DATA + assert record.producer_session_id == "activity-session" + assert record.sequence == 0 + assert placed.offset == Offset("1") + assert observed == [ + temporalio.converter.WorkflowSerializationContext( + namespace="ns", workflow_id="wf" + ) + ] + + +async def test_finish_waits_for_preceding_publish_and_closes_all_topic_handles() -> ( + None +): + backend = FakeOutputBackend() + backend.block_data = True + output = producer(backend) + topic = output.topic("events", type=str) + other_handle = output.topic("events", type=str) + + publish_task = asyncio.create_task(topic.publish("first")) + await backend.data_append_started.wait() + finish_task = asyncio.create_task(other_handle.finish_writing()) + backend.release_data_append.set() + await asyncio.gather(publish_task, finish_task) + + records = [record for _, record in backend.append_calls] + assert [record.kind for record in records] == [ + RecordKind.DATA, + RecordKind.FINISH, + ] + assert [record.sequence for record in records] == [0, 1] + assert records[1].payload == b"" + with pytest.raises(RuntimeError, match="already finished"): + await topic.publish("too late") + with pytest.raises(RuntimeError, match="already finished"): + await other_handle.finish_writing() + + +async def test_ambiguous_append_retries_exact_record_without_duplicate() -> None: + backend = FakeOutputBackend() + backend.fail_after_store_once = True + topic = producer(backend).topic("events", type=str) + + with pytest.raises(OutputAppendNotAcknowledgedError) as unknown: + await topic.publish("once") + with pytest.raises(OutputAppendNotAcknowledgedError) as blocked: + await topic.publish("must not draw another sequence") + assert blocked.value.record == unknown.value.record + + recovered = await topic.resolve_append(unknown.value.record) + following = await topic.publish("after recovery") + + assert recovered.offset == Offset("1") + assert following.offset == Offset("2") + assert len(backend.stored) == 2 + assert [record.sequence for _, record in backend.append_calls] == [0, 0, 1] + + +async def test_conflict_while_resolving_is_definitive_not_another_unknown() -> None: + backend = FakeOutputBackend() + backend.fail_before_store_once = True + topic = producer(backend).topic("events", type=str) + + with pytest.raises(OutputAppendNotAcknowledgedError) as unknown: + await topic.publish("original") + conflicting = dataclasses.replace(unknown.value.record, payload=b"different") + await backend.append_output(unknown.value.stream_key, conflicting) + + with pytest.raises(AppendConflictError): + await topic.resolve_append(unknown.value.record) + + # The provider answered definitively, so the failed recovery is no longer + # reported as an unknown operation on every later call. + placed = await topic.publish("next") + assert placed.offset == Offset("2") diff --git a/tests/contrib/external_workflow_streams/test_output_runtime.py b/tests/contrib/external_workflow_streams/test_output_runtime.py new file mode 100644 index 0000000000..1c58854a70 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_output_runtime.py @@ -0,0 +1,1066 @@ +"""Workflow-thread batching and staging for external output streams.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Sequence +from datetime import timedelta +from types import SimpleNamespace +from typing import Any, cast + +import pytest + +import temporalio.api.common.v1 +import temporalio.bridge.proto.external_data +import temporalio.bridge.proto.workflow_activation +import temporalio.bridge.proto.workflow_completion +import temporalio.converter +import temporalio.workflow +from temporalio.contrib.external_workflow_streams._backend import StreamKey +from temporalio.contrib.external_workflow_streams._errors import ( + ExternalStreamCapacityError, + StreamStorageError, +) +from temporalio.contrib.external_workflow_streams._manager import ( + ReadinessResult, + StreamSubscriptionManager, +) +from temporalio.contrib.external_workflow_streams._output_api import ( + ExternalOutputStreamOptions, +) +from temporalio.contrib.external_workflow_streams._output_backend import ( + OutputReadResult, + OutputStage, + OutputStageManifest, + OutputStageStatus, + OutputStreamBackend, + OutputStreamRecord, + StagedOutputRecord, +) +from temporalio.contrib.external_workflow_streams._output_codec import ( + fingerprint_logical_frames, +) +from temporalio.contrib.external_workflow_streams._record import ( + BEGINNING, + Cursor, + Offset, + RecordKind, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._runtime import ( + WorkflowStreamRuntime, +) +from temporalio.worker._workflow import _WorkflowWorker +from temporalio.worker._workflow_instance import _WorkflowInstanceImpl +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend + + +class _OutputMemoryBackend(MemoryStreamBackend, OutputStreamBackend): + provider_id = "output-memory" + provider_format_version = 1 + + def __init__(self) -> None: + super().__init__() + self.stages: dict[tuple[str, int], OutputStage] = {} + self.stage_calls = 0 + self.stage_failures = 0 + self.stage_payloads: list[tuple[bytes, ...]] = [] + + async def stage_output( + self, + manifest: OutputStageManifest, + records: Sequence[StagedOutputRecord], + ) -> OutputStage: + self.stage_calls += 1 + self.stage_payloads.append(tuple(record.payload for record in records)) + if self.stage_failures: + self.stage_failures -= 1 + raise OSError("output backend unavailable") + identity = (manifest.stage_token, manifest.sub_batch_id) + existing = self.stages.get(identity) + if existing is not None: + return existing + placed = tuple( + OutputStreamRecord( + record.kind, + record.payload, + Offset(f"{manifest.sub_batch_id + 1}-{record.publish_index}"), + ) + for record in records + ) + stage = OutputStage(manifest, placed, OutputStageStatus.PENDING) + self.stages[identity] = stage + return stage + + async def commit_output(self, manifest: OutputStageManifest) -> OutputStage: + return self._resolve(manifest, OutputStageStatus.COMMITTED) + + async def abort_output(self, manifest: OutputStageManifest) -> OutputStage: + return self._resolve(manifest, OutputStageStatus.ABORTED) + + def _resolve( + self, manifest: OutputStageManifest, status: OutputStageStatus + ) -> OutputStage: + existing = self.stages[(manifest.stage_token, manifest.sub_batch_id)] + resolved = OutputStage(existing.manifest, existing.records, status) + self.stages[(manifest.stage_token, manifest.sub_batch_id)] = resolved + return resolved + + async def output_stage(self, manifest: OutputStageManifest) -> OutputStage | None: + return self.stages.get((manifest.stage_token, manifest.sub_batch_id)) + + async def append_output( + self, key: StreamKey, record: StreamRecord + ) -> OutputStreamRecord: + placed = await self.append(key, record) + assert placed.offset is not None + return OutputStreamRecord(placed.kind, placed.payload, placed.offset) + + async def read_output_after( + self, + key: StreamKey, + after: Cursor, + *, + max_records: int, + block: timedelta | None = None, + ) -> OutputReadResult: + del key, after, max_records, block + return OutputReadResult(()) + + async def output_tail(self, key: StreamKey) -> Cursor: + del key + return BEGINNING + + +class _ChangingPayloadCodec(temporalio.converter.PayloadCodec): + """A valid codec whose wire bytes deliberately change on every call.""" + + def __init__(self) -> None: + self.encode_calls = 0 + + async def encode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + self.encode_calls += 1 + encoded = [] + for payload in payloads: + copied = temporalio.api.common.v1.Payload() + copied.CopyFrom(payload) + copied.data += f":encoding-{self.encode_calls}".encode() + encoded.append(copied) + return encoded + + async def decode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + return list(payloads) + + +async def _notify(run_id: str, wait_id: int, generation: int) -> str: + del run_id, wait_id, generation + return ReadinessResult.ACCEPTED + + +@pytest.fixture +def backend() -> _OutputMemoryBackend: + return _OutputMemoryBackend() + + +@pytest.fixture +def runtime(backend: _OutputMemoryBackend) -> WorkflowStreamRuntime: + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=1), + ) + return WorkflowStreamRuntime( + manager=manager, + backend=backend, + run_id="run", + namespace="namespace", + workflow_id="workflow", + first_execution_run_id="first-run", + data_converter=temporalio.converter.DataConverter.default, + default_idle_timeout=timedelta(seconds=1), + ) + + +def _successful_completion( + *variants: str, +) -> temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion: + completion = ( + temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion() + ) + completion.successful.SetInParent() + for variant in variants: + getattr(completion.successful.commands.add(), variant).SetInParent() + return completion + + +def _activation(*, finalize_output_latency: bool = False) -> Any: + activation = temporalio.bridge.proto.workflow_activation.WorkflowActivation( + run_id="run", + history_floor_event_id=7, + ) + if finalize_output_latency: + activation.jobs.add().finalize_external_streams.reason = ( + temporalio.bridge.proto.external_data.ParkReason.PARK_REASON_OUTPUT_LATENCY + ) + return activation + + +def _park_activation() -> Any: + activation = _activation() + activation.jobs.add().prepare_external_stream_park.quiescence_generation = 1 + return activation + + +class _ParkManager: + def __init__(self, *, became_ready: bool) -> None: + self.became_ready = became_ready + + async def prepare_park(self, *args: Any, **kwargs: Any) -> bool: + del args, kwargs + return self.became_ready + + +def _worker() -> _WorkflowWorker: + worker = object.__new__(_WorkflowWorker) + worker._pending_external_output_stages = {} + return worker + + +async def _publish_one( + runtime: WorkflowStreamRuntime, + *, + max_records: int = 10, +) -> None: + runtime.begin_activation(7) + await runtime.publish_output( + topic="events", + value="value", + value_type=str, + kind=RecordKind.DATA, + max_publish_latency=timedelta(milliseconds=250), + max_records=max_records, + max_logical_bytes=10_000, + ) + + +async def test_retainable_output_emits_buffered_without_provider_staging( + runtime: WorkflowStreamRuntime, + backend: _OutputMemoryBackend, +) -> None: + await _publish_one(runtime) + completion = _successful_completion( + "workflow_stream_progress", + "workflow_stream_quiescent", + ) + + await _worker()._stage_or_buffer_external_output(_activation(), completion, runtime) + + assert backend.stage_calls == 0 + assert runtime.has_output + assert [ + command.WhichOneof("variant") for command in completion.successful.commands + ] == [ + "workflow_stream_progress", + "workflow_output_stream_buffered", + "workflow_stream_quiescent", + ] + buffered = completion.successful.commands[1].workflow_output_stream_buffered + assert buffered.max_publish_latency.ToTimedelta() == timedelta(milliseconds=250) + + +async def test_output_latency_finalization_stages_and_answers_both_commands( + runtime: WorkflowStreamRuntime, + backend: _OutputMemoryBackend, +) -> None: + await _publish_one(runtime) + worker = _worker() + worker._external_stream_runtimes = {"run": runtime} + initial = _successful_completion("workflow_stream_quiescent") + await worker._stage_or_buffer_external_output(_activation(), initial, runtime) + assert backend.stage_calls == 0 + + activation = _activation(finalize_output_latency=True) + completion = await worker._handle_external_stream_jobs( + activation, + cast(Any, None), + ) + assert completion is not None + await worker._stage_or_buffer_external_output(activation, completion, runtime) + + assert backend.stage_calls == 1 + assert not runtime.has_output + assert [ + command.WhichOneof("variant") for command in completion.successful.commands + ] == [ + "workflow_output_stream_commit", + "external_stream_finalized", + ] + assert len(worker._pending_external_output_stages["run"]) == 1 + + +async def test_confirmed_park_stages_output_in_the_same_completion( + runtime: WorkflowStreamRuntime, + backend: _OutputMemoryBackend, +) -> None: + await _publish_one(runtime) + worker = _worker() + worker._external_stream_runtimes = {"run": runtime} + worker._external_stream_manager = _ParkManager(became_ready=False) + activation = _park_activation() + + completion = await worker._handle_external_stream_jobs( + activation, + cast(Any, None), + ) + assert completion is not None + await worker._stage_or_buffer_external_output(activation, completion, runtime) + + assert backend.stage_calls == 1 + assert not runtime.has_output + assert [ + command.WhichOneof("variant") for command in completion.successful.commands + ] == [ + "workflow_output_stream_commit", + "external_stream_park_result", + ] + assert completion.successful.commands[1].external_stream_park_result.HasField( + "confirmed" + ) + + +async def test_park_that_became_ready_keeps_output_buffered( + runtime: WorkflowStreamRuntime, + backend: _OutputMemoryBackend, +) -> None: + await _publish_one(runtime) + worker = _worker() + worker._external_stream_runtimes = {"run": runtime} + worker._external_stream_manager = _ParkManager(became_ready=True) + activation = _park_activation() + + completion = await worker._handle_external_stream_jobs( + activation, + cast(Any, None), + ) + assert completion is not None + await worker._stage_or_buffer_external_output(activation, completion, runtime) + + assert backend.stage_calls == 0 + assert runtime.has_output + assert [ + command.WhichOneof("variant") for command in completion.successful.commands + ] == [ + "workflow_output_stream_buffered", + "external_stream_park_result", + ] + assert completion.successful.commands[1].external_stream_park_result.HasField( + "became_ready" + ) + buffered = completion.successful.commands[0].workflow_output_stream_buffered + assert buffered.max_publish_latency.ToTimedelta() == timedelta(milliseconds=250) + + +async def test_output_without_quiescent_snapshot_stages_immediately( + runtime: WorkflowStreamRuntime, + backend: _OutputMemoryBackend, +) -> None: + await _publish_one(runtime) + completion = _successful_completion() + + await _worker()._stage_or_buffer_external_output(_activation(), completion, runtime) + + assert backend.stage_calls == 1 + assert [ + command.WhichOneof("variant") for command in completion.successful.commands + ] == ["workflow_output_stream_commit"] + + +async def test_capacity_rollover_bypasses_output_buffering( + runtime: WorkflowStreamRuntime, + backend: _OutputMemoryBackend, +) -> None: + await _publish_one(runtime, max_records=1) + blocked_publish = asyncio.create_task( + runtime.publish_output( + topic="events", + value="second", + value_type=str, + kind=RecordKind.DATA, + max_publish_latency=timedelta(milliseconds=250), + max_records=1, + max_logical_bytes=10_000, + ) + ) + await asyncio.sleep(0) + assert runtime.output_rollover_requested + completion = _successful_completion("workflow_stream_quiescent") + + try: + await _worker()._stage_or_buffer_external_output( + _activation(), completion, runtime + ) + assert backend.stage_calls == 1 + assert [ + command.WhichOneof("variant") for command in completion.successful.commands + ] == [ + "workflow_output_stream_commit", + "workflow_stream_quiescent", + ] + finally: + blocked_publish.cancel() + with pytest.raises(asyncio.CancelledError): + await blocked_publish + + +async def test_stage_uses_one_token_and_pre_codec_logical_manifests( + runtime: WorkflowStreamRuntime, backend: _OutputMemoryBackend +) -> None: + runtime.begin_activation(7) + await runtime.publish_output( + topic="events", + value={"value": 1}, + value_type=dict, + kind=RecordKind.DATA, + max_publish_latency=timedelta(milliseconds=100), + max_records=10, + max_logical_bytes=10_000, + ) + await runtime.publish_output( + topic="status", + value=None, + value_type=str, + kind=RecordKind.FINISH, + max_publish_latency=timedelta(milliseconds=50), + max_records=10, + max_logical_bytes=10_000, + ) + + staged = await runtime.stage_output(7) + + assert len(staged.topics) == 2 + assert {topic.manifest.stage_token for topic in staged.topics} == { + staged.stage_token + } + assert staged.segment_record_counts == ((1, 1),) + assert staged.topics[1].finished + assert all( + stage.status is OutputStageStatus.PENDING for stage in backend.stages.values() + ) + + +async def test_changed_codec_bytes_keep_logical_retry_identity_and_first_bytes( + backend: _OutputMemoryBackend, +) -> None: + codec = _ChangingPayloadCodec() + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=1), + ) + runtime = WorkflowStreamRuntime( + manager=manager, + backend=backend, + run_id="run", + namespace="namespace", + workflow_id="workflow", + first_execution_run_id="first-run", + data_converter=temporalio.converter.DataConverter(payload_codec=codec), + default_idle_timeout=timedelta(seconds=1), + ) + runtime.begin_activation(7) + await runtime.publish_output( + topic="events", + value="same logical value", + value_type=str, + kind=RecordKind.DATA, + max_publish_latency=timedelta(milliseconds=100), + max_records=10, + max_logical_bytes=10_000, + ) + + first = await runtime.stage_output(7) + repeated = await runtime.stage_output(7) + + assert codec.encode_calls == 2 + assert first == repeated + assert first.topics[0].manifest == repeated.topics[0].manifest + assert len(first.topics[0].manifest.fingerprint) == 32 + assert backend.stage_payloads[0] != backend.stage_payloads[1] + stored = backend.stages[(first.stage_token, 0)] + assert ( + tuple(record.payload for record in stored.records) == backend.stage_payloads[0] + ) + assert ( + tuple(record.payload for record in stored.records) != backend.stage_payloads[1] + ) + + +async def test_capacity_publish_waits_for_a_replacement_workflow_task( + runtime: WorkflowStreamRuntime, +) -> None: + options: dict[str, Any] = { + "value_type": str, + "kind": RecordKind.DATA, + "max_publish_latency": timedelta(milliseconds=100), + "max_records": 1, + "max_logical_bytes": 10_000, + } + runtime.begin_activation(7) + await runtime.publish_output(topic="events", value="one", **options) + second = asyncio.create_task( + runtime.publish_output(topic="events", value="two", **options) + ) + await asyncio.sleep(0) + assert runtime.output_rollover_requested + assert not second.done() + + await runtime.stage_output(7) + runtime.output_stage_recorded() + runtime.begin_activation(9) + await second + + staged = await runtime.stage_output(9) + assert staged.topics[0].manifest.record_count == 1 + + +async def test_stage_outage_is_storage_failure_and_retry_reuses_attempt( + runtime: WorkflowStreamRuntime, backend: _OutputMemoryBackend +) -> None: + runtime.begin_activation(7) + await runtime.publish_output( + topic="events", + value="value", + value_type=str, + kind=RecordKind.DATA, + max_publish_latency=timedelta(milliseconds=100), + max_records=10, + max_logical_bytes=10_000, + ) + backend.stage_failures = 1 + + with pytest.raises(StreamStorageError, match="events"): + await runtime.stage_output(7) + first_token = runtime._output_stage_token + + staged = await runtime.stage_output(7) + + assert staged.stage_token == first_token + assert backend.stage_calls == 2 + + +async def test_oversized_marker_manifest_is_rejected_before_external_io( + runtime: WorkflowStreamRuntime, backend: _OutputMemoryBackend +) -> None: + runtime.begin_activation(7) + await runtime.publish_output( + topic="t" * (70 * 1024), + value="value", + value_type=str, + kind=RecordKind.DATA, + max_publish_latency=timedelta(milliseconds=100), + max_records=10, + max_logical_bytes=100 * 1024, + ) + + with pytest.raises(ExternalStreamCapacityError, match="marker budget"): + await runtime.stage_output(7) + + assert backend.stage_calls == 0 + assert runtime._staged_output is None + + +async def test_replay_validates_the_recorded_shared_segment_schedule_without_io( + runtime: WorkflowStreamRuntime, backend: _OutputMemoryBackend +) -> None: + runtime.begin_output_replay( + SimpleNamespace( + schema_version=1, + fingerprint_version=1, + topics=[ + SimpleNamespace( + topic="events", + record_count=1, + logical_byte_count=0, + logical_fingerprint=b"", + finished=False, + ) + ], + segments=[SimpleNamespace(record_counts_by_topic=[1])], + ) + ) + runtime.begin_output_replay_segment() + await runtime.publish_output( + topic="events", + value="value", + value_type=str, + kind=RecordKind.DATA, + max_publish_latency=timedelta(milliseconds=100), + max_records=1, + max_logical_bytes=10_000, + ) + record = runtime._output_records[0] + fingerprint = fingerprint_logical_frames([record.frame]) + expected = runtime._output_replay_manifest.topics[0] + expected.logical_byte_count = fingerprint.logical_byte_count + expected.logical_fingerprint = fingerprint.digest + + runtime.verify_output_replay() + + assert backend.stage_calls == 0 + + +async def test_replay_ignores_live_capacity_and_does_not_leak_live_policy( + runtime: WorkflowStreamRuntime, +) -> None: + expected = SimpleNamespace( + schema_version=1, + fingerprint_version=1, + topics=[ + SimpleNamespace( + topic="events", + record_count=1, + logical_byte_count=0, + logical_fingerprint=b"", + finished=False, + ) + ], + segments=[SimpleNamespace(record_counts_by_topic=[1])], + ) + runtime.begin_output_replay(expected) + runtime.begin_output_replay_segment() + # The recorded value is larger than today's deliberately tiny live limit. + # Replay must validate it instead of moving or rejecting the old boundary. + await runtime.publish_output( + topic="events", + value="recorded before the limit changed", + value_type=str, + kind=RecordKind.DATA, + max_publish_latency=timedelta(milliseconds=1), + max_records=1, + max_logical_bytes=1, + ) + record = runtime._output_records[0] + fingerprint = fingerprint_logical_frames([record.frame]) + expected.topics[0].logical_byte_count = fingerprint.logical_byte_count + expected.topics[0].logical_fingerprint = fingerprint.digest + + assert runtime.output_max_publish_latency is None + runtime.verify_output_replay() + + runtime.begin_activation(11) + await runtime.publish_output( + topic="events", + value="live", + value_type=str, + kind=RecordKind.DATA, + max_publish_latency=timedelta(milliseconds=250), + max_records=10, + max_logical_bytes=10_000, + ) + assert runtime.output_max_publish_latency == timedelta(milliseconds=250) + + +async def test_output_replay_performs_no_io_token_mint_or_live_policy_split( + runtime: WorkflowStreamRuntime, + backend: _OutputMemoryBackend, + monkeypatch: pytest.MonkeyPatch, +) -> None: + expected = SimpleNamespace( + schema_version=1, + fingerprint_version=1, + topics=[ + SimpleNamespace( + topic="events", + record_count=2, + logical_byte_count=0, + logical_fingerprint=b"", + finished=False, + ) + ], + segments=[SimpleNamespace(record_counts_by_topic=[2])], + ) + runtime.begin_output_replay(expected) + runtime.begin_output_replay_segment() + monkeypatch.setattr( + "temporalio.contrib.external_workflow_streams._runtime.secrets.token_urlsafe", + lambda _size: pytest.fail("replay minted an output stage token"), + ) + + # Today's live policy would split after the first record and arm a 1 ms + # flush. Replay must instead reproduce the one recorded two-record segment. + for value in ("one", "two"): + await runtime.publish_output( + topic="events", + value=value, + value_type=str, + kind=RecordKind.DATA, + max_publish_latency=timedelta(milliseconds=1), + max_records=1, + max_logical_bytes=1, + ) + fingerprint = fingerprint_logical_frames( + record.frame for record in runtime._output_records + ) + expected.topics[0].logical_byte_count = fingerprint.logical_byte_count + expected.topics[0].logical_fingerprint = fingerprint.digest + + assert runtime._output_stage_token is None + assert not runtime.output_rollover_requested + assert runtime.output_max_publish_latency is None + runtime.verify_output_replay() + assert backend.stage_calls == 0 + + +class _SizedPayloadCodec(temporalio.converter.PayloadCodec): + def __init__(self, suffix_size: int) -> None: + self.suffix_size = suffix_size + self.encode_calls = 0 + + async def encode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + self.encode_calls += 1 + encoded: list[temporalio.api.common.v1.Payload] = [] + for payload in payloads: + copied = temporalio.api.common.v1.Payload() + copied.CopyFrom(payload) + copied.data += b"x" * self.suffix_size + encoded.append(copied) + return encoded + + async def decode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + return list(payloads) + + +def _runtime_with_codec( + backend: _OutputMemoryBackend, codec: temporalio.converter.PayloadCodec +) -> WorkflowStreamRuntime: + return WorkflowStreamRuntime( + manager=StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=1), + ), + backend=backend, + run_id="run", + namespace="namespace", + workflow_id="workflow", + first_execution_run_id="first-run", + data_converter=temporalio.converter.DataConverter(payload_codec=codec), + default_idle_timeout=timedelta(seconds=1), + ) + + +async def test_replay_keeps_recorded_segments_when_codec_wire_size_changes( + backend: _OutputMemoryBackend, +) -> None: + live_codec = _SizedPayloadCodec(1) + live = _runtime_with_codec(backend, live_codec) + for value in ("one", "two"): + live.begin_activation(7) + await live.publish_output( + topic="events", + value=value, + value_type=str, + kind=RecordKind.DATA, + max_publish_latency=timedelta(milliseconds=100), + max_records=10, + max_logical_bytes=10_000, + ) + staged = await live.stage_output(7) + assert staged.segment_record_counts == ((1,), (1,)) + assert live_codec.encode_calls == 2 + live_wire = backend.stage_payloads[-1] + + replay_codec = _SizedPayloadCodec(100_000) + replay = _runtime_with_codec(backend, replay_codec) + topic = staged.topics[0] + expected = SimpleNamespace( + schema_version=staged.schema_version, + fingerprint_version=staged.fingerprint_version, + topics=[ + SimpleNamespace( + topic=topic.manifest.stream_key.stream_name, + record_count=topic.manifest.record_count, + logical_byte_count=topic.manifest.logical_byte_count, + logical_fingerprint=topic.manifest.fingerprint, + finished=topic.finished, + ) + ], + segments=[ + SimpleNamespace(record_counts_by_topic=list(counts)) + for counts in staged.segment_record_counts + ], + ) + replay.begin_output_replay(expected) + for value in ("one", "two"): + replay.begin_output_replay_segment() + await replay.publish_output( + topic="events", + value=value, + value_type=str, + kind=RecordKind.DATA, + max_publish_latency=timedelta(milliseconds=1), + max_records=1, + max_logical_bytes=1, + ) + + replay.verify_output_replay() + assert replay_codec.suffix_size > sum(len(payload) for payload in live_wire) + assert replay_codec.encode_calls == 0 + assert backend.stage_calls == 1 + + +async def test_one_oversized_logical_record_is_refused_without_poisoning_batch( + runtime: WorkflowStreamRuntime, + backend: _OutputMemoryBackend, +) -> None: + runtime.begin_activation(7) + with pytest.raises(ExternalStreamCapacityError, match="one record"): + await runtime.publish_output( + topic="events", + value="too large", + value_type=str, + kind=RecordKind.DATA, + max_publish_latency=timedelta(milliseconds=100), + max_records=10, + max_logical_bytes=1, + ) + + assert not runtime.has_output + assert not runtime.output_rollover_requested + assert backend.stage_calls == 0 + await runtime.publish_output( + topic="events", + value="fits the next valid policy", + value_type=str, + kind=RecordKind.DATA, + max_publish_latency=timedelta(milliseconds=100), + max_records=10, + max_logical_bytes=10_000, + ) + assert runtime.has_output + + +async def test_abandon_output_replay_clears_partial_finish_and_expectations( + runtime: WorkflowStreamRuntime, +) -> None: + runtime.begin_output_replay( + SimpleNamespace( + schema_version=1, + fingerprint_version=1, + topics=[], + segments=[], + ) + ) + runtime.begin_output_replay_segment() + await runtime.publish_output( + topic="events", + value=None, + value_type=str, + kind=RecordKind.FINISH, + max_publish_latency=timedelta(milliseconds=1), + max_records=1, + max_logical_bytes=1, + ) + + runtime.abandon_output_replay() + runtime.begin_activation(11) + await runtime.publish_output( + topic="events", + value="live after failed replay", + value_type=str, + kind=RecordKind.DATA, + max_publish_latency=timedelta(milliseconds=250), + max_records=10, + max_logical_bytes=10_000, + ) + + assert runtime.has_output + assert runtime.output_max_publish_latency == timedelta(milliseconds=250) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("max_publish_latency", timedelta(0), "max_publish_latency"), + ("max_records", 0, "max_records"), + ("max_logical_bytes", 0, "max_logical_bytes"), + ], +) +def test_output_options_validate_direct_construction( + field: str, value: Any, message: str +) -> None: + kwargs = {field: value} + with pytest.raises(ValueError, match=message): + ExternalOutputStreamOptions(**kwargs) + + +async def test_workflow_output_rejects_read_only_context( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + temporalio.workflow.unsafe, + "is_read_only", + lambda: True, + ) + topic = ExternalOutputStreamOptions().topic("events", type=str) + + with pytest.raises( + temporalio.workflow.ReadOnlyContextError, match="external output" + ): + await topic.publish("not allowed from a query") + + +class _OutputOnlyReplayRuntime: + def __init__(self) -> None: + self.events: list[str] = [] + + def begin_output_replay(self, _manifest: Any) -> None: + self.events.append("begin-output") + + def take_replay_plan(self) -> None: + return None + + def begin_output_replay_segment(self) -> None: + self.events.append("segment") + + def resolve_all_pending(self) -> None: + self.events.append("resolve") + + +class _OutputOnlyReplayDriver: + def __init__(self, runtime: _OutputOnlyReplayRuntime) -> None: + self._external_stream_runtime = runtime + self._pending_output_replay_finish = False + + def _run_once(self, *, check_conditions: bool) -> None: + assert check_conditions + self._external_stream_runtime.events.append("drain") + + +def test_output_only_replay_drives_all_but_the_trailing_segment() -> None: + runtime = _OutputOnlyReplayRuntime() + driver = _OutputOnlyReplayDriver(runtime) + output = SimpleNamespace(segments=[object(), object(), object()]) + job = SimpleNamespace( + output=output, + HasField=lambda field: field == "output", + ) + + apply_replay = cast(Any, _WorkflowInstanceImpl._apply_replay_external_streams) + apply_replay(driver, job) + + assert runtime.events == [ + "begin-output", + "segment", + "resolve", + "drain", + "segment", + "resolve", + "drain", + "segment", + "resolve", + ] + assert driver._pending_output_replay_finish + + +class _CombinedReplayRuntime: + def __init__(self) -> None: + self.events: list[str] = [] + self._plan: Any = SimpleNamespace( + annotation=SimpleNamespace(header=SimpleNamespace(streams={1: object()})), + segments=tuple( + SimpleNamespace(deliveries=((1, index),)) for index in range(3) + ), + committed_boundaries={1: BEGINNING}, + ) + + def begin_output_replay(self, _manifest: Any) -> None: + self.events.append("begin-output") + + def take_replay_plan(self) -> Any: + plan, self._plan = self._plan, None + return plan + + def begin_replay(self, _streams: Any) -> None: + self.events.append("begin-input") + + def begin_output_replay_segment(self) -> None: + self.events.append("output-segment") + + def begin_replay_segment(self, _deliveries: Any) -> None: + self.events.append("input-segment") + + def resolve_all_pending(self) -> None: + self.events.append("resolve") + + def verify_replay_consumed(self) -> None: + self.events.append("verify-input") + + def reposition_after_replay(self, _boundaries: Any) -> None: + self.events.append("reposition-input") + + def verify_output_replay(self) -> None: + self.events.append("verify-output") + + def abandon_output_replay(self) -> None: + self.events.append("abandon-output") + + def end_replay(self) -> None: + self.events.append("end-input") + + +class _CombinedReplayDriver: + def __init__(self, runtime: _CombinedReplayRuntime) -> None: + self._external_stream_runtime = runtime + self._pending_replay_finish: Any = None + self._pending_output_replay_finish = False + + def _run_once(self, *, check_conditions: bool) -> None: + assert check_conditions + self._external_stream_runtime.events.append("drain") + + +def test_input_and_output_share_one_replay_segment_drain_schedule() -> None: + runtime = _CombinedReplayRuntime() + driver = _CombinedReplayDriver(runtime) + output = SimpleNamespace(segments=(object(), object(), object())) + job = SimpleNamespace( + output=output, + HasField=lambda field: field == "output", + ) + apply_replay = cast(Any, _WorkflowInstanceImpl._apply_replay_external_streams) + finish_replay = cast(Any, _WorkflowInstanceImpl._finish_replay_external_streams) + + apply_replay(driver, job) + driver._run_once(check_conditions=True) + finish_replay(driver) + + assert runtime.events == [ + "begin-output", + "begin-input", + "output-segment", + "input-segment", + "resolve", + "drain", + "output-segment", + "input-segment", + "resolve", + "drain", + "output-segment", + "input-segment", + "resolve", + "drain", + "verify-input", + "reposition-input", + "verify-output", + "abandon-output", + "end-input", + ] diff --git a/tests/contrib/external_workflow_streams/test_output_worker_integration.py b/tests/contrib/external_workflow_streams/test_output_worker_integration.py new file mode 100644 index 0000000000..668697f58f --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_output_worker_integration.py @@ -0,0 +1,1257 @@ +"""Real Worker/Core integration for Workflow-originated external output. + +PYTEST_DONT_REWRITE: the sandboxed Workflow re-imports this module, so pytest's +injected imports must not become part of sandbox validation. +""" + +from __future__ import annotations + +import asyncio +import time +import uuid +from collections.abc import Sequence +from dataclasses import replace +from datetime import timedelta +from typing import Any, cast + +import pytest + +import temporalio.api.enums.v1 +import temporalio.bridge.worker +from temporalio import workflow +from temporalio.client import Client, WorkflowExecutionStatus +from temporalio.contrib.external_workflow_streams import ( + ExternalOutputStreamClient, + WorkflowChainKey, +) +from temporalio.contrib.external_workflow_streams._backend import ( + ParkIntent, + ParkIntentRemoval, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec +from temporalio.contrib.external_workflow_streams._output_backend import ( + OutputReadResult, + OutputStage, + OutputStageConflictError, + OutputStageManifest, + OutputStageNotFoundError, + OutputStageResolutionError, + OutputStageStatus, + OutputStreamBackend, + OutputStreamRecord, + PendingOutputBarrier, + StagedOutputRecord, +) +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Cursor, + Offset, + RecordKind, + StreamRecord, +) +from temporalio.worker import Worker +from tests.contrib.external_workflow_streams.memory_backend import ( + MemoryStreamBackend, +) + +with workflow.unsafe.imports_passed_through(): + from temporalio.contrib.external_workflow_streams import ( + external_output_stream, + external_stream, + ) + + +@workflow.defn +class PublishExternalOutputWorkflow: + """Publishes a value without returning it through Temporal.""" + + @workflow.run + async def run(self) -> int: + topic = external_output_stream.topic("events", type=str) + await topic.publish("workflow-output-secret") + await topic.finish() + return 1 + + +@workflow.defn +class FinishOutputAcrossContinueAsNewWorkflow: + """Finishes in one Run and verifies the successor restores that fact.""" + + @workflow.run + async def run(self) -> str: + topic = external_output_stream.topic("events", type=str) + if workflow.info().continued_run_id is None: + await topic.finish() + workflow.continue_as_new() + try: + await topic.publish("must not be staged") + except workflow.NondeterminismError: + return "successor rejected finished topic" + return "successor incorrectly reopened finished topic" + + +@workflow.defn +class PublishOutputWhileInputIsRetainedWorkflow: + """Publishes output, then keeps the Workflow Task open on an input wait.""" + + @workflow.run + async def run(self) -> int: + output = external_output_stream.with_options( + max_publish_latency=timedelta(seconds=1) + ).topic("events", type=str) + input_topic = external_stream.with_options( + idle_timeout=timedelta(seconds=30) + ).topic("release", type=str) + await output.publish("visible-before-input") + await input_topic.subscribe().__aiter__().__anext__() + return 1 + + +@workflow.defn +class PublishAcrossRolloverAndContinueAsNewWorkflow: + """Crosses output capacity rollover, then continues as new.""" + + @workflow.run + async def run(self) -> str: + topic = external_output_stream.with_options(max_records=1).topic( + "events", type=str + ) + if workflow.info().continued_run_id is None: + await topic.publish("before-rollover") + await topic.publish("after-rollover") + workflow.continue_as_new() + await topic.publish("after-continue-as-new") + return "done" + + +@workflow.defn +class ConcurrentTurnOutputWorkflow: + """Returns stable turn IDs that clients correlate with output envelopes.""" + + def __init__(self) -> None: + self._done = False + + @workflow.run + async def run(self) -> str: + await workflow.wait_condition(lambda: self._done) + return "done" + + @workflow.update + async def start_turn(self, turn_id: str) -> str: + await external_output_stream.topic("events", type=str).publish( + f"turn_started:{turn_id}" + ) + return turn_id + + @workflow.update + async def finish(self) -> None: + await external_output_stream.topic("events", type=str).finish() + self._done = True + + +@workflow.defn +class PublishWhileTwoInputsParkWorkflow: + """Keeps two input intents outstanding while output latency wins.""" + + @workflow.run + async def run(self) -> None: + output = external_output_stream.with_options( + max_publish_latency=timedelta(milliseconds=500) + ).topic("events", type=str) + inputs = external_stream.with_options(idle_timeout=timedelta(milliseconds=50)) + first = inputs.topic("first", type=str).subscribe().__aiter__() + second = inputs.topic("second", type=str).subscribe().__aiter__() + await output.publish("latency-wins") + await asyncio.gather(first.__anext__(), second.__anext__()) + + +@workflow.defn +class ThreeRetainedOutputWindowsWorkflow: + """Publishes twice per retained latency window for three windows.""" + + @workflow.run + async def run(self) -> int: + output = external_output_stream.with_options( + max_publish_latency=timedelta(milliseconds=150) + ).topic("events", type=str) + triggers = external_stream.with_options( + idle_timeout=timedelta(seconds=30) + ).topic("triggers", type=str) + windows = 0 + async for trigger in triggers.subscribe(): + if trigger == "stop": + return windows + await output.publish(f"{trigger}:first") + await output.publish(f"{trigger}:second") + windows += 1 + return windows + + +class _OutputMemoryBackend(MemoryStreamBackend, OutputStreamBackend): + """In-memory output provider used to exercise the real Worker/Core path.""" + + provider_id = "output-memory" + provider_format_version = 1 + + def __init__(self) -> None: + super().__init__() + self.stages: dict[tuple[str, int], OutputStage] = {} + self._stage_by_offset: dict[tuple[StreamKey, Offset], tuple[str, int]] = {} + self.read_operations: list[str] = [] + self.commit_attempts = 0 + self.commit_failures_remaining = 0 + self.commit_failed = asyncio.Event() + self.stage_placed = asyncio.Event() + self.stage_release: asyncio.Event | None = None + self.pending_read = asyncio.Event() + self.expected_park_intents = 0 + self.all_park_intents_installed = asyncio.Event() + self.recheck_release: asyncio.Event | None = None + self.removed_park_intents: list[int] = [] + + async def stage_output( + self, + manifest: OutputStageManifest, + records: Sequence[StagedOutputRecord], + ) -> OutputStage: + identity = (manifest.stage_token, manifest.sub_batch_id) + existing = self.stages.get(identity) + if existing is not None: + if existing.manifest != manifest: + raise OutputStageConflictError(manifest) + return existing + if len(records) != manifest.record_count or any( + record.publish_index != index for index, record in enumerate(records) + ): + raise ValueError("staged output indexes do not match the manifest") + + placed: list[OutputStreamRecord] = [] + for record in records: + stored = await self.append( + manifest.stream_key, + StreamRecord( + kind=record.kind, + payload=record.payload, + producer_session_id=( + f"stage:{manifest.stage_token}:{manifest.sub_batch_id}" + ), + sequence=record.publish_index, + ), + ) + assert stored.offset is not None + placed.append( + OutputStreamRecord(stored.kind, stored.payload, stored.offset) + ) + self._stage_by_offset[(manifest.stream_key, stored.offset)] = identity + stage = OutputStage(manifest, tuple(placed), OutputStageStatus.PENDING) + self.stages[identity] = stage + self.stage_placed.set() + if self.stage_release is not None: + await self.stage_release.wait() + return stage + + async def commit_output(self, manifest: OutputStageManifest) -> OutputStage: + self.commit_attempts += 1 + if self.commit_failures_remaining: + self.commit_failures_remaining -= 1 + self.commit_failed.set() + raise OSError("injected post-report output commit failure") + return self._resolve(manifest, OutputStageStatus.COMMITTED) + + async def abort_output(self, manifest: OutputStageManifest) -> OutputStage: + return self._resolve(manifest, OutputStageStatus.ABORTED) + + def _resolve( + self, manifest: OutputStageManifest, status: OutputStageStatus + ) -> OutputStage: + identity = (manifest.stage_token, manifest.sub_batch_id) + existing = self.stages.get(identity) + if existing is None: + raise OutputStageNotFoundError(manifest) + if existing.manifest != manifest: + raise OutputStageConflictError(manifest) + if existing.status is not OutputStageStatus.PENDING: + if existing.status is status: + return existing + raise OutputStageResolutionError( + manifest, + current=existing.status, + requested=status, + ) + resolved = replace(existing, status=status) + self.stages[identity] = resolved + return resolved + + async def output_stage(self, manifest: OutputStageManifest) -> OutputStage | None: + self.read_operations.append("output_stage") + existing = self.stages.get((manifest.stage_token, manifest.sub_batch_id)) + if existing is not None and existing.manifest != manifest: + raise OutputStageConflictError(manifest) + return existing + + async def append_output( + self, key: StreamKey, record: StreamRecord + ) -> OutputStreamRecord: + placed = await self.append(key, record) + assert placed.offset is not None + return OutputStreamRecord(placed.kind, placed.payload, placed.offset) + + async def read_output_after( + self, + key: StreamKey, + after: Cursor, + *, + max_records: int, + block: timedelta | None = None, + ) -> OutputReadResult: + del block + self.read_operations.append("read_output_after") + visible: list[OutputStreamRecord] = [] + for record in self._after(key, after): + assert record.offset is not None + identity = self._stage_by_offset.get((key, record.offset)) + if identity is not None: + stage = self.stages[identity] + if stage.status is OutputStageStatus.PENDING: + self.pending_read.set() + return OutputReadResult( + tuple(visible), + PendingOutputBarrier(stage.manifest, stage.records[0].offset), + ) + if stage.status is OutputStageStatus.ABORTED: + continue + visible.append( + OutputStreamRecord(record.kind, record.payload, record.offset) + ) + if len(visible) >= max_records: + break + return OutputReadResult(tuple(visible)) + + async def output_tail(self, key: StreamKey) -> Cursor: + self.read_operations.append("output_tail") + result = await self.read_output_after( + key, BEGINNING, max_records=2**31 - 1, block=None + ) + if not result.records: + return BEGINNING + return AFTER(result.records[-1].offset) + + async def install_park_intent(self, key: StreamKey, intent: ParkIntent) -> None: + await super().install_park_intent(key, intent) + if ( + self.expected_park_intents + and len(self._intents) == self.expected_park_intents + ): + self.all_park_intents_installed.set() + + async def recheck(self, key: StreamKey, wait_id: int) -> bool: + if self.recheck_release is not None: + await self.recheck_release.wait() + return await super().recheck(key, wait_id) + + async def remove_park_intent_if_matches( + self, + key: StreamKey, + wait_id: int, + *, + run_id: str, + park_generation: int, + ) -> ParkIntentRemoval: + result = await super().remove_park_intent_if_matches( + key, + wait_id, + run_id=run_id, + park_generation=park_generation, + ) + if result in (ParkIntentRemoval.REMOVED, ParkIntentRemoval.ABSENT): + self.removed_park_intents.append(wait_id) + return result + + +class _UnavailableOutputMemoryBackend(_OutputMemoryBackend): + """Fails one stage, then holds its retry until a test restores storage.""" + + def __init__(self) -> None: + super().__init__() + self.failed_once = False + self.failed_stage_attempted = asyncio.Event() + self.recovery_release = asyncio.Event() + + async def stage_output( + self, + manifest: OutputStageManifest, + records: Sequence[StagedOutputRecord], + ) -> OutputStage: + if not self.failed_once: + self.failed_once = True + self.failed_stage_attempted.set() + raise OSError("injected output staging outage") + await self.recovery_release.wait() + return await super().stage_output(manifest, records) + + +async def _history(handle: Any) -> list[Any]: + return [event async for event in handle.fetch_history_events()] + + +def _output_marker_data(events: Sequence[Any]) -> list[tuple[Any, Any]]: + from temporalio.bridge.proto.external_data import ExternalStreamMarkerData + + markers: list[tuple[Any, Any]] = [] + for event in events: + if not event.HasField("marker_recorded_event_attributes"): + continue + attributes = event.marker_recorded_event_attributes + if attributes.marker_name != "core_external_stream": + continue + payloads = attributes.details["external_stream"].payloads + if len(payloads) != 1: + continue + marker = ExternalStreamMarkerData() + marker.ParseFromString(payloads[0].data) + if marker.HasField("output"): + markers.append((event, marker)) + return markers + + +def _output_marker_tokens(events: Sequence[Any]) -> list[str]: + return [marker.output.stage_token for _, marker in _output_marker_data(events)] + + +async def _first_failed_task(handle: Any, *, timeout: float = 30) -> Any: + deadline = asyncio.get_running_loop().time() + timeout + while asyncio.get_running_loop().time() < deadline: + for event in await _history(handle): + if event.HasField("workflow_task_failed_event_attributes"): + return event.workflow_task_failed_event_attributes + await asyncio.sleep(0.1) + raise AssertionError("the Workflow produced no failed Workflow Task") + + +def _failure_types(failure: Any) -> list[str]: + types: list[str] = [] + while True: + types.append(failure.application_failure_info.type) + if not failure.HasField("cause"): + return types + failure = failure.cause + + +async def test_workflow_output_round_trips_through_core_and_client( + client: Client, +) -> None: + """Stage, record compact proof, promote after reporting, and decode.""" + backend = _OutputMemoryBackend() + task_queue = f"tq-{uuid.uuid4()}" + workflow_id = f"wf-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[PublishExternalOutputWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + PublishExternalOutputWorkflow.run, + id=workflow_id, + task_queue=task_queue, + ) + assert await asyncio.wait_for(handle.result(), 30) == 1 + + description = await handle.describe() + chain = WorkflowChainKey( + client.namespace, + workflow_id, + description.raw_description.workflow_execution_info.first_run_id, + ) + assert len(backend.stages) == 1 + staged = next(iter(backend.stages.values())) + + # The server result and Worker's post-report History reconciliation + # race one another. Wait for the healthy-path promotion rather than + # requiring a reader to repair every successfully committed batch. + async def wait_until_committed() -> None: + while next(iter(backend.stages.values())).status is not ( + OutputStageStatus.COMMITTED + ): + await asyncio.sleep(0.01) + + await asyncio.wait_for(wait_until_committed(), 30) + + output_client = await ExternalOutputStreamClient.connect( + backend=backend, + workflow=chain, + client=client, + ) + + async def read_all() -> list[str]: + return [ + item.data + async for item in output_client.topic("events", type=str).subscribe() + ] + + assert await asyncio.wait_for(read_all(), 30) == ["workflow-output-secret"] + assert next(iter(backend.stages.values())).status is OutputStageStatus.COMMITTED + + events = [event async for event in handle.fetch_history_events()] + markers = [ + event.marker_recorded_event_attributes + for event in events + if event.HasField("marker_recorded_event_attributes") + and event.marker_recorded_event_attributes.marker_name + == "core_external_stream" + ] + assert len(markers) == 1 + envelope = markers[0].details["external_stream"] + assert len(envelope.payloads) == 1 + + from temporalio.bridge.proto.external_data import ExternalStreamMarkerData + + marker = ExternalStreamMarkerData() + marker.ParseFromString(envelope.payloads[0].data) + assert marker.HasField("output") + output = marker.output + assert output.stage_token == staged.manifest.stage_token + assert output.run_id == handle.result_run_id + assert output.provider_id == backend.provider_id + assert output.provider_format_version == backend.provider_format_version + assert output.history_floor_event_id == staged.manifest.history_floor_event_id + assert output.history_floor_event_id > 0 + by_id = {event.event_id: event for event in events} + assert output.history_floor_event_id in by_id + assert by_id[output.history_floor_event_id + 1].HasField( + "workflow_task_scheduled_event_attributes" + ) + assert [ + (topic.topic, topic.record_count, topic.finished) for topic in output.topics + ] == [("events", 2, True)] + assert [ + tuple(segment.record_counts_by_topic) for segment in output.segments + ] == [(2,)] + assert len(output.topics[0].logical_fingerprint) == 32 + + history_bytes = b"".join(event.SerializeToString() for event in events) + assert b"workflow-output-secret" not in history_bytes + assert b"workflow-output-secret" not in envelope.payloads[0].data + + +async def test_output_latency_flushes_a_retained_workflow_task( + client: Client, +) -> None: + """The latency deadline exposes output while the Workflow is still blocked.""" + backend = _OutputMemoryBackend() + task_queue = f"tq-{uuid.uuid4()}" + workflow_id = f"wf-{uuid.uuid4()}" + handle = None + async with Worker( + client, + task_queue=task_queue, + workflows=[PublishOutputWhileInputIsRetainedWorkflow], + external_stream_backend=backend, + ): + started_at = time.monotonic() + handle = await client.start_workflow( + PublishOutputWhileInputIsRetainedWorkflow.run, + id=workflow_id, + task_queue=task_queue, + ) + description = await handle.describe() + chain = WorkflowChainKey( + client.namespace, + workflow_id, + description.raw_description.workflow_execution_info.first_run_id, + ) + + try: + # The ordinary activation is retainable, so it must not stage on + # return. Core's one-second output deadline is the first legal + # boundary that can make this record visible. + await asyncio.sleep(0.1) + assert backend.stages == {} + + async def wait_until_committed() -> OutputStage: + while not backend.stages: + await asyncio.sleep(0.01) + stage = next(iter(backend.stages.values())) + while stage.status is not OutputStageStatus.COMMITTED: + await asyncio.sleep(0.01) + stage = next(iter(backend.stages.values())) + return stage + + staged = await asyncio.wait_for(wait_until_committed(), 30) + elapsed = time.monotonic() - started_at + assert elapsed >= 0.75 + assert staged.manifest.record_count == 1 + assert [record.kind for record in staged.records] == [RecordKind.DATA] + + # The output is externally readable even though Workflow code is + # still waiting for its input record. + assert (await handle.describe()).status == WorkflowExecutionStatus.RUNNING + output_client = await ExternalOutputStreamClient.connect( + backend=backend, + workflow=chain, + client=client, + ) + subscription = output_client.topic("events", type=str).subscribe() + item = await asyncio.wait_for(subscription.__anext__(), 30) + assert item.data == "visible-before-input" + + codec = StreamPayloadCodec(client.data_converter, str) + await backend.append( + chain.stream_key("release"), + StreamRecord( + RecordKind.DATA, + await codec.encode("release"), + "test", + 0, + ), + ) + assert await asyncio.wait_for(handle.result(), 30) == 1 + finally: + if (await handle.describe()).status == WorkflowExecutionStatus.RUNNING: + await handle.terminate() + + +@pytest.mark.timeout(90) +async def test_three_retained_output_windows_write_three_markers_and_wfts( + client: Client, +) -> None: + """Two publishes coalesce per window; three crossed windows stay distinct.""" + backend = _OutputMemoryBackend() + task_queue = f"tq-{uuid.uuid4()}" + workflow_id = f"wf-{uuid.uuid4()}" + handle = None + async with Worker( + client, + task_queue=task_queue, + workflows=[ThreeRetainedOutputWindowsWorkflow], + external_stream_backend=backend, + ): + try: + handle = await client.start_workflow( + ThreeRetainedOutputWindowsWorkflow.run, + id=workflow_id, + task_queue=task_queue, + ) + description = await handle.describe() + chain = WorkflowChainKey( + client.namespace, + workflow_id, + description.raw_description.workflow_execution_info.first_run_id, + ) + codec = StreamPayloadCodec(client.data_converter, str) + trigger_key = chain.stream_key("triggers") + + async def append_trigger(sequence: int, value: str) -> None: + await backend.append( + trigger_key, + StreamRecord( + RecordKind.DATA, + await codec.encode(value), + "latency-window-test", + sequence, + ), + ) + + async def wait_for_output_markers(count: int) -> list[tuple[Any, Any]]: + while True: + markers = _output_marker_data(await _history(handle)) + if len(markers) >= count: + return markers + await asyncio.sleep(0.01) + + for sequence, window in enumerate(("one", "two", "three")): + await append_trigger(sequence, window) + markers = await asyncio.wait_for( + wait_for_output_markers(sequence + 1), 30 + ) + assert len(markers) == sequence + 1 + + assert (await handle.describe()).status == WorkflowExecutionStatus.RUNNING + markers = _output_marker_data(await _history(handle)) + assert len(markers) == 3 + completed_event_ids = { + event.marker_recorded_event_attributes.workflow_task_completed_event_id + for event, _ in markers + } + assert len(completed_event_ids) == 3 + + from temporalio.bridge.proto.external_data import ParkReason + + for _, marker in markers: + assert marker.terminal_boundary == ParkReason.PARK_REASON_OUTPUT_LATENCY + assert len(marker.output.topics) == 1 + assert marker.output.topics[0].record_count == 2 + assert ( + sum( + segment.record_counts_by_topic[0] + for segment in marker.output.segments + ) + == 2 + ) + + events = await _history(handle) + assert completed_event_ids <= { + event.event_id + for event in events + if event.HasField("workflow_task_completed_event_attributes") + } + assert len(backend.stages) == 3 + assert all( + stage.status is OutputStageStatus.COMMITTED + for stage in backend.stages.values() + ) + + await append_trigger(3, "stop") + assert await asyncio.wait_for(handle.result(), 30) == 3 + assert len(_output_marker_data(await _history(handle))) == 3 + finally: + if ( + handle is not None + and (await handle.describe()).status == WorkflowExecutionStatus.RUNNING + ): + await handle.terminate() + + +async def test_finished_output_survives_continue_as_new_without_backend_reads( + client: Client, +) -> None: + backend = _OutputMemoryBackend() + task_queue = f"tq-{uuid.uuid4()}" + workflow_id = f"wf-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[FinishOutputAcrossContinueAsNewWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + FinishOutputAcrossContinueAsNewWorkflow.run, + id=workflow_id, + task_queue=task_queue, + ) + + assert await asyncio.wait_for(handle.result(), 30) == ( + "successor rejected finished topic" + ) + + # The first Run staged only FINISH. The successor rejected DATA from its + # reserved continuation header before consulting mutable provider state. + assert len(backend.stages) == 1 + [stage] = backend.stages.values() + assert [record.kind.name for record in stage.records] == ["FINISH"] + assert backend.read_operations == [] + + +async def test_cold_client_repairs_post_report_commit_failure( + client: Client, +) -> None: + """A marker stays authoritative when opportunistic promotion fails once.""" + backend = _OutputMemoryBackend() + backend.commit_failures_remaining = 1 + task_queue = f"tq-{uuid.uuid4()}" + workflow_id = f"wf-{uuid.uuid4()}" + chain: WorkflowChainKey + staged: OutputStage + + async with Worker( + client, + task_queue=task_queue, + workflows=[PublishExternalOutputWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + PublishExternalOutputWorkflow.run, + id=workflow_id, + task_queue=task_queue, + ) + assert await asyncio.wait_for(handle.result(), 30) == 1 + await asyncio.wait_for(backend.commit_failed.wait(), 30) + + description = await handle.describe() + chain = WorkflowChainKey( + client.namespace, + workflow_id, + description.raw_description.workflow_execution_info.first_run_id, + ) + assert len(backend.stages) == 1 + staged = next(iter(backend.stages.values())) + assert staged.status is OutputStageStatus.PENDING + assert backend.commit_attempts == 1 + + # The injected commit is reached only after the Worker found the token + # in exact-run History. Assert that proof at the failure boundary too, + # before shutting the Worker down and handing recovery to a new client. + events = [event async for event in handle.fetch_history_events()] + from temporalio.contrib.external_workflow_streams._output_client import ( + _event_has_output_stage_token, + ) + + assert any( + _event_has_output_stage_token(event, staged.manifest.stage_token) + for event in events + ) + + cold_client = await ExternalOutputStreamClient.connect( + backend=backend, + workflow=chain, + client=client, + ) + + async def read_all() -> list[str]: + return [ + item.data + async for item in cold_client.topic("events", type=str).subscribe() + ] + + assert await asyncio.wait_for(read_all(), 30) == ["workflow-output-secret"] + assert ( + backend.stages[ + (staged.manifest.stage_token, staged.manifest.sub_batch_id) + ].status + is OutputStageStatus.COMMITTED + ) + assert backend.commit_attempts == 2 + + +@pytest.mark.timeout(90) +async def test_reader_waits_at_pending_barrier_until_marker_commit( + client: Client, +) -> None: + """A staged record is physically present but unreadable before its proof.""" + backend = _OutputMemoryBackend() + backend.stage_release = asyncio.Event() + task_queue = f"tq-{uuid.uuid4()}" + workflow_id = f"wf-{uuid.uuid4()}" + handle = None + async with Worker( + client, + task_queue=task_queue, + workflows=[PublishExternalOutputWorkflow], + external_stream_backend=backend, + ): + try: + handle = await client.start_workflow( + PublishExternalOutputWorkflow.run, + id=workflow_id, + task_queue=task_queue, + ) + await asyncio.wait_for(backend.stage_placed.wait(), 30) + description = await handle.describe() + chain = WorkflowChainKey( + client.namespace, + workflow_id, + description.raw_description.workflow_execution_info.first_run_id, + ) + output_client = await ExternalOutputStreamClient.connect( + backend=backend, + workflow=chain, + client=client, + ) + + async def read_all() -> list[str]: + return [ + item.data + async for item in output_client.topic( + "events", type=str + ).subscribe() + ] + + read_task = asyncio.create_task(read_all()) + await asyncio.wait_for(backend.pending_read.wait(), 30) + await asyncio.sleep(0.1) + assert not read_task.done() + [stage] = backend.stages.values() + assert stage.status is OutputStageStatus.PENDING + assert _output_marker_tokens(await _history(handle)) == [] + + backend.stage_release.set() + assert await asyncio.wait_for(handle.result(), 30) == 1 + assert await asyncio.wait_for(read_task, 30) == ["workflow-output-secret"] + [stage] = backend.stages.values() + assert stage.status is OutputStageStatus.COMMITTED + assert _output_marker_tokens(await _history(handle)) == [ + stage.manifest.stage_token + ] + finally: + backend.stage_release.set() + if ( + handle is not None + and (await handle.describe()).status == WorkflowExecutionStatus.RUNNING + ): + await handle.terminate() + + +@pytest.mark.timeout(90) +async def test_rejected_post_stage_completion_never_exposes_phantom_output( + client: Client, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Reject the first report after staging; only its fresh retry is visible.""" + backend = _OutputMemoryBackend() + task_queue = f"tq-{uuid.uuid4()}" + workflow_id = f"wf-{uuid.uuid4()}" + original = temporalio.bridge.worker.Worker.complete_workflow_activation + rejected = asyncio.Event() + rejected_tokens: list[str] = [] + + async def reject_first_output_commit( + bridge: temporalio.bridge.worker.Worker, + completion: Any, + ) -> None: + if not rejected.is_set() and completion.HasField("successful"): + commits = [ + command.workflow_output_stream_commit + for command in completion.successful.commands + if command.HasField("workflow_output_stream_commit") + ] + if commits: + rejected_token = commits[0].manifest.stage_token + rejected_tokens.append(rejected_token) + matching = [ + stage + for stage in backend.stages.values() + if stage.manifest.stage_token == rejected_token + ] + assert len(matching) == 1 + assert matching[0].status is OutputStageStatus.PENDING + # Replace the successful report after provider staging. Core + # reports an ordinary failed WFT, then retries the Workflow. + completion.failed.failure.message = ( + "injected rejection after output stage" + ) + completion.failed.failure.application_failure_info.type = ( + "InjectedPostStageFailure" + ) + rejected.set() + await original(bridge, completion) + + monkeypatch.setattr( + temporalio.bridge.worker.Worker, + "complete_workflow_activation", + reject_first_output_commit, + ) + + async with Worker( + client, + task_queue=task_queue, + workflows=[PublishExternalOutputWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + PublishExternalOutputWorkflow.run, + id=workflow_id, + task_queue=task_queue, + ) + await asyncio.wait_for(rejected.wait(), 30) + assert await asyncio.wait_for(handle.result(), 30) == 1 + + [rejected_token] = rejected_tokens + events = await _history(handle) + assert any( + event.HasField("workflow_task_failed_event_attributes") for event in events + ) + stages = list(backend.stages.values()) + assert len(stages) == 2 + [abandoned] = [ + stage for stage in stages if stage.manifest.stage_token == rejected_token + ] + assert abandoned.status is OutputStageStatus.ABORTED + + marker_tokens = _output_marker_tokens(events) + assert len(marker_tokens) == 1 + assert marker_tokens != [rejected_token] + + description = await handle.describe() + chain = WorkflowChainKey( + client.namespace, + workflow_id, + description.raw_description.workflow_execution_info.first_run_id, + ) + output_client = await ExternalOutputStreamClient.connect( + backend=backend, + workflow=chain, + client=client, + ) + assert [ + item.data + async for item in output_client.topic("events", type=str).subscribe() + ] == ["workflow-output-secret"] + committed = [ + stage + for stage in backend.stages.values() + if stage.status is OutputStageStatus.COMMITTED + ] + assert len(committed) == 1 + assert marker_tokens == [committed[0].manifest.stage_token] + + +@pytest.mark.timeout(90) +async def test_concurrent_updates_preserve_their_own_turn_ids( + client: Client, +) -> None: + backend = _OutputMemoryBackend() + task_queue = f"tq-{uuid.uuid4()}" + workflow_id = f"wf-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[ConcurrentTurnOutputWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + ConcurrentTurnOutputWorkflow.run, + id=workflow_id, + task_queue=task_queue, + ) + description = await handle.describe() + chain = WorkflowChainKey( + client.namespace, + workflow_id, + description.raw_description.workflow_execution_info.first_run_id, + ) + output_client = await ExternalOutputStreamClient.connect( + backend=backend, + workflow=chain, + client=client, + ) + topic = output_client.topic("events", type=str) + boundary = await topic.tail() + first_id = f"turn-{uuid.uuid4()}" + second_id = f"turn-{uuid.uuid4()}" + + returned = await asyncio.gather( + handle.execute_update(ConcurrentTurnOutputWorkflow.start_turn, first_id), + handle.execute_update(ConcurrentTurnOutputWorkflow.start_turn, second_id), + ) + assert returned == [first_id, second_id] + await handle.execute_update(ConcurrentTurnOutputWorkflow.finish) + assert await asyncio.wait_for(handle.result(), 30) == "done" + + async def scan_for(own_id: str) -> list[str]: + seen: list[str] = [] + async for item in topic.subscribe(after=boundary): + seen.append(item.data) + if item.data == f"turn_started:{own_id}": + return seen + raise AssertionError(f"no output for {own_id}") + + first_seen, second_seen = await asyncio.gather( + scan_for(first_id), scan_for(second_id) + ) + assert first_seen[-1] == f"turn_started:{first_id}" + assert second_seen[-1] == f"turn_started:{second_id}" + # Both subscribers use the same boundary; one necessarily sees the + # other concurrent turn first, but still correlates its own envelope. + assert sorted((len(first_seen), len(second_seen))) == [1, 2] + + +@pytest.mark.timeout(90) +async def test_stage_outage_blocks_wft_and_reports_external_storage_cause( + client: Client, +) -> None: + backend = _UnavailableOutputMemoryBackend() + task_queue = f"tq-{uuid.uuid4()}" + workflow_id = f"wf-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[PublishExternalOutputWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + PublishExternalOutputWorkflow.run, + id=workflow_id, + task_queue=task_queue, + ) + result_task = asyncio.create_task(handle.result()) + await asyncio.wait_for(backend.failed_stage_attempted.wait(), 30) + await asyncio.sleep(0.05) + assert not result_task.done() + assert (await handle.describe()).status == WorkflowExecutionStatus.RUNNING + assert backend.stages == {} + + backend.recovery_release.set() + assert await asyncio.wait_for(result_task, 30) == 1 + failed = await _first_failed_task(handle) + assert ( + failed.cause + == temporalio.api.enums.v1.WorkflowTaskFailedCause.WORKFLOW_TASK_FAILED_CAUSE_EXTERNAL_STORAGE_FAILURE + ) + assert "StreamStorageError" in _failure_types(failed.failure) + assert len(_output_marker_tokens(await _history(handle))) == 1 + stages = list(backend.stages.values()) + assert len(stages) == 1 + assert stages[0].status is OutputStageStatus.COMMITTED + + description = await handle.describe() + chain = WorkflowChainKey( + client.namespace, + workflow_id, + description.raw_description.workflow_execution_info.first_run_id, + ) + output_client = await ExternalOutputStreamClient.connect( + backend=backend, + workflow=chain, + client=client, + ) + assert [ + item.data + async for item in output_client.topic("events", type=str).subscribe() + ] == ["workflow-output-secret"] + + +@pytest.mark.timeout(90) +async def test_cursor_resume_crosses_rollover_and_continue_as_new( + client: Client, + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend = _OutputMemoryBackend() + task_queue = f"tq-{uuid.uuid4()}" + workflow_id = f"wf-{uuid.uuid4()}" + original = temporalio.bridge.worker.Worker.complete_workflow_activation + rollover_requests: list[bool] = [] + + async def capture_rollover_request( + bridge: temporalio.bridge.worker.Worker, + completion: Any, + ) -> None: + if completion.HasField("successful"): + rollover_requests.extend( + command.workflow_output_stream_commit.request_rollover + for command in completion.successful.commands + if command.HasField("workflow_output_stream_commit") + ) + await original(bridge, completion) + + monkeypatch.setattr( + temporalio.bridge.worker.Worker, + "complete_workflow_activation", + capture_rollover_request, + ) + async with Worker( + client, + task_queue=task_queue, + workflows=[PublishAcrossRolloverAndContinueAsNewWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + PublishAcrossRolloverAndContinueAsNewWorkflow.run, + id=workflow_id, + task_queue=task_queue, + ) + description = await handle.describe() + first_run_id = description.raw_description.workflow_execution_info.first_run_id + + async def wait_for_replacement_stage() -> None: + while len(backend.stages) < 2: + await asyncio.sleep(0.01) + + try: + await asyncio.wait_for(wait_for_replacement_stage(), 10) + except asyncio.TimeoutError: + events = await _history(handle) + statuses = [stage.status.name for stage in backend.stages.values()] + marker_tokens = _output_marker_tokens(events) + status = (await handle.describe()).status + await handle.terminate() + pytest.fail( + "output capacity did not create its replacement Workflow Task: " + f"requests={rollover_requests}, stages={statuses}, " + f"markers={marker_tokens}, " + f"status={status.name if status is not None else 'NONE'}" + ) + assert await asyncio.wait_for(handle.result(), 30) == "done" + assert rollover_requests[:2] == [True, False] + + chain = WorkflowChainKey(client.namespace, workflow_id, first_run_id) + output_client = await ExternalOutputStreamClient.connect( + backend=backend, + workflow=chain, + client=client, + ) + topic = output_client.topic("events", type=str) + initial = topic.subscribe() + first = await asyncio.wait_for(initial.__anext__(), 30) + await cast(Any, initial).aclose() + assert first.data == "before-rollover" + + resumed = topic.subscribe(after=first.offset) + second = await asyncio.wait_for(resumed.__anext__(), 30) + third = await asyncio.wait_for(resumed.__anext__(), 30) + await cast(Any, resumed).aclose() + assert [second.data, third.data] == [ + "after-rollover", + "after-continue-as-new", + ] + assert len({first.offset, second.offset, third.offset}) == 3 + + stages_by_run: dict[str, int] = {} + for stage in backend.stages.values(): + stages_by_run[stage.manifest.run_id] = ( + stages_by_run.get(stage.manifest.run_id, 0) + 1 + ) + assert stage.status is OutputStageStatus.COMMITTED + assert stages_by_run[first_run_id] == 2 + successor_runs = { + run_id: count + for run_id, count in stages_by_run.items() + if run_id != first_run_id + } + assert list(successor_runs.values()) == [1] + + +@pytest.mark.timeout(90) +async def test_output_deadline_wins_park_race_and_removes_every_intent( + client: Client, +) -> None: + backend = _OutputMemoryBackend() + backend.expected_park_intents = 2 + backend.recheck_release = asyncio.Event() + task_queue = f"tq-{uuid.uuid4()}" + workflow_id = f"wf-{uuid.uuid4()}" + handle = None + async with Worker( + client, + task_queue=task_queue, + workflows=[PublishWhileTwoInputsParkWorkflow], + external_stream_backend=backend, + ): + try: + handle = await client.start_workflow( + PublishWhileTwoInputsParkWorkflow.run, + id=workflow_id, + task_queue=task_queue, + ) + await asyncio.wait_for(backend.all_park_intents_installed.wait(), 30) + assert len(backend._intents) == 2 + + # Keep PrepareExternalStreamPark outstanding until after the + # already-armed output deadline expires. Core must discard the + # losing park result and roll back both provider intents. + await asyncio.sleep(0.75) + backend.recheck_release.set() + + async def wait_for_flush_and_rollback() -> None: + while True: + committed = any( + stage.status is OutputStageStatus.COMMITTED + for stage in backend.stages.values() + ) + if committed and not backend._intents: + return + await asyncio.sleep(0.01) + + await asyncio.wait_for(wait_for_flush_and_rollback(), 30) + assert set(backend.removed_park_intents) == {1, 2} + assert len(backend.stages) == 1 + assert (await handle.describe()).status == WorkflowExecutionStatus.RUNNING + assert len(_output_marker_tokens(await _history(handle))) == 1 + finally: + backend.recheck_release.set() + if ( + handle is not None + and (await handle.describe()).status == WorkflowExecutionStatus.RUNNING + ): + await handle.terminate() diff --git a/tests/contrib/external_workflow_streams/test_parking_conformance.py b/tests/contrib/external_workflow_streams/test_parking_conformance.py new file mode 100644 index 0000000000..ad563d4914 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_parking_conformance.py @@ -0,0 +1,289 @@ +"""P2b — the parking conformance suite, and proof that it fails a broken stub.""" + +from __future__ import annotations + +import uuid +from datetime import timedelta + +import pytest + +from temporalio.contrib.external_workflow_streams._backend import ( + ParkIntent, + StreamBackend, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._record import BEGINNING +from tests.contrib.external_workflow_streams import conformance +from tests.contrib.external_workflow_streams.conformance import ( + PARKING_CONFORMANCE_CHECKS, + Check, +) +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend + + +@pytest.fixture +def stream_key() -> StreamKey: + return StreamKey("ns", "wf", uuid.uuid4().hex, "tokens") + + +@pytest.fixture +def backend() -> MemoryStreamBackend: + return MemoryStreamBackend() + + +# --- the reference backend conforms ----------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("check", PARKING_CONFORMANCE_CHECKS, ids=lambda c: c.__name__) +async def test_reference_backend_conforms( + check: Check, backend: MemoryStreamBackend, stream_key: StreamKey +) -> None: + await check(backend, stream_key) + + +def test_the_parking_suite_is_not_silently_empty() -> None: + assert len(PARKING_CONFORMANCE_CHECKS) >= 10 + assert len({c.__name__ for c in PARKING_CONFORMANCE_CHECKS}) == len( + PARKING_CONFORMANCE_CHECKS + ) + + +class ObserveOnlyBackend(MemoryStreamBackend): + """Cannot lease, and says so. Every producer signals idempotently.""" + + supports_leased_claims = False + + async def claim_park_generation( + self, + key: StreamKey, + wait_id: int, + park_generation: int, + *, + claimant: str, + lease: timedelta, + ) -> bool: + return True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("check", PARKING_CONFORMANCE_CHECKS, ids=lambda c: c.__name__) +async def test_an_observe_only_backend_conforms( + check: Check, stream_key: StreamKey +) -> None: + """The suite must accept the fallback it explicitly permits. + + A provider that cannot lease is not broken -- it is a different, weaker + contract the design allows for, and a suite that rejected it would push + implementers toward an unleased claim, which *is* broken. + """ + await check(ObserveOnlyBackend(), stream_key) + + +# --- deliberately broken backends ------------------------------------------- + + +class StreamKeyedIntentBackend(MemoryStreamBackend): + """Keys park intents by stream alone, ignoring the wait id. + + The realistic mistake: thinking of a park as a property of the stream rather + than of the subscription. + """ + + async def install_park_intent(self, key: StreamKey, intent: ParkIntent) -> None: + self._intents[(key, 0)] = intent + + async def park_intent(self, key: StreamKey, wait_id: int) -> ParkIntent | None: + return self._intents.get((key, 0)) + + +class NonExpiringClaimBackend(MemoryStreamBackend): + """Claims are exclusive forever. Nothing recovers a crashed producer's.""" + + async def claim_park_generation( + self, + key: StreamKey, + wait_id: int, + park_generation: int, + *, + claimant: str, + lease: timedelta, + ) -> bool: + held = self._claims.get((key, wait_id)) + if held is not None and held[0] != claimant and held[1] == park_generation: + return False + self._claims[(key, wait_id)] = (claimant, park_generation, float("inf")) + return True + + +class LastIntentOnlyBackend(MemoryStreamBackend): + """Enumerates only the most recently installed intent. + + A plausible implementation for a provider that tracks "the" park per stream, + and it looks correct until a Workflow subscribes to one stream twice: the + subscription it omits stays parked on records already sitting in the stream, + and nothing ever wakes it because no producer knows it exists. + """ + + async def parked_wait_ids(self, key: StreamKey) -> list[int]: + matching = [wait_id for (stored, wait_id) in self._intents if stored == key] + return matching[-1:] + + +class BlindRecheckBackend(MemoryStreamBackend): + """Rechecks against the stream's state when the intent was installed.""" + + async def recheck(self, key: StreamKey, wait_id: int) -> bool: + return False + + +class GenerationOutlivesRemovalBackend(MemoryStreamBackend): + """Answers the generation from a remembered value beside the intent. + + The plausible mistake: reading "the current park generation" as a fact about + the subscription's history rather than as "is a park outstanding right now". + It passes every other parking check -- the intent really is removed, and + really does stop being enumerable -- while the one call every wake path + actually asks keeps naming a park that is over. + """ + + def __init__(self) -> None: + super().__init__() + self._generations: dict[tuple[StreamKey, int], int] = {} + + async def install_park_intent(self, key: StreamKey, intent: ParkIntent) -> None: + await super().install_park_intent(key, intent) + self._generations[(key, intent.wait_id)] = intent.park_generation + + async def current_park_generation(self, key: StreamKey, wait_id: int) -> int | None: + return self._generations.get((key, wait_id)) + + +class DelimiterJoiningBackend(MemoryStreamBackend): + """Renders an identity by joining its fields, the way a first draft does. + + The natural implementation for any store with a flat keyspace, and the one + the Redis provider originally had. It looks correct for every identity that + contains no delimiter, which is every identity anyone writes by hand -- so + the collision needs a Workflow ID or a topic name carrying a `:`, which is + an ordinary character in both. + """ + + def _collapse(self, key: StreamKey) -> StreamKey: + joined = ":".join( + ( + key.namespace, + key.workflow_id, + key.first_execution_run_id, + key.stream_name, + ) + ) + # The remaining fields are placeholders: what matters is that two + # distinct identities collapse onto one, which is exactly what a + # delimiter-joined physical key does. They stay non-empty because + # StreamKey rejects empty components. + return StreamKey(joined, "-", "-", "-") + + async def append(self, key, record): # type: ignore[no-untyped-def] + return await super().append(self._collapse(key), record) + + async def read_after(self, key, after, *, max_records, block=None): # type: ignore[no-untyped-def] + return await super().read_after( + self._collapse(key), after, max_records=max_records, block=block + ) + + async def read_range(self, key, first, last): # type: ignore[no-untyped-def] + return await super().read_range(self._collapse(key), first, last) + + async def install_park_intent(self, key, intent): # type: ignore[no-untyped-def] + return await super().install_park_intent(self._collapse(key), intent) + + async def park_intent(self, key, wait_id): # type: ignore[no-untyped-def] + return await super().park_intent(self._collapse(key), wait_id) + + async def current_park_generation(self, key, wait_id): # type: ignore[no-untyped-def] + return await super().current_park_generation(self._collapse(key), wait_id) + + async def parked_wait_ids(self, key): # type: ignore[no-untyped-def] + return await super().parked_wait_ids(self._collapse(key)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("broken", "check", "reason"), + [ + pytest.param( + StreamKeyedIntentBackend, + conformance.check_intents_are_keyed_by_stream_and_wait_id, + r"keyed \(stream key, wait_id\)", + id="intents-keyed-by-stream-alone", + ), + pytest.param( + NonExpiringClaimBackend, + conformance.check_an_expired_claim_is_taken_over, + "expired claim was not taken over", + id="claim-never-expires", + ), + pytest.param( + DelimiterJoiningBackend, + conformance.check_distinct_identities_never_share_storage, + "share one physical stream", + id="identity-fields-joined-without-escaping", + ), + pytest.param( + LastIntentOnlyBackend, + conformance.check_every_parked_subscription_is_enumerable, + "expected both parked subscriptions", + id="enumerates-only-the-last-intent", + ), + pytest.param( + BlindRecheckBackend, + conformance.check_recheck_sees_an_append_past_the_cursor, + "append/park race is not closed", + id="recheck-cannot-see-a-late-append", + ), + pytest.param( + GenerationOutlivesRemovalBackend, + conformance.check_a_removed_intent_reports_no_generation, + "still reports a generation", + id="generation-outlives-its-intent", + ), + ], +) +async def test_a_broken_backend_fails_for_the_right_reason( + broken: type[StreamBackend], check: Check, reason: str, stream_key: StreamKey +) -> None: + with pytest.raises(AssertionError, match=reason): + await check(broken(), stream_key) + + +# --- the intent value type -------------------------------------------------- + + +def test_a_park_generation_of_zero_is_refused() -> None: + """0 is the reserved unparked-wake sentinel and can never be a real park.""" + with pytest.raises(ValueError, match="starts at 1"): + ParkIntent(1, BEGINNING, park_generation=0, run_id="run-a") + + +@pytest.mark.asyncio +async def test_two_same_stream_subscriptions_are_visible_separately( + backend: MemoryStreamBackend, stream_key: StreamKey +) -> None: + """The end state P21 asserts, checked here at the provider level. + + Inspecting both in the backend is what distinguishes "two subscriptions + parked independently" from "one overwrote the other and happens to look + parked". + """ + for wait_id in (1, 2): + await backend.install_park_intent( + stream_key, + ParkIntent(wait_id, BEGINNING, park_generation=5, run_id="run-a"), + ) + + assert await backend.current_park_generation(stream_key, 1) == 5 + assert await backend.current_park_generation(stream_key, 2) == 5 + assert (await backend.park_intent(stream_key, 1)) != ( + await backend.park_intent(stream_key, 2) + ) diff --git a/tests/contrib/external_workflow_streams/test_producer.py b/tests/contrib/external_workflow_streams/test_producer.py new file mode 100644 index 0000000000..6cea72eed3 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_producer.py @@ -0,0 +1,987 @@ +"""P6a/P6 — producer binding, and append plus write fence. + +PYTEST_DONT_REWRITE: sandboxed fixture Workflows re-import this module, so pytest's +injected imports would make sandbox validation depend on pytest's import locks. +""" + +from __future__ import annotations + +import asyncio +import dataclasses +import uuid + +import pytest + +import temporalio.converter +from temporalio import activity, workflow +from temporalio.client import Client +from temporalio.contrib.external_workflow_streams._backend import ( + AppendConflictError, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._producer import ( + AppendNotAcknowledgedError, + ChainKeyMismatchError, + ExternalStreamProducer, + PrecedingWriteFailedError, + WorkflowChainKey, + _default_session_id, +) +from temporalio.contrib.external_workflow_streams._record import ( + BEGINNING, + Offset, + RecordKind, + StreamRecord, +) +from temporalio.testing import ActivityEnvironment +from temporalio.worker import Worker +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend + + +@workflow.defn +class ChainWorkflow: + """Runs long enough for a producer to describe it.""" + + @workflow.run + async def run(self) -> None: + await workflow.wait_condition(lambda: False) + + +@pytest.fixture +def backend() -> MemoryStreamBackend: + return MemoryStreamBackend() + + +# --- the chain key ---------------------------------------------------------- + + +@pytest.mark.parametrize( + "missing", ["namespace", "workflow_id", "first_execution_run_id"] +) +def test_every_part_of_the_chain_key_is_required(missing: str) -> None: + parts = { + "namespace": "ns", + "workflow_id": "wf", + "first_execution_run_id": "run-1", + } + parts[missing] = "" + + with pytest.raises(ValueError, match=missing): + WorkflowChainKey(**parts) + + +def test_the_stream_name_appears_only_in_topic() -> None: + """No two arguments can disagree about the stream name. + + ``connect()`` takes the chain key and ``topic(name)`` completes it, so + there is exactly one place the name is written. + """ + key = WorkflowChainKey("ns", "wf", "run-1") + + assert key.stream_key("tokens") == StreamKey("ns", "wf", "run-1", "tokens") + assert key.stream_key("tokens") != key.stream_key("tool-events") + assert not hasattr(key, "stream_name") + + +# --- session IDs ------------------------------------------------------------ + + +def test_a_plain_process_must_supply_a_session_id() -> None: + """A random default would look correct until the first retry. + + Idempotency is on ``(session_id, sequence)``, so a fresh session per attempt + means every record appended twice -- and the failure shows up as duplicate + delivery to the Workflow, a long way from here. + """ + with pytest.raises(ValueError, match="no default outside an Activity"): + _default_session_id() + + +async def test_an_activity_derives_a_session_id_stable_across_attempts() -> None: + env = ActivityEnvironment() + base = dataclasses.replace( + ActivityEnvironment.default_info(), + activity_id="publish-tokens", + workflow_run_id="run-abc", + ) + + seen = [] + for attempt in (1, 2, 3): + env.info = dataclasses.replace(base, attempt=attempt) + + @activity.defn + async def capture() -> str: + return _default_session_id() + + seen.append(await env.run(capture)) + + assert len(set(seen)) == 1, ( + f"the session id must not vary by attempt, got {seen}; a retried attempt " + "must reuse its predecessor's idempotency keys" + ) + assert "run-abc" in seen[0] and "publish-tokens" in seen[0] + + +# --- binding and verification ----------------------------------------------- + + +async def test_a_wrong_first_execution_run_id_fails_loudly( + client: Client, backend: MemoryStreamBackend +) -> None: + task_queue = f"tq-{uuid.uuid4()}" + async with Worker(client, task_queue=task_queue, workflows=[ChainWorkflow]): + handle = await client.start_workflow( + ChainWorkflow.run, id=f"wf-{uuid.uuid4()}", task_queue=task_queue + ) + try: + with pytest.raises(ChainKeyMismatchError, match="first execution Run ID"): + await ExternalStreamProducer.connect( + backend=backend, + workflow=WorkflowChainKey( + client.namespace, handle.id, "definitely-not-the-run-id" + ), + client=client, + ) + finally: + await handle.terminate() + + +async def test_a_wrong_namespace_fails_before_any_server_call( + client: Client, backend: MemoryStreamBackend +) -> None: + with pytest.raises(ChainKeyMismatchError, match="namespace"): + await ExternalStreamProducer.connect( + backend=backend, + workflow=WorkflowChainKey("some-other-namespace", "wf", "run-1"), + client=client, + session_id="s", + ) + + +async def test_an_activity_and_a_plain_process_publish_under_one_verified_key( + client: Client, backend: MemoryStreamBackend +) -> None: + """Two very different producers, one stream, one verified chain key. + + The Activity gets its session id for free from its own identity; the plain + process supplies one. Neither can derive the chain key -- ``activity.Info`` + has no first execution Run ID -- so both are handed it by the Workflow. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker(client, task_queue=task_queue, workflows=[ChainWorkflow]): + handle = await client.start_workflow( + ChainWorkflow.run, id=f"wf-{uuid.uuid4()}", task_queue=task_queue + ) + try: + description = await handle.describe() + chain = WorkflowChainKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + ) + + @activity.defn + async def publish_from_activity() -> str: + producer = await ExternalStreamProducer.connect( + backend=backend, workflow=chain, client=client + ) + tokens = producer.topic("tokens", type=str) + await tokens.publish("from-activity", wake=False) + await tokens.finish_writing(wake=False) + return producer.session_id + + env = ActivityEnvironment() + env.info = dataclasses.replace( + ActivityEnvironment.default_info(), + activity_id="publish-tokens", + workflow_run_id=description.run_id, + ) + activity_session = await env.run(publish_from_activity) + assert activity_session.startswith("activity:"), ( + "an Activity must derive its session id from its own identity" + ) + + # And a plain, non-Temporal-hosted process against the same key. + plain = await ExternalStreamProducer.connect( + backend=backend, + workflow=chain, + client=client, + session_id="standalone-script", + ) + await plain.topic("tokens", type=str).publish("from-script", wake=False) + + records = backend.all_records(chain.stream_key("tokens")) + assert [r.kind for r in records] == [ + RecordKind.DATA, + RecordKind.WRITE_FENCE, + RecordKind.DATA, + ] + assert {r.producer_session_id for r in records} == { + activity_session, + "standalone-script", + } + finally: + await handle.terminate() + + +# --- append and the fence, without a server --------------------------------- + + +def _offline_producer(backend: MemoryStreamBackend, session: str = "s"): # type: ignore[no-untyped-def] + """A producer built directly, skipping the verification a bind performs. + + Verification is P6a's and is covered above against a real server; these + cases are about what append and the fence do once a producer is bound. + """ + import temporalio.converter + + return ExternalStreamProducer( + backend=backend, + workflow=WorkflowChainKey("ns", "wf", "run-1"), + data_converter=temporalio.converter.DataConverter.default, + session_id=session, + ) + + +def _placed_offsets(records: list[StreamRecord]) -> list[Offset]: + offsets: list[Offset] = [] + for record in records: + assert record.offset is not None + offsets.append(record.offset) + return offsets + + +async def test_records_are_appended_in_order(backend: MemoryStreamBackend) -> None: + tokens = _offline_producer(backend).topic("tokens", type=str) + + for value in ["a", "b", "c"]: + await tokens.publish(value, wake=False) + + records = backend.all_records(StreamKey("ns", "wf", "run-1", "tokens")) + assert [r.sequence for r in records] == [0, 1, 2] + assert backend.strictly_increasing(_placed_offsets(records)) + + +async def test_a_fence_takes_its_place_in_the_sequence( + backend: MemoryStreamBackend, +) -> None: + tokens = _offline_producer(backend).topic("tokens", type=str) + + await tokens.publish("a", wake=False) + await tokens.finish_writing(wake=False) + await tokens.publish("b", wake=False) + + records = backend.all_records(StreamKey("ns", "wf", "run-1", "tokens")) + assert [r.kind for r in records] == [ + RecordKind.DATA, + RecordKind.WRITE_FENCE, + RecordKind.DATA, + ] + + +async def test_a_retried_attempt_appends_no_duplicate( + backend: MemoryStreamBackend, +) -> None: + """The whole reason the session id must be stable across attempts.""" + key = StreamKey("ns", "wf", "run-1", "tokens") + + first = _offline_producer(backend, "attempt-stable").topic("tokens", type=str) + await first.publish("a", wake=False) + await first.publish("b", wake=False) + + # The Activity is retried: a fresh producer, the *same* session id, and the + # same values published again from the top. + retried = _offline_producer(backend, "attempt-stable").topic("tokens", type=str) + await retried.publish("a", wake=False) + await retried.publish("b", wake=False) + + assert len(backend.all_records(key)) == 2 + + +class GatedPayloadCodec(temporalio.converter.PayloadCodec): + """A codec whose ``encode`` completes only when the test releases that value. + + Stands in for what a real one does: an external payload store or a KMS round + trip, which is arbitrary I/O and completes in whatever order the service + answers. That order is not stable across Activity attempts, which is the whole + point -- so the test drives it explicitly instead of racing it. + """ + + def __init__(self) -> None: + self.arrived: dict[str, asyncio.Event] = {} + self.gates: dict[str, asyncio.Event] = {} + + def _for(self, marker: str) -> tuple[asyncio.Event, asyncio.Event]: + self.arrived.setdefault(marker, asyncio.Event()) + self.gates.setdefault(marker, asyncio.Event()) + return self.arrived[marker], self.gates[marker] + + async def wait_inside(self, marker: str) -> None: + """Blocks until this value's encode has started.""" + arrived, _ = self._for(marker) + await asyncio.wait_for(arrived.wait(), 2) + + def release(self, marker: str) -> None: + _, gate = self._for(marker) + gate.set() + + async def encode(self, payloads): # type: ignore[no-untyped-def] + # The default converter renders a `str` as JSON, so the value is legible + # in the payload's bytes and needs no side channel to identify it. + marker = payloads[0].data.decode() + arrived, gate = self._for(marker) + arrived.set() + await gate.wait() + return list(payloads) + + async def decode(self, payloads): # type: ignore[no-untyped-def] + return list(payloads) + + +def _gated_producer( + backend: MemoryStreamBackend, codec: GatedPayloadCodec, session: str +): # type: ignore[no-untyped-def] + import temporalio.converter + + return ExternalStreamProducer( + backend=backend, + workflow=WorkflowChainKey("ns", "wf", "run-1"), + data_converter=dataclasses.replace( + temporalio.converter.DataConverter.default, payload_codec=codec + ), + session_id=session, + ) + + +async def test_concurrent_publishes_take_their_sequence_in_invocation_order( + backend: MemoryStreamBackend, +) -> None: + """Idempotency may not depend on the order a payload codec answers in. + + ``(session_id, sequence)`` is the idempotency key and a retried Activity + reuses the session id on purpose, so the sequence a call draws has to be a + property of *the call*. Drawing it after awaiting the encode makes it a + property of the encode's completion order instead -- and a codec is allowed to + do real I/O, so that order is not stable across attempts. Two concurrent + publishes then exchange sequence numbers whenever the store answers the other + way round, the backend sees each stable key reused with different bytes, and + the retry raises ``AppendConflictError`` on both calls: a valid concurrent + Activity made permanently non-retryable by timing alone. + + The two attempts here are byte-identical in what they do and differ only in + which encode finishes first. + """ + key = StreamKey("ns", "wf", "run-1", "tokens") + + codec = GatedPayloadCodec() + topic = _gated_producer(backend, codec, "attempt-stable").topic("tokens", type=str) + first = asyncio.ensure_future(topic.publish("a", wake=False)) + second = asyncio.ensure_future(topic.publish("b", wake=False)) + await codec.wait_inside('"a"') + await codec.wait_inside('"b"') + codec.release('"a"') + codec.release('"b"') + offsets = [await first, await second] + + # The Activity is retried: a fresh producer, the *same* session id, the same + # calls in the same order -- and the store answers in the other order. + retry_codec = GatedPayloadCodec() + retried = _gated_producer(backend, retry_codec, "attempt-stable").topic( + "tokens", type=str + ) + first_again = asyncio.ensure_future(retried.publish("a", wake=False)) + second_again = asyncio.ensure_future(retried.publish("b", wake=False)) + await retry_codec.wait_inside('"a"') + await retry_codec.wait_inside('"b"') + retry_codec.release('"b"') + retry_codec.release('"a"') + + assert [await first_again, await second_again] == offsets, ( + "the retry appended the same values at different offsets, so the two " + "attempts disagree about what landed where" + ) + records = backend.all_records(key) + assert [r.sequence for r in records] == [0, 1] + assert len(records) == 2, "the retry was not idempotent; it appended duplicates" + + +async def test_reordered_encodes_cannot_duplicate_across_two_topics( + backend: MemoryStreamBackend, +) -> None: + """The same defect on the other side of the deduplication boundary. + + Deduplication is scoped by stream key, so two publishes to *different* topics + that exchange sequence numbers do not collide -- they each land under a key + the other stream has never seen, and the retry appends a second record + instead of raising. The conflict is the loud failure mode; this one is silent. + """ + tokens_key = StreamKey("ns", "wf", "run-1", "tokens") + events_key = StreamKey("ns", "wf", "run-1", "tool-events") + + codec = GatedPayloadCodec() + producer = _gated_producer(backend, codec, "attempt-stable") + tokens = producer.topic("tokens", type=str) + events = producer.topic("tool-events", type=str) + first = asyncio.ensure_future(tokens.publish("a", wake=False)) + second = asyncio.ensure_future(events.publish("b", wake=False)) + await codec.wait_inside('"a"') + await codec.wait_inside('"b"') + codec.release('"a"') + codec.release('"b"') + await first + await second + + retry_codec = GatedPayloadCodec() + retried = _gated_producer(backend, retry_codec, "attempt-stable") + retried_tokens = retried.topic("tokens", type=str) + retried_events = retried.topic("tool-events", type=str) + first_again = asyncio.ensure_future(retried_tokens.publish("a", wake=False)) + second_again = asyncio.ensure_future(retried_events.publish("b", wake=False)) + await retry_codec.wait_inside('"a"') + await retry_codec.wait_inside('"b"') + retry_codec.release('"b"') + retry_codec.release('"a"') + await first_again + await second_again + + assert len(backend.all_records(tokens_key)) == 1 + assert len(backend.all_records(events_key)) == 1, ( + "reversing the encode order appended a duplicate on a topic where the " + "swapped key could not collide, so nothing raised" + ) + + +# --- the fence's ordering claim --------------------------------------------- + + +TOKENS = StreamKey("ns", "wf", "run-1", "tokens") +EVENTS = StreamKey("ns", "wf", "run-1", "tool-events") + + +async def _let_every_ready_task_run() -> None: + """Runs every task that can run without waiting for anything. + + What makes "the fence has not appended" an assertion rather than a race. + Nothing here sleeps in wall-clock terms: an unordered ``finish_writing()`` + reaches the backend without awaiting anything that yields, so one turn of + the loop is enough for it to have finished, and the extra turns only make + that margin obvious. + """ + for _ in range(10): + await asyncio.sleep(0) + + +class GatedFailingCodec(GatedPayloadCodec): + """A gated codec whose release makes the encode *fail* rather than finish. + + The payload store rejecting a value, at the point where a concurrent fence + is already waiting behind the publish. + """ + + async def encode(self, payloads): # type: ignore[no-untyped-def] + await super().encode(payloads) + raise RuntimeError("the payload store rejected the value") + + +class CommittingThenFailingBackend(MemoryStreamBackend): + """Commits an append and then loses the answer, on cue. + + The unknown outcome the producer cannot see through: the record is durable + and its offset reached nobody, so it may not be treated as absent -- and a + fence appended in front of it would claim durability for a write whose place + in the stream is not yet fixed (ADR-038). + """ + + def __init__(self) -> None: + super().__init__() + self.committed = asyncio.Event() + self.lose_the_answer = asyncio.Event() + self.losing = True + + async def append(self, key, record): # type: ignore[no-untyped-def] + placed = await super().append(key, record) + if self.losing: + self.losing = False + self.committed.set() + await self.lose_the_answer.wait() + raise ConnectionError("connection reset") + return placed + + +async def test_a_fence_waits_for_a_publish_still_inside_its_codec( + backend: MemoryStreamBackend, +) -> None: + """The fence's claim is about invocation order, so it must hold under one. + + ``publish()`` draws its sequence before awaiting the codec -- which is what + keeps idempotency keys stable across attempts -- so an earlier publish can + still be encoding when ``finish_writing()`` is called. A fence appended there + says every preceding write in this session is in the stream while the first + one has not been sent yet, and a consumer that drains through it may park on + exactly that. The unsafe shape is this test's: a ``wake=False`` publish + batched behind a fence that carries the batch's only wake, where the wake is + spent before the data exists. + """ + codec = GatedPayloadCodec() + tokens = _gated_producer(backend, codec, "batch").topic("tokens", type=str) + + publishing = asyncio.ensure_future(tokens.publish("a", wake=False)) + await codec.wait_inside('"a"') + fencing = asyncio.ensure_future(tokens.finish_writing(wake=False)) + await _let_every_ready_task_run() + + assert not fencing.done(), ( + "the fence completed while the publish was still encoding" + ) + assert backend.all_records(TOKENS) == [], ( + "the fence reached the backend ahead of a publish invoked before it, so " + "its durability claim was false when a consumer could act on it" + ) + + codec.release('"a"') + await publishing + await fencing + + records = backend.all_records(TOKENS) + assert [(r.sequence, r.kind) for r in records] == [ + (0, RecordKind.DATA), + (1, RecordKind.WRITE_FENCE), + ], "backend order and invocation order disagree" + assert backend.strictly_increasing(_placed_offsets(records)) + + +async def test_two_handles_for_one_topic_share_the_fence_order( + backend: MemoryStreamBackend, +) -> None: + """``topic()`` returns a fresh handle per call; the stream is still one. + + So the order cannot live on the handle. A producer that publishes through + one handle and fences through another -- the shape a helper that takes a + topic name rather than a handle produces -- is the same stream and the same + claim. + """ + codec = GatedPayloadCodec() + producer = _gated_producer(backend, codec, "batch") + writing = producer.topic("tokens", type=str) + finishing = producer.topic("tokens", type=str) + assert writing is not finishing + + publishing = asyncio.ensure_future(writing.publish("a", wake=False)) + await codec.wait_inside('"a"') + fencing = asyncio.ensure_future(finishing.finish_writing(wake=False)) + await _let_every_ready_task_run() + + assert not fencing.done() and backend.all_records(TOKENS) == [], ( + "the second handle fenced without waiting for the first handle's publish" + ) + + codec.release('"a"') + await publishing + await fencing + + assert [r.kind for r in backend.all_records(TOKENS)] == [ + RecordKind.DATA, + RecordKind.WRITE_FENCE, + ] + + +async def test_a_fence_does_not_wait_for_another_stream( + backend: MemoryStreamBackend, +) -> None: + """The claim is per stream, so the ordering must be too. + + One connection serves several topics, and a publish blocked in a codec on + one of them says nothing about a fence on another -- stalling it there would + make an unrelated slow payload store hold up a finished topic. + """ + codec = GatedPayloadCodec() + producer = _gated_producer(backend, codec, "batch") + tokens = producer.topic("tokens", type=str) + events = producer.topic("tool-events", type=str) + + publishing = asyncio.ensure_future(tokens.publish("a", wake=False)) + await codec.wait_inside('"a"') + + await asyncio.wait_for(events.finish_writing(wake=False), 2) + + assert [r.kind for r in backend.all_records(EVENTS)] == [RecordKind.WRITE_FENCE] + assert backend.all_records(TOKENS) == [] + + codec.release('"a"') + await publishing + assert [r.kind for r in backend.all_records(TOKENS)] == [RecordKind.DATA] + + +async def test_a_fence_will_not_overtake_an_append_of_unknown_outcome() -> None: + """An append with no answer may still be in front of the fence. + + The record is durable and its offset reached nobody, so the one thing that + cannot be assumed is that it is absent. The fence is refused with *that* + operation's error rather than its own, which is the refusal that already + governs the stream while an append is unsettled: settle it with + ``resolve_append()``, then fence. + """ + backend = CommittingThenFailingBackend() + tokens = _offline_producer(backend, "batch").topic("tokens", type=str) + + publishing = asyncio.ensure_future(tokens.publish("a", wake=False)) + await asyncio.wait_for(backend.committed.wait(), 2) + fencing = asyncio.ensure_future(tokens.finish_writing(wake=False)) + await _let_every_ready_task_run() + + assert not fencing.done() + assert [r.sequence for r in backend.all_records(TOKENS)] == [0], ( + "the fixture must actually have committed, or this asserts nothing" + ) + + backend.lose_the_answer.set() + with pytest.raises(AppendNotAcknowledgedError) as lost: + await publishing + with pytest.raises(AppendNotAcknowledgedError) as refused: + await fencing + + assert refused.value.record.idempotency_key == lost.value.record.idempotency_key, ( + "the fence must report the unsettled operation's recovery, not invent one " + "for a fence that was never appended" + ) + assert [r.kind for r in backend.all_records(TOKENS)] == [RecordKind.DATA], ( + "a fence was appended in front of a record whose position was not settled" + ) + + await tokens.resolve_append(lost.value.record, wake=False) + await tokens.finish_writing(wake=False) + + records = backend.all_records(TOKENS) + assert [(r.sequence, r.kind) for r in records] == [ + (0, RecordKind.DATA), + (2, RecordKind.WRITE_FENCE), + ], "the settled record must be in the stream once, ahead of the fence" + + +async def test_a_fence_refuses_to_stand_in_for_a_failed_publish( + backend: MemoryStreamBackend, +) -> None: + """A failed earlier write is propagated, not passed over. + + The caller asked for that write before it asked for the fence, and a fence + appended over the gap tells a consumer the batch is complete when it is + short a record. Nothing is appended for the refused fence, so once the + failure has been dealt with -- by republishing the value or by accepting the + batch without it -- ``finish_writing()`` again appends one. + """ + codec = GatedFailingCodec() + tokens = _gated_producer(backend, codec, "batch").topic("tokens", type=str) + + publishing = asyncio.ensure_future(tokens.publish("a", wake=False)) + await codec.wait_inside('"a"') + fencing = asyncio.ensure_future(tokens.finish_writing(wake=False)) + await _let_every_ready_task_run() + codec.release('"a"') + + with pytest.raises(RuntimeError): + await publishing + with pytest.raises(PrecedingWriteFailedError) as caught: + await fencing + + assert caught.value.sequence == 0 and caught.value.stream_key == TOKENS + assert isinstance(caught.value.__cause__, RuntimeError), ( + "the fence must carry what actually failed, not describe it" + ) + assert backend.all_records(TOKENS) == [] + + await tokens.finish_writing(wake=False) + assert [(r.sequence, r.kind) for r in backend.all_records(TOKENS)] == [ + (2, RecordKind.WRITE_FENCE) + ], "the failed write is no longer outstanding, so a fence must go in now" + + +async def test_a_cancelled_fence_does_not_refuse_a_later_fence( + backend: MemoryStreamBackend, +) -> None: + """A failed *fence* is not a failed write, so it is not propagated as one. + + Two concurrent fences make independent claims about the publishes each was + invoked after, and neither is inside the other's claim. Holding them in one + append order made a fence that never reached the backend -- cancelled here, + but a refusal does the same -- look to a later fence exactly like a publish + whose record went missing, so a fence with every preceding write durable + behind it was refused with ``PrecedingWriteFailedError``. That error is + documented as reporting a failed ``publish()``, and nothing else. + """ + codec = GatedPayloadCodec() + tokens = _gated_producer(backend, codec, "batch").topic("tokens", type=str) + + publishing = asyncio.ensure_future(tokens.publish("a", wake=False)) + await codec.wait_inside('"a"') + abandoned = asyncio.ensure_future(tokens.finish_writing(wake=False)) + await _let_every_ready_task_run() + fencing = asyncio.ensure_future(tokens.finish_writing(wake=False)) + await _let_every_ready_task_run() + + abandoned.cancel() + with pytest.raises(asyncio.CancelledError): + await abandoned + assert not fencing.done(), ( + "the second fence appended while the publish it covers was still encoding" + ) + + codec.release('"a"') + await publishing + await fencing + + assert [(r.sequence, r.kind) for r in backend.all_records(TOKENS)] == [ + (0, RecordKind.DATA), + (2, RecordKind.WRITE_FENCE), + ], ( + "the surviving fence must append behind the data, and the cancelled one " + "must leave nothing in the stream" + ) + + +class RefusingRecoveryBackend(MemoryStreamBackend): + """Loses the answer to a data append, then refuses its recovery. + + Both halves are contract behaviour rather than a broken store: a backend + commits before it answers, so the answer can be lost, and + ``AppendConflictError`` is the one definite refusal the contract defines -- + the key it names holds different bytes, so the record being recovered did + not land and never will. Together they are the case where "no longer + unresolved" and "durable" are not the same thing, which is reachable from a + session-id collision or a retry that reused the sequence for other bytes. + """ + + def __init__(self) -> None: + super().__init__() + self.reached = asyncio.Event() + self.answer = asyncio.Event() + self.data_appends = 0 + + async def append(self, key, record): # type: ignore[no-untyped-def] + if record.kind is not RecordKind.DATA: + return await super().append(key, record) + self.data_appends += 1 + if self.data_appends == 1: + self.reached.set() + await self.answer.wait() + raise ConnectionError("connection reset") + # Refused without suspending, so nothing runs between the unresolved + # entry disappearing and the waiting fence reading the outcome. That is + # the window in which a stale "unknown" looked like durability. + raise AppendConflictError(record.idempotency_key) + + +async def test_a_fence_is_refused_when_recovery_proves_the_write_absent() -> None: + """A resolved append is not a durable one, so the fence may not infer it. + + An append whose answer was lost is unknown, not failed -- and the fence + waits rather than refusing outright. What ends the wait is + ``resolve_append()``, and it ends it *two* ways: the record is durable, or + the backend refuses the key and the record demonstrably never landed. + Deciding between them by whether the record is still in the producer's + unresolved set cannot work, because both outcomes remove it, and the fence + then appended over a hole while claiming the batch complete. + """ + backend = RefusingRecoveryBackend() + tokens = _offline_producer(backend, "batch").topic("tokens", type=str) + + async def publish_then_recover() -> None: + with pytest.raises(AppendNotAcknowledgedError) as lost: + await tokens.publish("a", wake=False) + with pytest.raises(AppendConflictError): + await tokens.resolve_append(lost.value.record, wake=False) + + publishing = asyncio.ensure_future(publish_then_recover()) + await asyncio.wait_for(backend.reached.wait(), 2) + fencing = asyncio.ensure_future(tokens.finish_writing(wake=False)) + await _let_every_ready_task_run() + + assert not fencing.done(), "the fence overtook an append of unknown outcome" + + backend.answer.set() + await publishing + with pytest.raises(PrecedingWriteFailedError) as refused: + await fencing + + assert refused.value.sequence == 0 and refused.value.stream_key == TOKENS + assert isinstance(refused.value.__cause__, AppendConflictError), ( + "the fence must carry the refusal recovery learned, not the unknown " + "outcome the interrupted call reported" + ) + assert backend.all_records(TOKENS) == [], ( + "a fence was appended claiming a write that recovery proved absent" + ) + + +async def test_a_fence_goes_in_once_recovery_makes_the_write_durable() -> None: + """The other half of the same decision, and the one that must not refuse. + + A recovery that finds the record already there resolves the unknown outcome + to *durable*, and the write the fence was waiting for is now in the stream + ahead of it. The fence has to be released, not refused: reporting + ``PrecedingWriteFailedError`` for a record that landed sends the caller to + republish a value the stream already holds. + """ + backend = CommittingThenFailingBackend() + tokens = _offline_producer(backend, "batch").topic("tokens", type=str) + + async def publish_then_recover() -> None: + with pytest.raises(AppendNotAcknowledgedError) as lost: + await tokens.publish("a", wake=False) + # The record did commit, so re-appending byte-identical content is the + # no-op the recovery depends on -- and it returns without suspending, so + # the fence's next turn is the first one after the resolution. + await tokens.resolve_append(lost.value.record, wake=False) + + publishing = asyncio.ensure_future(publish_then_recover()) + await asyncio.wait_for(backend.committed.wait(), 2) + fencing = asyncio.ensure_future(tokens.finish_writing(wake=False)) + await _let_every_ready_task_run() + + assert not fencing.done() + + backend.lose_the_answer.set() + await publishing + await fencing + + records = backend.all_records(TOKENS) + assert [(r.sequence, r.kind) for r in records] == [ + (0, RecordKind.DATA), + (1, RecordKind.WRITE_FENCE), + ], "the recovered record must be in the stream once, ahead of the fence" + + +class LosesOneAnswerThenRefusesIt(MemoryStreamBackend): + """Loses sequence 0's answer, refuses its recovery, and holds sequence 1. + + Two data writes precede one fence, and only the *second* is still in flight + when the first is recovered -- which is the schedule that decides whether the + fence reads an outcome or a memory of one. + """ + + def __init__(self) -> None: + super().__init__() + self.reached = asyncio.Event() + self.answer = asyncio.Event() + self.release_second = asyncio.Event() + + async def append(self, key, record): # type: ignore[no-untyped-def] + if record.kind is RecordKind.DATA and record.sequence == 0: + if not self.reached.is_set(): + self.reached.set() + await self.answer.wait() + raise ConnectionError("connection reset") + raise AppendConflictError(record.idempotency_key) + if record.kind is RecordKind.DATA and record.sequence == 1: + await self.release_second.wait() + return await super().append(key, record) + + +async def test_a_resolution_reached_while_the_fence_waits_on_a_later_write() -> None: + """An outcome read at wait time can still change; one read after cannot. + + The fence waits on each earlier publish in turn, so it can pass the first + one -- reading "unknown", which is not yet a failure -- and then sit on the + second while ``resolve_append()`` proves the first absent. An outcome + captured on the way past is the stale one, and stale here means permissive: + the fence skips a write recovery has refused and appends over the hole. + """ + backend = LosesOneAnswerThenRefusesIt() + tokens = _offline_producer(backend, "batch").topic("tokens", type=str) + + first = asyncio.ensure_future(tokens.publish("a", wake=False)) + await asyncio.wait_for(backend.reached.wait(), 2) + second = asyncio.ensure_future(tokens.publish("b", wake=False)) + await _let_every_ready_task_run() + fencing = asyncio.ensure_future(tokens.finish_writing(wake=False)) + await _let_every_ready_task_run() + + assert not fencing.done() + + backend.answer.set() + with pytest.raises(AppendNotAcknowledgedError) as lost: + await first + # Lets the fence take its turn on the *first* write and park on the second, + # which is the state this case is about. + await _let_every_ready_task_run() + assert not fencing.done(), "the fence must still be held by the second write" + + with pytest.raises(AppendConflictError): + await tokens.resolve_append(lost.value.record, wake=False) + + backend.release_second.set() + await second + + with pytest.raises(PrecedingWriteFailedError) as refused: + await fencing + + assert refused.value.sequence == 0, ( + "the fence must report the write recovery refused, not the one that landed" + ) + assert isinstance(refused.value.__cause__, AppendConflictError) + assert [(r.sequence, r.kind) for r in backend.all_records(TOKENS)] == [ + (1, RecordKind.DATA) + ], "only the second write may be in the stream, and no fence behind it" + + +async def test_republishing_different_content_under_one_key_is_an_error( + backend: MemoryStreamBackend, +) -> None: + """Silently accepting it would rewrite what a consumer may have delivered.""" + first = _offline_producer(backend, "same-session").topic("tokens", type=str) + await first.publish("a", wake=False) + + retried = _offline_producer(backend, "same-session").topic("tokens", type=str) + with pytest.raises(AppendConflictError): + await retried.publish("changed", wake=False) + + +async def test_sequences_are_per_connection_not_per_topic( + backend: MemoryStreamBackend, +) -> None: + """Two topics' first records must not collide on ``(session, 0)``.""" + producer = _offline_producer(backend) + tokens = producer.topic("tokens", type=str) + events = producer.topic("tool-events", type=str) + + await tokens.publish("a", wake=False) + await events.publish("b", wake=False) + + (token_record,) = backend.all_records(StreamKey("ns", "wf", "run-1", "tokens")) + (event_record,) = backend.all_records(StreamKey("ns", "wf", "run-1", "tool-events")) + assert token_record.sequence != event_record.sequence + + +async def test_a_topic_needs_a_name(backend: MemoryStreamBackend) -> None: + with pytest.raises(ValueError, match="non-empty name"): + _offline_producer(backend).topic("") + + +async def test_the_producer_handle_has_no_subscribe( + backend: MemoryStreamBackend, +) -> None: + """The two sides are mirror images, not one type used twice.""" + tokens = _offline_producer(backend).topic("tokens", type=str) + + assert not hasattr(tokens, "subscribe") + assert hasattr(tokens, "publish") + assert hasattr(tokens, "finish_writing") + + +async def test_a_standalone_read_back_sees_what_was_published( + backend: MemoryStreamBackend, +) -> None: + """The P6 criterion, without a Temporal server in the picture.""" + tokens = _offline_producer(backend).topic("tokens", type=str) + for value in ["a", "b", "c"]: + await tokens.publish(value, wake=False) + await tokens.finish_writing(wake=False) + + key = StreamKey("ns", "wf", "run-1", "tokens") + read_back = await backend.read_after(key, BEGINNING, max_records=100, block=None) + + import temporalio.converter + from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec + + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + values = [await codec.decode(r.payload) for r in read_back if not r.is_control] + assert values == ["a", "b", "c"] diff --git a/tests/contrib/external_workflow_streams/test_record.py b/tests/contrib/external_workflow_streams/test_record.py new file mode 100644 index 0000000000..680cd4632b --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_record.py @@ -0,0 +1,197 @@ +"""P1 — record model, offsets, cursor boundaries, control records.""" + +from __future__ import annotations + +import pytest + +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Cursor, + IdempotencyKey, + Offset, + RecordKind, + StreamRecord, +) + + +def redis_like_compare(left: Offset, right: Offset) -> int: + """The ``(ms, seq)`` numeric ordering, used here to stand in for a provider.""" + + def parts(offset: Offset) -> tuple[int, int]: + ms, _, seq = offset.token.partition("-") + return int(ms), int(seq or 0) + + return (parts(left) > parts(right)) - (parts(left) < parts(right)) + + +# --- offsets ---------------------------------------------------------------- + + +def test_offset_round_trips() -> None: + offset = Offset("1700000000000-3") + assert Offset.deserialize(offset.serialize()) == offset + + +def test_offset_is_hashable_and_equatable() -> None: + assert Offset("1-0") == Offset("1-0") + assert len({Offset("1-0"), Offset("1-0"), Offset("1-1")}) == 2 + + +def test_offset_rejects_an_empty_token() -> None: + with pytest.raises(ValueError, match="may not be empty"): + Offset("") + + +def test_offset_does_not_expose_lexical_ordering() -> None: + """The failure this prevents is invisible until a millisecond width change. + + ``"9-0" < "10-0"`` is false lexically and true numerically, so an offset + that supported ``<`` would silently order wrongly for one provider in ten. + """ + with pytest.raises(TypeError): + _ = Offset("9-0") < Offset("10-0") # type: ignore[operator] + + assert redis_like_compare(Offset("9-0"), Offset("10-0")) < 0 + + +# --- cursors ---------------------------------------------------------------- + + +def test_cursor_forms_round_trip() -> None: + for cursor in (BEGINNING, AFTER(Offset("1700000000000-0"))): + assert Cursor.deserialize(cursor.serialize()) == cursor + + +def test_beginning_is_distinguishable_from_a_provider_beginning_sentinel() -> None: + """Redis spells its beginning ``0-0``; that is an offset, not ``BEGINNING``.""" + assert BEGINNING != AFTER(Offset("0-0")) + assert BEGINNING.serialize() != AFTER(Offset("0-0")).serialize() + assert BEGINNING.is_beginning + assert not AFTER(Offset("0-0")).is_beginning + + +def test_cursor_deserialize_rejects_junk() -> None: + with pytest.raises(ValueError, match="not a serialized cursor"): + Cursor.deserialize("1700000000000-0") + + +def test_after_excludes_up_to_and_including_its_offset() -> None: + cursor = AFTER(Offset("100-0")) + + assert cursor.excludes(Offset("99-9"), redis_like_compare) + assert cursor.excludes(Offset("100-0"), redis_like_compare) + assert not cursor.excludes(Offset("100-1"), redis_like_compare) + + +def test_beginning_excludes_nothing() -> None: + assert not BEGINNING.excludes(Offset("0-1"), redis_like_compare) + + +def test_after_names_a_boundary_past_the_tail() -> None: + """A consumer parked at the tail names the last record, never the next one. + + Nothing here requires the record after ``100-0`` to exist -- which is the + whole point of a boundary cursor. + """ + assert str(AFTER(Offset("100-0"))) == "AFTER(100-0)" + + +# --- records ---------------------------------------------------------------- + + +def test_record_round_trips_through_fields() -> None: + record = StreamRecord( + kind=RecordKind.DATA, + payload=b"\x00\xffhello", + producer_session_id="session-a", + sequence=7, + ).placed_at(Offset("100-0")) + + assert StreamRecord.from_fields(Offset("100-0"), record.to_fields()) == record + + +def test_record_round_trips_from_bytes_keyed_fields() -> None: + """A provider that did not decode its responses must still round-trip.""" + record = StreamRecord(RecordKind.DATA, b"x", "session-a", 0).placed_at( + Offset("100-0") + ) + raw = {key.encode(): value for key, value in record.to_fields().items()} + + assert StreamRecord.from_fields(Offset("100-0"), raw) == record + + +def test_fence_round_trips_and_carries_no_payload() -> None: + fence = StreamRecord(RecordKind.WRITE_FENCE, b"", "session-a", 8).placed_at( + Offset("101-0") + ) + + assert fence.is_control + assert StreamRecord.from_fields(Offset("101-0"), fence.to_fields()) == fence + + +def test_a_fence_with_a_payload_is_rejected() -> None: + with pytest.raises(ValueError, match="carries no payload"): + StreamRecord(RecordKind.WRITE_FENCE, b"nope", "session-a", 0) + + +def test_from_fields_names_what_is_missing() -> None: + with pytest.raises(ValueError, match="missing field"): + StreamRecord.from_fields(Offset("1-0"), {"__tes_kind": b"1"}) + + +def test_data_and_control_share_one_offset_sequence() -> None: + """The fence sits between data records and advances the cursor like them.""" + placed = [ + StreamRecord(RecordKind.DATA, b"a", "s", 0).placed_at(Offset("100-0")), + StreamRecord(RecordKind.DATA, b"b", "s", 1).placed_at(Offset("100-1")), + StreamRecord(RecordKind.WRITE_FENCE, b"", "s", 2).placed_at(Offset("100-2")), + StreamRecord(RecordKind.DATA, b"c", "other", 0).placed_at(Offset("100-3")), + ] + + offsets = [r.offset for r in placed] + assert all( + redis_like_compare(a, b) < 0 # type: ignore[arg-type] + for a, b in zip(offsets, offsets[1:]) + ) + assert [r.is_control for r in placed] == [False, False, True, False] + + +def test_record_offset_is_absent_until_the_provider_places_it() -> None: + unplaced = StreamRecord(RecordKind.DATA, b"a", "s", 0) + + assert unplaced.offset is None + assert unplaced.placed_at(Offset("1-0")).offset == Offset("1-0") + + +def test_to_fields_excludes_the_offset() -> None: + """Two appends under one idempotency key must be byte-comparable. + + Including the offset would make every re-append look like different bytes, + turning the idempotent retry path into a spurious conflict. + """ + record = StreamRecord(RecordKind.DATA, b"a", "s", 0) + + assert record.to_fields() == record.placed_at(Offset("1-0")).to_fields() + + +# --- idempotency keys ------------------------------------------------------- + + +def test_idempotency_key_comes_from_session_and_sequence() -> None: + record = StreamRecord(RecordKind.DATA, b"a", "session-a", 4) + + assert record.idempotency_key == IdempotencyKey("session-a", 4) + + +def test_idempotency_key_rejects_a_missing_session() -> None: + with pytest.raises(ValueError, match="session id may not be empty"): + IdempotencyKey("", 0) + + +@pytest.mark.parametrize("bad", [-1, -100]) +def test_negative_sequences_are_rejected(bad: int) -> None: + with pytest.raises(ValueError, match="non-negative"): + IdempotencyKey("s", bad) + with pytest.raises(ValueError, match="non-negative"): + StreamRecord(RecordKind.DATA, b"a", "s", bad) diff --git a/tests/contrib/external_workflow_streams/test_redis_backend.py b/tests/contrib/external_workflow_streams/test_redis_backend.py new file mode 100644 index 0000000000..e1e0bcb792 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_redis_backend.py @@ -0,0 +1,480 @@ +"""P3/P3b — the Redis provider, against the conformance suite unmodified. + +"Unmodified" is the point: if a check had to be relaxed for Redis, it would be +encoding this provider's behaviour rather than the contract, and the next +provider would inherit the exemption. +""" + +from __future__ import annotations + +import uuid +from collections.abc import AsyncGenerator +from datetime import timedelta + +import pytest +import pytest_asyncio + +from temporalio.contrib.external_workflow_streams._backend import ParkIntent, StreamKey +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Offset, + RecordKind, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._redis import RedisStreamBackend +from tests.contrib.external_workflow_streams.conformance import ( + CONFORMANCE_CHECKS, + PARKING_CONFORMANCE_CHECKS, + Check, +) +from tests.contrib.external_workflow_streams.conftest import ( + KEY_NAMESPACE, + redis_url, +) + + +@pytest_asyncio.fixture # type: ignore[reportUntypedFunctionDecorator] +async def redis_backend( + redis_worker_id: str, +) -> AsyncGenerator[RedisStreamBackend, None]: + """A backend whose keys all sit under this test's own prefix. + + The prefix is the isolation: one Redis serves the whole suite, and two + xdist workers running this module must not see each other's streams. + """ + pytest.importorskip("redis.asyncio", reason="redis is not installed") + backend = RedisStreamBackend( + url=redis_url(), + key_prefix=f"{KEY_NAMESPACE}:{redis_worker_id}:{uuid.uuid4().hex}", + ) + try: + await backend._client.ping() + except Exception as err: + await backend.aclose() + pytest.skip(f"Redis is not reachable at {redis_url()}: {err}") + + try: + yield backend + finally: + keys = [ + k async for k in backend._client.scan_iter(match=f"{backend._key_prefix}*") + ] + if keys: + await backend._client.delete(*keys) + await backend.aclose() + + +@pytest.fixture +def stream_key() -> StreamKey: + return StreamKey("ns", "wf", uuid.uuid4().hex, "tokens") + + +# --- the conformance suites, unmodified ------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("check", CONFORMANCE_CHECKS, ids=lambda c: c.__name__) +async def test_redis_passes_the_core_conformance_suite( + check: Check, redis_backend: RedisStreamBackend, stream_key: StreamKey +) -> None: + await check(redis_backend, stream_key) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("check", PARKING_CONFORMANCE_CHECKS, ids=lambda c: c.__name__) +async def test_redis_passes_the_parking_conformance_suite( + check: Check, redis_backend: RedisStreamBackend, stream_key: StreamKey +) -> None: + await check(redis_backend, stream_key) + + +# --- Redis specifics the suite cannot state --------------------------------- + + +@pytest.mark.asyncio +async def test_offsets_are_real_redis_ids( + redis_backend: RedisStreamBackend, stream_key: StreamKey +) -> None: + placed = await redis_backend.append( + stream_key, StreamRecord(RecordKind.DATA, b"a", "s", 0) + ) + + assert placed.offset is not None + ms, _, seq = placed.offset.token.partition("-") + assert ms.isdigit() and seq.isdigit(), ( + f"a Redis offset is -, got {placed.offset.token!r}" + ) + + +@pytest.mark.asyncio +async def test_xrange_and_xread_disagree_exactly_as_the_contract_says( + redis_backend: RedisStreamBackend, stream_key: StreamKey +) -> None: + """The reason replay must not use XREAD, asserted against a real Redis. + + XRANGE includes the id it is given; XREAD starts strictly after it. A + provider that reached for XREAD to serve the replay read would drop the + first record of every recorded range -- and nothing before the first replay + would notice. + """ + first = await redis_backend.append( + stream_key, StreamRecord(RecordKind.DATA, b"a", "s", 0) + ) + second = await redis_backend.append( + stream_key, StreamRecord(RecordKind.DATA, b"b", "s", 1) + ) + assert first.offset is not None and second.offset is not None + + inclusive = await redis_backend.read_range(stream_key, first.offset, second.offset) + exclusive = await redis_backend.read_after( + stream_key, AFTER(first.offset), max_records=10, block=None + ) + + assert [r.payload for r in inclusive] == [b"a", b"b"] + assert [r.payload for r in exclusive] == [b"b"] + + +@pytest.mark.asyncio +async def test_beginning_reads_from_the_zero_sentinel( + redis_backend: RedisStreamBackend, stream_key: StreamKey +) -> None: + """`BEGINNING` is a boundary; `0-0` is how Redis spells that boundary. + + They are not the same thing, which is why the cursor type keeps them apart + -- but the provider has to know the translation. + """ + await redis_backend.append(stream_key, StreamRecord(RecordKind.DATA, b"a", "s", 0)) + + from_beginning = await redis_backend.read_after( + stream_key, BEGINNING, max_records=10, block=None + ) + assert [r.payload for r in from_beginning] == [b"a"] + + +@pytest.mark.asyncio +async def test_binary_payloads_survive_the_round_trip( + redis_backend: RedisStreamBackend, stream_key: StreamKey +) -> None: + """A decoding client would corrupt these, so the provider must not use one.""" + payload = bytes(range(256)) + placed = await redis_backend.append( + stream_key, StreamRecord(RecordKind.DATA, payload, "s", 0) + ) + + assert placed.offset is not None + (got,) = await redis_backend.read_range(stream_key, placed.offset, placed.offset) + assert got.payload == payload + + +@pytest.mark.asyncio +async def test_an_idempotent_reappend_writes_nothing_to_the_stream( + redis_backend: RedisStreamBackend, stream_key: StreamKey +) -> None: + """Not merely "returns the same offset" -- no second entry may exist. + + Asserted at the Redis level because a provider could return the original + offset while still having XADDed a duplicate, which replay would then find + inside a recorded range and count as an integrity failure. + """ + record = StreamRecord(RecordKind.DATA, b"payload", "retry", 0) + await redis_backend.append(stream_key, record) + await redis_backend.append(stream_key, record) + + length = await redis_backend._client.xlen(redis_backend.stream_key(stream_key)) + assert length == 1 + + +@pytest.mark.asyncio +async def test_a_blocking_read_returns_empty_rather_than_hanging( + redis_backend: RedisStreamBackend, stream_key: StreamKey +) -> None: + """`XREAD BLOCK 0` blocks forever, so a zero timeout must omit BLOCK.""" + got = await redis_backend.read_after( + stream_key, BEGINNING, max_records=10, block=timedelta(0) + ) + assert got == [] + + +@pytest.mark.asyncio +async def test_offsets_compare_numerically_not_lexically( + redis_backend: RedisStreamBackend, +) -> None: + """Real ids cross the width boundary; the comparator must not care.""" + assert redis_backend.compare_offsets(Offset("9-0"), Offset("10-0")) < 0 + assert ( + redis_backend.compare_offsets( + Offset("1700000000000-1"), Offset("1700000000000-2") + ) + < 0 + ) + assert redis_backend.compare_offsets(Offset("100-0"), Offset("100-0")) == 0 + + +# --- retention loss against real Redis --------------------------------------- + + +@pytest.mark.asyncio +async def test_a_trimmed_range_is_integrity_loss_not_a_storage_failure( + redis_backend: RedisStreamBackend, stream_key: StreamKey +) -> None: + """`XTRIM MAXLEN` is how a recorded range actually disappears in production. + + Not a hypothetical: a stream with a retention policy trims its own head + while a Workflow is still parked against it. That must read as integrity + loss -- an operator has to restore or terminate -- and never as a transient + storage failure, which would clear on its own and so would be retried + forever against data that is not coming back. + """ + from temporalio.contrib.external_workflow_streams._annotation import Run + from temporalio.contrib.external_workflow_streams._errors import ( + StreamIntegrityError, + ) + from temporalio.contrib.external_workflow_streams._replay import validate_run + + placed = [] + for i in range(5): + placed.append( + await redis_backend.append( + stream_key, StreamRecord(RecordKind.DATA, b"x", "producer", i) + ) + ) + run = Run(1, placed[0].offset, placed[-1].offset, 5) # type: ignore[arg-type] + + # The head of the recorded range is trimmed away, exactly as a retention + # policy would do it. + await redis_backend._client.xtrim( + redis_backend.stream_key(stream_key), maxlen=2, approximate=False + ) + survivors = await redis_backend.read_range( + stream_key, + placed[0].offset, # type: ignore[arg-type] + placed[-1].offset, # type: ignore[arg-type] + ) + + with pytest.raises(StreamIntegrityError) as caught: + validate_run(run, survivors, redis_backend) + + assert "missing from the stream" in str(caught.value) + + +@pytest.mark.asyncio +async def test_a_deleted_write_fence_is_integrity_loss( + redis_backend: RedisStreamBackend, stream_key: StreamKey +) -> None: + """A control record occupies an offset inside a run, so losing one is loss. + + It is never yielded to Workflow code, which is exactly why it could be + dismissed as harmless -- but the run's count includes it and its position is + recorded, so a range missing it no longer reads back as written. Letting it + pass would also shift every later control position by one. + """ + from temporalio.contrib.external_workflow_streams._annotation import Run + from temporalio.contrib.external_workflow_streams._errors import ( + StreamIntegrityError, + ) + from temporalio.contrib.external_workflow_streams._replay import validate_run + + first = await redis_backend.append( + stream_key, StreamRecord(RecordKind.DATA, b"a", "producer", 0) + ) + fence = await redis_backend.append( + stream_key, StreamRecord(RecordKind.WRITE_FENCE, b"", "producer", 1) + ) + last = await redis_backend.append( + stream_key, StreamRecord(RecordKind.DATA, b"b", "producer", 2) + ) + run = Run(1, first.offset, last.offset, 3, control_positions=(1,)) # type: ignore[arg-type] + + await redis_backend.delete_for_test(stream_key, fence.offset) # type: ignore[arg-type] + survivors = await redis_backend.read_range( + stream_key, + first.offset, # type: ignore[arg-type] + last.offset, # type: ignore[arg-type] + ) + + with pytest.raises(StreamIntegrityError) as caught: + validate_run(run, survivors, redis_backend) + + assert "contains 2 record(s)" in str(caught.value), ( + f"a deleted fence must fail the count check, got: {caught.value}" + ) + + +# --- key injectivity --------------------------------------------------------- + + +def _colliding_identities() -> tuple[StreamKey, StreamKey]: + """Two distinct identities a delimiter-joined key layout renders identically. + + Nothing here is exotic: a Workflow ID and a stream name are user-chosen + strings, and `:` is an ordinary character in both. + """ + first_run, second_run = uuid.uuid4().hex, uuid.uuid4().hex + return ( + StreamKey("ns", "wf", first_run, f"{second_run}:tokens"), + StreamKey("ns", f"wf:{first_run}", second_run, "tokens"), + ) + + +@pytest.mark.asyncio +async def test_distinct_identities_never_share_one_physical_key( + redis_backend: RedisStreamBackend, +) -> None: + """Every derivation, not only the stream: they all come from `stream_key`. + + Two `StreamKey`s that are not equal are two streams, and nothing they own -- + records, idempotency hashes, park intents, claims -- may land in one place. + Two Workflows sharing one data structure is a cross-Workflow leak, and it + reads as ordinary corruption rather than as a key-layout bug. + """ + first, second = _colliding_identities() + assert first != second + + for name, derive in ( + ("stream", redis_backend.stream_key), + ("idempotency", redis_backend._idempotency_key), + ("park intent", lambda k: redis_backend._intent_key(k, 1)), + ("claim", lambda k: redis_backend._claim_key(k, 1)), + ): + assert derive(first) != derive(second), ( + f"the {name} key is not injective: {first} and {second} both " + f"render as {derive(first)!r}" + ) + + +@pytest.mark.asyncio +async def test_colliding_identities_hold_isolated_records( + redis_backend: RedisStreamBackend, +) -> None: + """The consequence, asserted through the public operations rather than keys.""" + first, second = _colliding_identities() + + await redis_backend.append( + first, StreamRecord(RecordKind.DATA, b"first", "producer-first", 0) + ) + await redis_backend.append( + second, StreamRecord(RecordKind.DATA, b"second", "producer-second", 0) + ) + + from_first = await redis_backend.read_after( + first, BEGINNING, max_records=10, block=None + ) + from_second = await redis_backend.read_after( + second, BEGINNING, max_records=10, block=None + ) + + assert [r.payload for r in from_first] == [b"first"] + assert [r.payload for r in from_second] == [b"second"] + + +@pytest.mark.asyncio +async def test_colliding_identities_hold_isolated_park_state( + redis_backend: RedisStreamBackend, +) -> None: + """An intent and a claim installed on one must be invisible on the other. + + A shared claim key is worse than a shared stream: the second Workflow's + producer reads a claim it never made, concludes the wake is someone else's, + and stays silent -- so the Run is never woken at all. + """ + first, second = _colliding_identities() + + await redis_backend.install_park_intent( + first, ParkIntent(wait_id=1, cursor=BEGINNING, park_generation=3, run_id="r") + ) + + assert await redis_backend.park_intent(second, 1) is None + assert await redis_backend.current_park_generation(second, 1) is None + assert await redis_backend.parked_wait_ids(second) == [] + assert await redis_backend.parked_wait_ids(first) == [1] + + lease = timedelta(seconds=30) + assert await redis_backend.claim_park_generation( + first, 1, 3, claimant="producer-first", lease=lease + ) + assert await redis_backend.claim_park_generation( + second, 1, 3, claimant="producer-second", lease=lease + ), "a claim on one identity's generation must say nothing about another's" + + +@pytest.mark.parametrize( + ("metacharacter", "impostor_name"), + [ + # Each impostor name is a *different* stream whose rendering, read as a + # `SCAN MATCH` glob, matches the victim's park keys. The lengths are + # chosen so the suffix a naive implementation slices off still parses as + # a wait id -- otherwise the leak hides behind an incidental filter. + ("*", "token*"), + ("?", "token?"), + ("[", "toke[n]s"), + ], + ids=["star", "question", "bracket"], +) +@pytest.mark.asyncio +async def test_glob_metacharacters_cannot_enumerate_another_streams_intents( + redis_backend: RedisStreamBackend, + metacharacter: str, + impostor_name: str, +) -> None: + """`parked_wait_ids` must answer about one stream, whatever the name spells. + + The key goes into a `SCAN MATCH` pattern, so an unescaped `*`, `?` or `[` in + a user-chosen stream name turns the enumeration into a wildcard over the + keyspace -- and the producer then wakes, claims and rechecks waits that + belong to a different Workflow. + """ + assert metacharacter in impostor_name + run_id = uuid.uuid4().hex + victim = StreamKey("ns", "wf", run_id, "tokens") + impostor = StreamKey("ns", "wf", run_id, impostor_name) + + for wait_id in (101, 202): + await redis_backend.install_park_intent( + victim, + ParkIntent( + wait_id=wait_id, cursor=BEGINNING, park_generation=1, run_id="r" + ), + ) + await redis_backend.install_park_intent( + impostor, ParkIntent(wait_id=7, cursor=BEGINNING, park_generation=1, run_id="r") + ) + + assert await redis_backend.parked_wait_ids(impostor) == [7] + assert await redis_backend.parked_wait_ids(victim) == [101, 202] + + +@pytest.mark.asyncio +async def test_a_glob_metacharacter_in_the_key_prefix_does_not_widen_the_scan( + redis_backend: RedisStreamBackend, +) -> None: + """The prefix is the operator's, not the identity's, so encoding cannot fix it. + + `key_prefix` separates deployments sharing one Redis and is passed through + verbatim -- which is what makes it readable and what makes it the one place + a `SCAN MATCH` pattern can still be widened. The two prefixes here are the + same length so a leaked key's suffix still parses as a wait id; a shorter + impostor would be hidden by that filter rather than by the escaping. + """ + base = redis_backend._key_prefix # cleaned up by the fixture's scan + victim_backend = RedisStreamBackend( + client=redis_backend._client, key_prefix=f"{base}:px" + ) + impostor_backend = RedisStreamBackend( + client=redis_backend._client, key_prefix=f"{base}:p*" + ) + key = StreamKey("ns", "wf", uuid.uuid4().hex, "tokens") + + for wait_id in (101, 202): + await victim_backend.install_park_intent( + key, + ParkIntent( + wait_id=wait_id, cursor=BEGINNING, park_generation=1, run_id="r" + ), + ) + await impostor_backend.install_park_intent( + key, ParkIntent(wait_id=7, cursor=BEGINNING, park_generation=1, run_id="r") + ) + + assert await impostor_backend.parked_wait_ids(key) == [7] + assert await victim_backend.parked_wait_ids(key) == [101, 202] diff --git a/tests/contrib/external_workflow_streams/test_redis_fixture.py b/tests/contrib/external_workflow_streams/test_redis_fixture.py new file mode 100644 index 0000000000..afe0c0299e --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_redis_fixture.py @@ -0,0 +1,62 @@ +"""X2 — the Redis fixture itself, exercised through XADD/XREAD.""" + +from __future__ import annotations + +import pytest + +from tests.contrib.external_workflow_streams.conftest import ( + KEY_NAMESPACE, + RedisKeyspace, +) + + +@pytest.mark.asyncio +async def test_xadd_and_xread_through_the_fixture( + redis_keyspace: RedisKeyspace, +) -> None: + stream = redis_keyspace.key("tokens") + + first = await redis_keyspace.client.xadd(stream, {"payload": "a"}) # type: ignore[attr-defined] + second = await redis_keyspace.client.xadd(stream, {"payload": "b"}) # type: ignore[attr-defined] + + # XRANGE is inclusive of both endpoints -- this is the read replay uses. + entries = await redis_keyspace.client.xrange(stream, first, second) # type: ignore[attr-defined] + assert [fields["payload"] for _, fields in entries] == ["a", "b"] + + # XREAD is exclusive of the supplied id -- this is the read watching uses. + after_first = await redis_keyspace.client.xread({stream: first}) # type: ignore[attr-defined] + assert [fields["payload"] for _, fields in after_first[0][1]] == ["b"] + + +@pytest.mark.asyncio +async def test_keys_are_namespaced_to_this_test(redis_keyspace: RedisKeyspace) -> None: + assert redis_keyspace.key("tokens").startswith(f"{KEY_NAMESPACE}:") + assert redis_keyspace.prefix.endswith(":") + + +@pytest.mark.asyncio +async def test_two_keyspaces_do_not_collide( + redis_keyspace: RedisKeyspace, redis_client: object +) -> None: + """A second keyspace over the same server sees none of the first's keys. + + This is the property xdist parallelism depends on; asserting it here means + a regression shows up as a fixture failure rather than as a flaky suite. + """ + other = RedisKeyspace(client=redis_client, prefix=f"{KEY_NAMESPACE}:other:x:") + try: + await redis_keyspace.client.xadd(redis_keyspace.key("s"), {"v": "1"}) # type: ignore[attr-defined] + assert await redis_client.exists(other.key("s")) == 0 # type: ignore[attr-defined] + finally: + await other.cleanup() + + +@pytest.mark.asyncio +async def test_cleanup_removes_every_key( + redis_keyspace: RedisKeyspace, redis_client: object +) -> None: + stream = redis_keyspace.key("tokens") + await redis_keyspace.client.xadd(stream, {"payload": "a"}) # type: ignore[attr-defined] + + assert await redis_keyspace.cleanup() == 1 + assert await redis_client.exists(stream) == 0 # type: ignore[attr-defined] diff --git a/tests/contrib/external_workflow_streams/test_redis_output.py b/tests/contrib/external_workflow_streams/test_redis_output.py new file mode 100644 index 0000000000..9b0036a2f3 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_redis_output.py @@ -0,0 +1,284 @@ +"""Redis output staging, terminal transitions, and read-barrier behavior.""" + +from __future__ import annotations + +import dataclasses +import uuid +from collections.abc import AsyncGenerator + +import pytest +import pytest_asyncio + +from temporalio.contrib.external_workflow_streams._backend import ( + StreamDirection, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._errors import StreamIntegrityError +from temporalio.contrib.external_workflow_streams._output_backend import ( + OutputStageConflictError, + OutputStageManifest, + OutputStageResolutionError, + OutputStageStatus, + StagedOutputRecord, +) +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + RecordKind, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._redis import RedisStreamBackend +from tests.contrib.external_workflow_streams.conftest import ( + KEY_NAMESPACE, + redis_url, +) + + +@pytest_asyncio.fixture # type: ignore[reportUntypedFunctionDecorator] +async def output_backend( + redis_worker_id: str, +) -> AsyncGenerator[RedisStreamBackend, None]: + pytest.importorskip("redis.asyncio", reason="redis is not installed") + backend = RedisStreamBackend( + url=redis_url(), + key_prefix=f"{KEY_NAMESPACE}:{redis_worker_id}:{uuid.uuid4().hex}", + ) + try: + await backend._client.ping() + except Exception as err: + await backend.aclose() + pytest.skip(f"Redis is not reachable at {redis_url()}: {err}") + + try: + yield backend + finally: + keys = [ + key + async for key in backend._client.scan_iter(match=f"{backend._key_prefix}*") + ] + if keys: + await backend._client.delete(*keys) + await backend.aclose() + + +@pytest.fixture +def output_key() -> StreamKey: + return StreamKey( + "ns", + "wf", + uuid.uuid4().hex, + "events", + StreamDirection.OUTPUT, + ) + + +def manifest( + key: StreamKey, + *, + token: str = "stage-1", + sub_batch_id: int = 0, + fingerprint: bytes = b"f" * 32, + record_count: int = 2, +) -> OutputStageManifest: + return OutputStageManifest( + stream_key=key, + provider_id="redis-streams", + provider_format_version=1, + stage_token=token, + run_id="run-1", + history_floor_event_id=10, + sub_batch_id=sub_batch_id, + fingerprint_version=1, + fingerprint=fingerprint, + record_count=record_count, + logical_byte_count=20, + ) + + +def staged(*payloads: bytes) -> tuple[StagedOutputRecord, ...]: + return tuple( + StagedOutputRecord(index, RecordKind.DATA, payload) + for index, payload in enumerate(payloads) + ) + + +def direct(payload: bytes, *, sequence: int = 0) -> StreamRecord: + return StreamRecord( + RecordKind.DATA, + payload, + producer_session_id="activity", + sequence=sequence, + ) + + +@pytest.mark.parametrize( + ("terminal", "status"), + [ + ("commit", OutputStageStatus.COMMITTED), + ("abort", OutputStageStatus.ABORTED), + ], +) +async def test_exact_stage_retry_returns_actual_status_after_terminal( + output_backend: RedisStreamBackend, + output_key: StreamKey, + terminal: str, + status: OutputStageStatus, +) -> None: + batch = manifest(output_key) + records = staged(b"one", b"two") + + first = await output_backend.stage_output(batch, records) + repeated_pending = await output_backend.stage_output(batch, records) + resolved = await getattr(output_backend, f"{terminal}_output")(batch) + repeated_resolution = await getattr(output_backend, f"{terminal}_output")(batch) + repeated_terminal = await output_backend.stage_output( + batch, + # Encoded bytes may vary under a randomized PayloadCodec. The immutable + # logical manifest makes this the same operation, so Redis must retain + # and return the first successfully staged bytes. + staged(b"random-one", b"random-two"), + ) + + assert first.status is OutputStageStatus.PENDING + assert repeated_pending == first + assert resolved.status is status + assert repeated_resolution == resolved + assert repeated_terminal.status is status + assert repeated_terminal.records == first.records + assert [record.payload for record in repeated_terminal.records] == [ + b"one", + b"two", + ] + assert await output_backend.output_stage(batch) == repeated_terminal + + +async def test_reusing_stage_identity_with_another_manifest_conflicts_atomically( + output_backend: RedisStreamBackend, + output_key: StreamKey, +) -> None: + original = manifest(output_key) + first = await output_backend.stage_output(original, staged(b"one", b"two")) + conflicting = dataclasses.replace(original, fingerprint=b"x" * 32) + + with pytest.raises(OutputStageConflictError): + await output_backend.stage_output( + conflicting, staged(b"other-one", b"other-two") + ) + + assert await output_backend.output_stage(original) == first + entries = await output_backend._client.xrange(output_backend.stream_key(output_key)) + assert entries is not None + assert len(entries) == 2 + + +async def test_pending_stage_blocks_later_direct_output_until_commit( + output_backend: RedisStreamBackend, + output_key: StreamKey, +) -> None: + batch = manifest(output_key, record_count=1) + pending = await output_backend.stage_output(batch, staged(b"workflow")) + later = await output_backend.append_output(output_key, direct(b"activity")) + + blocked = await output_backend.read_output_after( + output_key, BEGINNING, max_records=10, block=None + ) + assert blocked.records == () + assert blocked.pending is not None + assert blocked.pending.manifest == batch + assert blocked.pending.offset == pending.records[0].offset + assert await output_backend.output_tail(output_key) == BEGINNING + + await output_backend.commit_output(batch) + visible = await output_backend.read_output_after( + output_key, BEGINNING, max_records=10, block=None + ) + assert [record.payload for record in visible.records] == [ + b"workflow", + b"activity", + ] + assert visible.records[-1].offset == later.offset + assert visible.pending is None + + +async def test_abort_skips_stage_and_cannot_be_reversed( + output_backend: RedisStreamBackend, + output_key: StreamKey, +) -> None: + batch = manifest(output_key, record_count=1) + staged_batch = await output_backend.stage_output(batch, staged(b"discarded")) + later = await output_backend.append_output(output_key, direct(b"visible")) + + aborted = await output_backend.abort_output(batch) + repeated = await output_backend.abort_output(batch) + with pytest.raises(OutputStageResolutionError): + await output_backend.commit_output(batch) + + visible = await output_backend.read_output_after( + output_key, BEGINNING, max_records=10, block=None + ) + assert aborted.status is OutputStageStatus.ABORTED + assert repeated == aborted + assert aborted.records == staged_batch.records + assert [record.payload for record in visible.records] == [b"visible"] + assert await output_backend.output_tail(output_key) == AFTER(later.offset) + + +async def test_missing_stage_metadata_is_an_integrity_failure( + output_backend: RedisStreamBackend, + output_key: StreamKey, +) -> None: + batch = manifest(output_key, record_count=1) + await output_backend.stage_output(batch, staged(b"workflow")) + stage_id = f"{len(batch.stage_token)}:{batch.stage_token}:{batch.sub_batch_id}" + await output_backend._client.hdel( + output_backend._output_status_key(output_key), stage_id + ) + + with pytest.raises(StreamIntegrityError, match="no stage status"): + await output_backend.read_output_after( + output_key, BEGINNING, max_records=10, block=None + ) + with pytest.raises(StreamIntegrityError, match="missing its status"): + await output_backend.output_stage(batch) + + +async def test_missing_staged_record_is_an_integrity_failure( + output_backend: RedisStreamBackend, + output_key: StreamKey, +) -> None: + batch = manifest(output_key, record_count=1) + stage = await output_backend.stage_output(batch, staged(b"workflow")) + await output_backend._client.xdel( + output_backend.stream_key(output_key), stage.records[0].offset.serialize() + ) + + with pytest.raises(StreamIntegrityError, match="record.*missing"): + await output_backend.output_stage(batch) + + +async def test_input_and_output_with_same_user_identity_are_physically_isolated( + output_backend: RedisStreamBackend, + output_key: StreamKey, +) -> None: + input_key = dataclasses.replace(output_key, direction=StreamDirection.INPUT) + input_record = direct(b"input") + output_record = direct(b"output") + + placed_input = await output_backend.append(input_key, input_record) + placed_output = await output_backend.append_output(output_key, output_record) + + input_records = await output_backend.read_after( + input_key, BEGINNING, max_records=10, block=None + ) + output_records = await output_backend.read_output_after( + output_key, BEGINNING, max_records=10, block=None + ) + assert output_backend.stream_key(input_key) != output_backend.stream_key(output_key) + assert [record.payload for record in input_records] == [b"input"] + assert placed_input.offset == input_records[0].offset + assert [record.payload for record in output_records.records] == [b"output"] + assert placed_output.offset == output_records.records[0].offset + with pytest.raises(ValueError, match="OUTPUT"): + await output_backend.append_output( + input_key, direct(b"wrong-direction", sequence=1) + ) diff --git a/tests/contrib/external_workflow_streams/test_registry.py b/tests/contrib/external_workflow_streams/test_registry.py new file mode 100644 index 0000000000..c368be89cd --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_registry.py @@ -0,0 +1,227 @@ +"""P17 — validation of the Worker's stream backend.""" + +from __future__ import annotations + +import uuid + +import pytest + +from temporalio import workflow +from temporalio.client import Client +from temporalio.contrib import external_workflow_streams +from temporalio.contrib.external_workflow_streams import ( + RedisStreamBackend, + StreamBackend, + external_stream, +) +from temporalio.contrib.external_workflow_streams._api import ( + external_stream as private_external_stream, +) +from temporalio.contrib.external_workflow_streams._backend import ( + _validate_backend, +) +from temporalio.contrib.external_workflow_streams._redis import ( + RedisStreamBackend as PrivateRedisStreamBackend, +) +from temporalio.worker import Worker +from temporalio.worker.workflow_sandbox._restrictions import SandboxRestrictions +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend + + +# These tests exercise Worker construction, not the sandbox. Keeping this +# unsandboxed avoids re-importing this pytest assertion-rewritten module while +# the constructor validates its registered workflows. +@workflow.defn(sandboxed=False) +class NoOpWorkflow: + """Workers need at least one registered workflow to be constructible.""" + + @workflow.run + async def run(self) -> None: + pass + + +class UndeclaredBackend(MemoryStreamBackend): + """Never considered whether its records can change.""" + + guarantees_immutability = None + + +class DeniedBackend(MemoryStreamBackend): + """Considered it and cannot make the guarantee.""" + + guarantees_immutability = False + + +class AnonymousBackend(MemoryStreamBackend): + """Conforming, but nameless in the annotation header.""" + + provider_id = "" + + +# --- public package --------------------------------------------------------- + + +def test_supported_entry_points_are_public() -> None: + assert external_stream is private_external_stream + assert RedisStreamBackend is PrivateRedisStreamBackend + assert issubclass(RedisStreamBackend, StreamBackend) + assert set(external_workflow_streams.__all__) == { + "AFTER", + "BEGINNING", + "AppendConflictError", + "AppendNotAcknowledgedError", + "ChainKeyMismatchError", + "ConcurrentStreamConsumerError", + "Cursor", + "ExternalStreamCapacityError", + "ExternalOutputStreamClient", + "ExternalOutputStreamClientTopic", + "ExternalOutputStreamItem", + "ExternalOutputStreamOptions", + "ExternalOutputStreamProducer", + "ExternalOutputStreamProducerTopic", + "ExternalOutputStreamTopic", + "ExternalStreamOptions", + "ExternalStreamProducer", + "ExternalStreamProducerTopic", + "ExternalStreamSubscription", + "ExternalStreamTopic", + "IdempotencyKey", + "Offset", + "OffsetComparator", + "OutputAppendNotAcknowledgedError", + "OutputReadResult", + "OutputStage", + "OutputStageConflictError", + "OutputStageManifest", + "OutputStageNotFoundError", + "OutputStageResolutionError", + "OutputStageStatus", + "OutputStreamBackend", + "OutputStreamRecord", + "ParkIntent", + "ParkIntentRemoval", + "PrecedingWriteFailedError", + "RecordKind", + "RedisStreamBackend", + "StreamBackend", + "StreamDecodeError", + "StreamError", + "StreamIntegrityError", + "StreamDirection", + "StreamKey", + "StreamPayloadCodec", + "StreamRecord", + "StreamStorageError", + "PendingOutputBarrier", + "StagedOutputRecord", + "WakeNotAcknowledgedError", + "WakeRequest", + "WorkflowChainKey", + "external_stream", + "external_output_stream", + "merge", + } + + +# --- backend validation ----------------------------------------------------- + + +def test_a_conforming_backend_is_accepted() -> None: + backend = MemoryStreamBackend() + + assert _validate_backend(backend) is backend + + +@pytest.mark.parametrize( + "backend_type", [UndeclaredBackend, DeniedBackend], ids=["undeclared", "denied"] +) +def test_a_backend_without_the_guarantee_is_rejected( + backend_type: type[MemoryStreamBackend], +) -> None: + """One message for "forgot" and "cannot": both break the same thing. + + Replay validates presence, count, order, and control positions only, and + that is sufficient *because* the bytes cannot change. A provider that + cannot promise it has not satisfied the contract, however it got there. + """ + with pytest.raises(ValueError) as caught: + _validate_backend(backend_type()) + + message = str(caught.value) + assert "guarantees_immutability" in message + assert "cannot change" in message + + +def test_a_backend_without_a_provider_id_is_rejected() -> None: + with pytest.raises(ValueError, match="no provider_id"): + _validate_backend(AnonymousBackend()) + + +def test_a_non_backend_is_rejected() -> None: + with pytest.raises(TypeError, match="not a StreamBackend"): + _validate_backend(object()) + + +# --- through Worker construction -------------------------------------------- + + +async def test_a_conforming_backend_registers_on_a_worker(client: Client) -> None: + async with Worker( + client, + task_queue=f"tq-{uuid.uuid4()}", + workflows=[NoOpWorkflow], + external_stream_backend=MemoryStreamBackend(), + ) as worker: + assert isinstance(worker._external_stream_backend, MemoryStreamBackend) + + +async def test_a_backend_without_the_guarantee_fails_worker_construction( + client: Client, +) -> None: + """Loudly, at construction -- not quietly, at replay. + + By replay time a cursor has already been committed against data the + provider could have changed, and there is nothing left to do about it. + """ + with pytest.raises(ValueError, match="guarantees_immutability"): + Worker( + client, + task_queue=f"tq-{uuid.uuid4()}", + workflows=[NoOpWorkflow], + external_stream_backend=DeniedBackend(), + ) + + +async def test_no_backend_is_the_default(client: Client) -> None: + async with Worker( + client, task_queue=f"tq-{uuid.uuid4()}", workflows=[NoOpWorkflow] + ) as worker: + assert worker._external_stream_backend is None + + +# --- the sandbox ------------------------------------------------------------ + + +def test_the_sandbox_refuses_a_direct_provider_import() -> None: + """Workflow code never constructs the configured backend. + + A provider reached from inside the sandbox would be a second, unregistered + instance -- its own connection, no watcher owning it, and none of the + registration checks applied to it. + """ + matcher = SandboxRestrictions.invalid_module_members_default + contrib = matcher.children["temporalio"].children["contrib"] + streams = contrib.children["external_workflow_streams"] + + assert streams.children["_redis"].match_self + assert "RedisStreamBackend" in streams.access + + from temporalio.worker.workflow_sandbox._restrictions import RestrictionContext + + context = RestrictionContext() + context.is_runtime = True + assert streams.children["_redis"].match_access(context) + assert "external_stream_backend=... instead" in ( + streams.children["_redis"].leaf_message or "" + ) diff --git a/tests/contrib/external_workflow_streams/test_replay.py b/tests/contrib/external_workflow_streams/test_replay.py new file mode 100644 index 0000000000..bc9ee978c5 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_replay.py @@ -0,0 +1,2186 @@ +"""P13 — the replay read path and its four range checks. + +The checks are sufficient *because* every provider guarantees a record's bytes +cannot change, so the only damage replay has to detect is a record that is no +longer there. Each test below deletes a record in a different position, and each +must fail a different one of the four -- which is what makes the set complete +rather than merely plausible. +""" + +# pyright: reportUnusedImport=false +from __future__ import annotations + +import asyncio +import uuid +from datetime import timedelta + +import pytest +import pytest_asyncio + +import temporalio.api.common.v1 +import temporalio.converter +import temporalio.workflow +from temporalio.contrib.external_workflow_streams._annotation import ( + Annotation, + AnnotationHeader, + Run, + Segment, + SegmentEndReason, + StreamBinding, + decode_annotation, + encode_annotation, +) +from temporalio.contrib.external_workflow_streams._api import ( + _install_runtime, + external_stream, + merge, +) +from temporalio.contrib.external_workflow_streams._backend import StreamKey +from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec +from temporalio.contrib.external_workflow_streams._errors import ( + StreamDecodeError, + StreamIntegrityError, + StreamStorageError, + classify_read_failure, +) +from temporalio.contrib.external_workflow_streams._manager import ( + ReadinessResult, + StreamSubscriptionManager, +) +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Cursor, + Offset, + RecordKind, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._replay import ( + ReplayPlan, + validate_run, +) +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend +from tests.contrib.external_workflow_streams.test_replay_end_to_end import publish + +RUN_ID = "run-1" + + +async def _notify( + run_id: str, # pyright: ignore[reportUnusedParameter] + wait_id: int, # pyright: ignore[reportUnusedParameter] + generation: int, # pyright: ignore[reportUnusedParameter] +) -> str: + return ReadinessResult.ACCEPTED + + +@pytest.fixture +def backend() -> MemoryStreamBackend: + return MemoryStreamBackend() + + +@pytest.fixture +def key() -> StreamKey: + return StreamKey("ns", "wf", uuid.uuid4().hex, "tokens") + + +@pytest_asyncio.fixture # type: ignore[reportUntypedFunctionDecorator] +async def manager(backend: MemoryStreamBackend): + mgr = StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=10), + ) + yield mgr + await mgr.shutdown() + + +async def append_five( + backend: MemoryStreamBackend, key: StreamKey +) -> list[StreamRecord]: + """Five data records, so first / middle / last deletions are all distinct.""" + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + placed = [] + for i in range(5): + placed.append( + await backend.append( + key, + StreamRecord( + RecordKind.DATA, await codec.encode(f"v{i}"), "producer", i + ), + ) + ) + return placed + + +def placed_offset(record: StreamRecord) -> Offset: + """Return the provider-assigned offset of a record already appended here.""" + assert record.offset is not None + return record.offset + + +def binding( + key: StreamKey, + start_cursor: Cursor = BEGINNING, + *, + provider_id: str = MemoryStreamBackend.provider_id, + provider_format_version: int = MemoryStreamBackend.provider_format_version, +) -> StreamBinding: + """A binding carrying the configured backend's provider identity.""" + return StreamBinding( + stream_key=key, + start_cursor=start_cursor, + provider_id=provider_id, + provider_format_version=provider_format_version, + ) + + +def annotation_for( + key: StreamKey, placed: list[StreamRecord], control_positions: tuple[int, ...] = () +) -> bytes: + """One marker recording all of ``placed`` as a single run.""" + first, last = placed[0].offset, placed[-1].offset + assert first is not None and last is not None + return encode_annotation( + Annotation( + header=AnnotationHeader(streams={1: binding(key)}), + segments=( + Segment( + runs=( + Run( + wait_id=1, + first_offset=first, + last_offset=last, + count=len(placed), + control_positions=control_positions, + ), + ), + end_reason=SegmentEndReason.NO_DATA_AVAILABLE, + ), + ), + terminal={1: AFTER(last)}, + ) + ) + + +# --- the happy path ---------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_an_intact_range_replays_in_recorded_order( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + placed = await append_five(backend, key) + + plan = await manager.prepare_replay(RUN_ID, annotation_for(key, placed)) + + assert plan.total_records == 5 + (segment,) = plan.segments + assert [r.offset for _, r in segment.deliveries] == [r.offset for r in placed] + assert all(wait_id == 1 for wait_id, _ in segment.deliveries) + + +@pytest.mark.asyncio +async def test_replay_performs_no_live_waiting( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """It reads recorded ranges; it never asks what comes next. + + A live watch here would block on an empty stream and, worse, could deliver a + record the original run never saw. + """ + placed = await append_five(backend, key) + + await manager.prepare_replay(RUN_ID, annotation_for(key, placed)) + + # `read_range` is the inclusive replay read; `read_after` is the exclusive + # live watch, and it must not have been used at all. + assert backend.range_reads, "replay must read the recorded range" + assert len(backend.range_reads) == 1, ( + "one read per run, not one per record -- replay I/O is a function of " + f"the consumed range, got {len(backend.range_reads)} reads for 5 records" + ) + + +@pytest.mark.asyncio +async def test_one_segment_per_recorded_activation( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """Segments stay apart so replay reproduces the same number of drains. + + Collapsing them would reproduce the record order while changing when + ``wait_condition`` predicates fire. + """ + placed = await append_five(backend, key) + first, second = placed[0].offset, placed[-1].offset + assert first is not None and second is not None + + annotation = encode_annotation( + Annotation( + header=AnnotationHeader({1: binding(key)}), + segments=( + Segment( + (Run(1, first, placed[1].offset, 2),), # type: ignore[arg-type] + SegmentEndReason.BATCH_LIMIT, + ), + Segment((), SegmentEndReason.NO_DATA_AVAILABLE), + Segment( + (Run(1, placed[2].offset, second, 3),), # type: ignore[arg-type] + SegmentEndReason.NO_DATA_AVAILABLE, + ), + ), + terminal={1: AFTER(second)}, + ) + ) + + plan = await manager.prepare_replay(RUN_ID, annotation) + + assert [len(s.deliveries) for s in plan.segments] == [2, 0, 3], ( + "the empty segment is a drain that happened and must survive replay" + ) + + +# --- integrity: a deletion in each position ---------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("position", [0, 2, 4], ids=["first", "middle", "last"]) +async def test_deleting_a_record_fails_as_integrity_loss( + manager: StreamSubscriptionManager, + backend: MemoryStreamBackend, + key: StreamKey, + position: int, +) -> None: + """And never resolves to an alternate stream result. + + Substituting a later record for a deleted one would hand Workflow code a + different history than the one its commands were derived from, and the + divergence would surface much later as an unrelated nondeterminism error. + """ + placed = await append_five(backend, key) + annotation = annotation_for(key, placed) + target = placed[position].offset + assert target is not None + await backend.delete_for_test(key, target) + + with pytest.raises(StreamIntegrityError) as caught: + await manager.prepare_replay(RUN_ID, annotation) + + assert str(target) in str(caught.value) or "contains" in str(caught.value) + + +@pytest.mark.asyncio +async def test_each_deletion_position_fails_a_different_check( + backend: MemoryStreamBackend, key: StreamKey +) -> None: + """What makes the set of four complete rather than merely plausible. + + If a middle deletion and an endpoint deletion tripped the same check, one of + the four would be redundant and some other damage would be going undetected. + """ + placed = await append_five(backend, key) + run = Run(1, placed[0].offset, placed[-1].offset, 5) # type: ignore[arg-type] + + messages = {} + for name, position in (("first", 0), ("middle", 2), ("last", 4)): + surviving = [r for i, r in enumerate(placed) if i != position] + with pytest.raises(StreamIntegrityError) as caught: + validate_run(run, surviving, backend) + messages[name] = str(caught.value) + + assert "is missing from the stream" in messages["first"] + assert "is missing from the stream" in messages["last"] + assert "contains 4 record(s)" in messages["middle"] + assert messages["first"] != messages["last"] + + +@pytest.mark.asyncio +async def test_a_wrong_control_position_fails_validation( + backend: MemoryStreamBackend, key: StreamKey +) -> None: + """A control record swapped for a data record would otherwise be yielded. + + Workflow code never saw the control record the first time, and delivering a + data record in its place would put a value in front of it that the original + run never produced. + """ + placed = await append_five(backend, key) + run = Run(1, placed[0].offset, placed[-1].offset, 5, control_positions=(2,)) # type: ignore[arg-type] + + with pytest.raises(StreamIntegrityError, match="control records at"): + validate_run(run, placed, backend) + + +@pytest.mark.asyncio +async def test_out_of_order_records_fail_validation( + backend: MemoryStreamBackend, key: StreamKey +) -> None: + placed = await append_five(backend, key) + reordered = [placed[0], placed[2], placed[1], placed[3], placed[4]] + run = Run(1, placed[0].offset, placed[-1].offset, 5) # type: ignore[arg-type] + + with pytest.raises(StreamIntegrityError, match="strictly increasing"): + validate_run(run, reordered, backend) + + +@pytest.mark.asyncio +async def test_a_marker_naming_an_unknown_wait_is_nondeterminism_not_integrity_loss( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """Row four of the taxonomy, and the row it must not be confused with. + + Nothing is wrong with the backend here: the recorded ranges are exactly + where they were written, and the Workflow code changed underneath them. + Reporting it as integrity loss would send an operator to repair a backend + that is fine, when the fix is to version the Workflow code -- the same + mistake in the opposite direction from calling an outage integrity loss. + + Reported rather than skipped either way, since skipping would silently + deliver a different stream result. + """ + placed = await append_five(backend, key) + annotation = encode_annotation( + Annotation( + header=AnnotationHeader({}), + segments=( + Segment( + (Run(7, placed[0].offset, placed[-1].offset, 5),), # type: ignore[arg-type] + SegmentEndReason.NO_DATA_AVAILABLE, + ), + ), + terminal={}, + ) + ) + + with pytest.raises( + temporalio.workflow.NondeterminismError, match="did not create" + ) as caught: + await manager.prepare_replay(RUN_ID, annotation) + + assert not isinstance(caught.value, StreamIntegrityError), ( + "this row must not be reachable through the storage-failure taxonomy" + ) + assert "workflow.patched" in str(caught.value), ( + "the message must name the remedy, which is versioning the Workflow " + "code rather than touching the backend" + ) + + +# --- storage failure is not integrity loss ----------------------------------- + + +class UnreachableBackend(MemoryStreamBackend): + """Stands in for a backend that is down rather than damaged.""" + + async def read_range(self, key, first, last): # type: ignore[no-untyped-def] + raise ConnectionError("connection refused") + + +@pytest.mark.asyncio +async def test_an_unreachable_backend_fails_as_transient_storage( + key: StreamKey, +) -> None: + """Distinguished from integrity loss by error type, never by retry behavior. + + The server retries a Workflow Task failure regardless of cause. An outage + clears on its own; integrity loss does not -- and an operator sent to repair + a backend that was merely unreachable would find nothing wrong with it. + """ + reachable = MemoryStreamBackend() + placed = await append_five(reachable, key) + annotation = annotation_for(key, placed) + + unreachable = UnreachableBackend() + manager = StreamSubscriptionManager( + backend=unreachable, + notify_ready=_notify, + watch_block=timedelta(milliseconds=10), + ) + try: + with pytest.raises(StreamStorageError, match="connection refused"): + await manager.prepare_replay(RUN_ID, annotation) + finally: + await manager.shutdown() + + +# --- decode failure is not integrity loss ------------------------------------ + + +@pytest.mark.asyncio +async def test_an_intact_but_undecodable_record_is_a_decode_error( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """The range validated, so the bytes are the bytes that were written. + + That makes this a configuration error on the *consumer* -- its converter or + codec does not match the producer's -- and reporting it as stream integrity + loss would send an operator to restore a backend that was never damaged. + """ + placed = [] + for i in range(3): + placed.append( + await backend.append( + key, + # Not a serialized Payload at all: intact bytes, undecodable. + StreamRecord(RecordKind.DATA, b"\\xff\\xfe raw", "producer", i), + ) + ) + annotation = annotation_for(key, placed) + + # The range itself validates -- nothing is missing. + plan = await manager.prepare_replay(RUN_ID, annotation) + assert plan.total_records == 3 + + # It is only decoding that fails, and only then is it classified. + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + _, record = plan.segments[0].deliveries[0] + try: + await codec.decode(record.payload) + raise AssertionError("expected the record to be undecodable") + except AssertionError: + raise + except Exception as err: + classified = classify_read_failure(range_validated=True, cause=err) + + assert isinstance(classified, StreamDecodeError) + assert not isinstance(classified, StreamIntegrityError) + + +class _FailingRetrieveDriver(temporalio.converter.StorageDriver): + """Stores payloads in memory and refuses to give them back once armed. + + Stands in for an external payload store that is briefly unreachable. It + raises the exception such a driver's client would raise -- a bare + ``ConnectionError`` -- rather than anything this module defines, which is the + whole point: the converter does not wrap it, so if no layer labels it the + record arrives at the consumer as an unclassified failure. + """ + + def __init__(self) -> None: + self._stored: dict[str, bytes] = {} + self.fail = False + self.retrieves = 0 + + def name(self) -> str: + return "failing-retrieve" + + async def store(self, context, payloads): # type: ignore[no-untyped-def] + claims = [] + for index, payload in enumerate(payloads): + key = f"k{len(self._stored) + index}" + self._stored[key] = payload.SerializeToString() + claims.append( + temporalio.converter.StorageDriverClaim(claim_data={"k": key}) + ) + return claims + + async def retrieve(self, context, claims): # type: ignore[no-untyped-def] + self.retrieves += 1 + if self.fail: + raise ConnectionError("the external payload store is unreachable") + out = [] + for claim in claims: + payload = temporalio.api.common.v1.Payload() + payload.ParseFromString(self._stored[claim.claim_data["k"]]) + out.append(payload) + return out + + +def _extstore_converter( + driver: _FailingRetrieveDriver, +) -> temporalio.converter.DataConverter: + return temporalio.converter.DataConverter( + external_storage=temporalio.converter.ExternalStorage( + drivers=[driver], payload_size_threshold=1 + ) + ) + + +@pytest.mark.parametrize("path", ["live", "replay"]) +@pytest.mark.asyncio +async def test_an_unreachable_payload_store_is_a_storage_failure_not_a_decode_one( + backend: MemoryStreamBackend, key: StreamKey, path: str +) -> None: + """Row one, on both delivery paths, whichever way the payload is fetched. + + A record's bytes can be a *reference*: with external storage configured, what + the stream holds is a claim and the value does not exist until the payload + store hands it over. A store that cannot be reached is therefore the same + condition as a stream backend that cannot be reached -- transient, clearing on + its own, with nothing for an operator to change. + + Left unlabelled it is not reported that way. The driver raises whatever its + client raises, the converter does not wrap it, and the consumer's + classification rule is "the range validated, so the bytes are the bytes that + were written" -- which turns it into a decode failure and tells an operator to + align a converter that was never wrong. Both paths are checked because they + prepare in different places: the watcher's loop for live delivery, the replay + read for a recorded range. + """ + driver = _FailingRetrieveDriver() + converter = _extstore_converter(driver) + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + data_converter=converter, + watch_block=timedelta(milliseconds=10), + ) + try: + codec = StreamPayloadCodec(converter, str) + placed = await backend.append( + key, StreamRecord(RecordKind.DATA, await codec.encode("alpha"), "p", 0) + ) + assert driver._stored, ( + "the payload was not externalized, so nothing here depends on the " + "payload store at all" + ) + + # The manager is what prepares, so the converter under test is the one it + # holds; the runtime here only needs to make the subscription. + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=key) + + driver.fail = True + if path == "live": + subscription = manager.subscription(RUN_ID, 1) + assert subscription is not None + await _until( + lambda: subscription.buffered == 1, + 5, + "the watcher never buffered the record", + ) + prepared = manager.drain(RUN_ID, 1)[0] + else: + plan = await manager.prepare_replay(RUN_ID, annotation_for(key, [placed])) + _, prepared = plan.segments[0].deliveries[0] + + carried = getattr(prepared, "prepare_error", None) + assert carried is not None, ( + f"the {path} path prepared the record without noticing that its " + "payload could not be fetched" + ) + assert isinstance(carried, StreamStorageError), ( + f"the {path} path carried {type(carried).__name__}, so the delivery " + "that raises it will classify a storage outage as a decode failure" + ) + assert not isinstance(carried, StreamDecodeError) + + # And the classification rule leaves it alone rather than relabelling it: + # a storage failure is about reaching the store at all, so "the range + # validated" says nothing about it. + assert classify_read_failure(range_validated=True, cause=carried) is carried + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_a_replay_plan_is_consumed_once( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """A second delivery would replay the same records twice.""" + placed = await append_five(backend, key) + await manager.prepare_replay(RUN_ID, annotation_for(key, placed)) + + assert manager.take_replay_plan(RUN_ID) is not None + assert manager.take_replay_plan(RUN_ID) is None + + +@pytest.mark.asyncio +async def test_eviction_discards_an_unconsumed_plan( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """A plan outliving its Run would be delivered to whatever came next.""" + placed = await append_five(backend, key) + await manager.prepare_replay(RUN_ID, annotation_for(key, placed)) + + await manager.evict_run(RUN_ID) + + assert manager.take_replay_plan(RUN_ID) is None + + +# --- the delivery driver (ADR-018) ------------------------------------------- + + +class DriverStub: + """The two things the replay driver is allowed to touch. + + Calling the real ``_apply_replay_external_streams`` against this asserts what + the driver does rather than what a Workflow happens to observe -- a Workflow + that never used ``wait_condition`` would pass either way. + """ + + def __init__(self, runtime) -> None: # type: ignore[no-untyped-def] + self._external_stream_runtime = runtime + self._pending_replay_finish: ReplayPlan | None = None + self.drains: list[int] = [] + + def _finish_replay_external_streams(self) -> None: + """The real one. A plan with no segments closes from inside the job. + + Present because a stub stands in for the instance, and the job calls this + on itself when there is no segment for the activation's drain to serve. + """ + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + + _WorkflowInstanceImpl._finish_replay_external_streams(self) # type: ignore[arg-type] + + def _run_once(self, *, check_conditions: bool) -> None: + assert check_conditions, ( + "a replay drain must check conditions -- otherwise `wait_condition` " + "predicates fire fewer times than they did live" + ) + self.drains.append(len(self.drains)) + + +def make_runtime(manager, backend): # type: ignore[no-untyped-def] + from temporalio.contrib.external_workflow_streams._runtime import ( + WorkflowStreamRuntime, + ) + + return WorkflowStreamRuntime( + manager=manager, + backend=backend, + run_id=RUN_ID, + namespace="ns", + workflow_id="wf", + first_execution_run_id=uuid.uuid4().hex, + data_converter=temporalio.converter.DataConverter.default, + default_idle_timeout=timedelta(seconds=1), + ) + + +class ConsumingStub(DriverStub): + """A stub that also takes what the segment holds, as Workflow code does. + + Every record in a run was handed to Workflow code in the activation that run + was recorded in, so a stub that drains nothing is not a stand-in for a + Workflow -- it is a stand-in for a Workflow that stopped subscribing, which + the driver now reports as nondeterminism. Tests that only care about how + many drains happened use this; tests that assert *what* each drain saw drain + for themselves. + + It drains only the waits the Run **registered**, for the same reason: a + Workflow has no handle to a subscription it never made, so a stub that + drains by whatever the segment holds reports a removed ``subscribe()`` as + consumed and hides the misrouting that goes with it. + """ + + def __init__(self, runtime) -> None: # type: ignore[no-untyped-def] + super().__init__(runtime) + self._runtime = runtime + + def _run_once(self, *, check_conditions: bool) -> None: + super()._run_once(check_conditions=check_conditions) + # Only the waits this Run actually registered. Draining whatever the + # segment happens to hold is something no Workflow can do -- it has no + # handle to a subscription it never made -- and a stub that does it + # hides exactly the failure the driver exists to catch: a removed + # `subscribe()` whose records are then taken by nobody's subscription + # and reported as consumed. + registered = set(self._runtime.subscriptions()) + for wait_id in sorted({w for w, _ in self._runtime._replay_ready or []}): + if wait_id in registered: + self._runtime.drain(wait_id) + + +def drive_with(stub): # type: ignore[no-untyped-def] + """Runs the whole activation shape around one replay job, not just `_apply`. + + The trailing drain is neither optional nor this helper's invention: the replay + job lands in the non-query job set, so the activation runs exactly one + `_run_once` for that set whatever `_apply` did with the job. Driving `_apply` + alone is what let the driver perform one drain more than the marker records + without any test noticing -- ADR-018 requires *k* drains for *k* segments and + it was performing *k + 1* -- so every test here drives the activation's drain + too, and the close that the driver now defers until after it. + """ + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + + try: + _WorkflowInstanceImpl._apply_replay_external_streams(stub, object()) # type: ignore[arg-type] + stub._run_once(check_conditions=True) + finally: + _WorkflowInstanceImpl._finish_replay_external_streams(stub) # type: ignore[arg-type] + return stub + + +def drive(runtime) -> DriverStub: # type: ignore[no-untyped-def] + return drive_with(ConsumingStub(runtime)) + + +@pytest.mark.asyncio +async def test_the_driver_drains_once_per_recorded_segment( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """Three recorded activations become three drains, not one. + + Delivering all five records in a single drain would reproduce the record + order while changing how many times the event loop ran, and a + ``wait_condition`` predicate would then see a different sequence of states + than it did live -- a nondeterminism that only shows up on replay. + """ + placed = await append_five(backend, key) + annotation = encode_annotation( + Annotation( + header=AnnotationHeader({1: binding(key)}), + segments=( + Segment( + (Run(1, placed_offset(placed[0]), placed_offset(placed[1]), 2),), + SegmentEndReason.BATCH_LIMIT, + ), # type: ignore[arg-type] + Segment( + (Run(1, placed_offset(placed[2]), placed_offset(placed[2]), 1),), + SegmentEndReason.BATCH_LIMIT, + ), # type: ignore[arg-type] + Segment( + (Run(1, placed_offset(placed[3]), placed_offset(placed[4]), 2),), + SegmentEndReason.NO_DATA_AVAILABLE, + ), # type: ignore[arg-type] + ), + terminal={1: AFTER(placed[4].offset)}, # type: ignore[arg-type] + ) + ) + await manager.prepare_replay(RUN_ID, annotation) + runtime = make_runtime(manager, backend) + # The `subscribe()` call the marker recorded. A stub stands in for Workflow + # code, so it has to make it: a replay that ends with a recorded binding + # never recreated is a removed subscription whatever is driving the drains. + runtime.register(wait_id=1, stream_key=key) + + stub = drive(runtime) + + assert len(stub.drains) == 3, ( + f"expected one drain per recorded segment, got {len(stub.drains)}" + ) + + +@pytest.mark.asyncio +async def test_a_drain_sees_only_its_own_segments_records( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """Records may not run ahead of the drain they were recorded in. + + Buffering all of them and letting the first drain take everything would put + values in front of Workflow code before the activation that produced them. + """ + placed = await append_five(backend, key) + annotation = encode_annotation( + Annotation( + header=AnnotationHeader({1: binding(key)}), + segments=( + Segment( + (Run(1, placed_offset(placed[0]), placed_offset(placed[1]), 2),), + SegmentEndReason.BATCH_LIMIT, + ), # type: ignore[arg-type] + Segment( + (Run(1, placed_offset(placed[2]), placed_offset(placed[4]), 3),), + SegmentEndReason.NO_DATA_AVAILABLE, + ), # type: ignore[arg-type] + ), + terminal={1: AFTER(placed[4].offset)}, # type: ignore[arg-type] + ) + ) + await manager.prepare_replay(RUN_ID, annotation) + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=key) + + seen_per_drain: list[list[Offset]] = [] + + class DrainingStub(DriverStub): + def _run_once(self, *, check_conditions: bool) -> None: + super()._run_once(check_conditions=check_conditions) + seen_per_drain.append([r.offset for r in runtime.drain(1)]) # type: ignore[misc] + + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + + drive_with(DrainingStub(runtime)) + + assert seen_per_drain == [ + [placed[0].offset, placed[1].offset], + [placed[2].offset, placed[3].offset, placed[4].offset], + ] + + +@pytest.mark.asyncio +async def test_the_driver_reads_nothing_and_ends_replay_mode( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """Delivery is memory-only, and the Run goes back to live drains afterwards. + + Leaving replay mode latched would make every later live drain read from an + empty recorded buffer, and the Workflow would block forever on records that + had already arrived. + """ + placed = await append_five(backend, key) + await manager.prepare_replay(RUN_ID, annotation_for(key, placed)) + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=key) + before = len(backend.range_reads) + + drive(runtime) + + assert len(backend.range_reads) == before, "the delivering pass must read nothing" + assert runtime._replay_ready is None, "replay mode must not outlive the job" + + +@pytest.mark.asyncio +async def test_a_replayed_drain_never_reaches_the_live_buffer( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """A watcher may have buffered past where the marker stops. + + That buffer is speculative -- reading a record is not consuming it -- so + delivering it would hand Workflow code records this Workflow Task never saw. + Replay would then be the thing introducing the nondeterminism it exists to + prevent, which is why the recorded segment is the *only* thing a drain can + see while one is being delivered. + """ + placed = await append_five(backend, key) + # The marker committed the first two records only. + await manager.prepare_replay(RUN_ID, annotation_for(key, placed[:2])) + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=key) + subscription = manager.subscription(RUN_ID, 1) + assert subscription is not None + subscription._append(placed[2:], subscription._prefetch_epoch) + + seen: list[list[Offset]] = [] + + class DrainingStub(DriverStub): + def _run_once(self, *, check_conditions: bool) -> None: + super()._run_once(check_conditions=check_conditions) + seen.append([r.offset for r in runtime.drain(1)]) # type: ignore[misc] + + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + + drive_with(DrainingStub(runtime)) + + assert seen == [[placed[0].offset, placed[1].offset]] + # And the buffered-past records were never handed to that drain. Closing the + # replay retracts them rather than leaving them in front of the next one: + # reading is not consuming, so the speculative buffer is discarded and the + # watcher re-reads from the boundary the marker committed. Withheld either + # way -- what changed is that the retraction is now certain by the time the + # replay returns rather than posted onto the manager's loop and raced. + assert subscription.committed_cursor == AFTER(placed_offset(placed[1])) + assert subscription.prefetch_cursor == AFTER(placed_offset(placed[1])), ( + "the watcher must resume reading at the marker's boundary, which is " + "what re-reads the records this retraction dropped" + ) + assert runtime.drain(1) == [], ( + "records read past the marker must not survive the replay in the " + "buffer, or the first live drain hands over what it never saw recorded" + ) + + +@pytest.mark.asyncio +async def test_two_markers_reassemble_in_workflow_task_order( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """A batch split by a rollover replays as one continuous consumption. + + The byte-budget high-water mark rolls the Workflow Task over rather than + growing the marker, so one logical batch ends up in *two* markers -- and + replay reaches them as two separate ``ReplayExternalStreams`` jobs, one per + replayed Workflow Task. Core's own test proves the two annotations + concatenate at the byte level; what only this side can show is that feeding + them through the manager in Workflow Task order reproduces the original + consumption exactly. + + Two failures this rules out, both silent: + + - a seam that loses or repeats a record, which is what an exclusive read or + an off-by-one at the split would produce -- the second marker resumes at + the first's terminal, so the record *at* that boundary must appear exactly + once across the pair; + - a second marker delivered against the first one's segmentation, which + would put records in front of Workflow code in a drain that never saw + them. + """ + placed = await append_five(backend, key) + split = 1 # The first marker stops after `placed[1]`. + + first_marker = encode_annotation( + Annotation( + header=AnnotationHeader({1: binding(key)}), + segments=( + Segment( + (Run(1, placed[0].offset, placed[0].offset, 1),), # type: ignore[arg-type] + SegmentEndReason.BATCH_LIMIT, + ), + Segment( + (Run(1, placed[1].offset, placed[1].offset, 1),), # type: ignore[arg-type] + # The annotation passed its high-water mark here, so this + # batch continues in the following marker. + SegmentEndReason.BUDGET_ROLLOVER, + ), + ), + terminal={1: AFTER(placed[split].offset)}, # type: ignore[arg-type] + ) + ) + # The replacement Workflow Task's annotation begins where the first one's + # terminal left the subscription: a rollover preserves the cursor rather + # than restarting it. + second_marker = encode_annotation( + Annotation( + header=AnnotationHeader( + {1: binding(key, AFTER(placed[split].offset))} # type: ignore[arg-type] + ), + segments=( + Segment( + (Run(1, placed[2].offset, placed[4].offset, 3),), # type: ignore[arg-type] + SegmentEndReason.NO_DATA_AVAILABLE, + ), + ), + terminal={1: AFTER(placed[4].offset)}, # type: ignore[arg-type] + ) + ) + + runtime = make_runtime(manager, backend) + # One subscription across both markers, as the rollover preserves it. + runtime.register(wait_id=1, stream_key=key) + per_marker: list[list[list[Offset]]] = [] + + for annotation in (first_marker, second_marker): + await manager.prepare_replay(RUN_ID, annotation) + drains: list[list[Offset]] = [] + + class DrainingStub(DriverStub): + def _run_once(self, *, check_conditions: bool) -> None: + super()._run_once(check_conditions=check_conditions) + drains.append([r.offset for r in runtime.drain(1)]) # type: ignore[misc] + + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + + drive_with(DrainingStub(runtime)) + per_marker.append(drains) + + assert per_marker == [ + [[placed[0].offset], [placed[1].offset]], + [[placed[2].offset, placed[3].offset, placed[4].offset]], + ], ( + "each marker must deliver its own segments and nothing else; a marker " + "reassembled against the wrong segmentation puts records in a drain " + f"that never saw them, got {per_marker}" + ) + + reassembled = [ + offset for drains in per_marker for drain in drains for offset in drain + ] + assert reassembled == [r.offset for r in placed], ( + "the two markers must reassemble into the original consumption, in " + f"Workflow Task order and with nothing lost or repeated at the seam: " + f"{reassembled}" + ) + + +@pytest.mark.asyncio +async def test_one_segment_delivers_each_waits_own_records( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """A segment can carry runs from several subscriptions at once. + + Two subscriptions are two independent cursors, so handing one's records to + the other is a different stream result rather than merely a different order. + """ + other = StreamKey( + key.namespace, key.workflow_id, key.first_execution_run_id, "tool-events" + ) + placed = await append_five(backend, key) + other_placed = await append_five(backend, other) + + annotation = encode_annotation( + Annotation( + header=AnnotationHeader({1: binding(key), 2: binding(other)}), + segments=( + Segment( + ( + Run(1, placed[0].offset, placed[4].offset, 5), # type: ignore[arg-type] + Run(2, other_placed[0].offset, other_placed[4].offset, 5), # type: ignore[arg-type] + ), + SegmentEndReason.NO_DATA_AVAILABLE, + ), + ), + terminal={1: AFTER(placed[4].offset), 2: AFTER(other_placed[4].offset)}, # type: ignore[arg-type] + ) + ) + await manager.prepare_replay(RUN_ID, annotation) + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=key) + runtime.register(wait_id=2, stream_key=other) + + drained: dict[int, list[Offset]] = {} + + class DrainingStub(DriverStub): + def _run_once(self, *, check_conditions: bool) -> None: + super()._run_once(check_conditions=check_conditions) + for wait_id in (1, 2): + drained[wait_id] = [r.offset for r in runtime.drain(wait_id)] # type: ignore[misc] + + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + + drive_with(DrainingStub(runtime)) + + assert drained[1] == [r.offset for r in placed] + assert drained[2] == [r.offset for r in other_placed] + + +# --- what the live buffer must look like afterwards --------------------------- + + +async def _until(predicate, timeout: float, message: str) -> None: # type: ignore[no-untyped-def] + """Polls rather than sleeping a fixed time, so a slow watcher is not a flake.""" + deadline = asyncio.get_running_loop().time() + timeout + while asyncio.get_running_loop().time() < deadline: + if predicate(): + return + await asyncio.sleep(0.01) + raise AssertionError(message) + + +@pytest.mark.asyncio +async def test_live_delivery_after_a_replay_resumes_past_the_marker( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """A replayed marker's records must not be handed over a second time, live. + + Replay delivers from the annotation, but the manager's watcher has been + reading the *same* records from the subscription's start cursor into the + live buffer the whole time -- nothing out there knows the marker exists. So + the first live drain after a replay re-delivers everything the marker + recorded: observed end-to-end as a Workflow that received + ``['alpha', 'alpha', 'beta', 'gamma']``. + + Repositioning to the marker's terminal is what makes live delivery resume + after the recorded range rather than in front of it. + """ + placed = await append_five(backend, key) + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=key) + + subscription = manager.subscription(RUN_ID, 1) + assert subscription is not None + await _until( + lambda: subscription.buffered == 5, + 5, + "the watcher never buffered the records, so nothing would be " + "re-delivered either way and this proves nothing", + ) + + # A marker recording only the first two of the five, so "resumed past it" + # and "delivered nothing at all" are distinguishable. + await manager.prepare_replay(RUN_ID, annotation_for(key, placed[:2])) + drive(runtime) + # The reposition is hopped onto this loop, exactly as the Workflow thread + # would hop it. + await asyncio.sleep(0) + + await _until( + lambda: subscription.buffered == 3, + 5, + "the buffer did not refill from the marker's boundary", + ) + assert [r.offset for r in manager.drain(RUN_ID, 1)] == [ + r.offset for r in placed[2:] + ], ( + "the live buffer still holds records the replayed marker already " + "delivered, so the Workflow receives them twice" + ) + assert subscription.committed_cursor == AFTER(placed_offset(placed[1])), ( + "the marker's boundary was never committed, so a later reset would " + "send the subscription back to the start cursor" + ) + + +@pytest.mark.asyncio +async def test_the_first_live_drain_after_a_replay_needs_no_loop_turn( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """Repositioning has to have *happened*, not merely been scheduled. + + The drain that follows replay is on the Workflow thread, and so is the replay + itself. A reposition posted to the manager's loop with `call_soon_threadsafe` + and then left to arrive whenever the loop next runs is ordered by nothing at + all: if Workflow code drains first -- which one ordinary scheduler pass is + enough to do -- the live buffer still holds every record the marker just + delivered and hands them over a second time. + + So this test deliberately gives the manager's loop **no** opportunity to run + between the replay and the drain: no `await`, no `sleep(0)`, nothing. That is + the whole point. A version of this test with a yield in the middle passes + against a posted reposition too, and proves only that the loop happened to + win. + """ + placed = await append_five(backend, key) + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=key) + + subscription = manager.subscription(RUN_ID, 1) + assert subscription is not None + await _until( + lambda: subscription.buffered == 5, + 5, + "the watcher never buffered the records, so nothing would be " + "re-delivered either way and this proves nothing", + ) + + # A marker covering the first two of the five only. + await manager.prepare_replay(RUN_ID, annotation_for(key, placed[:2])) + drive(runtime) + + # No yield here on purpose. Everything below runs in the same synchronous + # stretch of Workflow-thread work that the replay just finished in. + assert subscription.committed_cursor == AFTER(placed_offset(placed[1])), ( + "the marker's boundary must be committed by the time the replay returns, " + "not once some later loop turn gets round to it" + ) + assert manager.drain(RUN_ID, 1) == [], ( + "the drain immediately after a replay still saw the marker-covered " + "records in the live buffer, so the Workflow receives them twice" + ) + assert subscription.prefetch_cursor == AFTER(placed_offset(placed[1])), ( + "the watcher must have been moved to re-read from the marker's boundary" + ) + + # And the records past the marker are not lost: the watcher re-reads them + # from the boundary that was committed. + await _until( + lambda: subscription.buffered == 3, + 5, + "the buffer never refilled from the marker's boundary", + ) + assert [r.offset for r in manager.drain(RUN_ID, 1)] == [ + r.offset for r in placed[2:] + ] + + +@pytest.mark.asyncio +async def test_a_failing_activation_reports_its_own_error_not_the_replays( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """Closing a replay may not overwrite the error that stopped it. + + The close runs ``verify_replay_consumed``, which raises whenever a recorded + delivery is still armed -- and after a drain fails, one always is: the drain + that would have taken it is the one that failed. An exception raised while + another is propagating **replaces** it, so running the close from a bare + ``finally`` turned every failing replay activation into a nondeterminism error + blaming a ``subscribe()`` call nobody had touched. + + Two things are lost that way, not one. The diagnosis, and the *classification*: + a ``FailureError`` out of Workflow code fails the Workflow, while a + nondeterminism error fails the Workflow Task and is retried forever. + + Replay mode is still left -- a Run stuck in it has every later drain return + nothing -- so what the abandon path skips is only the checks and the cursor + move, and the cursor move has to be skipped anyway: an activation that failed + committed nothing. + """ + placed = await append_five(backend, key) + await manager.prepare_replay(RUN_ID, annotation_for(key, placed)) + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=key) + + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + + # A replay left mid-flight, exactly as a drain that raised leaves it: the + # segment armed, its deliveries untaken, the close still owed. + stub = DriverStub(runtime) + runtime.begin_replay(decode_annotation(annotation_for(key, placed)).header.streams) + plan = runtime.take_replay_plan() + assert plan is not None + runtime.begin_replay_segment(list(plan.segments[0].deliveries)) + stub._pending_replay_finish = plan + + # Closing from here is what a bare `finally` did, and it raises: the recorded + # deliveries are still armed, because the drain that would have taken them is + # the one that failed. Asserted so that what follows is a contrast and not a + # coincidence. + with pytest.raises(temporalio.workflow.NondeterminismError): + _WorkflowInstanceImpl._finish_replay_external_streams(stub) # type: ignore[arg-type] + + # Abandoning does not. Re-arm and take the other path. + runtime.begin_replay_segment(list(plan.segments[0].deliveries)) + stub._pending_replay_finish = plan + _WorkflowInstanceImpl._abandon_replay_external_streams(stub) # type: ignore[arg-type] + + assert runtime._replay_ready is None and runtime._replay_bindings is None, ( + "replay mode was left set, so every later drain on this Run returns nothing" + ) + assert stub._pending_replay_finish is None + subscription = manager.subscription(RUN_ID, 1) + assert subscription is not None + assert subscription.committed_cursor == BEGINNING, ( + "an activation that failed committed nothing, so the cursor may not have " + "moved to the marker's boundary" + ) + + +@pytest.mark.asyncio +async def test_a_marker_with_no_segments_closes_before_the_activations_own_drain( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """A marker that recorded no activation of its own defers nothing. + + A Workflow Task that subscribed to a quiet stream and blocked writes a marker + with a header and a terminal and **no segments** -- and its terminal names the + cursor nothing was ever delivered past. There is no recorded segment for the + activation's own drain to serve, so that drain is a *live* one: records that + arrived while the Run was evicted are in the buffer and it hands them over. + + Repositioning after that drain retracts what it just delivered. The cursor goes + back to the marker's boundary, the buffer is cleared, and the watcher re-reads + and re-delivers records Workflow code already has. So a plan with no segments + closes inside the job, before the drain, and only a plan with segments defers. + + This is not reachable by inspection alone: the drain only has something to + deliver if the watcher buffered first, so the defect it guards against + presents as an occasional end-to-end failure and passes in isolation. + """ + placed = await append_five(backend, key) + # The marker of a Workflow Task that saw nothing: header and terminal only, + # with the terminal at the start cursor. + annotation = encode_annotation( + Annotation( + header=AnnotationHeader({1: binding(key)}), + segments=(), + terminal={1: BEGINNING}, + ) + ) + await manager.prepare_replay(RUN_ID, annotation) + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=key) + + subscription = manager.subscription(RUN_ID, 1) + assert subscription is not None + await _until( + lambda: subscription.buffered == 5, + 5, + "the watcher never buffered anything, so the activation's drain has " + "nothing to deliver and this proves nothing", + ) + + taken: list[Offset] = [] + + class LiveDrainingStub(DriverStub): + def _run_once(self, *, check_conditions: bool) -> None: + super()._run_once(check_conditions=check_conditions) + taken.extend(r.offset for r in runtime.drain(1)) # type: ignore[misc] + + drive_with(LiveDrainingStub(runtime)) + + assert taken == [], ( + "the activation's drain handed over buffered records and the replay's " + "close then retracted the cursor past them, so the watcher re-reads and " + f"the Workflow receives every one of them twice: {taken}" + ) + assert subscription.committed_cursor == BEGINNING, ( + "the marker committed nothing, so its boundary is the start cursor" + ) + assert subscription.prefetch_cursor == BEGINNING, ( + "the watcher must be re-reading from the boundary the marker committed" + ) + + # And the records are not lost -- they arrive on the activation the re-read + # announces, which is the ordinary live path. + await _until( + lambda: subscription.buffered == 5, + 5, + "the buffer never refilled, so the records the retraction dropped are " + "gone rather than merely deferred", + ) + assert [r.offset for r in manager.drain(RUN_ID, 1)] == [r.offset for r in placed] + + +@pytest.mark.asyncio +async def test_a_read_in_flight_across_a_reposition_cannot_be_appended( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """Reposition and append must be atomic against the prefetch epoch. + + Making the reposition synchronous on the Workflow thread is only half a fix. + The watcher captures the epoch, reads, *awaits* a user codec, and appends -- so + a check made before the append can pass and then have the reposition land in + between, putting the retracted records straight back into the buffer and + re-advancing ``prefetch_cursor`` past them. Under the reposition's own lock, + the append is either cleared by it or rejected by it, with no third + interleaving. + + The watcher is held **between its read and its append**, which is the window + that matters and the one an epoch check placed outside the lock cannot cover. + """ + placed = await append_five(backend, key) + reached_prepare = asyncio.Event() + release_prepare = asyncio.Event() + real_prepare = manager._prepare + + async def paused_prepare(records): # type: ignore[no-untyped-def] + reached_prepare.set() + await release_prepare.wait() + return await real_prepare(records) + + manager._prepare = paused_prepare # type: ignore[assignment] + + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=key) + subscription = manager.subscription(RUN_ID, 1) + assert subscription is not None + + await asyncio.wait_for(reached_prepare.wait(), 5) + assert subscription.buffered == 0, "the watcher got past its append already" + + # The Workflow thread's half, while that read is still in flight. + manager.reposition_to_committed(RUN_ID, {1: AFTER(placed_offset(placed[1]))}) + assert subscription.committed_cursor == AFTER(placed_offset(placed[1])) + + release_prepare.set() + # Give the watcher every chance to append what it read from the old position. + for _ in range(20): + await asyncio.sleep(0.01) + if subscription.buffered: + break + + assert all(r.offset != placed[0].offset for r in list(subscription._buffer)), ( + "a read that started before the reposition was appended after it, so the " + "records the marker already accounts for are back in front of the Workflow" + ) + assert subscription.prefetch_cursor != BEGINNING, ( + "the watcher never resumed from the new boundary" + ) + + # And the contract that makes the above hold whatever the interleaving, which a + # single-threaded loop cannot demonstrate on its own: the epoch is compared + # under the same lock the reposition takes, so an append carrying a retracted + # epoch is *refused* rather than merely unlikely. In production the two run on + # different threads, where "there is no await between the check and the append" + # is not an argument. + with subscription._lock: + stale_epoch = subscription._prefetch_epoch - 1 + assert subscription._append(list(placed), stale_epoch) is False, ( + "an append from a retracted epoch was accepted, so a watcher that read " + "before the reposition can undo it" + ) + assert all(r.offset != placed[0].offset for r in list(subscription._buffer)) + + +@pytest.mark.asyncio +async def test_a_replay_activation_never_delivers_a_record_the_marker_omits( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """A record the annotation does not name may not reach Workflow code. + + The dangerous shape is a *coalesced* readiness job: a replayed Run registers a + live watcher while it replays, the stream already holds records, and Core can + resolve a stream wait in the same activation that carries the replay job. If + the activation's own drain reaches the live buffer, Workflow code receives a + record no marker records -- observed downstream as a replay that either + diverges or fails as nondeterminism. + + The record here is *poison*: it is in the buffer and in no run of the + annotation. Whatever the drains do, it must not come out. This drives the + activation shape rather than a real ``activate()``, so what it covers is the + Python-side invariant; the end-to-end case is + ``test_replay_end_to_end.py``'s parked-and-evicted test. + """ + placed = await append_five(backend, key) + # The marker records the first two records, across two segments. + annotation = encode_annotation( + Annotation( + header=AnnotationHeader({1: binding(key)}), + segments=( + Segment( + (Run(1, placed[0].offset, placed[0].offset, 1),), # type: ignore[arg-type] + SegmentEndReason.BATCH_LIMIT, + ), + Segment( + (Run(1, placed[1].offset, placed[1].offset, 1),), # type: ignore[arg-type] + SegmentEndReason.NO_DATA_AVAILABLE, + ), + ), + terminal={1: AFTER(placed[1].offset)}, # type: ignore[arg-type] + ) + ) + await manager.prepare_replay(RUN_ID, annotation) + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=key) + + subscription = manager.subscription(RUN_ID, 1) + assert subscription is not None + await _until( + lambda: subscription.buffered > 0, + 5, + "the watcher never filled the live buffer, so there is no poison record " + "for a stray drain to find", + ) + + taken: list[Offset] = [] + + class ResolvingStub(DriverStub): + """Resolves pending waits on every drain, as a coalesced readiness job does.""" + + def _run_once(self, *, check_conditions: bool) -> None: + super()._run_once(check_conditions=check_conditions) + runtime.resolve_all_pending() + taken.extend(r.offset for r in runtime.drain(1)) # type: ignore[misc] + + stub = drive_with(ResolvingStub(runtime)) + + assert len(stub.drains) == 2, ( + f"two recorded segments must produce two drains, got {len(stub.drains)}" + ) + assert taken == [placed[0].offset, placed[1].offset], ( + f"the replay activation delivered something the marker does not name: {taken}" + ) + # A delta is expected -- the `subscribe()` this Run replayed is itself + # replay-visible and puts the stream in the next annotation's header. What must + # not be in it is a *run*: every record this activation handed over came from + # the marker, and re-recording those would ask Core to write a second marker + # for observations already in History. + delta = runtime.take_observation_delta() + assert delta is not None + recorded = decode_annotation(delta) + assert all(not segment.runs for segment in recorded.segments), ( + "the replay activation recorded runs of its own, so a record it delivered " + "came from the live buffer rather than from the marker" + ) + + +@pytest.mark.asyncio +async def test_a_marker_drains_once_per_segment_including_the_activations_own( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend, key: StreamKey +) -> None: + """*k* segments produce *k* drains for the whole activation, not *k + 1*. + + The driver is not the only thing that drains. The replay job lands in the + non-query job set, so the activation runs one `_run_once` for that set after + every job in it has been applied -- meaning a driver that drained all *k* + segments itself produced one drain more than the marker records. ADR-018 + requires exactly *k*: under `_single_batch_activation` each live stream + activation drove exactly one condition-checking drain, so an extra one + evaluates `wait_condition` predicates at a point that did not exist live. + + Counted across the activation rather than inside `_apply`, because inside + `_apply` the count was always right and the defect was never visible there. + An empty segment is included: it records an activation that drained and found + nothing, and it must cost a drain like any other. + """ + placed = await append_five(backend, key) + annotation = encode_annotation( + Annotation( + header=AnnotationHeader({1: binding(key)}), + segments=( + Segment( + (Run(1, placed[0].offset, placed[1].offset, 2),), # type: ignore[arg-type] + SegmentEndReason.BATCH_LIMIT, + ), + # The empty one: an activation that ran a drain and found nothing. + Segment((), SegmentEndReason.NO_DATA_AVAILABLE), + Segment( + (Run(1, placed[2].offset, placed[4].offset, 3),), # type: ignore[arg-type] + SegmentEndReason.NO_DATA_AVAILABLE, + ), + ), + terminal={1: AFTER(placed[4].offset)}, # type: ignore[arg-type] + ) + ) + await manager.prepare_replay(RUN_ID, annotation) + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=key) + + predicate_evaluations = 0 + + class ConditionCountingStub(ConsumingStub): + def _run_once(self, *, check_conditions: bool) -> None: + nonlocal predicate_evaluations + super()._run_once(check_conditions=check_conditions) + if check_conditions: + # Stands in for a `wait_condition` predicate that stays false: + # what it counts is how many times the loop asked. + predicate_evaluations += 1 + + stub = drive_with(ConditionCountingStub(runtime)) + + assert len(stub.drains) == 3, ( + f"three recorded segments must produce three drains across the whole " + f"activation, got {len(stub.drains)}" + ) + assert predicate_evaluations == 3, ( + f"a still-false wait_condition predicate must be evaluated once per " + f"recorded segment, got {predicate_evaluations}" + ) + + +# --- the annotation must match the subscriptions the Workflow makes ---------- +# +# Delivery joins a recorded run to a subscription by `wait_id`, and an integer +# is not an identity. Everything below is row four of the failure taxonomy -- +# nondeterminism, remedied by versioning the Workflow -- and specifically *not* +# `StreamIntegrityError`, which would send an operator to repair a backend that +# holds exactly the bytes it was given. + + +async def _one_record(backend: MemoryStreamBackend, key: StreamKey) -> StreamRecord: + await publish(backend, key, [f"{key.stream_name}-value"]) + return backend.all_records(key)[-1] + + +def _single_run_annotation(bindings, runs) -> bytes: # type: ignore[no-untyped-def] + """One segment holding one run per wait, and a terminal that closes it.""" + return encode_annotation( + Annotation( + header=AnnotationHeader(bindings), + segments=(Segment(tuple(runs), SegmentEndReason.NO_DATA_AVAILABLE),), + terminal={run.wait_id: AFTER(run.last_offset) for run in runs}, + ) + ) + + +@pytest.mark.asyncio +async def test_a_wait_rebound_to_another_stream_is_nondeterminism( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend +) -> None: + """Moving wait 1 from one stream to another must fail, not deliver. + + This is the silent failure the whole check exists for: with the join done on + ``wait_id`` alone, the left stream's recorded bytes are handed to Workflow + code through the right stream's subscription. Nothing anywhere reports it -- + the range validates, because the records really are where the marker says -- + and the Workflow proceeds on data it never received. + + The taxonomy puts "annotation does not match the subscriptions made" in the + nondeterminism row precisely so the remedy is versioning the Workflow rather + than repairing a healthy backend. + """ + left = StreamKey("ns", "wf", uuid.uuid4().hex, "left") + right = StreamKey("ns", "wf", uuid.uuid4().hex, "right") + recorded = await _one_record(backend, left) + await _one_record(backend, right) + + await manager.prepare_replay( + RUN_ID, + _single_run_annotation( + {1: binding(left)}, + [Run(1, recorded.offset, recorded.offset, 1)], # type: ignore[arg-type] + ), + ) + runtime = make_runtime(manager, backend) + # The Workflow's `subscribe()` calls are unchanged in number and order, so + # wait 1 is still wait 1 -- it just reads a different stream now. + runtime.register(wait_id=1, stream_key=right) + + with pytest.raises(temporalio.workflow.NondeterminismError) as caught: + drive(runtime) + + assert not isinstance(caught.value, StreamIntegrityError), ( + "the backend holds exactly what it was given; this must not be " + "reachable through the storage-failure taxonomy" + ) + assert "workflow.patched" in str(caught.value) + assert runtime.drain(1) == [], ( + "the recorded record must never reach the rebound subscription" + ) + + +@pytest.mark.asyncio +async def test_a_rebinding_made_during_the_replay_activation_is_caught_too( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend +) -> None: + """The other order, which is the ordinary one for a Run's first marker. + + An activation carrying both ``InitializeWorkflow`` and + ``ReplayExternalStreams`` applies every job before any Workflow code runs, + so the ``subscribe()`` call happens *inside* the replay driver's first + drain. Checking only what was registered beforehand would find nothing to + check on exactly the activation where the first marker is replayed. + """ + left = StreamKey("ns", "wf", uuid.uuid4().hex, "left") + right = StreamKey("ns", "wf", uuid.uuid4().hex, "right") + recorded = await _one_record(backend, left) + await _one_record(backend, right) + + await manager.prepare_replay( + RUN_ID, + _single_run_annotation( + {1: binding(left)}, + [Run(1, recorded.offset, recorded.offset, 1)], # type: ignore[arg-type] + ), + ) + runtime = make_runtime(manager, backend) + + class LateRegisteringStub(DriverStub): + def _run_once(self, *, check_conditions: bool) -> None: + super()._run_once(check_conditions=check_conditions) + runtime.register(wait_id=1, stream_key=right) + + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + + with pytest.raises(temporalio.workflow.NondeterminismError, match="workflow"): + drive_with(LateRegisteringStub(runtime)) + + +@pytest.mark.asyncio +async def test_a_removed_subscription_leaves_recorded_deliveries_unconsumed( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend +) -> None: + """Records the marker recorded may not be quietly dropped. + + Deleting a ``subscribe()`` call leaves its wait unregistered, so nothing + drains it and the prepared deliveries are discarded when the segment is + replaced. The Workflow then reaches its next command having consumed less + than History says it consumed -- which surfaces later, somewhere else, as an + unrelated command mismatch. + """ + left = StreamKey("ns", "wf", uuid.uuid4().hex, "left") + right = StreamKey("ns", "wf", uuid.uuid4().hex, "right") + left_record = await _one_record(backend, left) + right_record = await _one_record(backend, right) + + await manager.prepare_replay( + RUN_ID, + _single_run_annotation( + {1: binding(left), 2: binding(right)}, + [ + Run(1, left_record.offset, left_record.offset, 1), # type: ignore[arg-type] + Run(2, right_record.offset, right_record.offset, 1), # type: ignore[arg-type] + ], + ), + ) + runtime = make_runtime(manager, backend) + # Wait 2's `subscribe()` call is gone; wait 1 is untouched. + runtime.register(wait_id=1, stream_key=left) + + class OnlyWaitOneStub(DriverStub): + def _run_once(self, *, check_conditions: bool) -> None: + super()._run_once(check_conditions=check_conditions) + runtime.drain(1) + + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + + with pytest.raises(temporalio.workflow.NondeterminismError) as caught: + drive_with(OnlyWaitOneStub(runtime)) + + assert "[2]" in str(caught.value), ( + f"the error must name the wait whose records went undelivered: {caught.value}" + ) + + +@pytest.mark.asyncio +async def test_a_removed_subscription_that_recorded_nothing_is_nondeterminism( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend +) -> None: + """A binding with no runs is still a subscription the Workflow must make. + + The first observation of a wait carries its provider, stream key, and start + cursor even when nothing was ever delivered, so a ``subscribe()`` on a + stream that stayed quiet for that Workflow Task records exactly this: a + binding and not one run. Every delivery-side check is then vacuous -- there + are no records left over to notice -- and the only thing that says the wait + existed at all is the marker's own header. + + Removing the *last* ``subscribe()`` renumbers nothing, so no other wait's + binding disagrees either. Without a check that every recorded binding was + recreated, the removal replays clean and the Workflow reaches its next + command having subscribed to one stream fewer than History says it did. + """ + left = StreamKey("ns", "wf", uuid.uuid4().hex, "left") + quiet = StreamKey("ns", "wf", uuid.uuid4().hex, "quiet") + left_record = await _one_record(backend, left) + + annotation = encode_annotation( + Annotation( + header=AnnotationHeader({1: binding(left), 2: binding(quiet)}), + segments=( + Segment( + (Run(1, left_record.offset, left_record.offset, 1),), # type: ignore[arg-type] + SegmentEndReason.NO_DATA_AVAILABLE, + ), + ), + # Wait 2 delivered nothing, so the terminal names only wait 1. + terminal={1: AFTER(left_record.offset)}, # type: ignore[arg-type] + ) + ) + await manager.prepare_replay(RUN_ID, annotation) + runtime = make_runtime(manager, backend) + # Wait 2's `subscribe()` call is gone. Wait 1 is untouched and consumes + # exactly what the marker recorded for it. + runtime.register(wait_id=1, stream_key=left) + + with pytest.raises(temporalio.workflow.NondeterminismError) as caught: + drive(runtime) + + assert not isinstance(caught.value, StreamIntegrityError), ( + "the backend holds exactly what it was given; this must not be " + "reachable through the storage-failure taxonomy" + ) + assert "[2]" in str(caught.value), ( + f"the error must name the wait the marker recorded and the code no " + f"longer creates: {caught.value}" + ) + assert "workflow.patched" in str(caught.value), ( + "the remedy is versioning the Workflow code, exactly as for a removed timer" + ) + + +@pytest.mark.asyncio +async def test_a_removed_middle_subscription_to_the_same_stream_is_nondeterminism( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend +) -> None: + """The renumbering case no binding comparison can see. + + Three waits on one stream and one backend, with the last of them quiet for + this Workflow Task. Deleting the *middle* ``subscribe()`` renumbers what was + wait 3 down to wait 2 -- and because all three name the same stream and the + same backend, the binding check compares equal for every wait that still + exists. The delivery-side check is quiet too: the records the marker + recorded for wait 2 are taken by the subscription that used to be wait 3, so + nothing is left over. + + That is the whole failure: the marker's records reach a *different* + subscription than the one that consumed them live, with its own cursor and + its own later reads, and every check that looks at what the code did agrees + with what the code did. Only the wait the marker recorded and the code no + longer creates gives it away. + """ + key = StreamKey("ns", "wf", uuid.uuid4().hex, "tokens") + placed = await append_five(backend, key) + + annotation = encode_annotation( + Annotation( + header=AnnotationHeader( + {1: binding(key), 2: binding(key), 3: binding(key)} + ), + segments=( + Segment( + ( + Run(1, placed[0].offset, placed[0].offset, 1), # type: ignore[arg-type] + Run(2, placed[1].offset, placed[1].offset, 1), # type: ignore[arg-type] + ), + SegmentEndReason.NO_DATA_AVAILABLE, + ), + ), + # Wait 3 was subscribed and quiet: a binding, no run, no terminal + # entry. + terminal={ + 1: AFTER(placed[0].offset), # type: ignore[arg-type] + 2: AFTER(placed[1].offset), # type: ignore[arg-type] + }, + ) + ) + await manager.prepare_replay(RUN_ID, annotation) + runtime = make_runtime(manager, backend) + # The middle `subscribe()` is gone, so the Workflow now makes two waits + # where History records three -- and the survivor that was wait 3 is + # registered as wait 2. + runtime.register(wait_id=1, stream_key=key) + runtime.register(wait_id=2, stream_key=key) + + with pytest.raises(temporalio.workflow.NondeterminismError) as caught: + drive(runtime) + + assert "[3]" in str(caught.value), ( + f"the error must name the wait whose subscribe() call went missing, " + f"which is the only trace the marker leaves of it: {caught.value}" + ) + assert "workflow.patched" in str(caught.value) + + +@pytest.mark.asyncio +async def test_a_subscription_the_marker_never_bound_is_not_a_removal( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend +) -> None: + """The check is one-directional, and that is not an oversight. + + Replay does not stop at the Workflow Task the marker covers: the Workflow + runs forward, and the ``subscribe()`` calls it makes past that point belong + to the *next* marker's header rather than to this one. Requiring every + registered wait to appear in the marker would report every one of them as + nondeterminism, and would turn adding a ``subscribe()`` at the end -- a + supported change, as adding a timer at the end is -- into a Workflow that + can never replay. + """ + left = StreamKey("ns", "wf", uuid.uuid4().hex, "left") + added = StreamKey("ns", "wf", uuid.uuid4().hex, "added") + left_record = await _one_record(backend, left) + await _one_record(backend, added) + + await manager.prepare_replay( + RUN_ID, + _single_run_annotation( + {1: binding(left)}, + [Run(1, left_record.offset, left_record.offset, 1)], # type: ignore[arg-type] + ), + ) + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=left) + # The new `subscribe()`, reached because replay keeps running after the + # marker's own Workflow Task. Nothing in this marker binds it, and nothing + # should. + runtime.register(wait_id=2, stream_key=added) + + stub = drive(runtime) + + assert len(stub.drains) == 1, ( + "the marker's one segment must still replay; a wait the marker never " + "bound is not a removed subscription" + ) + + +@pytest.mark.asyncio +async def test_a_subscription_made_during_the_replay_drive_is_not_missing( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend +) -> None: + """Which is why the check runs once at the end, not before each segment. + + An activation carrying both ``InitializeWorkflow`` and + ``ReplayExternalStreams`` applies every job before any Workflow code runs, + so on a Run's first marker the ``subscribe()`` call happens *inside* the + driver's first drain -- after the segment it belongs to has already been + made current. A binding check placed on the per-segment path would fail + every such replay, which is every Workflow that consumes a stream from its + first Workflow Task. + """ + key = StreamKey("ns", "wf", uuid.uuid4().hex, "tokens") + recorded = await _one_record(backend, key) + + await manager.prepare_replay( + RUN_ID, + _single_run_annotation( + {1: binding(key)}, + [Run(1, recorded.offset, recorded.offset, 1)], # type: ignore[arg-type] + ), + ) + runtime = make_runtime(manager, backend) + + taken: list[Offset] = [] + + class LateSubscribingStub(DriverStub): + def _run_once(self, *, check_conditions: bool) -> None: + super()._run_once(check_conditions=check_conditions) + if 1 not in runtime.subscriptions(): + runtime.register(wait_id=1, stream_key=key) + taken.extend(r.offset for r in runtime.drain(1)) # type: ignore[misc] + + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + + drive_with(LateSubscribingStub(runtime)) + + assert taken == [recorded.offset], ( + f"the subscription made inside the drive must satisfy the marker's " + f"binding and take its record: {taken}" + ) + + +@pytest.mark.asyncio +async def test_a_removed_subscriptions_records_are_taken_by_nobody( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend +) -> None: + """Nothing may consume on behalf of a subscription that does not exist. + + When the removed ``subscribe()`` did record records, the delivery check is + what reports it -- the records reach the end of the replay untaken, because + a Workflow holds no handle to a wait it never made and so cannot drain it. + Consuming them anyway, as a driver that drained by whatever the segment + holds would, reports the replay clean on that count and leaves the removal + to be caught, if at all, by something else. + """ + left = StreamKey("ns", "wf", uuid.uuid4().hex, "left") + right = StreamKey("ns", "wf", uuid.uuid4().hex, "right") + left_record = await _one_record(backend, left) + right_record = await _one_record(backend, right) + + await manager.prepare_replay( + RUN_ID, + _single_run_annotation( + {1: binding(left), 2: binding(right)}, + [ + Run(1, left_record.offset, left_record.offset, 1), # type: ignore[arg-type] + Run(2, right_record.offset, right_record.offset, 1), # type: ignore[arg-type] + ], + ), + ) + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=left) + + with pytest.raises( + temporalio.workflow.NondeterminismError, match="never took" + ) as caught: + drive(runtime) + + assert "[2]" in str(caught.value), ( + f"the error must name the wait whose recorded records nobody took: " + f"{caught.value}" + ) + + +# --- provider binding --------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_backend_declaring_another_provider_is_refused_before_any_read( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend +) -> None: + """A differently configured implementation is a deployment problem. + + Neither nondeterminism -- the Workflow is unchanged -- nor integrity loss: + the store that wrote the records is simply not the one configured here. It + is refused before the read, so an implementation that cannot + interpret those offsets never gets to try. + """ + left = StreamKey("ns", "wf", uuid.uuid4().hex, "tokens") + record = await _one_record(backend, left) + + with pytest.raises(StreamStorageError, match="declares"): + await manager.prepare_replay( + RUN_ID, + _single_run_annotation( + {1: binding(left, provider_id="redis-streams")}, + [Run(1, record.offset, record.offset, 1)], # type: ignore[arg-type] + ), + ) + + assert backend.range_reads == [], ( + "the recorded range must not be read through an implementation that " + "did not write it" + ) + + +@pytest.mark.asyncio +async def test_an_unreadable_provider_format_version_is_refused_before_any_read( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend +) -> None: + """The recorded format version is checked, not merely stored. + + A provider that keeps its id while changing how it lays records out is the + case a matching id cannot catch: the offsets read back mean something else, + and every one of replay's four range checks would be applied to the wrong + interpretation of the bytes. + """ + left = StreamKey("ns", "wf", uuid.uuid4().hex, "tokens") + record = await _one_record(backend, left) + + with pytest.raises(StreamStorageError, match="format version"): + await manager.prepare_replay( + RUN_ID, + _single_run_annotation( + { + 1: binding( + left, + provider_format_version=( + MemoryStreamBackend.provider_format_version + 1 + ), + ) + }, + [Run(1, record.offset, record.offset, 1)], # type: ignore[arg-type] + ), + ) + + assert backend.range_reads == [] + + +# --- a wait registered after the header, replayed ---------------------------- + + +class FakeInstance: + """Stands in for the Workflow object the per-Run subscription state hangs off.""" + + +@pytest.mark.asyncio +async def test_a_wait_registered_mid_annotation_replays( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend +) -> None: + """The whole round trip for a subscription created after the first delta. + + `register` accepts one at any activation of a retained Workflow Task, long + after the header frame went to Core -- and Core appends deltas rather than + rewriting them. A wait bound nowhere reaches the marker as a run and a + terminal entry alone, and `prepare_replay` then has no stream key and no + backend for it, so replay of *unchanged* code fails as a wait "this + Workflow did not create". + """ + from temporalio.contrib.external_workflow_streams._runtime import ( + WorkflowStreamRuntime, + ) + + chain = uuid.uuid4().hex + first_key = StreamKey("ns", "wf", chain, "tokens") + late_key = StreamKey("ns", "wf", chain, "tool-events") + first_record = (await append_five(backend, first_key))[0] + late_record = (await append_five(backend, late_key))[0] + + recorder = WorkflowStreamRuntime( + manager=manager, + backend=backend, + run_id="recording-run", + namespace="ns", + workflow_id="wf", + first_execution_run_id=chain, + data_converter=temporalio.converter.DataConverter.default, + default_idle_timeout=timedelta(seconds=1), + ) + recorder.register(wait_id=1, stream_key=first_key) + recorder.record_delivery(1, first_record) + deltas = [recorder.take_observation_delta()] + # The header is fixed from here on: Core already holds those bytes. + recorder.register(wait_id=2, stream_key=late_key) + recorder.record_delivery(2, late_record) + deltas.append(recorder.take_observation_delta()) + annotation = b"".join(d for d in deltas if d is not None) + recorder.add_terminal() + + assert set(decode_annotation(annotation).header.streams) == {1, 2} + + plan = await manager.prepare_replay(RUN_ID, annotation) + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=first_key) + runtime.register(wait_id=2, stream_key=late_key) + + drained: dict[int, list[Offset]] = {1: [], 2: []} + + class DrainingStub(DriverStub): + def _run_once(self, *, check_conditions: bool) -> None: + super()._run_once(check_conditions=check_conditions) + for wait_id in (1, 2): + drained[wait_id].extend(r.offset for r in runtime.drain(wait_id)) # type: ignore[misc] + + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + + drive_with(DrainingStub(runtime)) + + assert drained == {1: [first_record.offset], 2: [late_record.offset]} + assert plan.committed_boundaries == { + 1: AFTER(first_record.offset), # type: ignore[arg-type] + 2: AFTER(late_record.offset), # type: ignore[arg-type] + } + + +# --- a segment's global order is the order Workflow code receives in --------- + + +@pytest.mark.asyncio +async def test_a_segment_replays_in_its_recorded_cross_stream_order( + manager: StreamSubscriptionManager, + backend: MemoryStreamBackend, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A segment recorded as (wait 2, wait 1) must not replay as (wait 1, wait 2). + + The order is reachable live and `merge` is how: a pass asks in `wait_id` + order, finds wait 1's buffer empty and wait 2's holding a record, and wait + 1's record lands from the manager's loop before the next pass. That is one + activation, so it is one segment, and the segment records (2, 1). + + On replay `merge` asks in the same `wait_id` order -- but every record of + the segment is already in hand. A drain that searched the segment for its + own wait would answer the first ask with wait 1's record and reverse the + two, which is a different sequence of values reaching Workflow code and so + a different sequence of commands. Taking only from the front is what makes + the empty answer replay as empty. + """ + chain = uuid.uuid4().hex + left = StreamKey("ns", "wf", chain, "left") + right = StreamKey("ns", "wf", chain, "right") + left_record = (await append_five(backend, left))[0] + right_record = (await append_five(backend, right))[0] + + annotation = encode_annotation( + Annotation( + header=AnnotationHeader({1: binding(left), 2: binding(right)}), + segments=( + Segment( + ( + Run(2, right_record.offset, right_record.offset, 1), # type: ignore[arg-type] + Run(1, left_record.offset, left_record.offset, 1), # type: ignore[arg-type] + ), + SegmentEndReason.NO_DATA_AVAILABLE, + ), + ), + terminal={ + 1: AFTER(left_record.offset), # type: ignore[arg-type] + 2: AFTER(right_record.offset), # type: ignore[arg-type] + }, + ) + ) + plan = await manager.prepare_replay(RUN_ID, annotation) + assert [wait_id for wait_id, _ in plan.segments[0].deliveries] == [2, 1] + + runtime = make_runtime(manager, backend) + instance = FakeInstance() + monkeypatch.setattr(temporalio.workflow, "instance", lambda: instance) + _install_runtime(instance, runtime) + + runtime.begin_replay(plan.annotation.header.streams) + try: + first = external_stream.topic("left", type=str).subscribe() + second = external_stream.topic("right", type=str).subscribe() + assert (first.wait_id, second.wait_id) == (1, 2) + + runtime.begin_replay_segment(list(plan.segments[0].deliveries)) + merged = merge(first, second) + seen = [await merged.__anext__() for _ in range(2)] + await merged.aclose() + runtime.verify_replay_consumed() + finally: + runtime.end_replay() + + assert [(subscription.wait_id, value) for subscription, value in seen] == [ + (2, "v0"), + (1, "v0"), + ], "replay reordered a segment that recorded wait 2's record before wait 1's" + + +@pytest.mark.asyncio +async def test_a_replayed_drain_stops_at_another_waits_record( + manager: StreamSubscriptionManager, backend: MemoryStreamBackend +) -> None: + """The same rule at the drain, where a wait appears twice in one segment. + + Live, the second batch was not in this wait's buffer when the first drain + ran -- the record between them belongs to a drain that had not happened + yet. Handing both over at once would collapse two of the segment's drains + into one and let the values interleave differently. + """ + chain = uuid.uuid4().hex + left = StreamKey("ns", "wf", chain, "left") + right = StreamKey("ns", "wf", chain, "right") + left_records = await append_five(backend, left) + right_records = await append_five(backend, right) + + annotation = encode_annotation( + Annotation( + header=AnnotationHeader({1: binding(left), 2: binding(right)}), + segments=( + Segment( + ( + Run(1, left_records[0].offset, left_records[0].offset, 1), # type: ignore[arg-type] + Run(2, right_records[0].offset, right_records[0].offset, 1), # type: ignore[arg-type] + Run(1, left_records[1].offset, left_records[1].offset, 1), # type: ignore[arg-type] + ), + SegmentEndReason.NO_DATA_AVAILABLE, + ), + ), + terminal={ + 1: AFTER(left_records[1].offset), # type: ignore[arg-type] + 2: AFTER(right_records[0].offset), # type: ignore[arg-type] + }, + ) + ) + await manager.prepare_replay(RUN_ID, annotation) + runtime = make_runtime(manager, backend) + runtime.register(wait_id=1, stream_key=left) + runtime.register(wait_id=2, stream_key=right) + + taken: list[tuple[int, Offset]] = [] + + class InterleavingStub(DriverStub): + def _run_once(self, *, check_conditions: bool) -> None: + super()._run_once(check_conditions=check_conditions) + for wait_id in (1, 2, 1): + taken.extend((wait_id, r.offset) for r in runtime.drain(wait_id)) # type: ignore[misc] + + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + + drive_with(InterleavingStub(runtime)) + + assert taken == [ + (1, left_records[0].offset), + (2, right_records[0].offset), + (1, left_records[1].offset), + ] diff --git a/tests/contrib/external_workflow_streams/test_replay_end_to_end.py b/tests/contrib/external_workflow_streams/test_replay_end_to_end.py new file mode 100644 index 0000000000..279b9a5dba --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_replay_end_to_end.py @@ -0,0 +1,1098 @@ +"""Replay through the real path: a live run, then its own history replayed. + +The unit tests drive the replay driver directly, which proves the mechanism but +not that a history a Worker actually produced can be fed back through it. These +run a Workflow against a live server, fetch the history it wrote, and replay it +with the real ``Replayer`` -- the same tool a user reaches for to check new code +against old histories. + +PYTEST_DONT_REWRITE: sandboxed fixture Workflows re-import this module, so pytest's +injected imports would make sandbox validation depend on pytest's import locks. +""" + +# pyright: reportMissingParameterType=false +from __future__ import annotations + +import asyncio +import uuid +from datetime import timedelta + +import pytest + +import temporalio.api.enums.v1 +import temporalio.converter +from temporalio import workflow +from temporalio.client import Client +from temporalio.contrib.external_workflow_streams._annotation import decode_annotation +from temporalio.contrib.external_workflow_streams._backend import ( + DEFAULT_WATCH_BLOCK, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec +from temporalio.contrib.external_workflow_streams._record import ( + BEGINNING, + Cursor, + Offset, + RecordKind, + StreamRecord, +) +from temporalio.worker import Replayer, Worker +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend +from tests.contrib.external_workflow_streams.test_output_worker_integration import ( + _OutputMemoryBackend, +) + +with workflow.unsafe.imports_passed_through(): + from temporalio.contrib.external_workflow_streams import external_output_stream + from temporalio.contrib.external_workflow_streams._api import external_stream + from tests.contrib.external_workflow_streams import observations + + +EXTERNAL_STREAM_MARKER = "core_external_stream" +"""The marker name Core writes a replay annotation under.""" + + +# Unsandboxed on purpose: the observation sink below has to be the *same* +# object in Workflow code as in the test. A sandboxed Workflow re-imports every +# module it touches, including a passed-through one, so it would write its +# observations into a copy and the comparison would silently compare nothing. +# Nothing here depends on the sandbox; the sandbox itself is covered elsewhere. +@workflow.defn(sandboxed=False) +class ConditionWorkflow: + """Consumes records while a ``wait_condition`` watches the same state. + + The condition is the point. It is evaluated once per event-loop drain, so it + is a direct probe of activation segmentation: if replay collapses several + recorded segments into one, the predicate sees a different sequence of + states than it did live even though the records arrive in the same order. + """ + + def __init__(self) -> None: + self._seen: list[str] = [] + self._states_observed: list[int] = [] + + @workflow.run + async def run(self, expected: int) -> list[int]: + tokens = external_stream.with_options(idle_timeout=timedelta(seconds=30)).topic( + "tokens", type=str + ) + + async def watch() -> None: + # Records what the predicate saw on every evaluation. Returning this + # rather than the records is what makes a collapsed replay visible: + # the record order would look identical either way. + await workflow.wait_condition(self._observe) + + watcher = asyncio.ensure_future(watch()) + async for token in tokens.subscribe(): + self._seen.append(token) + if len(self._seen) >= expected: + break + await watcher + # Recorded outside the Workflow as well as returned: the return value is + # only visible for the live run, and the whole question here is whether + # the *replay* saw the same thing. + observations.record(workflow.info().run_id, self._states_observed) + return self._states_observed + + def _observe(self) -> bool: + self._states_observed.append(len(self._seen)) + return len(self._seen) >= 2 + + +@workflow.defn(sandboxed=False) +class CombinedInputOutputConditionWorkflow: + """Records the drain schedule while input and output share one retained WFT.""" + + def __init__(self) -> None: + self._seen: list[str] = [] + self._states_observed: list[int] = [] + + @workflow.run + async def run(self, expected: int) -> list[int]: + tokens = external_stream.with_options(idle_timeout=timedelta(seconds=30)).topic( + "tokens", type=str + ) + output = external_output_stream.with_options( + max_publish_latency=timedelta(seconds=30) + ).topic("events", type=str) + + async def watch() -> None: + await workflow.wait_condition(self._observe) + + watcher = asyncio.ensure_future(watch()) + async for token in tokens.subscribe(): + self._seen.append(token) + await output.publish(f"observed:{token}") + if len(self._seen) >= expected: + break + await watcher + observations.record(workflow.info().run_id, self._states_observed) + return self._states_observed + + def _observe(self) -> bool: + self._states_observed.append(len(self._seen)) + return len(self._seen) >= 2 + + +@workflow.defn +class EmptyStreamWorkflow: + """Subscribes to a stream nothing is ever written to, then gives up. + + The case an implicit "start wherever" would leave unrecorded: the marker has + to carry an explicit start cursor even though no record was ever delivered, + or replay has no boundary to reproduce. + """ + + @workflow.run + async def run(self) -> str: + tokens = external_stream.with_options(idle_timeout=timedelta(seconds=1)).topic( + "tokens", type=str + ) + subscription = tokens.subscribe() + iterator = subscription.__aiter__() + try: + await asyncio.wait_for(iterator.__anext__(), 3) + except asyncio.TimeoutError: + return "nothing arrived" + return "unexpected record" + + +@workflow.defn(name="QuietSubscriptionWorkflow") +class QuietSubscriptionWorkflow: + """Creates replay-visible stream state without reading a record.""" + + @workflow.run + async def run(self) -> str: + external_stream.topic("tokens", type=str).subscribe() + return "done" + + +@workflow.defn(name="QuietSubscriptionWorkflow") +class QuietSubscriptionRemovedWorkflow: + """The next version of ``QuietSubscriptionWorkflow``, with no stream call.""" + + @workflow.run + async def run(self) -> str: + return "done" + + +@pytest.fixture +def backend() -> MemoryStreamBackend: + return MemoryStreamBackend() + + +#: Per-stream sequence, so successive publishes do not collide. `(session_id, +#: sequence)` is the idempotency key: restarting the count re-uses a key with +#: different content, which the backend contract rejects outright -- correctly, +#: and it is the test that is wrong when it happens. +_sequences: dict[StreamKey, int] = {} + + +async def publish(backend: MemoryStreamBackend, key: StreamKey, values: list[str]): + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + start = _sequences.get(key, 0) + for i, value in enumerate(values, start=start): + await backend.append( + key, + StreamRecord(RecordKind.DATA, await codec.encode(value), "producer", i), + ) + _sequences[key] = start + len(values) + + +async def stream_key_for(client: Client, handle, name: str) -> StreamKey: # type: ignore[no-untyped-def] + description = await handle.describe() + return StreamKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + name, + ) + + +async def test_replaying_a_stream_history_reproduces_the_same_observations( + client: Client, backend: MemoryStreamBackend +) -> None: + """The end-to-end form of ADR-018, through the tool users actually run. + + A collapsed replay would deliver the same records in the same order and + still be wrong: the predicate would fire a different number of times, and + any Workflow whose control flow depends on a condition would diverge. + + So "replay did not fail" is not the assertion. Replay is only obliged to + match the *commands* in History, and this Workflow's commands say nothing + about how many times its predicate ran -- a replay that delivered both + records in one drain would produce the same completion and pass a + failure-only check. What is compared instead is the sequence the predicate + itself observed, live against replayed, for a marker that spans several + activations. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[ConditionWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + ConditionWorkflow.run, + 2, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + key = await stream_key_for(client, handle, "tokens") + await asyncio.sleep(1) + # Two separate arrivals, so the marker spans more than one activation. + await publish(backend, key, ["alpha"]) + await asyncio.sleep(0.5) + await publish(backend, key, ["beta"]) + + live = await asyncio.wait_for(handle.result(), 60) + history = await handle.fetch_history() + run_id = handle.first_execution_run_id or handle.result_run_id + + markers = [ + e + for e in history.events + if e.HasField("marker_recorded_event_attributes") + and e.marker_recorded_event_attributes.marker_name == EXTERNAL_STREAM_MARKER + ] + assert markers, "no stream marker was written, so there is nothing to replay" + + replayer = Replayer( + workflows=[ConditionWorkflow], + external_stream_backend=backend, + ) + result = await replayer.replay_workflow(history) + + assert result.replay_failure is None, ( + f"replaying the history the Worker just wrote failed: {result.replay_failure}" + ) + assert live, "the live run observed nothing, so this proves nothing" + assert len(live) > 1, ( + "the predicate ran only once live, so a collapsed replay would look " + f"identical and this proves nothing: {live}" + ) + + assert run_id is not None + runs = observations.executions(run_id) + assert len(runs) == 2, ( + "expected exactly two executions to be recorded -- the live run and the " + f"replay -- got {len(runs)}" + ) + live_observed, replayed_observed = runs + assert live_observed == live, ( + "the live run recorded something other than it returned" + ) + assert replayed_observed == live_observed, ( + "the replayed run's predicate saw a different sequence of states than " + "the live one did. The records arrive in the same order either way, so " + "this is the segmentation itself diverging: recorded activation " + f"boundaries were collapsed or re-cut. live={live_observed} " + f"replayed={replayed_observed}" + ) + + +async def test_live_and_replay_share_one_input_output_drain_schedule( + client: Client, +) -> None: + """A real shared marker reproduces one drain sequence, not two drivers.""" + backend = _OutputMemoryBackend() + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[CombinedInputOutputConditionWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + CombinedInputOutputConditionWorkflow.run, + 2, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + key = await stream_key_for(client, handle, "tokens") + await asyncio.sleep(0.5) + await publish(backend, key, ["alpha"]) + # The first readiness activation must finish and retain the same WFT + # before the second arrival, so the marker records two live drains. + await asyncio.sleep(0.25) + await publish(backend, key, ["beta"]) + + live = await asyncio.wait_for(handle.result(), 60) + history = await handle.fetch_history() + run_id = handle.first_execution_run_id or handle.result_run_id + + markers = stream_markers(history.events) + assert len(markers) == 1 + marker_event = markers[0] + + from temporalio.bridge.proto.external_data import ExternalStreamMarkerData + + marker = ExternalStreamMarkerData() + marker.ParseFromString( + marker_event.marker_recorded_event_attributes.details["external_stream"] + .payloads[0] + .data + ) + assert marker.HasField("output") + assert marker.waits + annotation = decode_annotation(marker.replay_annotation) + assert len(annotation.segments) >= 3 + assert len(marker.output.segments) == len(annotation.segments) + assert ( + sum(sum(segment.record_counts_by_topic) for segment in marker.output.segments) + == 2 + ) + assert sum(bool(segment.runs) for segment in annotation.segments) == 2 + + result = await Replayer( + workflows=[CombinedInputOutputConditionWorkflow], + external_stream_backend=backend, + ).replay_workflow(history) + assert result.replay_failure is None, ( + f"combined input/output replay failed: {result.replay_failure}" + ) + assert len(live) > 1 + assert run_id is not None + runs = observations.executions(run_id) + assert len(runs) == 2 + live_observed, replayed_observed = runs + assert live_observed == live + assert replayed_observed == live_observed, ( + "the shared input/output replay driver changed the live drain schedule: " + f"live={live_observed}, replay={replayed_observed}" + ) + assert len(backend.stages) == 1, "replay must not stage output again" + + +async def test_a_history_with_stream_markers_needs_its_backend_to_replay( + client: Client, backend: MemoryStreamBackend +) -> None: + """Replay re-reads the recorded ranges, so it needs the provider. + + Without the option there is no way to supply one, and the replay fails the + way a Worker with no backend would -- correct, but a configuration error + rather than a finding about the history. Asserted so the option cannot + quietly stop being threaded through. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[ConditionWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + ConditionWorkflow.run, + 2, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + key = await stream_key_for(client, handle, "tokens") + await asyncio.sleep(1) + await publish(backend, key, ["alpha", "beta"]) + await asyncio.wait_for(handle.result(), 60) + history = await handle.fetch_history() + + without = await Replayer(workflows=[ConditionWorkflow]).replay_workflow( + history, raise_on_replay_failure=False + ) + + assert without.replay_failure is not None + assert "external_stream_backend" in str(without.replay_failure), ( + "the failure must name the missing option rather than surface as an " + f"attribute error: {without.replay_failure}" + ) + + +async def test_a_no_backend_replayer_validates_a_removed_quiet_subscription( + client: Client, backend: MemoryStreamBackend +) -> None: + """History, not current Worker configuration, decides whether replay runs. + + A quiet subscription still writes its binding into the marker. The next + Workflow version removes the only ``subscribe()`` call while leaving its + ordinary Temporal command sequence unchanged, and the Replayer deliberately + has no backend configured. It must still decode the marker and apply the reverse + binding check rather than accepting the history because no live stream + runtime happened to be configured. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[QuietSubscriptionWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + QuietSubscriptionWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + assert await asyncio.wait_for(handle.result(), 30) == "done" + history = await handle.fetch_history() + + markers = [ + event + for event in history.events + if event.HasField("marker_recorded_event_attributes") + and event.marker_recorded_event_attributes.marker_name == EXTERNAL_STREAM_MARKER + ] + assert markers, "the quiet subscription produced no marker to validate" + + result = await Replayer( + workflows=[QuietSubscriptionRemovedWorkflow] + ).replay_workflow(history, raise_on_replay_failure=False) + + assert result.replay_failure is not None, ( + "removing the recorded subscription replayed successfully when the " + "Replayer had no backend configured" + ) + assert "never created" in str(result.replay_failure), ( + "the marker's reverse binding check did not report the removed quiet " + f"subscription: {result.replay_failure}" + ) + + +async def test_an_empty_stream_replays_from_its_recorded_boundary( + client: Client, backend: MemoryStreamBackend +) -> None: + """A subscription that never received a record still has a boundary. + + This is the case the explicit start cursor exists for: the marker records + where the subscription began even though nothing was delivered, so replay + reproduces an empty observation rather than resolving a position from + whatever the stream holds by then. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[EmptyStreamWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + EmptyStreamWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + assert await asyncio.wait_for(handle.result(), 60) == "nothing arrived" + history = await handle.fetch_history() + key = await stream_key_for(client, handle, "tokens") + + # Records land *after* the Run finished. Replay must not see them: the + # marker's boundary is where the subscription was, not where the stream got + # to afterwards. + await publish(backend, key, ["late-one", "late-two"]) + + result = await Replayer( + workflows=[EmptyStreamWorkflow], + external_stream_backend=backend, + ).replay_workflow(history, raise_on_replay_failure=False) + + assert result.replay_failure is None, ( + "replay resolved a position from live stream state rather than from the " + f"recorded boundary: {result.replay_failure}" + ) + + +#: ``run_id -> [whether each execution of the run body began while replaying]``. +#: +#: A Run that was evicted and came back executes its body from the top a second +#: time, replaying; one that merely stayed cached never does. Written from +#: Workflow code, which is why the Workflow below is unsandboxed: a sandboxed +#: one would write into its own copy of this module and the count would stay at +#: one however many times the Run was rebuilt. +STARTS: dict[str, list[bool]] = {} + + +@workflow.defn(sandboxed=False) +class EmptyStreamParkWorkflow: + """Subscribes to an empty stream, parks on it, and consumes what arrives later. + + Nothing else is pending -- no timer, no other command -- so the first + Workflow Task is retained until the idle timeout parks it. Both halves of + that matter here: the park is what writes the empty marker, and it is also + what leaves the Run resumable *only* by a wake Signal, so the records + published while it is evicted cannot be delivered by a watcher that happened + to survive. + """ + + def __init__(self) -> None: + self._seen: list[str] = [] + + @workflow.run + async def run(self, expected: int) -> list[str]: + STARTS.setdefault(workflow.info().run_id, []).append( + workflow.unsafe.is_replaying() + ) + tokens = external_stream.with_options(idle_timeout=timedelta(seconds=1)).topic( + "tokens", type=str + ) + async for token in tokens.subscribe(): + self._seen.append(token) + if len(self._seen) >= expected: + break + # Recorded outside the Workflow as well as returned: the return value + # exists only for the live run, and what the *replays* observed is half + # the question here. + observations.record(workflow.info().run_id, self._seen) + return self._seen + + +@workflow.defn +class CacheFillerWorkflow: + """Occupies the one cache slot, which is what evicts the Run under test.""" + + @workflow.run + async def run(self) -> str: + return "done" + + +class RecordedRangesOnlyBackend(MemoryStreamBackend): + """Serves recorded ranges faithfully and poisons every live read. + + Replay is obliged to read the ranges its markers name and nothing else, so + those ranges read back exactly as written while the *live* view of the same + stream answers with records no marker ever recorded. Any position or + delivery replay took from live state then shows up as a record the live Run + never saw, rather than as nothing -- which is the difference that matters + for an empty marker, whose live-resolved position would otherwise hand the + Workflow the very records it was about to receive anyway and look correct. + + The poison is served once. A live view that answered forever would fill the + buffer and say nothing more about it than the first answer already does. + + A watcher still runs during replay -- the Run goes on living once the last + marker is exhausted, and something has to be watching for that -- so + ``read_after`` is answered rather than refused. Writes and the park + handshake are recorded rather than refused for a reason worth keeping: an + exception raised inside a watcher is caught by the manager's own retry loop + and logged, so a refusal there would be asserted into a log line nobody + reads. + """ + + def __init__(self, source: MemoryStreamBackend, poison: list[StreamRecord]) -> None: + super().__init__() + # The same storage the live Run wrote, so the recorded ranges -- and + # only those -- still read back through `read_range`. + self._records = source._records + self._by_key = source._by_key + self._poison = poison + self._poison_served = False + #: Every provider call replay made that was not a recorded-range read. + self.live_calls: list[str] = [] + + async def read_after( # type: ignore[override] + self, + key: StreamKey, + after: Cursor, + *, + max_records: int, + block: timedelta | None = DEFAULT_WATCH_BLOCK, + ) -> list[StreamRecord]: + self.live_calls.append("read_after") + if not self._poison_served: + self._poison_served = True + return list(self._poison[:max_records]) + if block is not None: + # Blocked rather than returned empty immediately: a watcher whose + # read returns instantly spins, and the busy loop would be this + # test's own doing. + await asyncio.sleep(block.total_seconds()) + return [] + + async def append(self, key: StreamKey, record: StreamRecord) -> StreamRecord: + self.live_calls.append("append") + raise AssertionError("replay must never append to the stream") + + async def install_park_intent(self, key, intent) -> None: # type: ignore[no-untyped-def] + self.live_calls.append("install_park_intent") + + async def recheck(self, key: StreamKey, wait_id: int) -> bool: + self.live_calls.append("recheck") + return False + + +def stream_markers(events) -> list: # type: ignore[no-untyped-def] + return [ + e + for e in events + if e.HasField("marker_recorded_event_attributes") + and e.marker_recorded_event_attributes.marker_name == EXTERNAL_STREAM_MARKER + ] + + +def marker_annotation(marker_event): # type: ignore[no-untyped-def] + """The decoded annotation a marker carries.""" + from temporalio.bridge.proto.external_data import ExternalStreamMarkerData + + data = ExternalStreamMarkerData() + data.ParseFromString( + marker_event.marker_recorded_event_attributes.details["external_stream"] + .payloads[0] + .data + ) + return decode_annotation(data.replay_annotation) + + +async def wait_for_markers( + handle, count: int, message: str, timeout: float = 30 +) -> None: # type: ignore[no-untyped-def] + """Polls until History holds ``count`` stream markers. + + Polled rather than slept: a marker is written when the Workflow Task ends, + and waiting a fixed time instead would make every assertion after it a + statement about how fast this machine is. + """ + deadline = asyncio.get_running_loop().time() + timeout + while asyncio.get_running_loop().time() < deadline: + if ( + len(stream_markers([e async for e in handle.fetch_history_events()])) + >= count + ): + return + await asyncio.sleep(0.3) + raise AssertionError(message) + + +async def wait_until(predicate, message: str, timeout: float = 15) -> None: # type: ignore[no-untyped-def] + deadline = asyncio.get_running_loop().time() + timeout + while asyncio.get_running_loop().time() < deadline: + if predicate(): + return + await asyncio.sleep(0.1) + raise AssertionError(message) + + +async def _stall_diagnosis(client, handle, backend, key, run_id) -> str: # type: ignore[no-untyped-def] + """What the Run, its History, and the stream look like when the resume stalls. + + The failure is a 45-second timeout, and a timeout on its own cannot tell + "no Workflow Task was ever created" from "one was created and delivered + nothing" -- which are different defects in different components. Each line + below separates a pair of them. + """ + lines: list[str] = [] + try: + events = [event async for event in handle.fetch_history_events()] + kinds: dict[str, int] = {} + for event in events: + kinds[event.WhichOneof("attributes") or "?"] = ( + kinds.get(event.WhichOneof("attributes") or "?", 0) + 1 + ) + lines.append(f"history: {len(events)} event(s), {kinds}") + # A Workflow Task after the wake is the difference between "the wake never + # produced one" and "it did, and the Run delivered nothing". + started = [ + i + for i, e in enumerate(events) + if e.HasField("workflow_task_started_event_attributes") + ] + lines.append(f"workflow tasks started at event indices: {started}") + failures = [ + e.workflow_task_failed_event_attributes + for e in events + if e.HasField("workflow_task_failed_event_attributes") + ] + lines.append( + f"workflow task failures: {[(f.cause, f.failure.message[:200]) for f in failures]}" + ) + lines.append(f"stream markers in history: {len(stream_markers(events))}") + except Exception as err: # noqa: BLE001 -- diagnosis must not mask the failure + lines.append(f"history unavailable: {err!r}") + try: + description = await handle.describe() + lines.append(f"workflow status: {description.status}") + except Exception as err: # noqa: BLE001 + lines.append(f"describe failed: {err!r}") + try: + records = backend.all_records(key) + lines.append( + f"stream holds {len(records)} record(s) at " + f"{[str(r.offset) for r in records]}" + ) + parked = await backend.parked_wait_ids(key) + lines.append(f"parked wait ids: {parked}") + for wait_id in parked: + intent = await backend.park_intent(key, wait_id) + lines.append(f" wait {wait_id} intent: {intent}") + except Exception as err: # noqa: BLE001 + lines.append(f"backend state unavailable: {err!r}") + lines.append(f"workflow body starts (is_replaying per start): {STARTS.get(run_id)}") + lines.append(f"observations for this Run: {observations.executions(run_id)}") + # The readiness handshake, which is the only thing that can turn a buffered + # record into a Workflow Task. Its answers separate "the watcher never read + # the records" from "it read them and Core had nowhere to put them". + lines.append(f"readiness calls: {READINESS_TRACE}") + for manager, subscription in _live_subscriptions(): + lines.append( + f"subscription {subscription.run_id}/{subscription.wait_id}: " + f"buffered={subscription.buffered} " + f"committed={subscription.committed_cursor} " + f"delivery={subscription.delivery_cursor} " + f"prefetch={subscription.prefetch_cursor} " + f"generation={subscription.current_wait_generation()} " + f"wakes_owed={subscription.wakes_owed} " + f"cancelled={subscription._cancelled} " + f"watcher_done={None if subscription._watcher is None else subscription._watcher.done()}" + ) + del manager + return "Stall diagnosis:\n " + "\n ".join(lines) + + +#: Every readiness call the Worker's manager made, with what Core answered. +READINESS_TRACE: list[tuple[str, int, int, object]] = [] + +#: Managers alive in this process, so a stalled test can read their state. +LIVE_MANAGERS: list[object] = [] + + +def _live_subscriptions(): # type: ignore[no-untyped-def] + """Every subscription still registered on any manager in this process.""" + found = [] + for manager in LIVE_MANAGERS: + for run_subscriptions in getattr(manager, "_runs", {}).values(): + for subscription in run_subscriptions.values(): + found.append((manager, subscription)) + return found + + +@pytest.fixture +def trace_readiness(monkeypatch: pytest.MonkeyPatch): # type: ignore[no-untyped-def] + """Records what Core answered for every readiness call, and every manager. + + A stall in this file is always "a buffered record never became a Workflow + Task", and the readiness answers are where that happens. Without them the + failure cannot distinguish a watcher that never read from a Worker that read + and was told there was nowhere to deliver. + """ + from temporalio.contrib.external_workflow_streams._manager import ( + StreamSubscriptionManager, + ) + + READINESS_TRACE.clear() + LIVE_MANAGERS.clear() + + original_init = StreamSubscriptionManager.__init__ + original_notify = StreamSubscriptionManager._notify_ready_with_retries + + def spy_init(self, *args, **kwargs): # type: ignore[no-untyped-def] + original_init(self, *args, **kwargs) + LIVE_MANAGERS.append(self) + + async def spy_notify(self, subscription): # type: ignore[no-untyped-def] + result = await original_notify(self, subscription) + READINESS_TRACE.append( + ( + subscription.run_id[-8:], + subscription.wait_id, + subscription.current_wait_generation(), + result, + ) + ) + return result + + monkeypatch.setattr(StreamSubscriptionManager, "__init__", spy_init) + monkeypatch.setattr( + StreamSubscriptionManager, "_notify_ready_with_retries", spy_notify + ) + yield + LIVE_MANAGERS.clear() + + +async def test_an_empty_stream_parked_and_evicted_replays_from_the_recorded_cursor( + client: Client, + backend: MemoryStreamBackend, + monkeypatch: pytest.MonkeyPatch, + trace_readiness: None, # pyright: ignore[reportUnusedParameter] +) -> None: + """The empty boundary, carried across a park, an eviction, and two replays. + + Four things happen in order, and each is asserted rather than assumed -- + which is the whole difficulty of this case, because a test that quietly + skipped the park or the eviction would still pass on a Run that simply sat + in cache: + + 1. **Subscribe to an empty stream.** Nothing is ever published before the + eviction, so the subscription's entire first Workflow Task observes + nothing. + 2. **Park.** Asserted in the backend, where the park handshake installs its + intent, and not by waiting out the idle timeout and hoping: the intent's + own cursor is the empty boundary, ``BEGINNING``. + 3. **Evict.** One cache slot, so the filler Workflow pushes the Run out; the + manager's eviction of the Run is observed directly, and History is + checked at that boundary before a later wake can legitimately race the + terminal command. + 4. **Replay.** The records are published while the Run is *gone*, so the + cursor the resumed Run starts from cannot have come from anything still + running on this Worker. It starts where the marker says -- ``BEGINNING``, + the boundary the empty subscription recorded -- so it receives both, and + a start resolved from live backend state would have started at the tail + and received neither. + + Then the marker itself is decoded, and the finished history is replayed + twice: once against the real backend with records appended *after* the Run + finished, which replay must not see, and once against a backend whose only + working operation is the recorded-range read, which nothing but the marker's + own ranges can satisfy. + """ + from temporalio.contrib.external_workflow_streams._manager import ( + StreamSubscriptionManager, + ) + from temporalio.contrib.external_workflow_streams._wake import ( + UNPARKED_WAKE_GENERATION, + WakeRequest, + send_wake_signal, + ) + + #: Which Runs the Worker **finished** tearing down. `RemoveFromCache` is the + #: only thing that reaches this while a Worker is running, so it is the + #: eviction itself rather than a symptom of one. + #: + #: Recorded *after* the teardown is awaited, not before. Before, this said only + #: that teardown had **started**: the Run's watcher was still alive, and the + #: records this test publishes next could be read into a buffer that was about + #: to be discarded -- their readiness reported by a watcher with nothing left to + #: deliver it to. What the test needs to know is that nothing is watching any + #: more, which is the *post*-condition of this call. + evicted: list[str] = [] + evict_run = StreamSubscriptionManager.evict_run + + async def spy_on_eviction(self, run_id: str) -> None: # type: ignore[no-untyped-def] + await evict_run(self, run_id) + evicted.append(run_id) + + monkeypatch.setattr(StreamSubscriptionManager, "evict_run", spy_on_eviction) + + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[EmptyStreamParkWorkflow, CacheFillerWorkflow], + external_stream_backend=backend, + # One slot, so running anything else evicts the Run under test. + max_cached_workflows=1, + max_concurrent_workflow_tasks=2, + ): + handle = await client.start_workflow( + EmptyStreamParkWorkflow.run, + 2, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + key = await stream_key_for(client, handle, "tokens") + run_id = key.first_execution_run_id + + # --- 1 and 2: an empty subscription, parked ------------------------ + await wait_for_markers( + handle, + 1, + "the Run never ended a Workflow Task, so it cannot have parked on " + "the empty stream", + ) + parked = await backend.parked_wait_ids(key) + assert parked, ( + "no park intent was installed, so the Run never actually parked -- " + "it ended its Workflow Task some other way and the eviction below " + "would be evicting a Run in the wrong state" + ) + intent = await backend.park_intent(key, parked[0]) + assert intent is not None + assert intent.cursor == BEGINNING, ( + "the subscription parked at " + f"{intent.cursor} rather than at the beginning of the stream, so " + "the boundary it recorded is not the empty one this case is about" + ) + assert not backend.all_records(key), ( + "the stream was not empty while the Run was subscribed to it, so " + "this is no longer the empty-boundary case" + ) + + # --- 3: eviction --------------------------------------------------- + await client.execute_workflow( + CacheFillerWorkflow.run, + id=f"filler-{uuid.uuid4()}", + task_queue=task_queue, + ) + await wait_until( + lambda: run_id in evicted, + "the filler Workflow did not evict the Run: it is still cached, so " + "what follows would be a live resume rather than a replay", + ) + history_at_eviction = await handle.fetch_history() + assert not [ + event + for event in history_at_eviction.events + if event.HasField("workflow_task_failed_event_attributes") + ], ( + "a Workflow Task failed before the explicit cache eviction completed, " + "so the next execution would not prove that eviction caused the replay" + ) + assert STARTS.get(run_id) == [False], ( + "the Workflow body restarted before the explicit cache eviction boundary; " + f"the next replay would be ambiguous: {STARTS.get(run_id)}" + ) + + # --- 4: published while the Run is gone, then woken ---------------- + # Nothing on this Worker is watching for these: the Run is parked and + # evicted, its watcher torn down with it. The wake Signal is the only + # thing that can create a Workflow Task, and it carries no records. + await publish(backend, key, ["alpha", "beta"]) + # The unparked generation rather than the one the intent names, which is + # not the shape it looks like. A wake that names a park generation is + # identified *by* that generation -- its request ID deliberately ignores + # the sender, so every producer racing to wake one park is deduplicated + # to a single Workflow Task. That is correct for the producer's own wake + # and fatal for the Worker's follow-up: this Run's records land in the + # buffer just after the replayed Workflow Task closes, the Worker is + # answered `NoOpenWorkflowTask`, and the wake it then owes re-derives the + # very request ID the wake above already used. The server deduplicates + # it, no Workflow Task is created, and the Run waits forever on records + # it is already holding. Generation 0 keeps that follow-up wake + # distinguishable, which leaves this test about the cursor rather than + # about the wake protocol (ADR-023). + await send_wake_signal( + client, + WakeRequest( + namespace=client.namespace, + workflow_id=handle.id, + first_execution_run_id=run_id, + stream_name="tokens", + wait_id=intent.wait_id, + park_generation=UNPARKED_WAKE_GENERATION, + sender_identity=f"test-{uuid.uuid4()}", + wake_counter=1, + ), + ) + + try: + live = await asyncio.wait_for(handle.result(), 45) + except asyncio.TimeoutError: + # The same finding as the assertion below, which a Run that never + # finishes would otherwise report as a bare timeout -- reported with + # enough state to tell the two ways it can happen apart, because + # "it timed out" sends the next reader back to square one. + raise AssertionError( + "the resumed Run never received the records published while it " + "was evicted. They were in the stream before it came back and " + "its marker's boundary is the beginning of that stream, so a " + "Run starting where the marker says receives both; one that " + "resolved its position from live backend state starts at the " + "tail and waits for records that are already behind it.\n\n" + + await _stall_diagnosis(client, handle, backend, key, run_id) + ) from None + history = await handle.fetch_history() + + assert live == ["alpha", "beta"], ( + "the resumed Run did not start from the boundary its empty marker " + "recorded. Both records were published before it came back, so a Run " + "starting at BEGINNING sees both; one that resolved its position from " + f"live backend state starts at the tail and sees neither. Got {live}" + ) + + starts = STARTS.get(run_id, []) + assert len(starts) >= 2, ( + "the Workflow body ran only once, so the Run was never rebuilt and " + "nothing was replayed -- the eviction did not take effect" + ) + assert starts[0] is False and starts[1] is True, ( + "the second execution of the Workflow body did not begin in replay, so " + f"the Run came back some way other than by replaying its history: {starts}" + ) + unexpected_wft_failures = [ + event + for event in history.events + if event.HasField("workflow_task_failed_event_attributes") + and event.workflow_task_failed_event_attributes.cause + != temporalio.api.enums.v1.WorkflowTaskFailedCause.WORKFLOW_TASK_FAILED_CAUSE_UNHANDLED_COMMAND + ] + assert not unexpected_wft_failures, ( + "the resumed Run had a Workflow Task failure other than the expected " + "external-wake/terminal-command race: " + f"{unexpected_wft_failures}" + ) + + # --- the recorded empty boundary ----------------------------------------- + markers = stream_markers(history.events) + assert len(markers) >= 2, ( + "expected a marker for the parked Workflow Task and one for the " + f"Workflow Task that consumed, got {len(markers)}" + ) + empty = marker_annotation(markers[0]) + binding = empty.header.streams.get(intent.wait_id) + assert binding is not None, ( + "the empty subscription's marker records no stream in its header, so " + "replay of it has nothing to start from" + ) + assert binding.stream_key == key + assert binding.start_cursor == BEGINNING, ( + "the marker for a subscription that received nothing must still carry " + "an explicit start cursor, or replay resolves a position from whatever " + f"the stream holds by then; got {binding.start_cursor}" + ) + assert not [run for segment in empty.segments for run in segment.runs], ( + "the marker written for the empty subscription records deliveries, so " + "it is not the empty boundary" + ) + assert empty.terminal == {intent.wait_id: BEGINNING}, ( + "the parked Workflow Task's terminal must say the subscription stopped " + f"at the beginning of the stream; got {empty.terminal}" + ) + + # --- replay 1: the real backend, with records appended afterwards -------- + # They land after the Run finished. Replay must not see them: the marker's + # boundary is where the subscription was, not where the stream got to. + await publish(backend, key, ["late-one", "late-two"]) + against_live = await Replayer( + workflows=[EmptyStreamParkWorkflow], + external_stream_backend=backend, + ).replay_workflow(history, raise_on_replay_failure=False) + assert against_live.replay_failure is None, ( + f"replaying the history the Worker wrote failed: {against_live.replay_failure}" + ) + + # --- replay 2: recorded ranges and nothing else -------------------------- + # The late records are the poison: they are in the stream, they are in no + # recorded range, and the live Run never saw one. A replay that took a + # position or a delivery from live state hands them to Workflow code, and + # the observation below is then not the one the live Run recorded. + late = backend.all_records(key)[-2:] + recorded_only = RecordedRangesOnlyBackend(backend, late) + against_recorded = await Replayer( + workflows=[EmptyStreamParkWorkflow], + external_stream_backend=recorded_only, + ).replay_workflow(history, raise_on_replay_failure=False) + assert against_recorded.replay_failure is None, ( + "replay could not reproduce the Run from the ranges its markers name, " + "so something in it was reading the stream live: " + f"{against_recorded.replay_failure}" + ) + assert [call for call in recorded_only.live_calls if call != "read_after"] == [], ( + "replay wrote to the stream or ran the park handshake, neither of which " + f"it may do: {recorded_only.live_calls}" + ) + expected_ranges: list[tuple[Offset, Offset]] = [ + (run.first_offset, run.last_offset) + for marker in markers + for segment in marker_annotation(marker).segments + for run in segment.runs + ] + assert recorded_only.range_reads == expected_ranges, ( + "replay read ranges the markers do not name -- the empty marker names " + "none at all, so anything read for it came from live stream state: " + f"{recorded_only.range_reads} against {expected_ranges}" + ) + + runs = observations.executions(run_id) + assert len(runs) >= 3, ( + "expected at least the live Run and both explicit replays to record what " + f"they observed, got {len(runs)}" + ) + assert all(run == live for run in runs), ( + "a Workflow Task execution observed something other than the completed " + "Run did. An UnhandledCommand retry may add an execution, but records were " + "appended after the Run finished and every replay could reach them " + f"through the same provider: {runs}" + ) diff --git a/tests/contrib/external_workflow_streams/test_rollover_integration.py b/tests/contrib/external_workflow_streams/test_rollover_integration.py new file mode 100644 index 0000000000..6c156fb0b4 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_rollover_integration.py @@ -0,0 +1,619 @@ +"""Workflow Task rollover, against a live server. + +A retained Workflow Task is bounded by the server's Workflow Task timeout, and a +continuously fed stream whose inter-record gaps stay below the idle timeout never +reaches the idle parking path at all. Without rollover such a task does not +merely stay open too long -- it *fails*, and everything consumed on it is lost +with it. + +The Core tests prove the deadline exists. What only a live server can show is +that it fires *while a stream is being fed*: that a task genuinely held open by +arriving records is handed on before the server times it out, that the +subscription survives that hand-off, and that inputs queued behind the retained +task -- Signals especially -- are released no later than the deadline. + +Every test here stops short of the Workflow Task timeout on purpose. A retained +task that reaches it is not merely a failed assertion: the server schedules a +replacement, Core receives a Workflow Task for a run it still holds one for, and +the resulting ``dbg_panic`` takes the whole workflow-processing thread with it. + +PYTEST_DONT_REWRITE: sandboxed fixture Workflows re-import this module, so pytest's +injected imports would make sandbox validation depend on pytest's import locks. +""" + +from __future__ import annotations + +import asyncio +import uuid +from datetime import timedelta +from typing import Any + +import pytest + +from temporalio import workflow +from temporalio.client import Client +from temporalio.contrib.external_workflow_streams._backend import StreamKey +from temporalio.service import RPCError, RPCStatusCode +from temporalio.worker import Worker +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend +from tests.contrib.external_workflow_streams.test_worker_integration import publish + +with workflow.unsafe.imports_passed_through(): + from temporalio.contrib.external_workflow_streams._api import external_stream + +TASK_TIMEOUT = timedelta(seconds=10) +"""The bound a retained task must stay inside, and what the deadline derives +from. Long enough that the rollover deadline is unambiguous, short enough that +these tests do not have to run for a minute.""" + +ROLLOVER_FRACTION = 0.8 +"""Core arms the rollover deadline at this fraction of the Workflow Task +timeout, the same one the local-activity heartbeat uses.""" + +ROLLOVER_DEADLINE = TASK_TIMEOUT * ROLLOVER_FRACTION + +FEED_GAP_SECONDS = 0.3 +"""Well below the idle timeout, so the runtime never parks and the only thing +that can release the task is rollover.""" + +WAKE_SIGNAL_NAME = "__temporal_external_stream_wake" + +ROLLOVER_DELIVERY_TOLERANCE = timedelta(milliseconds=500) +"""How far past the deadline a released input may land and still count as bounded by it. + +The deadline is what Core aims the replacement Workflow Task at, not an instant +the server can hit. What a test can measure is the gap between two +``WorkflowTaskStarted`` event times, and between those sit Core's completion +RPC, the server scheduling the replacement, and a poll picking it up. Measured +overheads here are 16-65ms, so half a second is most of an order of magnitude of +headroom. + +It is deliberately far short of ``TASK_TIMEOUT - ROLLOVER_DEADLINE``, which is +two seconds: an input released because the *server* timed the retained task out +lands at about 10s and still fails the assertion, and telling that apart from a +rollover is the only thing the assertion is for. +""" + +POST_ROLLOVER_PAUSE = timedelta(seconds=10) +"""How long the post-rollover Workflow Task boundary is held open for. + +Long enough for an append and the wake it owes to be observed in History +without waiting on anything else, short enough that the test finishes. +""" + + +@workflow.defn +class SteadyStreamWorkflow: + """Consumes records until it has the number it was asked for. + + Nothing else it does can end a Workflow Task: no timers, no activities, and + the gaps the test feeds are far under the idle timeout, so retention is only + ever released by rollover. + """ + + def __init__(self) -> None: + self._seen = 0 + + @workflow.run + async def run(self, expected: int) -> int: + tokens = external_stream.topic("tokens", type=str) + async for _ in tokens.subscribe(): + self._seen += 1 + if self._seen >= expected: + break + return self._seen + + @workflow.query + def seen(self) -> int: + return self._seen + + +@workflow.defn +class SignalledStreamWorkflow: + """Consumes a stream until a Signal arrives, then reports when it did. + + The Signal is sent while the Workflow Task is retained, so the server has to + hold it until that task completes. What it returns is *Workflow* time, taken + the moment the handler ran, which is what makes "no later than the rollover + deadline" measurable from the outside rather than inferred from wall clock. + """ + + def __init__(self) -> None: + self._signalled: float | None = None + self._started: float | None = None + + @workflow.run + async def run(self) -> float: + tokens = external_stream.topic("tokens", type=str) + self._started = workflow.now().timestamp() + subscription = tokens.subscribe() + + async def consume() -> None: + async for _ in subscription: + pass + + consumer = asyncio.ensure_future(consume()) + await workflow.wait_condition(lambda: self._signalled is not None) + consumer.cancel() + assert self._signalled is not None and self._started is not None + return self._signalled - self._started + + @workflow.signal + def poke(self) -> None: + self._signalled = workflow.now().timestamp() + + +@workflow.defn +class RolloverThenPauseWorkflow: + """Consumes across a rollover, then pauses on a Timer and consumes once more. + + The pause is what makes the post-rollover window addressable. A completion + carrying a server-bound command cannot ask for retention, so the Workflow + Task ends there with the subscription still active and still unparked -- + which is the state a rollover leaves behind, held still. The rollover's own + window is not addressable from outside the process: Core completes a + rollover with ``force_new_wft``, and the server has a replacement Workflow + Task started within a millisecond of it. + """ + + def __init__(self) -> None: + self._seen = 0 + + @workflow.run + async def run(self, before_pause: int) -> int: + tokens = external_stream.topic("tokens", type=str) + subscription = tokens.subscribe() + iterator = subscription.__aiter__() + + while self._seen < before_pause: + await iterator.__anext__() + self._seen += 1 + + # A real Timer: the completion that carries it cannot be retained. + await asyncio.sleep(POST_ROLLOVER_PAUSE.total_seconds()) + + # Whatever was appended into that window. That it arrives at all proves + # only that it was not lost -- the Timer creates a Workflow Task of its + # own when it fires, and that task would find a buffered record anyway. + # What proves the wake is the Signal in History, which the test asserts. + await iterator.__anext__() + self._seen += 1 + return self._seen + + +@pytest.fixture +def backend() -> MemoryStreamBackend: + return MemoryStreamBackend() + + +async def stream_key_for(client: Client, handle: Any) -> StreamKey: + description = await handle.describe() + return StreamKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + "tokens", + ) + + +async def history(handle: Any) -> list[Any]: + return [e async for e in handle.fetch_history_events()] + + +def completed_tasks(events: list[Any]) -> list[Any]: + return [e for e in events if e.HasField("workflow_task_completed_event_attributes")] + + +def timed_out_tasks(events: list[Any]) -> list[Any]: + return [e for e in events if e.HasField("workflow_task_timed_out_event_attributes")] + + +async def stop_if_running(handle: Any) -> None: + """Terminates the Workflow, unless it has already finished on its own. + + Every test here terminates in a ``finally``, because a Workflow left running + keeps a Worker's Run alive past the test that owns it. A test whose whole + point is that the Workflow *finishes* has nothing left to terminate, though, + and the server answers that with ``NOT_FOUND`` -- an error about the + cleanup, raised after the assertions have all passed, which reads exactly + like the feature having failed. Only that one status is swallowed: anything + else is a real failure of the teardown and is left to surface. + """ + try: + await handle.terminate() + except RPCError as err: + if err.status is not RPCStatusCode.NOT_FOUND: + raise + + +def timers_started(events: list[Any]) -> list[Any]: + return [e for e in events if e.HasField("timer_started_event_attributes")] + + +def held_for(events: list[Any], completed: Any) -> timedelta: + """How long the Workflow Task that ``completed`` closes was held open.""" + started = { + e.event_id: e + for e in events + if e.HasField("workflow_task_started_event_attributes") + } + start = started[completed.workflow_task_completed_event_attributes.started_event_id] + return completed.event_time.ToDatetime() - start.event_time.ToDatetime() + + +def rollover_completion(events: list[Any]) -> Any | None: + """The first Workflow Task completion that can only be a rollover. + + Identified by how long the task it closes was held: an idle park closes one + an idle timeout after its last record, and a command-carrying completion + closes one as soon as the activation returns. Only rollover holds a task + open for the whole deadline, so a completion at the deadline is the + mechanism itself rather than a task that happened to take a while. + + The window is the deadline itself, less :data:`ROLLOVER_DELIVERY_TOLERANCE` + and up to the Workflow Task timeout. The same margin applies in this + direction for the same reason: what is measured is the distance between two + server event times, and Core arms its deadline from neither of them. + """ + for completed in completed_tasks(events): + held = held_for(events, completed) + if ROLLOVER_DEADLINE - ROLLOVER_DELIVERY_TOLERANCE <= held < TASK_TIMEOUT: + return completed + return None + + +def wake_signals(events: list[Any]) -> list[Any]: + """The reserved wake Signals in a Workflow's own History. + + The mechanism itself, rather than a record turning up -- which after a + rollover proves nothing, since the server has already scheduled a + replacement task that could have picked the record up on its own. + """ + return [ + e + for e in events + if e.HasField("workflow_execution_signaled_event_attributes") + and e.workflow_execution_signaled_event_attributes.signal_name + == WAKE_SIGNAL_NAME + ] + + +async def wait_for_history( + handle: Any, + predicate: Any, + *, + timeout: float, + message: str, +) -> list[Any]: + """Polls the Workflow's History until ``predicate`` holds, or fails saying why.""" + deadline = asyncio.get_running_loop().time() + timeout + while True: + events = await history(handle) + if predicate(events): + return events + if asyncio.get_running_loop().time() >= deadline: + raise AssertionError(message) + await asyncio.sleep(0.2) + + +async def feed( + backend: MemoryStreamBackend, + key: StreamKey, + session: str, + *, + count: int, + stop: asyncio.Event | None = None, +) -> None: + """Appends one record at a time, with gaps well under the idle timeout. + + One session with a continuous sequence, because ``(session_id, sequence)`` + is the append idempotency key: restarting the numbering re-uses a key with + different content, the backend rejects it, and the symptom is a Workflow + that appears to hang far from the cause. + """ + for i in range(count): + if stop is not None and stop.is_set(): + return + await publish(backend, key, [f"t{i}"], session=session) + await asyncio.sleep(FEED_GAP_SECONDS) + + +# --- case 19 ------------------------------------------------------------------ + + +@pytest.mark.timeout(180) +async def test_a_continuously_fed_stream_survives_a_rollover( + client: Client, backend: MemoryStreamBackend +) -> None: + """The case rollover exists for, fed for longer than one task may live. + + Three things have to hold at once, and each fails differently: + + - the deadline has to fire *while records are arriving*, which is the only + state in which the idle timer never will; + - the replacement task has to inherit the subscription and its cursor, so + the count the Workflow ends with is exactly what was published; + - no task may run past the Workflow Task timeout, asserted against the + server's own event timestamps rather than inferred from the Workflow + finishing. + + The check is made *before* the timeout could be reached, so a broken + rollover fails here rather than as a timed-out task and the Core panic that + follows one. + """ + total = 30 # 30 * 0.3s = 9s of feeding, past a deadline at 8s. + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[SteadyStreamWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + SteadyStreamWorkflow.run, + total, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + task_timeout=TASK_TIMEOUT, + ) + key = await stream_key_for(client, handle) + stop = asyncio.Event() + feeder = asyncio.ensure_future( + feed(backend, key, f"steady-{uuid.uuid4()}", count=total, stop=stop) + ) + try: + # A margin past the deadline, and still short of the timeout: by now + # a working rollover has completed at least one task. + await asyncio.sleep(ROLLOVER_DEADLINE.total_seconds() + 1) + events = await history(handle) + completed = completed_tasks(events) + assert completed, ( + "the stream was fed continuously for longer than the rollover " + f"deadline of {ROLLOVER_DEADLINE} and not one Workflow Task has " + "completed, so the retained task is on its way to being timed " + "out by the server -- which is the failure rollover exists to " + "prevent" + ) + + result = await asyncio.wait_for(handle.result(), 90) + assert result == total, ( + "the replacement task did not inherit the subscription's " + f"cursor: {result} of {total} records were consumed" + ) + + events = await history(handle) + assert not timed_out_tasks(events), ( + "a Workflow Task was held past the server's timeout, so " + "rollover did not bound it" + ) + for event in completed_tasks(events): + held = held_for(events, event) + assert held < TASK_TIMEOUT, ( + f"a Workflow Task was held for {held}, past the " + f"{TASK_TIMEOUT} it must stay inside" + ) + assert rollover_completion(events) is not None, ( + "no Workflow Task was held to the rollover deadline, so the " + f"{total} records were consumed without one -- and this test " + "asserts nothing about rollover unless one fired" + ) + finally: + stop.set() + feeder.cancel() + await stop_if_running(handle) + + +# --- case 21 ------------------------------------------------------------------ + + +@pytest.mark.timeout(180) +async def test_a_signal_into_a_retained_task_lands_by_the_rollover_deadline( + client: Client, backend: MemoryStreamBackend +) -> None: + """Retention delays Signals, and rollover is the only thing bounding that. + + While a Workflow Task is retained the server cannot start another one, so + Signals, Updates, and non-legacy Queries queue behind it -- and a stream can + hold a task open far longer than an outstanding local activity ever would. + Retention latency for those inputs is therefore bounded by the rollover + deadline and by nothing else. + + Both halves are asserted, because either alone is satisfiable by accident: + that the task really was retained when the Signal was sent -- no Workflow + Task had completed, so the Signal had nothing to be delivered on -- and that + the Workflow saw it within the deadline, measured in *Workflow* time from + the Run's own start. + + The second half carries :data:`ROLLOVER_DELIVERY_TOLERANCE`, because what it + measures is the distance between two ``WorkflowTaskStarted`` event times and + the deadline can only be the moment Core *decides* to hand the task on. A + completion RPC, the server scheduling the replacement, and a poll picking it + up all land between the two, so an exact bound is unreachable by + construction rather than merely flaky. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[SignalledStreamWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + SignalledStreamWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + task_timeout=TASK_TIMEOUT, + ) + started_at = asyncio.get_running_loop().time() + key = await stream_key_for(client, handle) + stop = asyncio.Event() + feeder = asyncio.ensure_future( + feed(backend, key, f"signalled-{uuid.uuid4()}", count=40, stop=stop) + ) + try: + # Long enough for the stream to be driving the Workflow, far short + # of the deadline. + await asyncio.sleep(2) + assert not completed_tasks(await history(handle)), ( + "a Workflow Task had already completed, so the Signal below " + "would not be sent into a retained one and this proves nothing" + ) + + await handle.signal(SignalledStreamWorkflow.poke) + # Waited out only as far as the deadline and its tolerance, never as + # far as the Workflow Task timeout: a retained task that reaches that + # is a panicking Core, not a failed assertion. + bound = (ROLLOVER_DEADLINE + ROLLOVER_DELIVERY_TOLERANCE).total_seconds() + waited = asyncio.get_running_loop().time() - started_at + try: + elapsed = await asyncio.wait_for(handle.result(), bound + 0.5 - waited) + except asyncio.TimeoutError: + raise AssertionError( + "the Signal had still not been delivered by the rollover " + f"deadline of {ROLLOVER_DEADLINE}: the Workflow Task is " + "still retained, and nothing but that deadline bounds how " + "long an input queues behind it" + ) from None + + assert elapsed <= bound, ( + f"the Signal was delivered {elapsed:.3f}s into the Run, past the " + f"rollover deadline of {ROLLOVER_DEADLINE} and the " + f"{ROLLOVER_DELIVERY_TOLERANCE} allowed for getting the " + "replacement task started. Retention latency is bounded by that " + "deadline and by nothing else" + ) + events = await history(handle) + assert not timed_out_tasks(events), ( + "the Signal was released by the server timing the retained task " + "out rather than by a rollover" + ) + finally: + stop.set() + feeder.cancel() + await stop_if_running(handle) + + +# --- case 23 ------------------------------------------------------------------ + + +@pytest.mark.timeout(180) +async def test_an_append_after_a_rollover_completion_wakes_the_subscription( + client: Client, backend: MemoryStreamBackend +) -> None: + """A subscription that has been through a rollover is still woken by a Signal. + + A rollover ends a Workflow Task with the subscription active and unparked: + there is no retained task to notify locally and no park generation for a + producer to observe, so the consumer's own watcher is what covers the + window -- it observes ``NoOpenWorkflowTask`` and sends the reserved wake + Signal itself. + + That window cannot be appended into directly. Core completes a rollover with + ``force_new_wft``, so the server starts the replacement Workflow Task within + a millisecond, and an append made from a test lands on the replacement task + instead: readiness is answered ``Accepted``, the record is delivered live, + and no wake is owed or sent -- which is correct behaviour and not the + mechanism this case is about. The Workflow therefore *holds* the same state + open on the far side of the rollover with a Timer, whose completion is + server-bound and equally cannot be retained, and the append goes in there. + + Both halves are asserted, since either alone is satisfiable without the + other: that a rollover really happened, by a Workflow Task held to the + deadline and closed inside the timeout; and that the append was answered by + a wake, by ``__temporal_external_stream_wake`` reaching the Workflow's own + History *after* the completion that opened the window. The record arriving + proves nothing on its own -- the Timer creates a Workflow Task when it + fires, which would have picked the record up anyway. + """ + before_pause = 30 # 30 * 0.3s = 9s of feeding, past a deadline at 8s. + task_queue = f"tq-{uuid.uuid4()}" + session = f"rollover-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[RolloverThenPauseWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + RolloverThenPauseWorkflow.run, + before_pause, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + task_timeout=TASK_TIMEOUT, + ) + key = await stream_key_for(client, handle) + stop = asyncio.Event() + feeder = asyncio.ensure_future( + feed(backend, key, session, count=before_pause, stop=stop) + ) + try: + # The Timer is started only once every fed record has been consumed, + # which takes longer than the deadline -- so its appearance is also + # the point past which the rollover has already happened. + events = await wait_for_history( + handle, + timers_started, + timeout=60, + message=( + f"the Workflow never consumed the {before_pause} records fed " + "to it and reached its pause, so there is no unparked " + "Workflow Task boundary to append into" + ), + ) + stop.set() + feeder.cancel() + + rollover = rollover_completion(events) + assert rollover is not None, ( + "no Workflow Task was held to the rollover deadline of " + f"{ROLLOVER_DEADLINE}, so the subscription this appends to has " + "not been through a rollover and the case is untested" + ) + assert not timed_out_tasks(events), ( + "the task was released by the server's timeout rather than by a " + "rollover" + ) + pause = timers_started(events)[0] + assert pause.event_id > rollover.event_id, ( + "the pause began before the rollover, so the window appended " + "into below is not a post-rollover one" + ) + + # Counted by name, and by the *same* name the check below uses: a + # baseline that counted every Signal would be compared against a + # total that counts only wake Signals, and any unrelated Signal in + # this Workflow's History would then make the comparison meaningless + # in whichever direction it happened to fall. + before = len(wake_signals(events)) + + await publish(backend, key, ["after-rollover"], session=session) + + def woken(events: list[Any]) -> bool: + wakes = wake_signals(events) + # Position as well as count: a Signal sent while a Workflow Task + # was open reaches History only when that task completes, so + # counting alone would accept one that was owed long before this + # append and merely landed late. + return len(wakes) > before and wakes[-1].event_id > pause.event_id + + await wait_for_history( + handle, + woken, + timeout=POST_ROLLOVER_PAUSE.total_seconds(), + message=( + "no wake Signal followed the append: the Workflow Task that " + "carried the pause left the subscription active and " + "unparked, so there was no retained task to notify and no " + "park generation to observe, and the watcher owes one" + ), + ) + + assert await asyncio.wait_for(handle.result(), 60) == before_pause + 1, ( + "the record appended into the unparked window never reached the " + "Workflow, so the wake Signal did not carry the subscription " + "over the boundary" + ) + finally: + stop.set() + feeder.cancel() + await stop_if_running(handle) diff --git a/tests/contrib/external_workflow_streams/test_runtime.py b/tests/contrib/external_workflow_streams/test_runtime.py new file mode 100644 index 0000000000..021637b192 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_runtime.py @@ -0,0 +1,1386 @@ +"""P10a/P10b — the observation delta and the quiescent snapshot. + +Driven directly against the runtime handle: what is under test is what gets +*recorded*, and threading it through a real Worker would obscure that behind +scheduling. +""" + +from __future__ import annotations + +import asyncio +import uuid +from datetime import timedelta + +import pytest +import pytest_asyncio + +import temporalio.converter +from temporalio.contrib.external_workflow_streams._annotation import ( + ROLLOVER_HIGH_WATER, + AnnotationHeader, + SegmentEndReason, + StreamBinding, + decode_annotation, + encode_header, + encode_terminal, +) +from temporalio.contrib.external_workflow_streams._backend import StreamKey +from temporalio.contrib.external_workflow_streams._errors import ( + ExternalStreamCapacityError, +) +from temporalio.contrib.external_workflow_streams._manager import ( + ReadinessResult, + StreamSubscriptionManager, +) +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Offset, + RecordKind, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._runtime import ( + _RUN_COST_FLOOR, + WorkflowStreamRuntime, +) +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend + +RUN_ID = "run-1" + + +async def _notify( + run_id: str, # pyright: ignore[reportUnusedParameter] + wait_id: int, # pyright: ignore[reportUnusedParameter] + generation: int, # pyright: ignore[reportUnusedParameter] +) -> str: + return ReadinessResult.ACCEPTED + + +@pytest.fixture +def backend() -> MemoryStreamBackend: + return MemoryStreamBackend() + + +@pytest_asyncio.fixture # type: ignore[reportUntypedFunctionDecorator] +async def runtime(backend: MemoryStreamBackend): + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=10), + ) + yield WorkflowStreamRuntime( + manager=manager, + backend=backend, + run_id=RUN_ID, + namespace="ns", + workflow_id="wf", + first_execution_run_id=uuid.uuid4().hex, + data_converter=temporalio.converter.DataConverter.default, + default_idle_timeout=timedelta(seconds=1), + ) + # Watchers outlive the runtime unless the manager is told to stop, and a + # leaked watcher keeps polling a backend nothing is reading any more. + await manager.shutdown() + + +def subscribe( + runtime: WorkflowStreamRuntime, + wait_id: int, + name: str = "tokens", + idle_timeout: timedelta | None = None, +) -> StreamKey: + key = runtime.stream_key(name) + runtime.register( + wait_id=wait_id, + stream_key=key, + idle_timeout=idle_timeout, + ) + return key + + +def data(offset: str, session: str = "s", seq: int = 0) -> StreamRecord: + return StreamRecord(RecordKind.DATA, b"x", session, seq).placed_at(Offset(offset)) + + +def fence(offset: str, session: str = "s", seq: int = 0) -> StreamRecord: + return StreamRecord(RecordKind.WRITE_FENCE, b"", session, seq).placed_at( + Offset(offset) + ) + + +def annotation_of(runtime: WorkflowStreamRuntime) -> object: + """The annotation Core would hold, decoded. + + Core accumulates deltas by byte concatenation, so this joins them the same + way rather than reaching for anything the runtime keeps internally. + """ + parts = [] + delta = runtime.take_observation_delta() + if delta is not None: + parts.append(delta) + parts.append(runtime.add_terminal()) + return decode_annotation(b"".join(parts)) + + +# --- emission is not conditional on records --------------------------------- + + +async def test_a_first_subscription_to_an_empty_stream_still_emits( + runtime: WorkflowStreamRuntime, +) -> None: + """Without this, replay of an empty stream has nowhere to begin. + + The binding carries the stream key, an explicit start cursor, and the + configured provider identity. None is derivable from mutable backend state. + """ + key = subscribe(runtime, 1) + + decoded = annotation_of(runtime) + + binding = decoded.header.streams[1] # type: ignore[attr-defined] + assert binding.provider_id == "memory" + assert decoded.header.streams[1].stream_key == key # type: ignore[attr-defined] + assert decoded.header.streams[1].start_cursor == BEGINNING # type: ignore[attr-defined] + assert decoded.terminal == {1: BEGINNING} # type: ignore[attr-defined] + + +async def test_an_activation_that_drained_nothing_emits_an_empty_segment( + runtime: WorkflowStreamRuntime, +) -> None: + """The drain happened, so replay must reproduce it. + + Reproducing only the record order while changing how many drains occurred + would change when ``wait_condition`` predicates fire. + """ + subscribe(runtime, 1) + + decoded = annotation_of(runtime) + + assert len(decoded.segments) == 1 # type: ignore[attr-defined] + assert decoded.segments[0].runs == () # type: ignore[attr-defined] + + +async def test_an_activation_that_observed_nothing_at_all_emits_nothing( + runtime: WorkflowStreamRuntime, +) -> None: + """The one case where a completion legitimately carries no progress.""" + assert runtime.take_observation_delta() is None + + +# --- runs ------------------------------------------------------------------- + + +async def test_consecutive_deliveries_from_one_stream_collapse_into_one_run( + runtime: WorkflowStreamRuntime, +) -> None: + """The claim the whole design rests on: no field in a run is per-record.""" + subscribe(runtime, 1) + for i in range(1, 1001): + runtime.record_delivery(1, data(f"{i}-0")) + + decoded = annotation_of(runtime) + + (run,) = decoded.segments[0].runs # type: ignore[attr-defined] + assert run.count == 1000 + assert run.first_offset == Offset("1-0") + assert run.last_offset == Offset("1000-0") + + +async def test_alternating_streams_produce_one_run_per_delivery( + runtime: WorkflowStreamRuntime, +) -> None: + """The honest worst case, recorded as such rather than assumed away. + + Per-stream ranges alone are insufficient: concurrent coroutines could + otherwise observe a different order than the one that actually happened. + """ + subscribe(runtime, 1) + subscribe(runtime, 2, name="tool-events") + for i, wait_id in enumerate([1, 2, 1, 2], start=1): + runtime.record_delivery(wait_id, data(f"{i}-0")) + + decoded = annotation_of(runtime) + + runs = decoded.segments[0].runs # type: ignore[attr-defined] + assert [r.wait_id for r in runs] == [1, 2, 1, 2] + assert all(r.count == 1 for r in runs) + + +async def test_a_run_resumes_after_the_other_stream_interrupts( + runtime: WorkflowStreamRuntime, +) -> None: + """Maximal, not merged: the interruption is itself part of the schedule.""" + subscribe(runtime, 1) + subscribe(runtime, 2, name="tool-events") + runtime.record_delivery(1, data("1-0")) + runtime.record_delivery(1, data("2-0")) + runtime.record_delivery(2, data("3-0")) + runtime.record_delivery(1, data("4-0")) + + runs = annotation_of(runtime).segments[0].runs # type: ignore[attr-defined] + + assert [(r.wait_id, r.count) for r in runs] == [(1, 2), (2, 1), (1, 1)] + + +# --- control records -------------------------------------------------------- + + +async def test_control_records_are_counted_and_their_positions_recorded( + runtime: WorkflowStreamRuntime, +) -> None: + """They occupy offsets, so replay's range read will find them. + + Leaving them out of ``count`` would make every such range read find more + records than the marker claims and fail as integrity loss -- for a stream + that is perfectly intact. + """ + subscribe(runtime, 1) + runtime.record_delivery(1, data("1-0")) + runtime.record_delivery(1, fence("2-0")) + runtime.record_delivery(1, data("3-0")) + + (run,) = annotation_of(runtime).segments[0].runs # type: ignore[attr-defined] + + assert run.count == 3 + assert run.control_positions == (1,) + assert run.last_offset == Offset("3-0") + + +async def test_control_positions_are_sparse(runtime: WorkflowStreamRuntime) -> None: + """One entry per fence, not one kind tag per record.""" + subscribe(runtime, 1) + for i in range(1, 101): + runtime.record_delivery(1, data(f"{i}-0")) + runtime.record_delivery(1, fence("101-0")) + + (run,) = annotation_of(runtime).segments[0].runs # type: ignore[attr-defined] + + assert run.count == 101 + assert run.control_positions == (100,) + + +# --- the terminal ----------------------------------------------------------- + + +async def test_the_terminal_is_where_delivery_stopped( + runtime: WorkflowStreamRuntime, +) -> None: + subscribe(runtime, 1) + subscribe(runtime, 2, name="tool-events") + runtime.record_delivery(1, data("5-0")) + + decoded = annotation_of(runtime) + + assert decoded.terminal == { # type: ignore[attr-defined] + 1: AFTER(Offset("5-0")), + 2: BEGINNING, + } + + +async def test_the_terminal_reads_no_backend( + runtime: WorkflowStreamRuntime, backend: MemoryStreamBackend +) -> None: + """Finalization performs no backend I/O (ADR-010). + + The boundary is where this Workflow Task's deliveries stopped, already + fixed. Refreshing it against the stream could name a position replay must + not reproduce. + """ + subscribe(runtime, 1) + runtime.record_delivery(1, data("5-0")) + reads_before = len(backend.range_reads) + + runtime.add_terminal() + + assert len(backend.range_reads) == reads_before + + +# --- deltas concatenate ----------------------------------------------------- + + +async def test_deltas_across_activations_concatenate_into_one_annotation( + runtime: WorkflowStreamRuntime, +) -> None: + """What Core holds is the byte concatenation, so that is what is decoded.""" + subscribe(runtime, 1) + runtime.record_delivery(1, data("1-0")) + first = runtime.take_observation_delta() + + runtime.record_delivery(1, data("2-0")) + second = runtime.take_observation_delta() + + terminal = runtime.add_terminal() + + assert first is not None and second is not None + decoded = decode_annotation(first + second + terminal) + assert len(decoded.segments) == 2, "one segment per activation, not one per task" + assert [r.count for s in decoded.segments for r in s.runs] == [1, 1] + assert decoded.terminal == {1: AFTER(Offset("2-0"))} + + +async def test_terminating_an_already_closed_annotation_adds_nothing( + runtime: WorkflowStreamRuntime, +) -> None: + """A second terminal for one Workflow Task must not start a second annotation. + + Core can ask for a terminal for a boundary it decided -- a rollover + deadline, a shutdown -- on a task whose last completion had already closed + the annotation. Creating a fresh header and terminal for that would append a + complete second annotation to the one Core is about to write, and the marker + would decode as far as the first terminal and fail on the frame after it. + """ + subscribe(runtime, 1) + runtime.record_delivery(1, data("1-0")) + delta = runtime.take_observation_delta() + assert delta is not None + + closed = delta + runtime.add_terminal() + again = runtime.add_terminal() + + assert again == b"", "the annotation was already closed; there was nothing to add" + decoded = decode_annotation(closed + again) + assert decoded.terminal == {1: AFTER(Offset("1-0"))} + + +async def test_a_terminal_after_new_observations_opens_a_fresh_annotation( + runtime: WorkflowStreamRuntime, +) -> None: + """The suppression above is about a *closed* annotation, not about the Run. + + Anything observed after the close begins the next annotation, and that one + needs its own header and its own terminal like any other. + """ + subscribe(runtime, 1) + runtime.record_delivery(1, data("1-0")) + first = runtime.take_observation_delta() + assert first is not None + first += runtime.add_terminal() + + runtime.record_delivery(1, data("2-0")) + second = runtime.take_observation_delta() + assert second is not None + second += runtime.add_terminal() + + assert decode_annotation(first).terminal == {1: AFTER(Offset("1-0"))} + assert decode_annotation(second).terminal == {1: AFTER(Offset("2-0"))} + + +async def test_a_new_annotation_starts_from_the_current_cursors( + runtime: WorkflowStreamRuntime, +) -> None: + """The next Workflow Task continues where this one stopped. + + Restarting the header at the original start cursor would make replay of the + second marker re-deliver everything the first one already recorded. + """ + subscribe(runtime, 1) + runtime.record_delivery(1, data("9-0")) + runtime.take_observation_delta() + runtime.add_terminal() + + runtime.start_new_annotation() + runtime.record_delivery(1, data("10-0")) + decoded = annotation_of(runtime) + + assert decoded.header.streams[1].start_cursor == AFTER(Offset("9-0")) # type: ignore[attr-defined] + + +# --- a subscription created after the header went out ----------------------- + + +async def test_a_wait_registered_after_the_first_delta_reaches_the_header( + runtime: WorkflowStreamRuntime, +) -> None: + """`register` accepts a subscription at any activation of a retained task. + + The header frame is emitted with the first delta and Core appends rather + than rewrites, so a wait that joined later cannot be added to it in place. + Without a binding of its own that wait reaches the marker as a run and a + terminal entry with no stream key, no backend, and no start cursor, and + replay of *unchanged* code fails as a wait "this Workflow did not create". + """ + subscribe(runtime, 1, "tokens") + runtime.record_delivery(1, data("1-0")) + first = runtime.take_observation_delta() + assert first is not None + assert set(decode_annotation(first).header.streams) == {1}, ( + "the first delta's header is what this test is about; if it already " + "carried wait 2 the fix under test is not being exercised" + ) + + key = runtime.stream_key("tokens") + runtime.register(wait_id=2, stream_key=key) + runtime.record_delivery(2, data("2-0")) + second = runtime.take_observation_delta() + assert second is not None + terminal = runtime.add_terminal() + + decoded = decode_annotation(first + second + terminal) + + assert set(decoded.header.streams) == {1, 2}, ( + "wait 2 has a run and a terminal entry but no binding, so replay has " + "no stream key, no backend, and no start cursor for it" + ) + assert decoded.header.streams[2].stream_key == key + assert decoded.terminal == {1: AFTER(Offset("1-0")), 2: AFTER(Offset("2-0"))} + + +async def test_a_late_wait_is_bound_before_the_segment_that_records_it( + runtime: WorkflowStreamRuntime, +) -> None: + """Order within the delta, not merely presence somewhere in the marker. + + A decoder that met the run first would have no binding to attach it to, and + a marker truncated at a rollover boundary would carry the run without the + binding at all. + """ + subscribe(runtime, 1) + runtime.record_delivery(1, data("1-0")) + first = runtime.take_observation_delta() + assert first is not None + + subscribe(runtime, 2) + runtime.record_delivery(2, data("2-0")) + second = runtime.take_observation_delta() + assert second is not None + + # 0x04 is the bindings frame tag. + assert second[0] == 0x04, ( + "the delta must open with the binding for the wait it goes on to " + f"record a run for, got {second.hex()}" + ) + + +async def test_a_late_wait_starts_at_its_own_cursor( + runtime: WorkflowStreamRuntime, +) -> None: + """Not wherever the waits already in the header have reached. + + The annotation-wide position belongs to the waits that were there when it + opened. Recording it for a wait that joined afterwards would start replay + of that wait past records it in fact received. + """ + subscribe(runtime, 1) + runtime.record_delivery(1, data("9-0")) + first = runtime.take_observation_delta() + assert first is not None + + subscribe(runtime, 2) + runtime.record_delivery(2, data("3-0")) + second = runtime.take_observation_delta() + assert second is not None + + decoded = decode_annotation(first + second + runtime.add_terminal()) + + assert decoded.header.streams[2].start_cursor == BEGINNING + assert decoded.header.streams[1].start_cursor == BEGINNING + + +async def test_a_wait_is_bound_once_per_annotation( + runtime: WorkflowStreamRuntime, +) -> None: + """Three activations, one binding each for the two waits that exist. + + A binding re-emitted every activation would grow the marker with the + activation count and be rejected on decode as a wait bound twice. + """ + subscribe(runtime, 1) + runtime.record_delivery(1, data("1-0")) + parts = [runtime.take_observation_delta()] + + subscribe(runtime, 2) + runtime.record_delivery(2, data("2-0")) + parts.append(runtime.take_observation_delta()) + + runtime.record_delivery(2, data("3-0")) + parts.append(runtime.take_observation_delta()) + + joined = b"".join(p for p in parts if p is not None) + + # A re-emitted binding is not merely wasteful: the decoder refuses a wait + # bound twice, so this decoding at all is the assertion. + decoded = decode_annotation(joined + runtime.add_terminal()) + assert set(decoded.header.streams) == {1, 2} + + +async def test_the_next_annotation_binds_every_wait_in_its_own_header( + runtime: WorkflowStreamRuntime, +) -> None: + """A late binding belongs to the annotation it was emitted into. + + The next Workflow Task writes a header from scratch, so the wait that + joined late must appear in *that* header rather than being remembered as + already announced -- a second marker missing it is the same defect one + Workflow Task later. + """ + subscribe(runtime, 1) + runtime.record_delivery(1, data("1-0")) + runtime.take_observation_delta() + subscribe(runtime, 2) + runtime.record_delivery(2, data("2-0")) + runtime.take_observation_delta() + runtime.add_terminal() + + runtime.start_new_annotation() + runtime.record_delivery(2, data("3-0")) + decoded = annotation_of(runtime) + + assert set(decoded.header.streams) == {1, 2} # type: ignore[attr-defined] + assert decoded.header.streams[2].start_cursor == AFTER(Offset("2-0")) # type: ignore[attr-defined] + + +# --- quiescence (P10a) ------------------------------------------------------ + + +async def test_the_quiescent_snapshot_is_the_complete_blocked_set( + runtime: WorkflowStreamRuntime, +) -> None: + """A partial set would be worse than none. + + It would let one idle stream park a Workflow Task another stream is still + driving, which is the whole reason the timeout is a property of the set. + """ + subscribe(runtime, 1) + subscribe(runtime, 2, name="tool-events") + subscribe(runtime, 3, name="more") + + snapshot = runtime.quiescent_snapshot() + + assert snapshot is not None + assert [w.wait_id for w in snapshot] == [1, 2, 3] + + +async def test_nothing_blocked_asks_for_no_retention( + runtime: WorkflowStreamRuntime, +) -> None: + subscribe(runtime, 1) + runtime.note_blocked(1, False) + + assert runtime.quiescent_snapshot() is None + + +async def test_a_fenced_subscription_reports_immediately_parkable( + runtime: WorkflowStreamRuntime, +) -> None: + """One fenced stream does not park the task; Core decides that from the set.""" + subscribe(runtime, 1) + subscribe(runtime, 2, name="tool-events") + runtime.record_delivery(1, fence("1-0")) + + snapshot = runtime.quiescent_snapshot() + + assert snapshot is not None + assert [(w.wait_id, w.immediately_parkable) for w in snapshot] == [ + (1, True), + (2, False), + ] + + +async def test_a_later_record_clears_the_fence(runtime: WorkflowStreamRuntime) -> None: + """A fence asserts nothing about other producers; a later record just wakes.""" + subscribe(runtime, 1) + runtime.record_delivery(1, fence("1-0")) + runtime.record_delivery(1, data("2-0")) + + snapshot = runtime.quiescent_snapshot() + + assert snapshot is not None + assert snapshot[0].immediately_parkable is False + + +async def test_differing_idle_timeouts_reduce_to_their_min( + runtime: WorkflowStreamRuntime, +) -> None: + """The set shares one timer, so the configured values must reduce to one. + + ``min`` in ``wait_id`` order over the blocked set and nothing else, so the + result is deterministic and reproduces on replay. + """ + subscribe(runtime, 1, idle_timeout=timedelta(seconds=5)) + subscribe(runtime, 2, name="b", idle_timeout=timedelta(seconds=2)) + subscribe(runtime, 3, name="c", idle_timeout=timedelta(seconds=9)) + + assert runtime.effective_idle_timeout() == timedelta(seconds=2) + + +async def test_only_blocked_subscriptions_contribute_to_the_reduction( + runtime: WorkflowStreamRuntime, +) -> None: + """The inputs are the quiescent set's values and nothing else.""" + subscribe(runtime, 1, idle_timeout=timedelta(seconds=5)) + subscribe(runtime, 2, name="b", idle_timeout=timedelta(milliseconds=1)) + runtime.note_blocked(2, False) + + assert runtime.effective_idle_timeout() == timedelta(seconds=5) + + +async def test_re_entering_the_blocked_state_bumps_the_generation( + runtime: WorkflowStreamRuntime, +) -> None: + """This is what makes a readiness notification for the old block stale.""" + subscribe(runtime, 1) + snapshot = runtime.quiescent_snapshot() + assert snapshot is not None and snapshot[0].generation == 0 + + runtime.note_blocked(1, False) + runtime.note_blocked(1, True) + + snapshot = runtime.quiescent_snapshot() + assert snapshot is not None and snapshot[0].generation == 1 + + +async def test_readiness_after_a_re_block_names_the_current_generation( + backend: MemoryStreamBackend, +) -> None: + """Bumping the generation is only half of it; the manager has to be told. + + Core compares the generation a readiness notification carries against the + one lang put in the quiescent snapshot. The manager is what makes that call + and only the runtime knows when a wait re-blocks, so a manager left to + itself reports 0 for the life of the subscription. Every notification after + the first block is then answered ``Stale``, which the manager treats as + "re-probe later" -- but the watcher's prefetch cursor is already past the + record and no second notification is coming, so an append after a confirmed + park never wakes the Workflow at all. + """ + reported: list[int] = [] + notified = asyncio.Event() + + async def notify( + run_id: str, # pyright: ignore[reportUnusedParameter] + wait_id: int, # pyright: ignore[reportUnusedParameter] + generation: int, + ) -> str: + reported.append(generation) + notified.set() + return ReadinessResult.ACCEPTED + + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=notify, + watch_block=timedelta(milliseconds=10), + ) + runtime = WorkflowStreamRuntime( + manager=manager, + backend=backend, + run_id=RUN_ID, + namespace="ns", + workflow_id="wf", + first_execution_run_id=uuid.uuid4().hex, + data_converter=temporalio.converter.DataConverter.default, + default_idle_timeout=timedelta(seconds=1), + ) + try: + key = subscribe(runtime, 1) + # One resolved block, so the wait is on its second one -- the state + # every subscription is in for every record after its first. + runtime.note_blocked(1, False) + runtime.note_blocked(1, True) + + await backend.append(key, StreamRecord(RecordKind.DATA, b"x", "s", 0)) + await asyncio.wait_for(notified.wait(), 5) + + snapshot = runtime.quiescent_snapshot() + assert snapshot is not None + assert reported == [snapshot[0].generation] == [1], ( + "the generation reported to Core is not the one lang put in the " + f"quiescent snapshot, so Core answers Stale: reported={reported} " + f"snapshot={snapshot[0].generation}" + ) + finally: + await manager.shutdown() + + +async def test_a_closed_wait_can_never_re_enter_the_blocked_set( + runtime: WorkflowStreamRuntime, +) -> None: + """Closing keeps the wait's state, so the state has to refuse to block again. + + The subscription is deliberately not forgotten -- its binding is what a + replay of unchanged code reads, and its cursor is what stops a + Continue-As-New successor restarting the stream from the beginning -- which + leaves it reachable by ``wait_id`` long after the coroutine that was + awaiting it is gone. A wait that could re-enter the blocked set from there + would be named by the quiescent snapshot, and Core would go back to + retaining -- and eventually parking -- the Workflow Task for nobody. + """ + subscribe(runtime, 1) + runtime.note_blocked(1, False) + runtime.unsubscribe(1) + + runtime.note_blocked(1, True) + + assert runtime.quiescent_snapshot() is None, ( + "a closed wait re-entered the quiescent set, so Core is asked to hold " + "the Workflow Task open for a wait nothing is awaiting" + ) + # And its binding is still there to be read, which is the whole reason + # closing does not simply drop the state. + assert set(annotation_of(runtime).header.streams) == {1} # type: ignore[attr-defined] + + +# --- the byte budget -------------------------------------------------------- + + +async def test_rollover_is_requested_before_the_budget_is_reached( + backend: MemoryStreamBackend, +) -> None: + """The runtime asks Core to roll the task over rather than grow the marker.""" + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=10), + ) + runtime = WorkflowStreamRuntime( + manager=manager, + backend=backend, + run_id=RUN_ID, + namespace="ns", + workflow_id="wf", + first_execution_run_id="first", + data_converter=temporalio.converter.DataConverter.default, + default_idle_timeout=timedelta(seconds=1), + max_annotation_bytes=2048, + ) + try: + subscribe(runtime, 1) + subscribe(runtime, 2, name="tool-events") + + assert not runtime.request_rollover + + delivered = 0 + while not runtime.request_rollover: + # Alternating, so every delivery is its own run -- the one workload + # that cannot be range-compressed and therefore the one that reaches + # the cap. + for wait_id in (1, 2): + delivered += 1 + runtime.record_delivery(wait_id, data(f"{delivered}-0")) + runtime.take_observation_delta() + + assert delivered > 0 + finally: + await manager.shutdown() + + +def _budget_runtime( + backend: MemoryStreamBackend, + manager: StreamSubscriptionManager, + *, + max_annotation_bytes: int, +) -> WorkflowStreamRuntime: + return WorkflowStreamRuntime( + manager=manager, + backend=backend, + run_id=RUN_ID, + namespace="ns", + workflow_id="wf", + first_execution_run_id="first", + data_converter=temporalio.converter.DataConverter.default, + default_idle_timeout=timedelta(seconds=1), + max_annotation_bytes=max_annotation_bytes, + ) + + +async def test_the_segment_that_crosses_the_mark_asks_for_rollover_itself( + backend: MemoryStreamBackend, +) -> None: + """The crossing segment's own completion carries the request, not the next one. + + ``take_observation_delta`` is what *closes* the activation's segment, so a + rollover flag sampled before it describes the annotation as it stood one + activation ago. The segment that crossed the high-water mark then went out + with ``request_rollover = false``, and the activation after it was free to add + another frame -- and to overflow -- before Core had ever been asked to roll + over. One activation of delay is the whole margin the high-water mark exists + to provide. + """ + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + from tests.contrib.external_workflow_streams.test_delivery_budget import ( + _CompletionStub, + ) + + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=10), + ) + runtime = _budget_runtime(backend, manager, max_annotation_bytes=2048) + try: + subscribe(runtime, 1) + subscribe(runtime, 2, name="tool-events") + + # Fill to just short of the mark, one flushed activation at a time, so + # that the *next* activation is the one that crosses it. + accumulated: list[bytes] = [] + delivered = 0 + high_water = int(2048 * ROLLOVER_HIGH_WATER) + while True: + for wait_id in (1, 2): + delivered += 1 + runtime.record_delivery(wait_id, data(f"{delivered}-0")) + delta = runtime.take_observation_delta() + if delta is not None: + accumulated.append(delta) + assert runtime._accumulator is not None + # Stops with a gap the next activation's segment is bigger than, so + # that segment is unambiguously the frame that crosses the mark. + if runtime._accumulator.size >= high_water - 400: + break + assert not runtime.request_rollover, ( + "the mark was crossed by a flushed activation, so nothing here " + "is testing the activation that crosses it" + ) + + # One more activation, whose closing segment is what crosses the mark -- + # and which is therefore the completion that has to carry the request. + for _ in range(15): + for wait_id in (1, 2): + delivered += 1 + runtime.record_delivery(wait_id, data(f"{delivered}-0")) + assert not runtime.request_rollover, ( + "the crossing segment is still open, so the accumulator cannot know " + "about it yet -- which is the entire reason the flag must be read " + "after the delta is taken and not before" + ) + + stub = _CompletionStub(runtime) + _WorkflowInstanceImpl._emit_external_stream_commands(stub) # type: ignore[arg-type] + + commands = stub._current_completion.successful.commands + progress = commands[0].workflow_stream_progress + assert progress.request_rollover, ( + "the completion that carried the crossing segment did not ask for the " + "rollover, so an entire further activation may append to an " + "annotation Core has not been told to close" + ) + accumulated.append(progress.observation_delta) + annotation = decode_annotation(b"".join(accumulated)) + assert annotation.terminal is not None, ( + "Core issues no finalization job for a rollover it was asked for, so " + "the terminal has to ride this same delta" + ) + assert annotation.segments[-1].end_reason in ( + SegmentEndReason.BUDGET_ROLLOVER, + SegmentEndReason.NO_DATA_AVAILABLE, + ) + finally: + await manager.shutdown() + + +async def test_a_frame_larger_than_the_slack_rolls_over_instead_of_raising( + backend: MemoryStreamBackend, +) -> None: + """An indivisible frame bigger than the remaining budget is not a failure. + + The high-water mark is a *fraction* of the budget, and a segment frame can be + larger than the fraction that is left: a run costs two provider-supplied + offset strings, and nothing bounds their length from this side. Checking at + encode time and raising is what ADR-007 rejects -- the encoding that + overflowed overflows again on the retry, so the Workflow Task fails forever + with no marker, no terminal, and no rollover ever asked for. + + The runtime instead stops handing records over while the annotation can still + record them, and asks for the rollover on that boundary. Long offsets are what + make one activation's segment expensive here; that is exactly the shape the + old check-and-fail path could not survive. + """ + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=10), + ) + runtime = _budget_runtime(backend, manager, max_annotation_bytes=1024) + long = "o" * 300 + try: + subscribe(runtime, 1) + subscribe(runtime, 2, name="tool-events") + + # Each record is its own run, and each run costs two 300-byte offsets -- + # far more than the 25% of 1024 bytes that the high-water margin leaves. + delivered = 0 + while runtime.delivery_budget_remaining() > 0 and delivered < 50: + wait_id = 1 + (delivered % 2) + delivered += 1 + runtime.record_delivery(wait_id, data(f"{long}{delivered}-0")) + + assert delivered < 50, ( + "the annotation budget never stopped delivery, so this test is not " + "exercising the boundary it exists for" + ) + assert runtime.request_rollover, ( + "the runtime stopped delivering because the annotation could record " + "no more, and said nothing about it -- Core is never asked to roll " + "the task over and the next activation overflows" + ) + assert runtime.annotation_budget_exhausted + assert runtime.delivery_budget_exhausted(), ( + "records left buffered by the annotation budget need their readiness " + "re-reported exactly as the record cap's do" + ) + + # And the annotation still closes, which is what raising would have + # prevented: a marker with no terminal is durable and undecodable past + # the frame after it. + delta = runtime.take_observation_delta() + assert delta is not None + terminal = runtime.add_terminal() + annotation = decode_annotation(delta + terminal) + assert annotation.terminal is not None + assert annotation.segments[-1].end_reason == SegmentEndReason.BUDGET_ROLLOVER, ( + "a segment cut by the byte budget has to say so: the batch continues " + "in the following marker, which no other end reason implies" + ) + finally: + await manager.shutdown() + + +async def test_an_activation_the_annotation_budget_stopped_is_not_wedged_by_it( + backend: MemoryStreamBackend, +) -> None: + """Stopping delivery obliges the same completion to ask for the rollover. + + These two are one mechanism and each is useless alone. The annotation budget + stops delivery so a segment can never overflow; the rollover is what gives the + next Workflow Task a fresh annotation to deliver into. Stop without asking and + the Workflow is **wedged**: the next activation begins against the same full + annotation, is handed a delivery budget of zero, delivers nothing, and so + observes nothing -- and a rollover condition that depended on having observed + something would then never become true again. + + That is not hypothetical. The flag is read after `take_observation_delta`, + which it must be, since that call is what closes the crossing segment -- and + that call also clears the observed flag. The frame here is deliberately larger + than the slack the *fractional* high-water mark leaves but small enough that + the mark itself is not crossed, so the mark cannot cover for it. + """ + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + from tests.contrib.external_workflow_streams.test_delivery_budget import ( + _CompletionStub, + ) + + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=10), + ) + runtime = _budget_runtime(backend, manager, max_annotation_bytes=1024) + long = "o" * 300 + try: + subscribe(runtime, 1) + subscribe(runtime, 2, name="tool-events") + + delivered = 0 + while runtime.delivery_budget_remaining() > 0 and delivered < 50: + wait_id = 1 + (delivered % 2) + delivered += 1 + runtime.record_delivery(wait_id, data(f"{long}{delivered}-0")) + + assert delivered < 50, "the annotation budget never stopped delivery" + assert runtime._accumulator is None or not ( + runtime._accumulator.request_rollover + ), ( + "the accumulator's own high-water mark is asking for this rollover, " + "so it would mask the condition under test -- the frame has to be " + "larger than the slack the fractional mark leaves without reaching " + "the mark itself" + ) + + stub = _CompletionStub(runtime) + _WorkflowInstanceImpl._emit_external_stream_commands(stub) # type: ignore[arg-type] + + commands = stub._current_completion.successful.commands + assert commands[0].workflow_stream_progress.request_rollover, ( + "delivery stopped for the annotation budget and the completion asked " + "for no rollover, so the next Workflow Task inherits a full " + "annotation and can never deliver again" + ) + assert ( + decode_annotation( + commands[0].workflow_stream_progress.observation_delta + ).terminal + is not None + ) + + # And the next Workflow Task can in fact deliver: the annotation was + # closed and a fresh one begins from its own header. + assert runtime.delivery_budget_remaining() > 0, ( + "the Workflow Task after the rollover is still unable to take a " + "single record" + ) + assert not runtime.annotation_budget_exhausted + finally: + await manager.shutdown() + + +async def test_the_first_record_of_an_activation_is_priced_from_a_measurement( + backend: MemoryStreamBackend, +) -> None: + """A per-segment maximum prices every activation's first record at the floor. + + ``close_segment`` empties the open segment at the end of each activation, so a + price taken from *that* segment's runs is back to the bare floor every time -- + and a real run costs two provider-chosen offset strings, which can be many + times the floor. Delivery then goes ahead on a price that is wrong by a factor, + the closing segment no longer fits, and what surfaces is a byte-budget error + that fails the Workflow. + + The measurement therefore belongs to the *annotation*, not the segment: the + largest run encoded since the header went out is what the next record costs + until a larger one appears. + """ + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=10), + ) + runtime = _budget_runtime(backend, manager, max_annotation_bytes=8192) + long = "9" * 300 + try: + subscribe(runtime, 1) + subscribe(runtime, 2, name="tool-events") + + # One activation's worth of long-offset runs, then the flush that closes + # its segment. The price has to survive that flush: `close_segment` empties + # the open segment, and a price read from it is back to the floor. + for i in range(4): + runtime.record_delivery(1 + (i % 2), data(f"{long}{i}-0")) + measured = runtime._max_run_bytes + assert measured > _RUN_COST_FLOOR, ( + "the long offsets did not even cost more than the floor, so this test " + "cannot tell a measurement from a guess" + ) + + runtime.take_observation_delta() + + assert runtime._run_sizes == [], "the open segment must have been emptied" + assert runtime._max_run_bytes == measured, ( + "the price fell back to the floor when the segment was closed, so the " + "first record of the next activation is priced at " + f"{_RUN_COST_FLOOR} instead of {measured} -- and delivering on that " + "price hands over a record the closing segment cannot record" + ) + + # And a fresh annotation *does* start from the floor again, which is right: + # it has measured nothing, and the measurement is a property of the + # annotation rather than of the Run. + runtime.add_terminal() + assert runtime._max_run_bytes == 0 + + # Driven to the end, nothing escapes and every annotation closes: the + # arithmetic never hands over a record it cannot then record. + total = 0 + for _ in range(40): + while runtime.delivery_budget_remaining() > 0 and total < 400: + total += 1 + runtime.record_delivery(1 + (total % 2), data(f"{long}{total}-0")) + delta = runtime.take_observation_delta() + if delta is not None and runtime.request_rollover: + assert decode_annotation(delta + runtime.add_terminal()).terminal + assert total > 0, "nothing was ever delivered, so this asserts nothing" + finally: + await manager.shutdown() + + +async def test_a_run_too_large_for_any_annotation_says_so_and_fails_the_workflow( + backend: MemoryStreamBackend, +) -> None: + """The one boundary no rollover can move, reported as itself. + + A record is priced before its offsets are seen, so a provider whose offsets are + far longer than anything measured can make one run cost more than the whole + budget has left -- and a fresh annotation has to carry that same run, so + rolling over changes nothing. That is a genuine capacity limit, and the two + things that matter are how it is *reported*: not as an internal byte-budget + error, which names nothing an operator can act on, and not as a Workflow Task + failure, which the server retries on an encoding that cannot succeed (ADR-007). + """ + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=10), + ) + runtime = _budget_runtime(backend, manager, max_annotation_bytes=2048) + try: + subscribe(runtime, 1) + with pytest.raises(ExternalStreamCapacityError) as caught: + for i in range(10): + runtime.record_delivery(1, data(f"{'9' * 5000}{i}-0")) + + assert "offsets" in str(caught.value), ( + f"the error must name what an operator can change: {caught.value}" + ) + assert caught.value.non_retryable, ( + "a retryable capacity limit is a Workflow Task the server retries " + "against an encoding that can never fit" + ) + finally: + await manager.shutdown() + + +async def test_a_subscription_set_too_large_to_record_is_refused_at_subscribe( + backend: MemoryStreamBackend, +) -> None: + """A header that cannot fit is rejected where the Workflow can still act on it. + + A header is one indivisible frame and a rollover writes a fresh one, so an + oversized header is not a rollover problem -- every annotation would be the + same size. Discovered while encoding a completion it fails the Workflow Task, + the server retries Workflow Task failures regardless of cause, and the retry + encodes the identical bytes: a Workflow stuck permanently with nothing durable + to say why. + + Raised from the ``subscribe()`` call instead, it is deterministic, reproduces + under replay, and names what to change. + """ + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=10), + ) + runtime = _budget_runtime(backend, manager, max_annotation_bytes=512) + try: + subscribe(runtime, 1) + with pytest.raises(ExternalStreamCapacityError) as caught: + subscribe(runtime, 2, name="s" * 600) + + assert "512" in str(caught.value) + assert "2" in str(caught.value), "the error must name the refused wait" + # Refused *before* it became replay-visible: a wait left half-registered + # would reach the next header as a binding nothing is watching. + assert set(runtime.subscriptions()) == {1} + assert manager.subscription(RUN_ID, 2) is None + + # And the Workflow Task that survives the refusal still completes: the + # annotation for the subscription it does hold encodes and closes. + runtime.record_delivery(1, data("1-0")) + annotation = annotation_of(runtime) + assert set(annotation.header.streams) == {1} # type: ignore[attr-defined] + assert annotation.terminal is not None # type: ignore[attr-defined] + finally: + await manager.shutdown() + + +async def test_the_capacity_floor_covers_everything_an_empty_annotation_carries( + backend: MemoryStreamBackend, +) -> None: + """Clearing header-plus-terminal is not clearing the floor. + + Every annotation also carries at least one **segment frame** -- an activation + that drained and observed nothing still encodes one, and the empty segment is + meaningful (ADR-018) -- and the **spill margin** a mispriced record overruns + into. A check that priced only the header and the terminal accepted a + subscription set that cleared it by a byte or two and then could not encode its + very first completion, which relocates the failure rather than preventing it. + """ + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=10), + ) + try: + # One wait whose stream name makes its header large enough that the + # remaining segment and spill reserve cross the budget. + # The budget below fits both with room to spare -- and does not fit the + # 3-byte segment frame and the 64-byte margin behind them. + name = "s" * 200 + runtime = _budget_runtime(backend, manager, max_annotation_bytes=292) + header_and_terminal = len( + encode_header( + AnnotationHeader( + { + 1: StreamBinding( + stream_key=runtime.stream_key(name), + start_cursor=BEGINNING, + provider_id=MemoryStreamBackend.provider_id, + provider_format_version=( + MemoryStreamBackend.provider_format_version + ), + ) + } + ) + ) + ) + len(encode_terminal({1: BEGINNING})) + assert header_and_terminal <= 292, ( + "this set clears header-plus-terminal, which is the whole point: a " + "floor priced on those two alone accepts it" + ) + + with pytest.raises(ExternalStreamCapacityError) as caught: + runtime.register(wait_id=1, stream_key=runtime.stream_key(name)) + assert "margin" in str(caught.value), ( + f"the error must say what the floor covers: {caught.value}" + ) + assert set(runtime.subscriptions()) == set() + finally: + await manager.shutdown() + + +# --- the completion path closes every annotation it ends -------------------- + + +def _progress_delta(stub: object) -> bytes | None: + """The observation delta a completion is carrying, if any.""" + for command in stub._current_completion.successful.commands: # type: ignore[attr-defined] + if command.HasField("workflow_stream_progress"): + return command.workflow_stream_progress.observation_delta + return None + + +async def test_a_command_producing_completion_closes_its_annotation( + runtime: WorkflowStreamRuntime, +) -> None: + """Two Workflow Tasks in a row, each ended by the Workflow's own command. + + Core writes and clears a marker on each of them, so each delta has to be a + whole annotation: its own header, and its own terminal. Carrying the header + across the first would leave the second starting at whatever frame came + first, and the leading byte of an annotation is read as its schema version + -- so the failure is reported as an unsupported version rather than as + anything to do with framing. + """ + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + from tests.contrib.external_workflow_streams.test_delivery_budget import ( + _CompletionStub, + ) + + subscribe(runtime, 1) + runtime.record_delivery(1, data("1-0")) + first = _CompletionStub(runtime) + first._add_command() # a timer, an activity: server-bound, so the task ends + _WorkflowInstanceImpl._emit_external_stream_commands(first) # type: ignore[arg-type] + + runtime.record_delivery(1, data("2-0")) + second = _CompletionStub(runtime) + second._add_command() + _WorkflowInstanceImpl._emit_external_stream_commands(second) # type: ignore[arg-type] + + deltas = {"first": _progress_delta(first), "second": _progress_delta(second)} + # Decoded before anything is asserted about them: it is the *second* delta + # that loses its header, and asserting per delta as each is decoded would + # stop on the first one's own shortcoming instead. + decoded = {} + for name, delta in deltas.items(): + assert delta is not None, f"the {name} completion reported no progress" + decoded[name] = decode_annotation(delta) + + for name, annotation in decoded.items(): + assert annotation.header.streams, f"the {name} marker records no stream" + assert annotation.terminal is not None, ( + f"the {name} marker has no terminal, so nothing in it says where " + "that Workflow Task's deliveries stopped" + ) + + assert decoded["second"].terminal == {1: AFTER(Offset("2-0"))} + + +async def test_a_rollover_request_closes_the_annotation_it_splits( + backend: MemoryStreamBackend, +) -> None: + """A budget rollover ends the task too, and needs no finalization round trip. + + Core takes ``request_rollover`` as authoritative over the retention the same + completion asks for, writes the marker, and forces a replacement task -- + without asking for a terminal, because the progress command carrying the + request is supposed to have carried one. A completion that asked for the + split without closing the annotation would produce a marker with no terminal + *and* leave the next one headerless. + """ + from temporalio.worker._workflow_instance import _WorkflowInstanceImpl + from tests.contrib.external_workflow_streams.test_delivery_budget import ( + _CompletionStub, + ) + + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=10), + ) + runtime = WorkflowStreamRuntime( + manager=manager, + backend=backend, + run_id=RUN_ID, + namespace="ns", + workflow_id="wf", + first_execution_run_id="first", + data_converter=temporalio.converter.DataConverter.default, + default_idle_timeout=timedelta(seconds=1), + max_annotation_bytes=2048, + ) + try: + subscribe(runtime, 1) + subscribe(runtime, 2, name="tool-events") + + # Accumulated the way Core accumulates, by byte append: the marker + # carries every delta of the Workflow Task, not just its last one. + accumulated: list[bytes] = [] + delivered = 0 + while not runtime.request_rollover: + # Alternating, so every delivery is its own run -- the one workload + # that cannot be range-compressed and therefore the one that reaches + # the cap. + for wait_id in (1, 2): + delivered += 1 + runtime.record_delivery(wait_id, data(f"{delivered}-0")) + earlier = runtime.take_observation_delta() + if earlier is not None: + accumulated.append(earlier) + + # The Workflow is still blocked on both streams, so this completion asks + # for retention as well -- and the rollover wins over it. + stub = _CompletionStub(runtime) + _WorkflowInstanceImpl._emit_external_stream_commands(stub) # type: ignore[arg-type] + + delta = _progress_delta(stub) + assert delta is not None, "the completion carried no progress to roll over" + accumulated.append(delta) + commands = stub._current_completion.successful.commands + assert commands[0].workflow_stream_progress.request_rollover, ( + "the completion did not ask Core to roll the task over" + ) + assert decode_annotation(b"".join(accumulated)).terminal is not None, ( + "the annotation Core is about to write has no terminal, and the " + "rollover path asks for none" + ) + + # And the next annotation begins from a header of its own. + runtime.record_delivery(1, data("9999-0")) + following = runtime.take_observation_delta() + assert following is not None + assert decode_annotation(following).header.streams + finally: + await manager.shutdown() + + +async def test_a_fence_ends_the_segment_as_fence_reached( + runtime: WorkflowStreamRuntime, +) -> None: + """How replay knows the batch stopped at a fence rather than running dry.""" + subscribe(runtime, 1) + runtime.record_delivery(1, fence("1-0")) + + decoded = annotation_of(runtime) + + assert decoded.segments[0].end_reason == SegmentEndReason.FENCE_REACHED # type: ignore[attr-defined] + + +async def test_running_dry_ends_the_segment_as_no_data_available( + runtime: WorkflowStreamRuntime, +) -> None: + subscribe(runtime, 1) + runtime.record_delivery(1, data("1-0")) + + decoded = annotation_of(runtime) + + assert decoded.segments[0].end_reason == SegmentEndReason.NO_DATA_AVAILABLE # type: ignore[attr-defined] + + +async def test_the_stream_key_comes_from_the_chain_not_the_run( + runtime: WorkflowStreamRuntime, +) -> None: + """The stream spans the whole Continue-As-New chain.""" + key = runtime.stream_key("tokens") + + assert key.workflow_id == "wf" + assert key.first_execution_run_id != RUN_ID + assert key.stream_name == "tokens" diff --git a/tests/contrib/external_workflow_streams/test_runtime_only_jobs.py b/tests/contrib/external_workflow_streams/test_runtime_only_jobs.py new file mode 100644 index 0000000000..51ead11e8a --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_runtime_only_jobs.py @@ -0,0 +1,346 @@ +"""P19 — the jobs answered outside the synchronous Workflow thread. + +Two of the four stream activation jobs are *themselves* backend operations, and +a third has to be prepared by one. Routing them through ``activate()`` would put +a multi-second transaction inside a call running under a 2-second deadlock +timeout: it would fail the Workflow Task for a perfectly healthy backend, and get +worse the more records replay had to validate. +""" + +from __future__ import annotations + +import asyncio +import uuid +from datetime import timedelta + +import pytest + +import temporalio.converter +from temporalio.contrib.external_workflow_streams._backend import StreamKey +from temporalio.contrib.external_workflow_streams._manager import ( + ReadinessResult, + StreamSubscriptionManager, +) +from temporalio.contrib.external_workflow_streams._record import ( + Offset, + RecordKind, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._runtime import WorkflowStreamRuntime +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend +from tests.contrib.external_workflow_streams.test_replay import ( + annotation_for, + append_five, +) + +RUN_ID = "run-1" +DEADLOCK_TIMEOUT_SECONDS = 2 +"""What `_WorkflowWorker` gives an activation before declaring a deadlock.""" + + +async def _notify( + run_id: str, # pyright: ignore[reportUnusedParameter] + wait_id: int, # pyright: ignore[reportUnusedParameter] + generation: int, # pyright: ignore[reportUnusedParameter] +) -> str: + return ReadinessResult.ACCEPTED + + +class HostileBackend(MemoryStreamBackend): + """Raises on every provider method. + + Registering one of these is how "this path performs no backend I/O" is + asserted rather than trusted: if any layer reached for the provider, the + call would explode instead of quietly succeeding. + """ + + def _refuse(self, name: str): # type: ignore[no-untyped-def] + raise AssertionError(f"{name} must not be called on this path") + + async def append(self, key, record): # type: ignore[no-untyped-def] + self._refuse("append") + + async def read_range(self, key, first, last): # type: ignore[no-untyped-def] + self._refuse("read_range") + + async def read_after(self, key, after, *, max_records, block=None): # type: ignore[no-untyped-def] + self._refuse("read_after") + + async def install_park_intent(self, key, intent): # type: ignore[no-untyped-def] + self._refuse("install_park_intent") + + async def recheck(self, key, wait_id): # type: ignore[no-untyped-def] + self._refuse("recheck") + + +class SlowParkBackend(MemoryStreamBackend): + """A park handshake that takes longer than the deadlock timeout.""" + + async def install_park_intent(self, key, intent): # type: ignore[no-untyped-def] + await asyncio.sleep(DEADLOCK_TIMEOUT_SECONDS + 0.5) + return await super().install_park_intent(key, intent) + + +class SlowRangeBackend(MemoryStreamBackend): + """A recorded-range read that takes longer than the deadlock timeout. + + Not a contrived latency: an inclusive read over a whole recorded batch is + the operation that gets *slower* the more records the marker committed, so + this is the shape a large, healthy replay actually has. + """ + + async def read_range(self, key, first, last): # type: ignore[no-untyped-def] + await asyncio.sleep(DEADLOCK_TIMEOUT_SECONDS + 0.5) + return await super().read_range(key, first, last) + + +def make_runtime(backend, manager): # type: ignore[no-untyped-def] + return WorkflowStreamRuntime( + manager=manager, + backend=backend, + run_id=RUN_ID, + namespace="ns", + workflow_id="wf", + first_execution_run_id=uuid.uuid4().hex, + data_converter=temporalio.converter.DataConverter.default, + default_idle_timeout=timedelta(seconds=1), + ) + + +def make_manager(backend): # type: ignore[no-untyped-def] + return StreamSubscriptionManager( + backend=backend, + notify_ready=_notify, + watch_block=timedelta(milliseconds=10), + ) + + +# --- finalization performs no backend I/O (ADR-010) -------------------------- + + +@pytest.mark.asyncio +async def test_finalization_touches_no_provider_at_all() -> None: + """Asserted against a provider that raises, rather than trusted. + + The boundary is not "wherever the stream is now"; it is where this Workflow + Task's deliveries stopped, which was fixed the moment the last activation + returned. Refreshing it against the backend would be actively wrong -- it + could name a position replay must not reproduce -- so the only correct + number of provider calls here is zero. + """ + backend = HostileBackend() + manager = make_manager(backend) + runtime = make_runtime(backend, manager) + try: + # Register without going through the manager's watcher, which would call + # `read_after` -- what is under test is the finalization path only. + runtime._subscriptions.clear() + runtime.register( + wait_id=1, + stream_key=runtime.stream_key("tokens"), + ) + + terminal = runtime.add_terminal() + + assert terminal, "finalization must produce a terminal" + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_finalization_reports_where_delivery_stopped_not_where_the_stream_is() -> ( + None +): + """A record arriving mid-finalization changes nothing about the terminal. + + It belongs to the *next* Workflow Task and reaches Core through the normal + readiness path -- or, if none is open by then, through the wake Signal. + """ + backend = MemoryStreamBackend() + manager = make_manager(backend) + runtime = make_runtime(backend, manager) + try: + key = runtime.stream_key("tokens") + runtime.register(wait_id=1, stream_key=key) + runtime.record_delivery( + 1, + StreamRecord(RecordKind.DATA, b"a", "s", 0).placed_at(Offset("5-0")), + ) + + # Something lands in the stream while finalization is being answered. + await backend.append(key, StreamRecord(RecordKind.DATA, b"later", "s", 9)) + + assert runtime.blocked_snapshot()[1].offset == Offset("5-0"), ( + "the terminal must name where deliveries stopped, not where the " + "stream has since got to" + ) + finally: + await manager.shutdown() + + +# --- the park handshake outlives the deadlock timeout ------------------------ + + +@pytest.mark.asyncio +async def test_a_park_slower_than_the_deadlock_timeout_is_still_answered() -> None: + """It is answered off the Workflow thread, so its duration is irrelevant. + + Were this routed through `activate()`, a healthy-but-slow backend would fail + the Workflow Task -- and it would get worse under exactly the conditions + parking exists for. + """ + backend = SlowParkBackend() + manager = make_manager(backend) + runtime = make_runtime(backend, manager) + try: + key = runtime.stream_key("tokens") + runtime.register(wait_id=1, stream_key=key) + + started = asyncio.get_running_loop().time() + became_ready = await asyncio.wait_for( + manager.prepare_park(RUN_ID, 1, runtime.blocked_snapshot()), 30 + ) + elapsed = asyncio.get_running_loop().time() - started + + assert elapsed > DEADLOCK_TIMEOUT_SECONDS, ( + "this test is only meaningful if the handshake really did outlast " + f"the deadlock timeout, took {elapsed:.2f}s" + ) + assert became_ready is False, "nothing was appended, so nothing became ready" + assert await backend.park_intent(key, 1) is not None + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_a_recheck_that_finds_records_abandons_the_whole_park() -> None: + """All-or-nothing across the set, and every intent installed comes back out. + + A park confirmed for a set with a ready member would lose that member's + record until a producer happened to signal it awake. + """ + backend = MemoryStreamBackend() + manager = make_manager(backend) + runtime = make_runtime(backend, manager) + try: + first = runtime.stream_key("tokens") + second = runtime.stream_key("tool-events") + runtime.register(wait_id=1, stream_key=first) + runtime.register(wait_id=2, stream_key=second) + + # A record on *one* of the two streams. + await backend.append(first, StreamRecord(RecordKind.DATA, b"a", "s", 0)) + + became_ready = await manager.prepare_park(RUN_ID, 1, runtime.blocked_snapshot()) + + assert became_ready is True + assert await backend.park_intent(first, 1) is None + assert await backend.park_intent(second, 2) is None, ( + "the other subscription's intent must be removed too -- parking is " + "all-or-nothing across the complete set" + ) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_intents_are_installed_before_anything_is_rechecked() -> None: + """This ordering is what closes the append/park race. + + A producer appends *before* it observes the park generation, so an append is + either seen by the recheck or paired with a wake Signal. Rechecking before + every intent was installed would leave a window where it is neither. + """ + backend = MemoryStreamBackend() + order: list[str] = [] + + class RecordingBackend(MemoryStreamBackend): + async def install_park_intent(self, key, intent): # type: ignore[no-untyped-def] + order.append(f"install:{intent.wait_id}") + return await super().install_park_intent(key, intent) + + async def recheck(self, key, wait_id): # type: ignore[no-untyped-def] + order.append(f"recheck:{wait_id}") + return await super().recheck(key, wait_id) + + backend = RecordingBackend() + manager = make_manager(backend) + runtime = make_runtime(backend, manager) + try: + for wait_id, name in ((1, "tokens"), (2, "tool-events")): + runtime.register( + wait_id=wait_id, + stream_key=runtime.stream_key(name), + ) + + await manager.prepare_park(RUN_ID, 1, runtime.blocked_snapshot()) + + installs = [i for i, step in enumerate(order) if step.startswith("install")] + rechecks = [i for i, step in enumerate(order) if step.startswith("recheck")] + assert installs and rechecks + assert max(installs) < min(rechecks), ( + f"every intent must be installed before any recheck, got {order}" + ) + finally: + await manager.shutdown() + + +# --- replay's recorded reads outlive the deadlock timeout too ---------------- + + +@pytest.mark.asyncio +async def test_a_replay_read_slower_than_the_deadlock_timeout_is_still_answered() -> ( + None +): + """The reason the replay job is partitioned rather than answered in `_apply`. + + Routing it through ``activate()`` would put the recorded range reads inside + a call running under a 2-second deadlock timeout -- failing the Workflow + Task for a perfectly healthy backend, and getting worse the more records + replay had to validate. Out here the duration is simply irrelevant. + """ + backend = SlowRangeBackend() + manager = make_manager(backend) + key = StreamKey("ns", "wf", uuid.uuid4().hex, "tokens") + try: + placed = await append_five(backend, key) + + started = asyncio.get_running_loop().time() + plan = await asyncio.wait_for( + manager.prepare_replay(RUN_ID, annotation_for(key, placed)), 30 + ) + elapsed = asyncio.get_running_loop().time() - started + + assert elapsed > DEADLOCK_TIMEOUT_SECONDS, ( + "this test is only meaningful if the read really did outlast the " + f"deadlock timeout, took {elapsed:.2f}s" + ) + assert plan.total_records == len(placed) + finally: + await manager.shutdown() + + +@pytest.mark.asyncio +async def test_the_delivering_activation_reads_nothing_itself() -> None: + """Everything a replay activation needs is in memory before it starts. + + Preparation and delivery are two different moments, and the second one must + perform no I/O at all -- which is asserted here by taking the prepared plan + from a provider that refuses every call. + """ + backend = MemoryStreamBackend() + manager = make_manager(backend) + key = StreamKey("ns", "wf", uuid.uuid4().hex, "tokens") + try: + placed = await append_five(backend, key) + await manager.prepare_replay(RUN_ID, annotation_for(key, placed)) + + # From here on the provider is hostile. Delivery still has to work. + manager._backend = HostileBackend() + + plan = manager.take_replay_plan(RUN_ID) + + assert plan is not None + assert plan.total_records == len(placed) + finally: + await manager.shutdown() diff --git a/tests/contrib/external_workflow_streams/test_serialization_context.py b/tests/contrib/external_workflow_streams/test_serialization_context.py new file mode 100644 index 0000000000..4b1e768920 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_serialization_context.py @@ -0,0 +1,746 @@ +"""A stream record is converted with the same context as every other payload. + +A ``DataConverter`` may carry a codec or a payload converter that derives what +it does from *which Workflow* the payload belongs to -- an encryption key per +Workflow ID is the reason ``WithSerializationContext`` exists at all. The Worker +already binds a :py:class:`~temporalio.converter.WorkflowSerializationContext` +around every payload an activation carries, on both halves: ``decode_activation`` +on the Worker's loop and the Workflow instance's own converters on the Workflow +thread. + +A stream record is a payload like any other and has to arrive the same way. If +it does not, a per-Workflow key decrypts nothing: the record fails to decode, +the error travels with it, and the Workflow Task is retried forever against a +record that will never decode -- reported as a converter mismatch when the +configuration is in fact correct. + +The producer is half of the same statement, and the more dangerous half. Both +sides being context-*free* is self-consistent and works; fixing only the +consumer would break every deployment that works today, because the producer +would encrypt with no context while the consumer decrypts with a +Workflow-derived key. These tests therefore pin both sides to the same context. + +PYTEST_DONT_REWRITE: sandboxed fixture Workflows re-import this module, so pytest's +injected imports would make sandbox validation depend on pytest's import locks. +""" + +# pyright: reportUnusedImport=false +from __future__ import annotations + +import asyncio +import uuid +from collections.abc import Sequence +from datetime import timedelta + +import pytest + +import temporalio.api.common.v1 +import temporalio.converter +from temporalio import workflow +from temporalio.client import Client +from temporalio.contrib.external_workflow_streams._backend import StreamKey +from temporalio.contrib.external_workflow_streams._producer import ( + ExternalStreamProducer, + WorkflowChainKey, +) +from temporalio.contrib.external_workflow_streams._record import ( + BEGINNING, + RecordKind, + StreamRecord, +) +from temporalio.converter import ( + DataConverter, + DefaultPayloadConverter, + PayloadCodec, + SerializationContext, + WithSerializationContext, + WorkflowSerializationContext, +) +from temporalio.worker import Replayer, Worker +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend + +with workflow.unsafe.imports_passed_through(): + from temporalio.contrib.external_workflow_streams._api import external_stream + + +#: Marks the one payload these tests care about. Every other payload an +#: activation carries -- the Workflow's own argument, its result -- travels +#: through the same converter, and the Worker's handling of those is not what is +#: under test here. +_SENTINEL = b"context-probe" + +#: Contexts the codec saw, in the order it saw them. Module-level because the +#: converter is constructed by the Worker, not by the test. +_codec_contexts: list[SerializationContext | None] = [] + +#: Contexts the payload converter saw while converting a stream record. Written +#: from the Workflow thread, which is the half a codec-only probe cannot reach. +_converter_contexts: list[SerializationContext | None] = [] + +#: The same two observations, from the chain-keyed probes an offline replay is +#: driven through. Kept apart from the pair above so the replay assertions can +#: be exact rather than a suffix of everything the process has ever converted. +_chain_codec_contexts: list[SerializationContext | None] = [] +_chain_converter_contexts: list[SerializationContext | None] = [] + + +class ContextRequiredCodec(PayloadCodec, WithSerializationContext): + """A codec that cannot decode a stream record without a context. + + Deliberately *fails* rather than merely recording: a codec keyed on the + Workflow ID has nothing to decrypt with when it is handed no context, and a + probe that succeeded either way would pass against a converter that was + never bound at all. + + Only the sentinel payload is treated specially, so the Worker's own + ``decode_activation`` work -- which is separately context-bound already -- + contributes no observations of its own. + """ + + def __init__(self, context: SerializationContext | None = None) -> None: + self.context = context + + def with_context(self, context: SerializationContext) -> ContextRequiredCodec: + return ContextRequiredCodec(context) + + async def encode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + return list(payloads) + + async def decode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + for payload in payloads: + if _SENTINEL not in payload.data: + continue + # Recorded before the refusal, so the assertion below can name the + # context that was actually supplied instead of only reporting that + # nothing worked. + _codec_contexts.append(self.context) + if self.context is None: + raise RuntimeError( + "a stream record was decoded with no serialization context" + ) + return list(payloads) + + +class ContextRequiredPayloadConverter( + DefaultPayloadConverter, WithSerializationContext +): + """The same refusal, on the synchronous half that runs on the Workflow thread. + + The codec probe cannot reach here: ``prepare()`` runs the codec on the + Worker's loop and ``convert()`` runs the payload converter inside + ``activate()``, and the two halves are handed converters from two different + places. A fix applied to only one of them leaves the other context-free. + """ + + def __init__(self, context: SerializationContext | None = None) -> None: + super().__init__() + self.context = context + + def with_context( + self, context: SerializationContext + ) -> ContextRequiredPayloadConverter: + return ContextRequiredPayloadConverter(context) + + def from_payloads( + self, + payloads: Sequence[temporalio.api.common.v1.Payload], + type_hints: list[type] | None = None, + ) -> list[object]: + for payload in payloads: + if _SENTINEL not in payload.data: + continue + _converter_contexts.append(self.context) + if self.context is None: + raise RuntimeError( + "a stream record was converted with no serialization context" + ) + return super().from_payloads(payloads, type_hints) + + +def _chain_tag(context: SerializationContext | None) -> bytes: + """The Workflow identity a chain-keyed probe writes into the bytes. + + Both the namespace and the Workflow id, because an offline ``Replayer`` + keeps the Workflow id and substitutes only the namespace: a probe keyed on + the id alone reports success against a converter that never saw the + recorded namespace at all. + """ + assert isinstance(context, WorkflowSerializationContext), ( + f"a stream record's payload was handled with no Workflow context: {context}" + ) + return f"|{context.namespace}|{context.workflow_id}".encode() + + +def _chain_apply( + payloads: Sequence[temporalio.api.common.v1.Payload], + context: SerializationContext | None, + seen: list[SerializationContext | None], + encoding: bool, +) -> list[temporalio.api.common.v1.Payload]: + """Appends or strips the tag, refusing bytes written under another context. + + A probe that only *recorded* the context would pass against a converter + bound to the wrong Workflow, because a wrong context still decodes -- to a + different value, silently. Tagging makes the mismatch the failure it would + be for a real per-Workflow key. + + Only the sentinel payload is touched, so the Workflow's own argument and + result -- which an offline replay legitimately handles under the replay + harness's own namespace -- travel through untouched. + """ + out: list[temporalio.api.common.v1.Payload] = [] + for payload in payloads: + if _SENTINEL not in payload.data: + out.append(payload) + continue + if not encoding: + seen.append(context) + tag = _chain_tag(context) + copied = temporalio.api.common.v1.Payload() + copied.CopyFrom(payload) + if encoding: + copied.data = payload.data + tag + else: + if not payload.data.endswith(tag): + raise RuntimeError( + "a stream record written under one Workflow context was " + f"read back under {context}" + ) + copied.data = payload.data[: -len(tag)] + out.append(copied) + return out + + +class ChainKeyedCodec(PayloadCodec, WithSerializationContext): + """The asynchronous half, keyed on the Workflow the record belongs to.""" + + def __init__(self, context: SerializationContext | None = None) -> None: + self.context = context + + def with_context(self, context: SerializationContext) -> ChainKeyedCodec: + return ChainKeyedCodec(context) + + async def encode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + return _chain_apply(payloads, self.context, _chain_codec_contexts, True) + + async def decode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + return _chain_apply(payloads, self.context, _chain_codec_contexts, False) + + +class ChainKeyedPayloadConverter(DefaultPayloadConverter, WithSerializationContext): + """The synchronous half, keyed the same way and run on the Workflow thread. + + The half an offline ``Replayer`` gets wrong on its own: the manager binds + the codec to the stream key the marker recorded, while the converter that + runs inside ``activate()`` belongs to a runtime the harness built under its + own namespace. + """ + + def __init__(self, context: SerializationContext | None = None) -> None: + super().__init__() + self.context = context + + def with_context(self, context: SerializationContext) -> ChainKeyedPayloadConverter: + return ChainKeyedPayloadConverter(context) + + def to_payloads( + self, values: Sequence[object] + ) -> list[temporalio.api.common.v1.Payload]: + return _chain_apply( + super().to_payloads(values), self.context, _chain_converter_contexts, True + ) + + def from_payloads( + self, + payloads: Sequence[temporalio.api.common.v1.Payload], + type_hints: list[type] | None = None, + ) -> list[object]: + return super().from_payloads( + _chain_apply(payloads, self.context, _chain_converter_contexts, False), + type_hints, + ) + + +@workflow.defn +class CountProbeTokensWorkflow: + """Consumes records and returns only how many, never the values.""" + + @workflow.run + async def run(self, expected: int) -> int: + tokens = external_stream.with_options(idle_timeout=timedelta(seconds=30)).topic( + "tokens", type=str + ) + + seen = 0 + async for _ in tokens.subscribe(): + seen += 1 + if seen >= expected: + break + return seen + + +@pytest.fixture +def backend() -> MemoryStreamBackend: + return MemoryStreamBackend() + + +async def _await_observation(seen: list[SerializationContext | None]) -> None: + """Waits for the probe to fire, so a missing context is not read as a hang. + + A context-free decode raises, the error travels with the record, and the + Workflow Task is retried forever -- so waiting on the Workflow's *result* + would turn a wrong context into a timeout minutes later. The context is + observable before any of that, which is what makes the assertion sharp. + """ + for _ in range(200): + if seen: + return + await asyncio.sleep(0.1) + + +async def _publish_probe_record( + backend: MemoryStreamBackend, key: StreamKey, converter: DataConverter +) -> None: + from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec + + codec = StreamPayloadCodec(converter, str) + await backend.append( + key, + StreamRecord(RecordKind.DATA, await codec.encode(_SENTINEL.decode()), "p", 0), + ) + + +async def _stream_key(client: Client, handle) -> StreamKey: # type: ignore[no-untyped-def] + description = await handle.describe() + return StreamKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + "tokens", + ) + + +# --- the consumer ----------------------------------------------------------- + + +async def test_a_stream_records_codec_runs_with_the_consuming_workflows_context( + client: Client, backend: MemoryStreamBackend +) -> None: + """The asynchronous half, on the Worker's loop. + + The manager prepares records for every Run on the Worker, so it cannot hold + a bound converter of its own; the context has to come from the record's own + subscription. Getting it from the manager's construction instead would + decode one Workflow's records under another Workflow's key. + """ + _codec_contexts.clear() + config = client.config() + config["data_converter"] = DataConverter(payload_codec=ContextRequiredCodec()) + probe_client = Client(**config) + + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + probe_client, + task_queue=task_queue, + workflows=[CountProbeTokensWorkflow], + external_stream_backend=backend, + ): + handle = await probe_client.start_workflow( + CountProbeTokensWorkflow.run, + 1, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + key = await _stream_key(client, handle) + await _publish_probe_record(backend, key, DataConverter.default) + await _await_observation(_codec_contexts) + + assert _codec_contexts == [ + WorkflowSerializationContext( + namespace=client.namespace, workflow_id=handle.id + ) + ], ( + "the stream record's codec ran with the wrong serialization " + "context; a codec keyed on the Workflow decrypts nothing here, " + "while every other payload in the same activation is bound" + ) + assert await asyncio.wait_for(handle.result(), 30) == 1 + + +async def test_a_stream_records_conversion_runs_with_the_consuming_workflows_context( + client: Client, backend: MemoryStreamBackend +) -> None: + """The synchronous half, inside ``activate()``. + + ``convert()`` runs on the Workflow thread through the runtime's converter, + which is a different object from the manager's. The Workflow's own argument + is converted with a bound converter on this very thread, so a record that is + not is inconsistent with the activation it arrived in. + """ + _converter_contexts.clear() + config = client.config() + config["data_converter"] = DataConverter( + payload_converter_class=ContextRequiredPayloadConverter + ) + probe_client = Client(**config) + + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + probe_client, + task_queue=task_queue, + workflows=[CountProbeTokensWorkflow], + external_stream_backend=backend, + ): + handle = await probe_client.start_workflow( + CountProbeTokensWorkflow.run, + 1, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + key = await _stream_key(client, handle) + await _publish_probe_record(backend, key, DataConverter.default) + await _await_observation(_converter_contexts) + + assert _converter_contexts[:1] == [ + WorkflowSerializationContext( + namespace=client.namespace, workflow_id=handle.id + ) + ], ( + "the stream record was converted on the Workflow thread with the " + "wrong serialization context, although the Workflow's own argument " + "was converted with the right one on the same thread" + ) + assert await asyncio.wait_for(handle.result(), 30) == 1 + + +async def test_a_replayed_record_is_prepared_with_the_recorded_streams_context( + backend: MemoryStreamBackend, +) -> None: + """Replay prepares from the annotation, and must reach the same context. + + A replayed record never passes a subscription: the Workflow has not run far + enough to have called ``subscribe()`` again, so the live path's source for + the context does not exist yet. The annotation's own header carries the + stream key, which is why it records it -- and a replay prepared without it + would decode nothing on exactly the Runs that a Worker restart produces, + while the live path worked. + """ + from temporalio.contrib.external_workflow_streams._annotation import ( + Annotation, + AnnotationHeader, + Run, + Segment, + SegmentEndReason, + StreamBinding, + encode_annotation, + ) + from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec + from temporalio.contrib.external_workflow_streams._manager import ( + ReadinessResult, + StreamSubscriptionManager, + ) + from temporalio.contrib.external_workflow_streams._record import AFTER + + _codec_contexts.clear() + key = StreamKey("ns", "wf-replayed", uuid.uuid4().hex, "tokens") + placed = await backend.append( + key, + StreamRecord( + RecordKind.DATA, + await StreamPayloadCodec(DataConverter.default, str).encode( + _SENTINEL.decode() + ), + "p", + 0, + ), + ) + assert placed.offset is not None + + async def notify( + run_id: str, # pyright: ignore[reportUnusedParameter] + wait_id: int, # pyright: ignore[reportUnusedParameter] + generation: int, # pyright: ignore[reportUnusedParameter] + ) -> str: + return ReadinessResult.ACCEPTED + + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=notify, + data_converter=DataConverter(payload_codec=ContextRequiredCodec()), + watch_block=timedelta(milliseconds=10), + ) + try: + await manager.prepare_replay( + "run-1", + encode_annotation( + Annotation( + header=AnnotationHeader( + streams={ + 1: StreamBinding( + stream_key=key, + start_cursor=BEGINNING, + provider_id=MemoryStreamBackend.provider_id, + provider_format_version=( + MemoryStreamBackend.provider_format_version + ), + ) + } + ), + segments=( + Segment( + runs=( + Run( + wait_id=1, + first_offset=placed.offset, + last_offset=placed.offset, + count=1, + ), + ), + end_reason=SegmentEndReason.NO_DATA_AVAILABLE, + ), + ), + terminal={1: AFTER(placed.offset)}, + ) + ), + ) + finally: + await manager.shutdown() + + assert _codec_contexts == [ + WorkflowSerializationContext(namespace="ns", workflow_id="wf-replayed") + ], ( + "the replayed record was prepared with the wrong serialization " + "context, so a Worker restart would fail to decode records the live " + "path decodes without complaint" + ) + + +async def test_an_offline_replay_converts_a_record_under_the_recorded_context( + client: Client, backend: MemoryStreamBackend +) -> None: + """Both halves of one record, through the real ``Replayer``. + + A ``Replayer`` runs under its own ``ReplayNamespace`` -- deliberately, and + :meth:`_verify_binding` accommodates it by comparing only the stream name + from a recorded binding. The record's *conversion* has to make the same + accommodation from the other side: the manager prepares a replayed record + under the stream key the marker recorded, so a converter left on the + harness's namespace would see two different Workflow identities while + decoding one payload -- and a namespace-keyed converter then refuses a + history that is entirely valid. + + Live execution cannot show this. There the recorded key and the runtime's + own identity are the same Workflow, so both halves agree whether or not + anything binds them separately. Only a harness that supplies its own + namespace pulls them apart, which is why this runs the history back through + the tool a user would. + """ + assert client.namespace != "ReplayNamespace", ( + "the Replayer's default namespace matches the live one, so this test " + "cannot tell a bound converter from an unbound one" + ) + _chain_codec_contexts.clear() + _chain_converter_contexts.clear() + converter = DataConverter( + payload_converter_class=ChainKeyedPayloadConverter, + payload_codec=ChainKeyedCodec(), + ) + config = client.config() + config["data_converter"] = converter + probe_client = Client(**config) + + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + probe_client, + task_queue=task_queue, + workflows=[CountProbeTokensWorkflow], + external_stream_backend=backend, + ): + handle = await probe_client.start_workflow( + CountProbeTokensWorkflow.run, + 1, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + key = await _stream_key(client, handle) + # Written under the *producing* Workflow's context, which is what the + # marker goes on to record and therefore what replay has to reproduce. + await _publish_probe_record( + backend, + key, + converter.with_context( + WorkflowSerializationContext( + namespace=key.namespace, workflow_id=key.workflow_id + ) + ), + ) + assert await asyncio.wait_for(handle.result(), 30) == 1 + history = await handle.fetch_history() + + assert _chain_codec_contexts and _chain_converter_contexts, ( + "the live run decoded no stream record, so there is nothing recorded " + "for the replay to disagree with" + ) + # Only the replay's observations are asserted on; the live run's are the + # subject of the two tests above. + _chain_codec_contexts.clear() + _chain_converter_contexts.clear() + + result = await Replayer( + workflows=[CountProbeTokensWorkflow], + data_converter=converter, + external_stream_backend=backend, + ).replay_workflow(history, raise_on_replay_failure=False) + + assert result.replay_failure is None, ( + "replaying a valid history failed because the record was converted " + "under the replay harness's namespace rather than the one the marker " + f"recorded: {result.replay_failure}" + ) + recorded = WorkflowSerializationContext( + namespace=client.namespace, workflow_id=handle.id + ) + assert _chain_codec_contexts == [recorded], ( + "the asynchronous half of the replayed record's decoding did not run " + f"with the recorded context: {_chain_codec_contexts}" + ) + assert _chain_converter_contexts == [recorded], ( + "the synchronous half ran with a different context than the " + "asynchronous half of the same record; one payload was decoded under " + f"two Workflow identities: {_chain_converter_contexts}" + ) + + +# --- the producer ----------------------------------------------------------- + + +async def test_the_producer_encodes_with_the_consuming_workflows_context( + backend: MemoryStreamBackend, +) -> None: + """The half that must not be forgotten. + + Producer and consumer share one ``DataConverter``, so they must also share + one context. Two context-*free* sides are self-consistent and work; one side + bound and the other not is a mismatch that appears only on the far side, as + a decode failure long after the append reported success. + + The chain key is the only Workflow identity a producer has, and it is the + right one: it names the whole Continue-As-New chain, which is what the + stream spans and what ``workflow_id`` means on the consumer too. + """ + recorded: list[SerializationContext | None] = [] + + class RecordingCodec(PayloadCodec, WithSerializationContext): + def __init__(self, context: SerializationContext | None = None) -> None: + self.context = context + + def with_context(self, context: SerializationContext) -> RecordingCodec: + return RecordingCodec(context) + + async def encode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + recorded.append(self.context) + return list(payloads) + + async def decode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + return list(payloads) + + producer = ExternalStreamProducer( + backend=backend, + workflow=WorkflowChainKey("ns", "wf", "run-1"), + data_converter=DataConverter(payload_codec=RecordingCodec()), + session_id="s", + ) + await producer.topic("tokens", type=str).publish("a", wake=False) + + assert recorded == [ + WorkflowSerializationContext(namespace="ns", workflow_id="wf") + ], ( + "the producer encoded the record with a different serialization " + "context than the consuming Worker decodes it with; the two sides " + "share one converter and must share one context" + ) + + +async def test_a_producers_record_decodes_on_a_context_bound_consumer( + backend: MemoryStreamBackend, +) -> None: + """The round trip, with a context that actually changes the bytes. + + A codec that transforms differently per Workflow is the only probe that can + tell "both sides bound to the same context" from "both sides bound to + nothing": the first survives this, and so does today's context-free pair, + but a fix applied to one side alone does not. + """ + + class PerWorkflowCodec(PayloadCodec, WithSerializationContext): + """XORs with a key derived from the Workflow ID, as a real one would.""" + + def __init__(self, context: SerializationContext | None = None) -> None: + self.context = context + + def with_context(self, context: SerializationContext) -> PerWorkflowCodec: + return PerWorkflowCodec(context) + + def _mask(self) -> int: + assert isinstance(self.context, WorkflowSerializationContext), ( + "this codec has no key without a Workflow context" + ) + return sum(self.context.workflow_id.encode()) % 251 or 7 + + def _apply( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + mask = self._mask() + out: list[temporalio.api.common.v1.Payload] = [] + for payload in payloads: + copied = temporalio.api.common.v1.Payload() + copied.CopyFrom(payload) + copied.data = bytes(b ^ mask for b in payload.data) + out.append(copied) + return out + + async def encode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + return self._apply(payloads) + + async def decode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + return self._apply(payloads) + + from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec + + converter = DataConverter(payload_codec=PerWorkflowCodec()) + producer = ExternalStreamProducer( + backend=backend, + workflow=WorkflowChainKey("ns", "wf", "run-1"), + data_converter=converter, + session_id="s", + ) + await producer.topic("tokens", type=str).publish("a", wake=False) + + key = StreamKey("ns", "wf", "run-1", "tokens") + records = await backend.read_after(key, BEGINNING, max_records=10, block=None) + data = [r for r in records if not r.is_control] + assert len(data) == 1 + + consumer = StreamPayloadCodec( + converter.with_context( + WorkflowSerializationContext(namespace="ns", workflow_id="wf") + ), + str, + ) + assert await consumer.decode(data[0].payload) == "a" diff --git a/tests/contrib/external_workflow_streams/test_shutdown_sweep.py b/tests/contrib/external_workflow_streams/test_shutdown_sweep.py new file mode 100644 index 0000000000..b735f32df2 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_shutdown_sweep.py @@ -0,0 +1,1250 @@ +"""P20 — the Worker shutdown wake sweep. + +Two obligations Core cannot discharge. Teardown ordering, so a finalization in +flight is always answered before the Run's state disappears; and the sweep +itself, because an idle cached Run gets no eviction activation at shutdown at +all -- and that is exactly the Run that most needs one, since its records are +buffered in a process about to exit. +""" + +# pyright: reportMissingParameterType=false +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from datetime import timedelta + +import pytest +import pytest_asyncio + +from temporalio.contrib.external_workflow_streams._backend import ( + ParkIntentRemoval, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._manager import ( + SHUTDOWN_WAKE_ATTEMPTS, + ReadinessResult, + RunStatus, + StreamSubscriptionManager, +) +from temporalio.contrib.external_workflow_streams._record import ( + BEGINNING, + RecordKind, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._wake import ( + WakeRequest, + wake_request_id, +) +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend + +RUN_ID = "run-1" + + +class Harness: + """A manager wired to recording stand-ins for Core and the Signal path.""" + + def __init__( + self, + backend: MemoryStreamBackend, + status: str, + *, + wake_fails: bool = False, + readiness: str = ReadinessResult.ACCEPTED, + ) -> None: + self.backend = backend + self.status = status + self.wake_fails = wake_fails + self.readiness = readiness + self.probes: list[str] = [] + self.wakes: list[tuple[str, int]] = [] + self.metric: list[str] = [] + self.manager = StreamSubscriptionManager( + backend=backend, + notify_ready=self._notify, + send_wake=self._wake, + run_status=self._probe, + shutdown_wake_failed_metric=self._on_metric, + watch_block=timedelta(milliseconds=10), + ) + + async def _notify( + self, + run_id: str, # pyright: ignore[reportUnusedParameter] + wait_id: int, # pyright: ignore[reportUnusedParameter] + generation: int, # pyright: ignore[reportUnusedParameter] + ) -> str: + return self.readiness + + async def _probe(self, run_id: str) -> str: + self.probes.append(run_id) + return self.status + + async def _wake(self, subscription) -> None: # type: ignore[no-untyped-def] + if self.wake_fails: + raise ConnectionError("service unavailable") + self.wakes.append((subscription.run_id, subscription.wait_id)) + + def _on_metric(self, subscription) -> None: # type: ignore[no-untyped-def] + self.metric.append(subscription.stream_key.stream_name) + + def register(self, wait_id: int = 1, stream_name: str = "tokens") -> StreamKey: + key = StreamKey("ns", "wf", "first-run", stream_name) + self.manager.register( + run_id=RUN_ID, + wait_id=wait_id, + stream_key=key, + start_cursor=BEGINNING, + ) + return key + + +def _swept_request_id(harness: Harness, subscription) -> str: # type: ignore[no-untyped-def] + """What the Worker's real sender derives for one swept subscription. + + Built from the manager's own sender identity rather than a literal, because + that identity is exactly what has to be fixed across one sweep's retries and + distinct between two Workers. + """ + key = subscription.stream_key + return wake_request_id( + WakeRequest( + namespace=key.namespace, + workflow_id=key.workflow_id, + first_execution_run_id=key.first_execution_run_id, + stream_name=key.stream_name, + wait_id=subscription.wait_id, + park_generation=0, + sender_identity=harness.manager.wake_sender_identity, + wake_counter=subscription.wake_counter, + ) + ) + + +@pytest.fixture +def backend() -> MemoryStreamBackend: + return MemoryStreamBackend() + + +@pytest_asyncio.fixture # type: ignore[reportUntypedFunctionDecorator] +async def harness_factory(backend: MemoryStreamBackend): + made: list[Harness] = [] + + def make(status: str, **kwargs) -> Harness: # type: ignore[no-untyped-def] + harness = Harness(backend, status, **kwargs) + made.append(harness) + return harness + + yield make + for harness in made: + if not harness.manager._shutting_down: + await harness.manager.shutdown() + + +# --- the probe ---------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_the_sweep_probes_rather_than_asserting_readiness( + harness_factory, +) -> None: + """Readiness means "a record is buffered" and would be a lie here. + + Probing with it would manufacture a spurious Workflow Task on the way out of + a Worker that is shutting down -- for a Run that may have nothing waiting at + all. + """ + harness = harness_factory(RunStatus.NO_OPEN_WORKFLOW_TASK) + harness.register() + + await harness.manager.shutdown() + + assert harness.probes == [RUN_ID], "the sweep must use the read-only probe" + + +@pytest.mark.asyncio +async def test_a_run_with_no_subscriptions_is_not_probed(harness_factory) -> None: + """Nothing is owed for a Run holding nothing.""" + harness = harness_factory(RunStatus.NO_OPEN_WORKFLOW_TASK) + + await harness.manager.shutdown() + + assert harness.probes == [] + + +# --- the four states ---------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_shutdown_in_the_no_open_task_window_sends_an_unparked_wake( + harness_factory, +) -> None: + """The window the whole sweep exists for. + + The Run is cached with no open Workflow Task, its records are buffered in a + process about to exit, and nothing else will ever tell the Workflow they + arrived. Waiting for an unrelated Workflow event is not a plan. + """ + harness = harness_factory(RunStatus.NO_OPEN_WORKFLOW_TASK) + harness.register() + + await harness.manager.shutdown() + + assert harness.wakes == [(RUN_ID, 1)] + assert harness.metric == [] + + +@pytest.mark.asyncio +async def test_shutdown_with_a_workflow_task_open_sends_nothing( + harness_factory, +) -> None: + """Core owns that transition; a wake here would race it. + + The result would be a second Workflow Task for a Run already being attended + to, which is waste at best and a duplicate resolve at worst. + """ + harness = harness_factory(RunStatus.WFT_OPEN) + harness.register() + + await harness.manager.shutdown() + + assert harness.wakes == [], "a Run with an open task must be left to Core" + + +@pytest.mark.asyncio +async def test_shutdown_on_a_parked_run_sends_nothing(harness_factory) -> None: + """A producer's next append wakes it through the ordinary path.""" + harness = harness_factory(RunStatus.PARKED) + harness.register() + + await harness.manager.shutdown() + + assert harness.wakes == [] + + +@pytest.mark.asyncio +async def test_a_run_core_no_longer_knows_still_gets_its_wake( + harness_factory, +) -> None: + """Evicted and cached-with-no-task differ in what happens next, not in what + is owed now. + + The Run is gone from *this* Worker, but the Workflow still exists and its + records are still in the stream; another Worker will pick it up from the + marker. + """ + harness = harness_factory(RunStatus.RUN_NOT_FOUND) + harness.register() + + await harness.manager.shutdown() + + assert harness.wakes == [(RUN_ID, 1)] + + +@pytest.mark.asyncio +async def test_every_subscription_of_a_swept_run_is_woken(harness_factory) -> None: + """Each is an independent wait with its own park intent. + + Waking one and not the other would leave the second waiting on records that + were already in its buffer. + """ + harness = harness_factory(RunStatus.NO_OPEN_WORKFLOW_TASK) + harness.register(wait_id=1, stream_name="a") + harness.register(wait_id=2, stream_name="b") + + await harness.manager.shutdown() + + assert sorted(harness.wakes) == [(RUN_ID, 1), (RUN_ID, 2)] + + +# --- failure is reported, never dropped --------------------------------------- + + +@pytest.mark.asyncio +async def test_an_unacknowledged_wake_surfaces_on_the_metric( + harness_factory, +) -> None: + """A dropped wake is silent by nature. + + The Workflow simply waits, and nothing distinguishes that from a producer + having nothing to say -- so it must be counted, not merely logged. + """ + harness = harness_factory(RunStatus.NO_OPEN_WORKFLOW_TASK, wake_fails=True) + harness.register() + + await harness.manager.shutdown() + + assert harness.metric == ["tokens"] + assert harness.manager.shutdown_wake_failures == 1 + assert harness.wakes == [], "a failed wake must never be counted as delivered" + + +@pytest.mark.asyncio +async def test_a_wake_with_no_sender_configured_is_reported_not_ignored( + backend: MemoryStreamBackend, +) -> None: + """Silently skipping it would report a clean shutdown that lost a record.""" + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=lambda *_: _accepted(), + run_status=lambda _: _no_open_task(), + watch_block=timedelta(milliseconds=10), + ) + manager.register( + run_id=RUN_ID, + wait_id=1, + stream_key=StreamKey("ns", "wf", "first-run", "tokens"), + start_cursor=BEGINNING, + ) + + await manager.shutdown() + + assert manager.shutdown_wake_failures == 1 + + +async def _accepted() -> str: + return ReadinessResult.ACCEPTED + + +async def _no_open_task() -> str: + return RunStatus.NO_OPEN_WORKFLOW_TASK + + +@pytest.mark.asyncio +async def test_a_probe_failure_does_not_stop_the_shutdown( + backend: MemoryStreamBackend, +) -> None: + """Shutdown is never blocked by a server that has stopped answering.""" + + async def failing_probe( + run_id: str, # pyright: ignore[reportUnusedParameter] + ) -> str: + raise ConnectionError("service unavailable") + + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=lambda *_: _accepted(), + run_status=failing_probe, + watch_block=timedelta(milliseconds=10), + ) + manager.register( + run_id=RUN_ID, + wait_id=1, + stream_key=StreamKey("ns", "wf", "first-run", "tokens"), + start_cursor=BEGINNING, + ) + + await asyncio.wait_for(manager.shutdown(), 5) + + assert manager._runs == {} + + +@pytest.mark.asyncio +async def test_shutdown_is_never_blocked_past_the_grace_period( + backend: MemoryStreamBackend, +) -> None: + """A wake that has not been acknowledged by now is better reported than + waited on indefinitely. + + Without the bound, a Worker that could not reach the server would hang on + the way out -- turning a recoverable delivery problem into an unrecoverable + process one. + """ + + async def hanging_probe( + run_id: str, # pyright: ignore[reportUnusedParameter] + ) -> str: + await asyncio.sleep(60) + return RunStatus.NO_OPEN_WORKFLOW_TASK + + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=lambda *_: _accepted(), + run_status=hanging_probe, + watch_block=timedelta(milliseconds=10), + ) + manager.register( + run_id=RUN_ID, + wait_id=1, + stream_key=StreamKey("ns", "wf", "first-run", "tokens"), + start_cursor=BEGINNING, + ) + + started = asyncio.get_running_loop().time() + await manager.shutdown(grace=timedelta(milliseconds=200)) + elapsed = asyncio.get_running_loop().time() - started + + assert elapsed < 5, f"shutdown waited {elapsed:.1f}s past its grace period" + assert manager._runs == {} + + +@pytest.mark.asyncio +async def test_a_grace_period_expiry_counts_every_wake_it_abandons( + backend: MemoryStreamBackend, +) -> None: + """A Worker may give up on a wake; it may not do so quietly. + + The grace period cancels the sweep wherever it happens to be, and + ``_send_owed_wake`` re-raises ``CancelledError`` by design -- so the failure + accounting that sits after it was never reached, and neither was any + subscription later in the serial loop. Shutdown then reported + ``shutdown_wake_failures == 0`` for a Worker that had just abandoned every + one of its handoffs, which is precisely the silence this counter exists to + break: a dropped wake looks exactly like a producer with nothing to say. + """ + hanging = asyncio.Event() + + class HangingWakeHarness(Harness): + async def _wake(self, subscription) -> None: # type: ignore[no-untyped-def] + hanging.set() + await asyncio.sleep(60) + + harness = HangingWakeHarness(backend, RunStatus.NO_OPEN_WORKFLOW_TASK) + for wait_id in (1, 2): + harness.manager.register( + run_id=RUN_ID, + wait_id=wait_id, + stream_key=StreamKey("ns", "wf", "first-run", f"tokens-{wait_id}"), + start_cursor=BEGINNING, + ) + + started = asyncio.get_running_loop().time() + await harness.manager.shutdown(grace=timedelta(milliseconds=200)) + elapsed = asyncio.get_running_loop().time() - started + + assert elapsed < 5, f"shutdown waited {elapsed:.1f}s past its grace period" + assert hanging.is_set(), "the wake never got as far as hanging" + assert harness.wakes == [], "no wake was acknowledged, so none may be reported" + assert harness.manager.shutdown_wake_failures == 2, ( + "both handoffs were abandoned -- the one the cancellation landed inside " + f"and the one never reached -- and {harness.manager.shutdown_wake_failures} " + "were counted" + ) + assert sorted(harness.metric) == ["tokens-1", "tokens-2"], ( + f"the metric must fire once per abandoned subscription: {harness.metric}" + ) + + +@pytest.mark.asyncio +async def test_a_probe_that_cannot_answer_is_not_reported_as_nothing_owed( + backend: MemoryStreamBackend, +) -> None: + """ "We could not tell" is not "nothing was owed". + + A Run whose status cannot be read may be holding a buffered record with + nowhere to announce it. Sending a wake anyway would race a Workflow Task that + might be open, so the sweep sends nothing -- but passing over the Run without + counting it tears the Run down and reports a clean shutdown, which is the same + silent loss by a different route. + """ + + class FailingProbeHarness(Harness): + async def _probe(self, run_id: str) -> str: + self.probes.append(run_id) + raise ConnectionError("service unavailable") + + harness = FailingProbeHarness(backend, RunStatus.NO_OPEN_WORKFLOW_TASK) + for wait_id in (1, 2): + harness.manager.register( + run_id=RUN_ID, + wait_id=wait_id, + stream_key=StreamKey("ns", "wf", "first-run", f"tokens-{wait_id}"), + start_cursor=BEGINNING, + ) + + await asyncio.wait_for(harness.manager.shutdown(grace=timedelta(seconds=2)), 5) + + assert harness.wakes == [] + assert harness.manager.shutdown_wake_failures == 2, ( + "a Run this Worker could say nothing about was torn down reporting a " + "clean shutdown" + ) + assert sorted(harness.metric) == ["tokens-1", "tokens-2"] + assert harness.manager._runs == {}, "shutdown must still complete" + + +@pytest.mark.asyncio +async def test_a_hanging_probe_counts_the_runs_it_never_answered_for( + backend: MemoryStreamBackend, +) -> None: + """The grace period expiring inside the probe is the same loss as inside the wake.""" + + class HangingProbeHarness(Harness): + async def _probe(self, run_id: str) -> str: + self.probes.append(run_id) + await asyncio.sleep(60) + return RunStatus.NO_OPEN_WORKFLOW_TASK + + harness = HangingProbeHarness(backend, RunStatus.NO_OPEN_WORKFLOW_TASK) + harness.manager.register( + run_id=RUN_ID, + wait_id=1, + stream_key=StreamKey("ns", "wf", "first-run", "tokens-1"), + start_cursor=BEGINNING, + ) + + await harness.manager.shutdown(grace=timedelta(milliseconds=200)) + + assert harness.wakes == [] + assert harness.manager.shutdown_wake_failures == 1 + assert harness.metric == ["tokens-1"] + + +@pytest.mark.asyncio +async def test_a_run_with_nothing_owed_is_not_counted_as_a_failure( + backend: MemoryStreamBackend, +) -> None: + """The counter has to stay quiet where it should, or it says nothing at all. + + A parked Run is woken by a producer's append through the ordinary path, and a + Run with an open Workflow Task is Core's to finish. Counting either as an + abandoned handoff would make the metric fire on every clean shutdown and + become unalertable. + """ + for status in (RunStatus.PARKED, RunStatus.WFT_OPEN): + harness = Harness(backend, status) + harness.manager.register( + run_id=RUN_ID, + wait_id=1, + stream_key=StreamKey("ns", "wf", "first-run", "tokens"), + start_cursor=BEGINNING, + ) + + await asyncio.wait_for(harness.manager.shutdown(grace=timedelta(seconds=2)), 5) + + assert harness.manager.shutdown_wake_failures == 0, ( + f"a Run reported as {status} owes no wake, so nothing was abandoned" + ) + assert harness.metric == [] + + +@pytest.mark.asyncio +async def test_a_wake_the_live_path_delivered_is_not_counted_as_abandoned( + backend: MemoryStreamBackend, +) -> None: + """The sweep's set is a snapshot, and the live path empties it underneath. + + Readiness that comes back `RunNotFound` owes a wake, sends it, and *then* drops + the subscription -- and `RunNotFound` is a likely answer during shutdown, while + watchers keep running for the whole grace window and the sweep awaits inside + itself for them to interleave with. A subscription that leaves the Run that way + has had its handoff made, so counting it reports a loss that did not happen on + the counter operators are told to alert on. + """ + other_run = "run-2" + harness = Harness(backend, RunStatus.PARKED) + for run_id in (RUN_ID, other_run): + harness.manager.register( + run_id=run_id, + wait_id=1, + stream_key=StreamKey("ns", "wf", run_id, "tokens"), + start_cursor=BEGINNING, + ) + + # While the sweep is busy with the first Run, the live path takes the second + # Run's subscription away -- which is what it does on `RunNotFound`, and only + # once its own owed wake has been acknowledged. The sweep then reaches that Run, + # finds no subscriptions, and moves on without saying anything about it. + # + # Both halves of that drop are reproduced, in the order the production path + # uses them: the flag first, then the pop. Popping alone takes the + # subscription out of the sweep's reach while leaving its watcher looping on + # `while not subscription._cancelled`, so the watcher outlives the test's + # event loop and reads from the backend after it closes -- an unraisable + # `RuntimeError: Event loop is closed` that this test's own shortcut, not the + # manager, is responsible for. + async def drop_the_other(run_id: str) -> str: + if run_id == RUN_ID: + dropped = harness.manager._runs.get(other_run, {}).pop(1, None) + if dropped is not None: + dropped._cancelled = True + return RunStatus.PARKED + + harness.manager._run_status = drop_the_other # type: ignore[assignment] + + await asyncio.wait_for(harness.manager.shutdown(grace=timedelta(seconds=2)), 5) + + assert harness.manager.shutdown_wake_failures == 0, ( + "a subscription whose wake the live path had already delivered was counted " + "as an abandoned handoff, on the counter operators are told to alert on" + ) + assert harness.metric == [] + + +@pytest.mark.asyncio +async def test_a_manager_with_no_probe_owes_nothing_and_reports_nothing( + backend: MemoryStreamBackend, +) -> None: + """No probe wired means no sweep, which means nothing to have failed at. + + The sweep is defined entirely in terms of what Core answers, so a manager with + no run-status probe has no handoff obligation at all. Counting its + subscriptions as abandoned wakes would make the metric fire wherever the + mechanism simply is not configured -- which is noise in exactly the series an + operator is expected to alert on. + """ + manager = StreamSubscriptionManager( + backend=backend, + notify_ready=lambda *_: _accepted(), + watch_block=timedelta(milliseconds=10), + ) + manager.register( + run_id=RUN_ID, + wait_id=1, + stream_key=StreamKey("ns", "wf", "first-run", "tokens"), + start_cursor=BEGINNING, + ) + + await asyncio.wait_for(manager.shutdown(grace=timedelta(seconds=2)), 5) + + assert manager.shutdown_wake_failures == 0 + assert manager._runs == {} + + +# --- teardown ordering -------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_the_sweep_runs_before_any_subscription_is_torn_down( + harness_factory, +) -> None: + """A subscription torn down first has nothing left to name in its wake. + + Teardown resets cursors to the committed boundary and cancels the watcher; + sweeping afterwards would send a wake for a wait whose buffered records had + already been discarded. + """ + harness = harness_factory(RunStatus.NO_OPEN_WORKFLOW_TASK) + key = harness.register() + await harness.backend.append(key, StreamRecord(RecordKind.DATA, b"x", "p", 0)) + await asyncio.sleep(0.1) + + await harness.manager.shutdown() + + assert harness.wakes == [(RUN_ID, 1)] + assert harness.manager._runs == {}, "teardown still happens, just afterwards" + + +@pytest.mark.asyncio +async def test_teardown_removes_every_run(harness_factory) -> None: + """The backstop half of the obligation: no Run outlives the manager. + + A Run whose watchers survived shutdown would hold a backend connection open + for the life of the process. + """ + harness = harness_factory(RunStatus.PARKED) + harness.register(wait_id=1, stream_name="a") + harness.register(wait_id=2, stream_name="b") + + await harness.manager.shutdown() + + assert harness.manager._runs == {} + + +@pytest.mark.asyncio +async def test_eviction_remains_the_normal_teardown_path(harness_factory) -> None: + """The sweep is a backstop, not a replacement. + + Per-Run teardown is driven by ``RemoveFromCache``, which is what guarantees a + finalization in flight is answered before the Run's state disappears. A + shutdown hook that tore Runs down itself would have no such ordering. + """ + harness = harness_factory(RunStatus.NO_OPEN_WORKFLOW_TASK) + harness.register() + + await harness.manager.evict_run(RUN_ID) + await harness.manager.shutdown() + + assert harness.probes == [], ( + "a Run already evicted has nothing left to sweep, and probing it would " + "ask Core about a Run this Worker has finished with" + ) + assert harness.wakes == [] + + +# --- the retry within the grace period ---------------------------------------- + + +class FlakyWakeHarness(Harness): + """Fails a fixed number of attempts, then succeeds.""" + + def __init__(self, backend: MemoryStreamBackend, failures: int) -> None: + super().__init__(backend, RunStatus.NO_OPEN_WORKFLOW_TASK) + self.remaining_failures = failures + self.attempts = 0 + + async def _wake(self, subscription) -> None: # type: ignore[no-untyped-def] + self.attempts += 1 + if self.remaining_failures > 0: + self.remaining_failures -= 1 + raise ConnectionError("service unavailable") + self.wakes.append((subscription.run_id, subscription.wait_id)) + + +@pytest.mark.asyncio +async def test_a_momentarily_failing_wake_is_retried_and_succeeds( + backend: MemoryStreamBackend, +) -> None: + """A Worker shutting down is often shutting down because something is unhealthy. + + That makes the first attempt the one most likely to land in the middle of + it, and giving up there would report a clean shutdown that lost a record. + """ + harness = FlakyWakeHarness(backend, failures=1) + harness.register() + + await harness.manager.shutdown() + + assert harness.attempts == 2, f"expected one retry, got {harness.attempts} attempts" + assert harness.wakes == [(RUN_ID, 1)] + assert harness.metric == [], "a wake that eventually succeeded is not a failure" + + +@pytest.mark.asyncio +async def test_the_retry_is_bounded_and_then_reported( + backend: MemoryStreamBackend, +) -> None: + """Retrying forever would trade a lost record for a Worker that never exits. + + The metric is what makes giving up visible, which is the only reason giving + up is acceptable. + """ + harness = FlakyWakeHarness(backend, failures=99) + harness.register() + + await harness.manager.shutdown() + + assert harness.attempts == SHUTDOWN_WAKE_ATTEMPTS + assert harness.metric == ["tokens"] + assert harness.wakes == [] + + +@pytest.mark.asyncio +async def test_the_retry_is_the_same_wake_not_a_second_one( + backend: MemoryStreamBackend, +) -> None: + """Derived from the wake's identity, so the server deduplicates it. + + This is what makes retrying safe at all: without a stable request ID the + retry would ask for a second Workflow Task, and an attempt that had in fact + arrived would be duplicated rather than resolved. + """ + harness = FlakyWakeHarness(backend, failures=1) + subscription = harness.register() + del subscription + + seen: list[str] = [] + + async def recording_wake(sub) -> None: # type: ignore[no-untyped-def] + seen.append(_swept_request_id(harness, sub)) + if len(seen) == 1: + raise ConnectionError("service unavailable") + + harness.manager._send_wake = recording_wake + + await harness.manager.shutdown() + + assert len(seen) == 2 + assert seen[0] == seen[1], ( + "the retry must derive the same request ID; a fresh one would ask for a " + "second Workflow Task rather than resolve the first attempt" + ) + + +@pytest.mark.asyncio +async def test_two_workers_sweeps_do_not_deduplicate_each_other( + backend: MemoryStreamBackend, +) -> None: + """The opposite half of the same requirement, and the reason it is per instance. + + Two Workers shutting down at different times are two separate asks, and they + share a ``Client`` -- so a sender identity taken from the client identity + would give both first unparked wakes the same request ID. The server would + deduplicate the second, and the Run the surviving Worker picked up would + never get its Workflow Task. + """ + harnesses = [Harness(backend, RunStatus.NO_OPEN_WORKFLOW_TASK) for _ in range(2)] + seen: list[str] = [] + + for harness in harnesses: + harness.register() + + async def recording_wake(sub, harness=harness) -> None: # type: ignore[no-untyped-def] + seen.append(_swept_request_id(harness, sub)) + + harness.manager._send_wake = recording_wake + + for harness in harnesses: + await harness.manager.shutdown() + + assert len(seen) == 2 + assert seen[0] != seen[1], ( + "both Workers' sweeps derived the same request ID, so the second " + "Worker's wake would be deduplicated away and its Run would stall" + ) + + +# --- the sweep driven through the Worker's real wake callback ----------------- + + +def _worker_wake_callback(manager): # type: ignore[no-untyped-def] + """The Worker's own sender, bound to this manager and nothing else. + + The retry and the metric are the manager's, but whether they can ever run is + the callback's: a callback that returns normally after a failed Signal makes + an unacknowledged wake indistinguishable from a delivered one, and the loop + below exits after its first attempt. Driving the real method is the only way + that shows. + """ + from temporalio.worker._workflow import _WorkflowWorker + + worker = object.__new__(_WorkflowWorker) + worker._client = object() # only ever handed to the patched sender + worker._external_stream_manager = manager + return worker._send_external_stream_wake + + +@pytest.fixture +def failing_signals(monkeypatch): # type: ignore[no-untyped-def] + """Makes the raw Signal call fail a fixed number of times, and counts it.""" + import temporalio.contrib.external_workflow_streams._wake as wake_module + + attempts: list[str] = [] + + def install(failures: int) -> list[str]: + remaining = [failures] + + async def fake_send( + client, # pyright: ignore[reportUnusedParameter] + wake_request, + *, + producer_session_id: str = "", # pyright: ignore[reportUnusedParameter] + ) -> str: + request_id = wake_request_id(wake_request) + attempts.append(request_id) + if remaining[0] > 0: + remaining[0] -= 1 + raise ConnectionError("service unavailable") + return request_id + + monkeypatch.setattr(wake_module, "send_wake_signal", fake_send) + return attempts + + return install + + +@pytest.mark.asyncio +async def test_a_signal_failure_reaches_the_sweeps_retry( + backend: MemoryStreamBackend, failing_signals +) -> None: + """The retry is only real if the callback tells the sweep it failed. + + A callback that logs and returns reports success, the three-attempt loop + ends after one call, and the record the Worker was holding is never + announced -- while shutdown reports itself clean. + """ + attempts = failing_signals(failures=2) + harness = Harness(backend, RunStatus.NO_OPEN_WORKFLOW_TASK) + harness.register() + harness.manager._send_wake = _worker_wake_callback(harness.manager) + + await harness.manager.shutdown() + + assert len(attempts) == 3, ( + f"the sweep made {len(attempts)} Signal attempts; a failure the callback " + "swallows leaves the retry loop with nothing to retry" + ) + assert len(set(attempts)) == 1, ( + "the retries must be the same wake, not three separate asks" + ) + assert harness.metric == [], "a wake that eventually succeeded is not a failure" + assert harness.manager.shutdown_wake_failures == 0 + + +@pytest.mark.asyncio +async def test_a_wake_that_never_lands_is_retried_then_counted( + backend: MemoryStreamBackend, failing_signals +) -> None: + """And the metric fires, which is the only thing that makes giving up acceptable.""" + attempts = failing_signals(failures=99) + harness = Harness(backend, RunStatus.NO_OPEN_WORKFLOW_TASK) + harness.register() + harness.manager._send_wake = _worker_wake_callback(harness.manager) + + await harness.manager.shutdown() + + assert len(attempts) == SHUTDOWN_WAKE_ATTEMPTS + assert harness.metric == ["tokens"], ( + "the wake was never acknowledged and shutdown reported nothing" + ) + assert harness.manager.shutdown_wake_failures == 1 + + +class _RecordedSignals(list): # type: ignore[type-arg] + """The wake requests each Signal would carry, and what a send costs. + + The request ID alone is not enough for these tests: what is under test is + *which park* the wake names, and an obsolete generation reaches the service + exactly as happily as a current one. + + ``delay`` is what makes a send a round trip rather than a function call. A + shutdown step that merely *schedules* a Signal passes every assertion a + zero-cost send can make, because the task it left behind runs on the way out + of the next await either way. + """ + + delay: float = 0.0 + + +@pytest.fixture +def recorded_signals(monkeypatch) -> _RecordedSignals: # type: ignore[no-untyped-def] + import temporalio.contrib.external_workflow_streams._wake as wake_module + + requests = _RecordedSignals() + + async def fake_send( + client, # pyright: ignore[reportUnusedParameter] + wake_request, + *, + producer_session_id: str = "", # pyright: ignore[reportUnusedParameter] + ) -> str: + if requests.delay: + await asyncio.sleep(requests.delay) + requests.append(wake_request) + return wake_request_id(wake_request) + + monkeypatch.setattr(wake_module, "send_wake_signal", fake_send) + return requests + + +def _removals_failing_until(harness: Harness, resume: Callable[[], bool]) -> None: + """Fails every park-intent removal until ``resume`` says the backend recovered.""" + original = harness.backend.remove_park_intent_if_matches + + async def flaky_remove( + key: StreamKey, + wait_id: int, + *, + run_id: str, + park_generation: int, + ) -> ParkIntentRemoval: + if not resume(): + raise ConnectionError("backend unavailable") + return await original( + key, + wait_id, + run_id=run_id, + park_generation=park_generation, + ) + + harness.backend.remove_park_intent_if_matches = ( # type: ignore[method-assign] + flaky_remove + ) + + +@pytest.mark.asyncio +async def test_the_sweep_does_not_wake_on_an_intent_it_has_decided_to_remove( + backend: MemoryStreamBackend, recorded_signals: _RecordedSignals +) -> None: + """A Signal the service accepts is not the same as a wake Core accepts. + + The sweep composed its wake from whatever ``current_park_generation`` + answered, and a stale intent is precisely a generation Core has already + discarded. So the send succeeded, the subscription was resolved as a handoff + made, and ``shutdown_wake_failures`` reported a clean exit for a Worker whose + only wake was ignored on arrival. The manager holds the ledger that says the + intent is stale, so it is the only thing that can tell the two apart -- and + the honest composition for an intent already decided against is the unparked + wake, which carries no generation to be discarded. + """ + harness = Harness(backend, RunStatus.NO_OPEN_WORKFLOW_TASK) + key = harness.register() + harness.manager._send_wake = _worker_wake_callback(harness.manager) + # Never recovers, so the intent is still installed when the sweep runs and + # the only Signal of this shutdown is the sweep's own. + _removals_failing_until(harness, lambda: False) + await harness.manager.prepare_park(RUN_ID, 4, {1: BEGINNING}) + await harness.manager.resolve_park(RUN_ID) + assert await harness.backend.current_park_generation(key, 1) == 4, ( + "the failed removal must leave the stale intent installed" + ) + + await harness.manager.shutdown(grace=timedelta(seconds=1)) + + assert [request.park_generation for request in recorded_signals] == [0], ( + "the sweep named the generation behind an intent it had already decided " + "to remove, so Core discards the wake while the send reports success" + ) + assert harness.manager.shutdown_wake_failures == 0 + assert harness.metric == [] + + +@pytest.mark.asyncio +async def test_a_removal_the_final_pass_confirms_carries_its_own_handoff( + backend: MemoryStreamBackend, recorded_signals: _RecordedSignals +) -> None: + """The last removal of a Worker's life ends a suppression window too. + + An evicted Run has no subscription for the sweep to wake and none for the + cleanup to announce through, which is the whole reason the ledger is not + kept on the Subscription. When the backend recovers only after the retry + loops are cancelled -- the case :meth:`_final_owed_removal_pass` exists for + -- that pass is what ends the suppression, and it has to carry the handoff + itself: nothing runs after it, and a wake merely *scheduled* as the process + exits is a wake never sent. + + Which is what the send delay pins. The cleanup handoff is posted as a task, + because every drain reaches it holding the Run's park lock under a budget + meant for one backend call rather than for a Signal with retries behind it -- + so with a send that costs nothing, a shutdown that never waits for that task + still looks like one that did. + """ + harness = Harness(backend, RunStatus.NO_OPEN_WORKFLOW_TASK) + key = harness.register() + harness.manager._send_wake = _worker_wake_callback(harness.manager) + recorded_signals.delay = 0.05 + # Recovers exactly when the autonomous retries have been cancelled, so the + # final pass is the attempt that succeeds rather than the loop. + _removals_failing_until( + harness, lambda: harness.manager._stopping_owed_removal_retries + ) + await harness.manager.prepare_park(RUN_ID, 4, {1: BEGINNING}) + await harness.manager.resolve_park(RUN_ID) + await harness.manager.evict_run(RUN_ID) + assert harness.manager.subscription(RUN_ID, 1) is None, ( + "eviction must drop the Subscription -- that is the premise" + ) + assert await harness.backend.current_park_generation(key, 1) == 4 + + await harness.manager.shutdown(grace=timedelta(seconds=1)) + + assert await harness.backend.current_park_generation(key, 1) is None, ( + "the final pass did not retire the intent, so this proves nothing" + ) + assert [request.park_generation for request in recorded_signals] == [0], ( + "shutdown reported itself clean while the record the retired intent had " + "silenced was announced to nobody" + ) + assert harness.manager.shutdown_wake_failures == 0 + + +@pytest.mark.asyncio +async def test_the_sweeps_wake_is_not_sent_twice_by_the_final_pass( + backend: MemoryStreamBackend, recorded_signals: _RecordedSignals +) -> None: + """One obligation, not two. + + The sweep already composed this wait's wake as unparked *because* the ledger + said its intent was stale, which makes it the same generation-0 handoff the + removal owes once it finally goes through. Sending it again asks the server + for a second, empty Workflow Task. + """ + harness = Harness(backend, RunStatus.NO_OPEN_WORKFLOW_TASK) + key = harness.register() + harness.manager._send_wake = _worker_wake_callback(harness.manager) + _removals_failing_until( + harness, lambda: harness.manager._stopping_owed_removal_retries + ) + await harness.manager.prepare_park(RUN_ID, 4, {1: BEGINNING}) + await harness.manager.resolve_park(RUN_ID) + + await harness.manager.shutdown(grace=timedelta(seconds=1)) + + assert await harness.backend.current_park_generation(key, 1) is None + assert [request.park_generation for request in recorded_signals] == [0], ( + "the sweep's wake and the cleanup's are one handoff; sending both costs " + "an empty Workflow Task" + ) + + +# --- the park intent's lifetime ----------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_resolved_park_leaves_no_intent_behind(harness_factory) -> None: + """A park intent exists only while that park is actually outstanding. + + An aborted park was never the only park that ends. A **confirmed** one ends + too -- when a wake Signal or a fresh quiescent snapshot clears Core's + ``park_generation`` -- and nothing about that is visible in the backend, so + only the manager can take the intent back out. + + What a left-behind intent costs is not tidiness. It *is* the answer + ``current_park_generation`` gives everyone who asks: a producer choosing what + its wake Signal names gets a generation Core has already discarded and its + wake is ignored as stale, and the shutdown sweep gets the same answer and + sends a parked wake whose request ID -- deliberately independent of the + sender -- is byte-identical to the wake that resolved the park in the first + place, so the server deduplicates it and no Workflow Task is ever created. + """ + harness = harness_factory(RunStatus.NO_OPEN_WORKFLOW_TASK) + key = harness.register() + + confirmed = not await harness.manager.prepare_park(RUN_ID, 1, {1: BEGINNING}) + assert confirmed, "nothing was appended, so this park must have confirmed" + assert await harness.backend.current_park_generation(key, 1) == 1 + + await harness.manager.resolve_park(RUN_ID) + + assert await harness.backend.park_intent(key, 1) is None, ( + "the intent of a park that is over is still installed" + ) + assert await harness.backend.current_park_generation(key, 1) is None, ( + "a resolved park still reports a generation, so the next wake -- a " + "producer's or the shutdown sweep's -- will name a park Core no longer " + "recognises instead of the unparked wake it owes" + ) + + +@pytest.mark.asyncio +async def test_a_resolve_with_no_park_installed_touches_nothing( + harness_factory, +) -> None: + """A resolve is the ordinary delivery path, not a rare one. + + Every record delivered live arrives on a resolve activation, so removing + intents there unconditionally would put a backend write on the hot path for + a park that never existed. The manager mirrors what it installed and removes + only that. + """ + harness = harness_factory(RunStatus.NO_OPEN_WORKFLOW_TASK) + key = harness.register() + removals: list[int] = [] + + original = harness.backend.remove_park_intent_if_matches + + async def recording_remove( + removed_key: StreamKey, + wait_id: int, + *, + run_id: str, + park_generation: int, + ) -> ParkIntentRemoval: + removals.append(wait_id) + return await original( + removed_key, + wait_id, + run_id=run_id, + park_generation=park_generation, + ) + + harness.backend.remove_park_intent_if_matches = ( # type: ignore[method-assign] + recording_remove + ) + + await harness.manager.resolve_park(RUN_ID) + assert removals == [], "a Run that never parked owes the backend nothing" + + await harness.manager.prepare_park(RUN_ID, 1, {1: BEGINNING}) + await harness.manager.resolve_park(RUN_ID) + await harness.manager.resolve_park(RUN_ID) + + assert removals == [1], ( + "the intent must come out exactly once: once for the park that was " + "installed, and never again for resolves with nothing outstanding" + ) + assert await harness.backend.park_intent(key, 1) is None + + +@pytest.mark.asyncio +async def test_a_backend_failure_during_removal_leaves_it_owed( + harness_factory, +) -> None: + """Cleanup does not get to fail a Workflow Task, and does not give up either. + + A resolve is a delivery activation; raising here would trade a stale intent + for a repeated Workflow Task on the healthy path. Swallowing the failure + *and* forgetting the intent would be the original defect back again, so the + removal stays owed and the next resolve retries it. + """ + harness = harness_factory(RunStatus.NO_OPEN_WORKFLOW_TASK) + key = harness.register() + await harness.manager.prepare_park(RUN_ID, 1, {1: BEGINNING}) + + original = harness.backend.remove_park_intent_if_matches + failures = [True] + + async def flaky_remove( + removed_key: StreamKey, + wait_id: int, + *, + run_id: str, + park_generation: int, + ) -> ParkIntentRemoval: + if failures: + failures.pop() + raise ConnectionError("backend unavailable") + return await original( + removed_key, + wait_id, + run_id=run_id, + park_generation=park_generation, + ) + + harness.backend.remove_park_intent_if_matches = ( # type: ignore[method-assign] + flaky_remove + ) + + await harness.manager.resolve_park(RUN_ID) + assert await harness.backend.current_park_generation(key, 1) == 1 + + await harness.manager.resolve_park(RUN_ID) + assert await harness.backend.current_park_generation(key, 1) is None, ( + "a removal that failed once was never retried, so the intent is stale " + "for the rest of the Run's life" + ) + + +# --- when the probe is asked --------------------------------------------------- + + +@pytest.mark.asyncio +async def test_the_sweep_acts_on_what_the_probe_heard(harness_factory) -> None: + """The probe and the wakes belong at different moments, so they are separate. + + Core keeps the state the probe reports only until the Worker's shutdown is + initiated: an idle cached Run has no pending work, so ``shutdown_done`` is + satisfied by the first input after the shutdown token is cancelled and the + workflow-state lane ends there. Everything afterwards answers + ``RunNotFound``. + + That answer owes a wake, which is what makes the degradation invisible: a + sweep that asks too late still sends its wake and still looks correct, while + the two answers that mean *do not* send one stop occurring. This Run is + parked -- a producer's append reaches it through the ordinary path and it + needs nothing from the sweep -- and the sweep must still know that after + Core has forgotten. + """ + harness = harness_factory(RunStatus.PARKED) + harness.register() + + await harness.manager.probe_runs() + harness.status = RunStatus.RUN_NOT_FOUND + + await harness.manager.shutdown() + + assert harness.wakes == [], ( + "a parked Run was woken on the way out. Its state was read after Core " + "had dropped it, so the Parked branch could not be taken" + ) + + +@pytest.mark.asyncio +async def test_a_run_the_probe_never_reached_is_still_swept(harness_factory) -> None: + """Probing early must not become a way to miss a Run. + + A Workflow Task already in flight when the probe ran can still cache a Run + afterwards. Nothing is known about it, so it is asked -- and both answers + Core has left to give owe a wake. + """ + harness = harness_factory(RunStatus.NO_OPEN_WORKFLOW_TASK) + + await harness.manager.probe_runs() + harness.register() + + await harness.manager.shutdown() + + assert harness.wakes == [(RUN_ID, 1)] + assert harness.probes == [RUN_ID], ( + "a Run the probe phase never saw must still be asked about" + ) diff --git a/tests/contrib/external_workflow_streams/test_wake.py b/tests/contrib/external_workflow_streams/test_wake.py new file mode 100644 index 0000000000..97ec27d1ed --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_wake.py @@ -0,0 +1,1815 @@ +"""P14 — the producer wake-signal path. + +The wakeup used whenever no open Workflow Task can accept local readiness. Two of +its properties are the reason it does not reuse the public Signal API, and both +are asserted here rather than assumed: the envelope bypasses the user's +``DataConverter``, and the request ID is a pure function of the wake's identity. + +PYTEST_DONT_REWRITE: sandboxed fixture Workflows re-import this module, so pytest's +injected imports would make sandbox validation depend on pytest's import locks. +""" + +# pyright: reportUnusedImport=false, reportMissingParameterType=false +from __future__ import annotations + +import asyncio +import dataclasses +import pathlib +import re +import uuid +from datetime import timedelta + +import pytest + +import temporalio.api.common.v1 +import temporalio.bridge +import temporalio.converter +from temporalio import workflow +from temporalio.bridge.proto.external_stream.external_stream_pb2 import WakeSignal +from temporalio.client import Client +from temporalio.contrib.external_workflow_streams._backend import ( + AppendConflictError, + ParkIntent, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._producer import ( + AppendNotAcknowledgedError, + ExternalStreamProducer, + WakeNotAcknowledgedError, + WorkflowChainKey, +) +from temporalio.contrib.external_workflow_streams._record import ( + BEGINNING, + Offset, + RecordKind, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._wake import ( + UNPARKED_WAKE_GENERATION, + WAKE_SIGNAL_ENCODING, + WAKE_SIGNAL_ENVELOPE_VERSION, + WAKE_SIGNAL_MESSAGE_TYPE, + WAKE_SIGNAL_NAME, + WakeRequest, + build_signal_request, + send_wake_signal, + wake_request_id, +) +from temporalio.worker import Worker +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend + +CHAIN = WorkflowChainKey("ns", "wf-1", "first-run-1") + + +def request(**overrides) -> WakeRequest: # type: ignore[no-untyped-def] + fields = dict( + namespace="ns", + workflow_id="wf-1", + first_execution_run_id="first-run-1", + stream_name="tokens", + wait_id=1, + park_generation=7, + ) + fields.update(overrides) + return WakeRequest(**fields) # type: ignore[arg-type] + + +# --- the stable request ID ---------------------------------------------------- + + +def test_two_producers_retrying_one_wake_derive_the_identical_request_id() -> None: + """Which is what lets the server deduplicate them into a single wake. + + A generation is woken once. Producers racing to wake it must not each + produce a Workflow Task, and the public Signal path -- which draws a fresh + UUID per attempt -- would give exactly that. + """ + first_producer = request() + second_producer = request() + + assert wake_request_id(first_producer) == wake_request_id(second_producer) + + +def test_the_request_id_is_stable_across_processes() -> None: + """Pinned, because a derivation that drifted would silently double every wake. + + Hashing anything process-local -- ``hash()``, an object id, a random seed -- + would still pass a same-process equality test and fail in production. + """ + assert wake_request_id(request()) == "f5d6b411-3047-5b8e-a755-46cb74ad0ed8" + + +@pytest.mark.parametrize( + "field", + ["namespace", "workflow_id", "first_execution_run_id", "stream_name"], +) +def test_every_identity_field_changes_the_request_id(field: str) -> None: + """Otherwise two different wakes would deduplicate into one, losing one.""" + assert wake_request_id(request()) != wake_request_id(request(**{field: "other"})) + + +def test_the_wait_id_and_generation_change_the_request_id() -> None: + assert wake_request_id(request()) != wake_request_id(request(wait_id=2)) + assert wake_request_id(request()) != wake_request_id(request(park_generation=8)) + + +def test_fields_cannot_be_confused_with_each_other() -> None: + """A stream named with a separator must not collide with a different tuple. + + Delimiter-joining the identity would make ``("a", "b")`` and ``("a:b", "")`` + the same wake. + """ + left = request(workflow_id="a", stream_name="b") + right = request(workflow_id="a\x00b", stream_name="") + + assert wake_request_id(left) != wake_request_id(right) + + +# --- unparked wakes ----------------------------------------------------------- + + +def test_two_unparked_wakes_from_different_senders_differ() -> None: + """Generation 0 carries no attempt identity, so the sender must supply one. + + Without this, two Workers shutting down at different times would derive the + same request ID and the server would deduplicate the second wake away -- + turning a correct retry mechanism into silent loss. + """ + first = request(park_generation=0, sender_identity="worker-a", wake_counter=1) + second = request(park_generation=0, sender_identity="worker-b", wake_counter=1) + + assert wake_request_id(first) != wake_request_id(second) + + +def test_one_senders_successive_unparked_wakes_differ() -> None: + """They are two separate asks, not a retry of one.""" + first = request(park_generation=0, sender_identity="worker-a", wake_counter=1) + second = request(park_generation=0, sender_identity="worker-a", wake_counter=2) + + assert wake_request_id(first) != wake_request_id(second) + + +def test_retrying_one_unparked_wake_keeps_its_request_id() -> None: + """The counter is per attempt, not per send -- that is what makes it a retry.""" + attempt = request(park_generation=0, sender_identity="worker-a", wake_counter=3) + + assert wake_request_id(attempt) == wake_request_id( + request(park_generation=0, sender_identity="worker-a", wake_counter=3) + ) + + +def test_an_unparked_wake_without_a_sender_identity_is_refused() -> None: + """Refused at construction rather than producing a colliding request ID.""" + with pytest.raises(ValueError, match="sender identity"): + request(park_generation=UNPARKED_WAKE_GENERATION) + + +def test_a_parked_wake_needs_no_sender_identity() -> None: + """Its generation already identifies the attempt, for every sender alike.""" + assert wake_request_id(request(park_generation=7)) == wake_request_id( + request(park_generation=7, sender_identity="", wake_counter=0) + ) + + +# --- the Worker's sender identity --------------------------------------------- + + +class _StubBackend: + """Reports no confirmed park, so every wake here is an unparked one.""" + + async def current_park_generation(self, stream_key, wait_id): # type: ignore[no-untyped-def] + return None + + +class _StubSubscription: + def __init__(self, wake_counter: int = 1) -> None: + #: The manager keys its owed-removal ledger by Run, and the sender asks + #: it what generation this wake must name, so a target without one is + #: not a target the real sender can compose from. + self.run_id = "run-1" + self.stream_key = StreamKey("ns", "wf-1", "first-run-1", "tokens") + self.wait_id = 1 + #: The sender's sequence number for this wake, which is what the + #: unparked request ID is derived from. Held by the subscription but + #: drawn from the manager, so it does not restart when a Run comes back. + self.wake_counter = wake_counter + self.wakes_owed = wake_counter + self.backend = _StubBackend() + + +def _worker_sending_wakes(client_identity: str): # type: ignore[no-untyped-def] + """One Worker's real wake sender, wired to a manager and nothing else. + + Built around the actual ``_WorkflowWorker`` method rather than a + reimplementation of it, because the defect this guards was in *which* + identity that method passed, not in the derivation it passed it to. + """ + from temporalio.contrib.external_workflow_streams._manager import ( + StreamSubscriptionManager, + ) + from temporalio.worker._workflow import _WorkflowWorker + + async def notify_ready( + run_id: str, # pyright: ignore[reportUnusedParameter] + wait_id: int, # pyright: ignore[reportUnusedParameter] + generation: int, # pyright: ignore[reportUnusedParameter] + ) -> str: + raise AssertionError("readiness is not part of this path") + + worker = object.__new__(_WorkflowWorker) + worker._client = object() # only ever handed to the patched sender + worker._external_stream_manager = StreamSubscriptionManager( + backend=None, + notify_ready=notify_ready, + client_identity=client_identity, + ) + return worker + + +@pytest.fixture +def sent_request_ids(monkeypatch): # type: ignore[no-untyped-def] + """Records the request ID each wake would actually be sent under.""" + import temporalio.contrib.external_workflow_streams._wake as wake_module + + recorded: list[str] = [] + + async def fake_send( + client, # pyright: ignore[reportUnusedParameter] + wake_request, + *, + producer_session_id: str = "", # pyright: ignore[reportUnusedParameter] + ) -> str: + request_id = wake_request_id(wake_request) + recorded.append(request_id) + return request_id + + monkeypatch.setattr(wake_module, "send_wake_signal", fake_send) + return recorded + + +@pytest.mark.asyncio +async def test_two_workers_sharing_one_client_derive_different_request_ids( + sent_request_ids: list[str], +) -> None: + """The client identity is shared; the sender identity must not be. + + Two Workers in one process share a ``Client`` and so share its identity, and + each one's counter restarts at 1 -- so deriving from the client identity + gives both first unparked wakes the same request ID. The server + deduplicates the second, no Workflow Task is created, and the Run the + surviving Worker picked up stalls. + """ + first = _worker_sending_wakes("one-shared-client") + second = _worker_sending_wakes("one-shared-client") + + await first._send_external_stream_wake(_StubSubscription(wake_counter=1)) + await second._send_external_stream_wake(_StubSubscription(wake_counter=1)) + + assert sent_request_ids[0] != sent_request_ids[1], ( + "both Workers derived the same request ID, so the server would " + "deduplicate the second wake away and the Run would never be woken" + ) + + +@pytest.mark.asyncio +async def test_one_workers_retry_of_an_unparked_wake_keeps_its_request_id( + sent_request_ids: list[str], +) -> None: + """The identity is fixed for the Worker's lifetime, which is what makes it a retry. + + The shutdown sweep re-sends one owed wake within its grace period. If the + sender identity were redrawn per attempt the retry would ask for a second + Workflow Task instead of resolving the attempt that may in fact have + arrived. + """ + worker = _worker_sending_wakes("one-shared-client") + subscription = _StubSubscription(wake_counter=1) + + await worker._send_external_stream_wake(subscription) + await worker._send_external_stream_wake(subscription) + + assert len(sent_request_ids) == 2 + assert sent_request_ids[0] == sent_request_ids[1], ( + "the retry derived a fresh request ID, so it would wake the Workflow a " + "second time rather than deduplicate against the first attempt" + ) + + +@pytest.mark.asyncio +async def test_one_worker_retries_an_unacknowledged_wake_with_the_same_request_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A failed send remains retriable under the same request identity.""" + import temporalio.contrib.external_workflow_streams._wake as wake_module + + recorded: list[str] = [] + + async def fail_once( + client, # pyright: ignore[reportUnusedParameter] + wake_request, + *, + producer_session_id: str = "", # pyright: ignore[reportUnusedParameter] + ) -> str: + request_id = wake_request_id(wake_request) + recorded.append(request_id) + if len(recorded) == 1: + raise ConnectionError("service unavailable") + return request_id + + monkeypatch.setattr(wake_module, "send_wake_signal", fail_once) + worker = _worker_sending_wakes("one-shared-client") + subscription = _StubSubscription(wake_counter=1) + + with pytest.raises(ConnectionError, match="service unavailable"): + await worker._send_external_stream_wake(subscription) + await worker._send_external_stream_wake(subscription) + + assert len(recorded) == 2 + assert recorded[0] == recorded[1] + + +@pytest.mark.asyncio +async def test_a_workers_sender_identity_still_names_its_client() -> None: + """So a request ID stays traceable to a client in server-side logs. + + The per-instance part is what makes two Workers distinct; the client + identity is what makes either of them identifiable afterwards. + """ + worker = _worker_sending_wakes("client-identity-here") + + assert worker._external_stream_manager.wake_sender_identity.startswith( + "client-identity-here" + ) + + +# --- the envelope ------------------------------------------------------------- + + +def test_the_signal_bypasses_the_user_data_converter() -> None: + """Core must read this Signal, and Core has no access to a user codec. + + Asserted structurally: the payload's encoding is the protocol's own, and its + bytes parse as the protobuf envelope. A ``DataConverter`` that encrypted + payloads would make it unreadable to the only reader that matters. + """ + signal = build_signal_request(request(), identity="producer-1") + + (payload,) = signal.input.payloads + assert payload.metadata["encoding"] == WAKE_SIGNAL_ENCODING + assert payload.metadata["messageType"].decode() == WAKE_SIGNAL_MESSAGE_TYPE + + envelope = WakeSignal() + envelope.ParseFromString(payload.data) + assert envelope.envelope_version == WAKE_SIGNAL_ENVELOPE_VERSION + assert envelope.stream_name == "tokens" + assert envelope.wait_id == 1 + assert envelope.park_generation == 7 + assert envelope.first_execution_run_id == "first-run-1" + + +def test_the_signal_name_is_the_reserved_one_and_not_the_other_features() -> None: + """``workflow_streams`` is a different feature that coexists with this one.""" + signal = build_signal_request(request(), identity="producer-1") + + assert signal.signal_name == WAKE_SIGNAL_NAME == "__temporal_external_stream_wake" + assert not signal.signal_name.startswith("__temporal_workflow_stream") + + +def test_the_signal_carries_no_run_id() -> None: + """So it always lands on the current Run of the chain. + + A wake sent while a Continue-As-New is in flight must reach the successor, + not fail against a Run that has already closed. Chain identity travels inside + the envelope instead, which is what lets Core reject a reused Workflow ID. + """ + signal = build_signal_request(request(), identity="producer-1") + + assert signal.workflow_execution.workflow_id == "wf-1" + assert signal.workflow_execution.run_id == "" + + +def test_the_envelope_carries_no_stream_payload() -> None: + """The record's bytes stay in the backend; the Signal is pure identity.""" + signal = build_signal_request( + request(), identity="producer-1", producer_session_id="session-1" + ) + + raw = signal.SerializeToString() + envelope = WakeSignal() + envelope.ParseFromString(signal.input.payloads[0].data) + assert set(envelope.DESCRIPTOR.fields_by_name) == { + "envelope_version", + "stream_name", + "wait_id", + "park_generation", + "first_execution_run_id", + "producer_session_id", + }, "a new envelope field must be reviewed for whether it can carry user data" + assert b"session-1" in raw # diagnostics only, and the only free-form string + + +# --- the producer send sequence ----------------------------------------------- + + +class RecordingClient: + """Captures the raw service call without a server. + + The wake path is a raw ``SignalWorkflowExecution``, so what needs asserting + is the request itself -- not that some client method was reached. + """ + + def __init__(self, fail: bool = False) -> None: + self.sent: list = [] + self.fail = fail + self.namespace = "ns" + self.data_converter = temporalio.converter.DataConverter.default + self.service_client = self + self.config = self + self.identity = "producer-identity" + self.workflow_service = self + + async def signal_workflow_execution(self, request) -> None: # type: ignore[no-untyped-def] + if self.fail: + raise ConnectionError("service unavailable") + self.sent.append(request) + + +def make_producer(backend, client) -> ExternalStreamProducer: # type: ignore[no-untyped-def] + return ExternalStreamProducer( + backend=backend, + workflow=CHAIN, + data_converter=temporalio.converter.DataConverter.default, + session_id="session-1", + client=client, # type: ignore[arg-type] + ) + + +async def park(backend: MemoryStreamBackend, key: StreamKey, wait_id: int, gen: int): + await backend.install_park_intent( + key, + ParkIntent( + wait_id=wait_id, cursor=BEGINNING, park_generation=gen, run_id="run-1" + ), + ) + + +@pytest.mark.asyncio +async def test_a_producer_that_finds_a_wakeable_generation_sends_exactly_one_signal() -> ( + None +): + backend = MemoryStreamBackend() + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + await park(backend, topic.stream_key, wait_id=1, gen=4) + + await topic.publish("a", wake=False) + sent = await topic.wake() + + assert len(sent) == 1 + assert len(client.sent) == 1 + envelope = WakeSignal() + envelope.ParseFromString(client.sent[0].input.payloads[0].data) + assert envelope.park_generation == 4 + assert envelope.wait_id == 1 + + +@pytest.mark.asyncio +async def test_a_producer_that_loses_the_claim_sends_the_same_wake_anyway() -> None: + """A lost claim says someone *intends* to send. It is not evidence one did. + + The Signal goes out either way, and it costs no second Workflow Task: a + parked wake's request ID is derived from the generation and ignores sender + identity, so both producers issue byte-identical requests that the server + deduplicates into one wake. Staying silent here is what strands a parked + Workflow when the claim holder crashed between claiming and signalling. + """ + backend = MemoryStreamBackend() + key = StreamKey("ns", "wf-1", "first-run-1", "tokens") + await park(backend, key, wait_id=1, gen=4) + await backend.claim_park_generation( + key, 1, 4, claimant="someone-else", lease=timedelta(seconds=30) + ) + + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + await topic.publish("a", wake=False) + + assert len(await topic.wake()) == 1 + assert len(client.sent) == 1 + assert client.sent[0].request_id == wake_request_id(request(park_generation=4)), ( + "the claim holder's wake and this one must be the same request, or the " + "duplicate the claim exists to avoid becomes a second Workflow Task" + ) + + +@pytest.mark.asyncio +async def test_an_expired_claim_is_taken_over_rather_than_stranding_the_wake() -> None: + """A producer crashing between claim and Signal must not park the Workflow forever. + + Without lease expiry every other producer concludes the wake is already + handled, and the Workflow waits with data sitting in the stream. + """ + backend = MemoryStreamBackend() + key = StreamKey("ns", "wf-1", "first-run-1", "tokens") + await park(backend, key, wait_id=1, gen=4) + # The crashed producer's claim, already expired. + await backend.claim_park_generation( + key, 1, 4, claimant="crashed", lease=timedelta(milliseconds=1) + ) + await backend.expire_claims_for_test() + + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + await topic.publish("a", wake=False) + + assert len(await topic.wake()) == 1, "an expired claim must be takeable" + + +@pytest.mark.asyncio +async def test_a_stream_with_no_park_intent_gets_an_unparked_wake() -> None: + """Cached-with-no-open-task and evicted are invisible from out here (ADR-023). + + Staying silent because no park was observed loses the record until something + else happens to wake the Workflow. + """ + backend = MemoryStreamBackend() + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + + await topic.publish("a", wake=False) + assert len(await topic.wake()) == 1 + + envelope = WakeSignal() + envelope.ParseFromString(client.sent[0].input.payloads[0].data) + assert envelope.park_generation == UNPARKED_WAKE_GENERATION + + +@pytest.mark.asyncio +async def test_two_subscriptions_to_one_stream_are_woken_independently() -> None: + """They are two waits with their own cursors, generations, and claims. + + Waking only one leaves the other parked on a record it can already see. + """ + backend = MemoryStreamBackend() + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + await park(backend, topic.stream_key, wait_id=1, gen=4) + await park(backend, topic.stream_key, wait_id=2, gen=4) + + await topic.publish("a", wake=False) + await topic.wake() + + woken = [] + for signal in client.sent: + envelope = WakeSignal() + envelope.ParseFromString(signal.input.payloads[0].data) + woken.append(envelope.wait_id) + assert sorted(woken) == [1, 2] + + +@pytest.mark.asyncio +async def test_a_failed_wake_reports_unacknowledged_and_the_retry_completes_it() -> ( + None +): + """The record is durable, the wake is not, and only one of them may be retried. + + Reporting success here would make the "durable producer" row of the + wakeup-durability boundary false. + """ + backend = MemoryStreamBackend() + client = RecordingClient(fail=True) + topic = make_producer(backend, client).topic("tokens", type=str) + await park(backend, topic.stream_key, wait_id=1, gen=4) + offset = await topic.publish("a", wake=False) + + with pytest.raises(WakeNotAcknowledgedError) as caught: + await topic.wake() + + assert caught.value.pending, "the retry must know what is still owed" + + # The record stayed put: the append is not what failed. + assert await backend.read_range(topic.stream_key, offset, offset) + + client.fail = False + assert len(await topic.retry_wake(caught.value.pending)) == 1 + + +@pytest.mark.asyncio +async def test_a_retried_wake_reuses_the_original_request_id() -> None: + """Recomputing it would draw a fresh counter and defeat the deduplication.""" + backend = MemoryStreamBackend() + client = RecordingClient(fail=True) + topic = make_producer(backend, client).topic("tokens", type=str) + await topic.publish("a", wake=False) + + with pytest.raises(WakeNotAcknowledgedError) as caught: + await topic.wake() + expected = wake_request_id(caught.value.pending[0]) + + client.fail = False + assert await topic.retry_wake(caught.value.pending) == [expected] + + +class FlakyCoordinationBackend(MemoryStreamBackend): + """A backend whose append works and whose coordination calls do not. + + One named method fails once. Appending keeps working throughout, which is the + whole point: the record is durable before the failure happens, so what the + caller has to be told is which half of ``publish()`` failed. + """ + + def __init__(self, failing: str) -> None: + super().__init__() + self.failing = failing + self.failures = 0 + + def _maybe_fail(self, name: str) -> None: + if name == self.failing: + self.failing = "" + self.failures += 1 + raise ConnectionError(f"{name} is unavailable") + + async def parked_wait_ids(self, key): # type: ignore[no-untyped-def] + self._maybe_fail("parked_wait_ids") + return await super().parked_wait_ids(key) + + async def current_park_generation(self, key, wait_id): # type: ignore[no-untyped-def] + self._maybe_fail("current_park_generation") + return await super().current_park_generation(key, wait_id) + + async def claim_park_generation( + self, + key: StreamKey, + wait_id: int, + park_generation: int, + *, + claimant: str, + lease: timedelta, + ) -> bool: + self._maybe_fail("claim_park_generation") + return await super().claim_park_generation( + key, + wait_id, + park_generation, + claimant=claimant, + lease=lease, + ) + + +@pytest.mark.parametrize( + "failing", + ["parked_wait_ids", "current_park_generation", "claim_park_generation"], +) +@pytest.mark.asyncio +async def test_a_coordination_failure_after_the_append_is_still_unacknowledged( + failing: str, +) -> None: + """The wake is three steps, and all three are after a durable append. + + Only the Signal used to be inside the guarantee. The observe and claim steps + raised whatever the provider raised -- a bare ``ConnectionError`` -- which + passed straight through ``publish()``, because ``publish()`` catches only the + durable-but-unacknowledged error. The caller then had neither the offset nor + any statement about what had already landed, and its only obvious move -- + retry ``publish()`` -- appends the record a *second* time, under a new + sequence number and therefore a new idempotency key. + """ + backend = FlakyCoordinationBackend(failing) + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + await park(backend, topic.stream_key, wait_id=1, gen=4) + + with pytest.raises(WakeNotAcknowledgedError) as caught: + await topic.publish("a") + + assert backend.failures == 1, "the test did not exercise the failure it names" + assert caught.value.offset is not None, ( + "the caller has no other way to learn that the append it must not retry " + "did succeed" + ) + assert caught.value.restart, ( + "no Signal was composed, so `pending` is empty and only a fresh wake() " + "recovers -- a caller told to retry_wake([]) would get a no-op that " + "looks like success" + ) + assert not caught.value.pending + with pytest.raises(ValueError, match=r"call wake\(\) again"): + await topic.retry_wake(caught.value.pending) + + # Recovering the wake alone leaves exactly one record in the stream. + assert await topic.wake() + records = await backend.read_after(topic.stream_key, BEGINNING, max_records=10) + assert [r.sequence for r in records if r.kind == RecordKind.DATA] == [0], ( + "the stream holds more than the one record that was published" + ) + + +@pytest.mark.asyncio +async def test_a_coordination_failure_after_a_fence_is_still_unacknowledged() -> None: + """A fence is the record most likely to find the Workflow parked. + + Same guarantee as ``publish()``, and the duplicate it prevents is worse: + retrying ``finish_writing()`` appends a second fence, which reads back as a + producer session that ended twice. + """ + backend = FlakyCoordinationBackend("parked_wait_ids") + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + await park(backend, topic.stream_key, wait_id=1, gen=4) + + with pytest.raises(WakeNotAcknowledgedError) as caught: + await topic.finish_writing() + + assert caught.value.offset is not None + assert caught.value.restart + + assert await topic.wake() + records = await backend.read_after(topic.stream_key, BEGINNING, max_records=10) + assert len([r for r in records if r.kind == RecordKind.WRITE_FENCE]) == 1, ( + "the stream holds more than the one fence that was written" + ) + + +class BlockingCoordinationBackend(MemoryStreamBackend): + """Appends, then blocks once in one named coordination call. + + One-shot, so the recovery the test performs afterwards can get through. The + append is never blocked: the record is durable before the block is reached, + which is the state the whole contract is about. + """ + + def __init__(self, blocking: str) -> None: + super().__init__() + self.blocking = blocking + self.reached = asyncio.Event() + + async def _maybe_block(self, name: str) -> None: + if name != self.blocking: + return + self.blocking = "" + self.reached.set() + await asyncio.Event().wait() + + async def parked_wait_ids(self, key): # type: ignore[no-untyped-def] + await self._maybe_block("parked_wait_ids") + return await super().parked_wait_ids(key) + + async def current_park_generation(self, key, wait_id): # type: ignore[no-untyped-def] + await self._maybe_block("current_park_generation") + return await super().current_park_generation(key, wait_id) + + async def claim_park_generation(self, key, wait_id, generation, **kwargs): # type: ignore[no-untyped-def,override] + await self._maybe_block("claim_park_generation") + return await super().claim_park_generation(key, wait_id, generation, **kwargs) + + +class BlockingClient(RecordingClient): + """A client whose first Signal send never completes.""" + + def __init__(self) -> None: + super().__init__() + self.reached = asyncio.Event() + self.blocking = True + + async def signal_workflow_execution(self, request) -> None: # type: ignore[no-untyped-def] + if self.blocking: + self.blocking = False + self.reached.set() + await asyncio.Event().wait() + await super().signal_workflow_execution(request) + + +async def _recover(topic, error: WakeNotAcknowledgedError) -> None: # type: ignore[no-untyped-def] + """The recovery the error names, and only that one.""" + assert bool(error.pending) != bool(error.restart), ( + "the error must name exactly one recovery: `pending` to re-send, or " + "`restart` to compose a fresh wake. Both or neither leaves the caller " + "guessing about a record that is already durable" + ) + if error.pending: + await topic.retry_wake(error.pending) + else: + await topic.wake() + + +@pytest.mark.parametrize( + "stage", + ["parked_wait_ids", "current_park_generation", "claim_park_generation", "signal"], +) +@pytest.mark.asyncio +async def test_cancellation_after_the_append_is_an_unacknowledged_wake( + stage: str, +) -> None: + """Cancellation is not a third outcome, and it must not read like the first. + + ``CancelledError`` is a ``BaseException``, so it passed through every handler + on this path -- the coordination block re-raised it deliberately, the Signal + loop caught only ``Exception``, and ``publish()`` catches only the + durable-but-unacknowledged error. What escaped was a bare ``CancelledError`` + with no offset, no ``pending`` and no ``restart``: indistinguishable from + cancellation *before* the append, and so unrecoverable in both directions. The + caller could neither wake the record it did not know about nor re-publish the + value safely -- retrying ``publish()`` draws a new sequence number and appends + a second record. + + Parameterized over all four post-append stages, because the state they leave + is the same and only the recovery differs. + """ + backend = BlockingCoordinationBackend("" if stage == "signal" else stage) + signalling = BlockingClient() if stage == "signal" else None + client: RecordingClient = signalling or RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + await park(backend, topic.stream_key, wait_id=1, gen=4) + reached = signalling.reached if signalling is not None else backend.reached + + task = asyncio.ensure_future(topic.publish("a")) + await asyncio.wait_for(reached.wait(), 2) + task.cancel() + + with pytest.raises(WakeNotAcknowledgedError) as caught: + await task + + assert caught.value.cancelled, ( + "the cancellation is reported rather than swallowed: a caller that wants " + "to honour it after recovering the wake still has to know it was asked " + "to stop" + ) + assert caught.value.offset is not None, ( + "a cancellation carrying no offset cannot be told apart from one before " + "the append, so the caller cannot know whether the value is published" + ) + await _recover(topic, caught.value) + + records = await backend.read_after(topic.stream_key, BEGINNING, max_records=10) + assert [r.sequence for r in records if r.kind == RecordKind.DATA] == [0], ( + "recovering the wake alone left more than the one record that was published" + ) + assert client.sent, "the recovery sent no Signal, so the record stays unannounced" + + +@pytest.mark.asyncio +async def test_cancellation_after_a_fence_is_an_unacknowledged_wake() -> None: + """A fence carries the identical contract, and a worse duplicate. + + Retrying ``finish_writing()`` appends a second fence, which reads back as a + producer session that ended twice. + """ + backend = BlockingCoordinationBackend("parked_wait_ids") + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + await park(backend, topic.stream_key, wait_id=1, gen=4) + + task = asyncio.ensure_future(topic.finish_writing()) + await asyncio.wait_for(backend.reached.wait(), 2) + task.cancel() + + with pytest.raises(WakeNotAcknowledgedError) as caught: + await task + + assert caught.value.cancelled and caught.value.offset is not None + await _recover(topic, caught.value) + + records = await backend.read_after(topic.stream_key, BEGINNING, max_records=10) + assert len([r for r in records if r.kind == RecordKind.WRITE_FENCE]) == 1, ( + "the stream holds more than the one fence that was written" + ) + + +class BlockingEncodeCodec(temporalio.converter.PayloadCodec): + """A codec whose first ``encode`` never completes. + + The one place a cancellation is *knowably* pre-append: the record has not + been built, let alone handed to the backend, so no provider anywhere can have + committed it. + """ + + def __init__(self) -> None: + self.reached = asyncio.Event() + + async def encode(self, payloads): # type: ignore[no-untyped-def] + self.reached.set() + await asyncio.Event().wait() + return list(payloads) + + async def decode(self, payloads): # type: ignore[no-untyped-def] + return list(payloads) + + +@pytest.mark.asyncio +async def test_cancellation_before_the_append_stays_a_cancellation() -> None: + """The other half of the distinction, without which the first is meaningless. + + Nothing was sent to the backend, so there is nothing to say about a durable + record and nothing to recover: reporting an unacknowledged wake here would + send the caller to wake a record that does not exist, and refusing a + cancellation that costs nothing to honour buys nothing. + + "Before the append" has to mean *before the backend was called*, not merely + before it answered. Cancellation while the call is in flight is a different + state, and it is the one below (ADR-038). + """ + backend = MemoryStreamBackend() + client = RecordingClient() + codec = BlockingEncodeCodec() + producer = ExternalStreamProducer( + backend=backend, + workflow=CHAIN, + data_converter=dataclasses.replace( + temporalio.converter.DataConverter.default, payload_codec=codec + ), + session_id="session-1", + client=client, # type: ignore[arg-type] + ) + topic = producer.topic("tokens", type=str) + + task = asyncio.ensure_future(topic.publish("a")) + await asyncio.wait_for(codec.reached.wait(), 2) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert not await backend.read_after(topic.stream_key, BEGINNING, max_records=10) + assert not client.sent + assert not producer._unresolved, ( + "an append that was never made owes no recovery, and holding one would " + "block the stream for a record that cannot exist" + ) + + +# --- the append acknowledgement window ---------------------------------------- + + +class CommittingThenBlockingBackend(MemoryStreamBackend): + """A backend whose first ``append`` **commits** and then loses its answer. + + This is the shape a remote provider actually has: the Redis backend runs an + atomic script server-side and receives its result in a separate client-side + step, so a cancellation or a dropped connection between the two leaves a + durable record that the caller was never told about. A backend that blocks + *before* storing is indistinguishable from this one from the producer's side + -- which is the whole reason the window needs an outcome of its own. + """ + + def __init__(self, fail: BaseException | None = None) -> None: + super().__init__() + self.committed = asyncio.Event() + self.blocking = True + self._fail = fail + + async def append(self, key, record): # type: ignore[no-untyped-def] + placed = await super().append(key, record) + if self.blocking: + self.blocking = False + self.committed.set() + if self._fail is not None: + raise self._fail + await asyncio.Event().wait() + return placed + + +async def _cancel_after_commit(backend, call): # type: ignore[no-untyped-def] + """Runs ``call`` and cancels it once the backend has committed but not answered.""" + task = asyncio.ensure_future(call()) + await asyncio.wait_for(backend.committed.wait(), 2) + task.cancel() + return task + + +@pytest.mark.asyncio +async def test_cancellation_after_backend_commit_before_append_ack_is_recoverable() -> ( + None +): + """The record is durable and the caller was never told. Both must be true. + + An ``append()`` that does not return is not an append that did not happen: a + backend commits on its own side and only then answers. A bare + ``CancelledError`` here says nothing about the record, and the caller's + obvious move -- publishing the same value again -- draws a fresh sequence + number and puts it in the stream twice. + + So the window gets its own outcome, carrying the exact record, and the + recovery re-appends *that* record: byte-identical under the same + ``(session_id, sequence)``, which the backend contract makes a no-op + returning the original offset. + """ + backend = CommittingThenBlockingBackend() + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + + task = await _cancel_after_commit(backend, lambda: topic.publish("a", wake=False)) + + with pytest.raises(AppendNotAcknowledgedError) as caught: + await task + + error = caught.value + assert error.cancelled, ( + "a caller that wants to honour the cancellation after settling the " + "append still has to know it was asked to stop" + ) + assert error.record.sequence == 0 and error.record.kind == RecordKind.DATA + assert error.record.offset is None, ( + "the record is carried in the state it was sent in; one carrying an " + "offset would mean the append had been acknowledged after all" + ) + + committed = await backend.read_after(topic.stream_key, BEGINNING, max_records=10) + assert [r.sequence for r in committed] == [0], ( + "the fixture must actually have committed, or this asserts nothing" + ) + + offset = await topic.resolve_append(error.record, wake=error.wake) + + assert offset == committed[0].offset, ( + "the recovery must recover the original offset, not mint a second " + "position for a record that already had one" + ) + records = await backend.read_after(topic.stream_key, BEGINNING, max_records=10) + assert [(r.sequence, r.kind) for r in records] == [(0, RecordKind.DATA)], ( + "settling the append appended a second copy of a record that had already landed" + ) + + +@pytest.mark.asyncio +async def test_a_lost_append_response_is_recoverable_the_same_way() -> None: + """Cancellation is not what makes the window ambiguous -- the commit is. + + A provider exception raised after the server committed carries exactly as + little information as a cancellation does, so it must not reach the caller as + an ordinary failure either. + """ + backend = CommittingThenBlockingBackend(fail=ConnectionError("connection reset")) + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + + with pytest.raises(AppendNotAcknowledgedError) as caught: + await topic.publish("a", wake=False) + + assert not caught.value.cancelled + assert isinstance(caught.value.__cause__, ConnectionError) + + await topic.resolve_append(caught.value.record, wake=False) + + records = await backend.read_after(topic.stream_key, BEGINNING, max_records=10) + assert [(r.sequence, r.kind) for r in records] == [(0, RecordKind.DATA)] + + +@pytest.mark.asyncio +async def test_an_unresolved_append_refuses_the_publish_that_would_duplicate_it() -> ( + None +): + """The failure the outcome exists to prevent, prevented rather than described. + + A caller that treats the unknown outcome as a failure retries `publish()`, + which draws sequence 1 and appends the value a second time. The stream will + not take it. + """ + backend = CommittingThenBlockingBackend() + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + + task = await _cancel_after_commit(backend, lambda: topic.publish("a", wake=False)) + with pytest.raises(AppendNotAcknowledgedError) as caught: + await task + + with pytest.raises(AppendNotAcknowledgedError) as refused: + await topic.publish("a", wake=False) + assert refused.value.record.idempotency_key == caught.value.record.idempotency_key + assert "resolve_append" in str(refused.value) + + with pytest.raises(AppendNotAcknowledgedError): + await topic.finish_writing(wake=False) + + records = await backend.read_after(topic.stream_key, BEGINNING, max_records=10) + assert [(r.sequence, r.kind) for r in records] == [(0, RecordKind.DATA)], ( + "the refused publish still appended, which is the duplicate this whole " + "outcome exists to prevent" + ) + + await topic.resolve_append(caught.value.record, wake=False) + assert await topic.publish("b", wake=False), ( + "once the append is settled the stream has to be usable again, or the " + "refusal is a permanent wedge rather than a guard" + ) + + +@pytest.mark.asyncio +async def test_settling_an_unacknowledged_append_wakes_exactly_once() -> None: + """The recovery finishes the call, rather than half of it. + + The wake-enabled form owes both halves: the record's offset and one Signal. + Recovering only the append would leave a parked Workflow on a durable record, + which is the state the acknowledged-wake contract exists to refuse. + """ + backend = CommittingThenBlockingBackend() + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + await park(backend, topic.stream_key, wait_id=1, gen=4) + + task = await _cancel_after_commit(backend, lambda: topic.publish("a")) + with pytest.raises(AppendNotAcknowledgedError) as caught: + await task + + assert caught.value.wake, "the recovery has to know the call owed a wake" + assert not client.sent, "nothing can have been signalled before the append landed" + + await topic.resolve_append(caught.value.record) + + assert len(client.sent) == 1, ( + "the settled record was announced once, or not at all -- the recovery " + "defaults to what the interrupted call was doing rather than to a wake " + "policy of its own" + ) + records = await backend.read_after(topic.stream_key, BEGINNING, max_records=10) + assert [(r.sequence, r.kind) for r in records] == [(0, RecordKind.DATA)] + + +@pytest.mark.asyncio +async def test_an_unacknowledged_fence_settles_to_exactly_one_fence() -> None: + """A duplicate fence reads back as a producer session that ended twice.""" + backend = CommittingThenBlockingBackend() + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + + task = await _cancel_after_commit(backend, lambda: topic.finish_writing(wake=False)) + with pytest.raises(AppendNotAcknowledgedError) as caught: + await task + + assert caught.value.record.kind == RecordKind.WRITE_FENCE + await topic.resolve_append(caught.value.record) + + records = await backend.read_after(topic.stream_key, BEGINNING, max_records=10) + assert len([r for r in records if r.kind == RecordKind.WRITE_FENCE]) == 1, ( + "the stream holds more than the one fence that was written" + ) + assert not client.sent, ( + "the fence was appended with wake=False and the recovery invented a " + "Signal for it; the recovery finishes the interrupted call, it does not " + "choose a different one" + ) + + +@pytest.mark.asyncio +async def test_a_refused_append_preserves_the_unresolved_operations_recovery() -> None: + """The refusal reports the *unsettled* operation, not the call it refused. + + An append outcome is more than its record: whether a wake was owed, under + which lease, and whether cancellation is still to be honoured once the state + is settled. Filling those from whichever later call happened to be refused + makes the error's own instructions wrong -- a caller following them settles a + record that owed a Signal with ``wake=False``, and the parked Workflow stays + parked on a durable record nobody announced. + """ + backend = CommittingThenBlockingBackend() + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + await park(backend, topic.stream_key, wait_id=1, gen=4) + lease = timedelta(seconds=97) + + task = await _cancel_after_commit( + backend, lambda: topic.publish("a", wake=True, lease=lease) + ) + with pytest.raises(AppendNotAcknowledgedError) as caught: + await task + assert (caught.value.wake, caught.value.lease, caught.value.cancelled) == ( + True, + lease, + True, + ) + + with pytest.raises(AppendNotAcknowledgedError) as refused: + await topic.publish("b", wake=False, lease=timedelta(seconds=5)) + + assert refused.value.record == caught.value.record, ( + "the refusal must name the append that is actually unsettled" + ) + assert refused.value.wake is True, ( + "the refusal reported the refused call's wake policy beside the older " + "call's record, so recovering by its own fields drops a Signal" + ) + assert refused.value.lease == lease + assert refused.value.cancelled is True, ( + "the caller still has to know it was asked to stop; the refusal is not " + "a fresh, uncancelled attempt" + ) + assert refused.value.stream_key == topic.stream_key + + # Recovering by the refusal's own fields, which is what its message says. + await topic.resolve_append( + refused.value.record, wake=refused.value.wake, lease=refused.value.lease + ) + + assert len(client.sent) == 1, "the record is durable and nobody was told" + records = await backend.read_after(topic.stream_key, BEGINNING, max_records=10) + assert [(r.sequence, r.kind) for r in records] == [(0, RecordKind.DATA)] + + +@pytest.mark.asyncio +async def test_a_reinterrupted_append_recovery_preserves_its_latest_wake() -> None: + """A second lost response must not restore the first attempt's wake policy. + + ``resolve_append`` may deliberately override the unresolved operation's wake + and lease. If that append also commits without returning, the error describes + the recovery attempt that was interrupted. Retaining the older operation + beside that newer error makes the next defaulted recovery silently use stale + instructions -- here it returns success without the Signal the error says is + owed. + """ + + class LoseThreeResponsesAfterCommit(MemoryStreamBackend): + def __init__(self) -> None: + super().__init__() + self.attempt = 0 + self.second_committed = asyncio.Event() + + async def append(self, key, record): # type: ignore[no-untyped-def] + placed = await super().append(key, record) + self.attempt += 1 + if self.attempt == 1: + raise ConnectionError("the append committed but its response was lost") + if self.attempt == 2: + self.second_committed.set() + await asyncio.Event().wait() + if self.attempt == 3: + raise ConnectionError("the recovery response was also lost") + return placed + + backend = LoseThreeResponsesAfterCommit() + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + await park(backend, topic.stream_key, wait_id=1, gen=4) + + with pytest.raises(AppendNotAcknowledgedError) as first: + await topic.publish("a", wake=False) + + recovery_lease = timedelta(seconds=97) + recovery = asyncio.ensure_future( + topic.resolve_append(first.value.record, wake=True, lease=recovery_lease) + ) + await asyncio.wait_for(backend.second_committed.wait(), 2) + recovery.cancel() + with pytest.raises(AppendNotAcknowledgedError) as second: + await recovery + + assert second.value.record == first.value.record + assert second.value.wake is True + assert second.value.lease == recovery_lease + assert second.value.cancelled is True + assert not client.sent + + with pytest.raises(AppendNotAcknowledgedError) as third: + await topic.resolve_append(second.value.record) + + assert third.value.record == second.value.record + assert third.value.wake is True + assert third.value.lease == recovery_lease + assert third.value.cancelled is True, ( + "a later connection failure erased the cancellation the caller still " + "has to honour after settling the operation" + ) + + with pytest.raises(AppendNotAcknowledgedError) as refused: + await topic.publish("b", wake=False, lease=timedelta(seconds=5)) + + assert refused.value.record == third.value.record + assert refused.value.wake is True + assert refused.value.lease == recovery_lease + assert refused.value.cancelled is True + + offset = await topic.resolve_append(refused.value.record) + + assert len(client.sent) == 1, ( + "the second recovery defaulted to the original publish's wake=False " + "instead of the interrupted resolve_append's wake=True" + ) + records = await backend.read_after(topic.stream_key, BEGINNING, max_records=10) + assert [(r.sequence, r.offset) for r in records] == [(0, offset)], ( + "repeated recovery appended more than the one record being settled" + ) + + +@pytest.mark.asyncio +async def test_an_unknown_append_can_only_be_resolved_on_its_originating_topic() -> ( + None +): + """Idempotency is scoped to the stream, so the recovery has to be too. + + A ``StreamRecord`` names its producer session and sequence but not its + stream. On any *other* topic that key has never been used, so settling there + does not deduplicate against the copy that may already exist -- it appends + the value a second time, onto a stream no consumer of it is watching, and + leaves the real stream still blocked. + """ + backend = CommittingThenBlockingBackend() + client = RecordingClient() + producer = make_producer(backend, client) + tokens = producer.topic("tokens", type=str) + events = producer.topic("events", type=str) + + task = await _cancel_after_commit(backend, lambda: tokens.publish("a", wake=False)) + with pytest.raises(AppendNotAcknowledgedError) as caught: + await task + assert caught.value.stream_key == tokens.stream_key, ( + "the error has to name where it must be settled, since the record cannot" + ) + + with pytest.raises(ValueError, match="not on"): + await events.resolve_append(caught.value.record) + assert not await backend.read_after(events.stream_key, BEGINNING, max_records=10), ( + "the wrong topic took the record, so one logical input is now on two streams" + ) + + # Nor do different bytes under the outstanding key count as that record. + tampered = dataclasses.replace(caught.value.record, payload=b'"z"') + with pytest.raises(ValueError, match="does not hold these bytes"): + await tokens.resolve_append(tampered) + assert producer._unresolved, ( + "a rejected recovery must not clear the state that is still unsettled" + ) + + offset = await tokens.resolve_append(caught.value.record) + records = await backend.read_after(tokens.stream_key, BEGINNING, max_records=10) + assert [(r.sequence, r.offset) for r in records] == [(0, offset)] + assert await tokens.publish("b", wake=False), "the settled stream stays usable" + + +@pytest.mark.asyncio +async def test_recovery_is_bound_to_the_producer_instance_that_made_the_append() -> ( + None +): + """A same-session replacement producer is not the one that owes the append. + + Its sequence and wake counters start at zero, so settling there leaves both + invalid: the next publish reuses a sequence number the recovered record + already holds, and the next unparked wake re-derives a request ID an earlier, + *different* wake already used -- which the server deduplicates away, leaving + a durable record unannounced. + + The recovery for a producer that is gone is the one the Activity retry + already performs: rebuild with the same session id and re-run the same calls + in the same order. That re-derives the same sequences, so the backend + deduplicates the appends, and it leaves the new session's counters correct. + """ + backend = CommittingThenBlockingBackend() + backend.blocking = False # the first publish has to land normally + client = RecordingClient() + first = make_producer(backend, client) + topic = first.topic("tokens", type=str) + + first_offset = await topic.publish("a") + assert len(client.sent) == 1 + first_request_id = client.sent[0].request_id + + backend.blocking = True + task = await _cancel_after_commit(backend, lambda: topic.publish("b")) + with pytest.raises(AppendNotAcknowledgedError) as caught: + await task + + replacement = make_producer(backend, client) + assert replacement.session_id == first.session_id + with pytest.raises(ValueError, match="no unsettled append"): + await replacement.topic("tokens", type=str).resolve_append(caught.value.record) + + # The recovery that refusal names, and what makes it the right one. + replacement_topic = replacement.topic("tokens", type=str) + assert await replacement_topic.publish("a") == first_offset, ( + "re-running the same call has to deduplicate onto the original record" + ) + assert client.sent[1].request_id == first_request_id, ( + "re-running one attempt is a retry of its wake, which the server " + "collapses rather than turning into a second Workflow Task" + ) + + await replacement_topic.publish("b") + assert client.sent[2].request_id != first_request_id, ( + "the second value's wake collided with the first value's, so the server " + "deduplicates it away and the record stays unannounced" + ) + records = await backend.read_after(topic.stream_key, BEGINNING, max_records=10) + assert [r.sequence for r in records] == [0, 1], ( + "each value must appear exactly once, under the sequence its call drew" + ) + + +@pytest.mark.asyncio +async def test_a_conflicting_append_is_a_refusal_not_an_unknown_outcome() -> None: + """The one failure the contract lets the producer read as "did not land". + + ``AppendConflictError`` says the key was used with *different* bytes, so this + record definitely did not land and re-appending it would raise the identical + error. Folding it into the unknown outcome would send the caller to settle an + append that has no settlement, and would wedge the stream behind it. + """ + + class ConflictingBackend(MemoryStreamBackend): + async def append(self, key, record): # type: ignore[no-untyped-def] + raise AppendConflictError(record.idempotency_key) + + backend = ConflictingBackend() + producer = make_producer(backend, RecordingClient()) + topic = producer.topic("tokens", type=str) + + with pytest.raises(AppendConflictError): + await topic.publish("a", wake=False) + + assert not producer._unresolved + with pytest.raises(AppendConflictError): + await topic.publish("b", wake=False) + + +@pytest.mark.asyncio +async def test_settling_an_append_that_never_landed_appends_it_once() -> None: + """The other history the same call has to be right for. + + The producer cannot tell the two apart -- that is the premise -- so the + recovery has to leave one record whichever one happened. + """ + + class LosingBackend(MemoryStreamBackend): + """Loses its first append entirely, then behaves.""" + + def __init__(self) -> None: + super().__init__() + self.reached = asyncio.Event() + + async def append(self, key, record): # type: ignore[no-untyped-def] + if not self.reached.is_set(): + self.reached.set() + raise ConnectionError("connection reset before the write") + return await super().append(key, record) + + backend = LosingBackend() + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + + with pytest.raises(AppendNotAcknowledgedError) as caught: + await topic.publish("a", wake=False) + + assert not await backend.read_after(topic.stream_key, BEGINNING, max_records=10) + + offset = await topic.resolve_append(caught.value.record, wake=False) + + records = await backend.read_after(topic.stream_key, BEGINNING, max_records=10) + assert [(r.sequence, r.offset) for r in records] == [(0, offset)], ( + "settling an append that had not landed has to land it, under the " + "sequence number the interrupted call drew rather than a fresh one" + ) + + +@pytest.mark.asyncio +async def test_settling_a_record_from_another_session_is_refused() -> None: + """Identity is the whole recovery; a foreign record is a different record.""" + backend = MemoryStreamBackend() + topic = make_producer(backend, RecordingClient()).topic("tokens", type=str) + foreign = StreamRecord( + kind=RecordKind.DATA, + payload=b'"a"', + producer_session_id="someone-else", + sequence=0, + ) + + with pytest.raises(ValueError, match="not by this one"): + await topic.resolve_append(foreign, wake=False) + with pytest.raises(ValueError, match="already carries offset"): + await topic.resolve_append( + StreamRecord( + kind=RecordKind.DATA, + payload=b'"a"', + producer_session_id="session-1", + sequence=0, + ).placed_at(Offset("1-1")), + wake=False, + ) + + +@pytest.mark.asyncio +async def test_waking_without_a_client_says_so() -> None: + backend = MemoryStreamBackend() + producer = ExternalStreamProducer( + backend=backend, + workflow=CHAIN, + data_converter=temporalio.converter.DataConverter.default, + session_id="session-1", + ) + + with pytest.raises(RuntimeError, match="requires a Temporal client"): + await producer.topic("tokens", type=str).wake() + + +# --- the two halves of the envelope live in different repos ------------------- + + +def test_the_python_constants_match_cores() -> None: + """Core is the reader; a drifted constant here is an unreadable Signal there. + + Nothing else would catch it: Python's own tests would still pass, Core's own + tests would still pass, and the wake would simply stop working -- appearing + as a Workflow that waits out its idle timeout for no visible reason. + """ + core = ( + pathlib.Path(temporalio.bridge.__file__).parent + / "sdk-core" + / "crates" + / "protos" + / "src" + / "protos" + / "mod.rs" + ) + if not core.exists(): + pytest.skip("the Core submodule is not checked out") + source = core.read_text() + + for name, value in ( + ("WAKE_SIGNAL_NAME", f'"{WAKE_SIGNAL_NAME}"'), + ("WAKE_SIGNAL_MESSAGE_TYPE", f'"{WAKE_SIGNAL_MESSAGE_TYPE}"'), + ("WAKE_SIGNAL_ENVELOPE_VERSION", str(WAKE_SIGNAL_ENVELOPE_VERSION)), + ("UNPARKED_WAKE_GENERATION", str(UNPARKED_WAKE_GENERATION)), + ): + assert re.search(rf"{name}:\s*\S+\s*=\s*{re.escape(value)}\s*;", source), ( + f"{name} is {value} in Python but Core declares something else" + ) + + +# --- P6b: publish()'s acknowledged-wake contract ------------------------------ + + +@pytest.mark.asyncio +async def test_publish_completes_only_once_its_wake_is_acknowledged() -> None: + """Returning means durable *and* signalled, which is the whole contract. + + An append that lands but is never signalled leaves the Workflow parked on + data already sitting in the stream -- and a publish() that returned success + there would be reporting a delivery that never happened. + """ + backend = MemoryStreamBackend() + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + await park(backend, topic.stream_key, wait_id=1, gen=4) + + offset = await topic.publish("a") + + assert offset is not None + assert len(client.sent) == 1, "publish must not return before it has signalled" + + +@pytest.mark.asyncio +async def test_publish_acknowledges_the_wake_even_when_another_producer_claimed_it() -> ( + None +): + """A claim held by a producer that crashed must not report a wake nobody sent. + + The lease permits takeover *after* it expires; it schedules nobody to take + over. If this publish is the last producer action -- and here it is the only + one -- a claim-based silence leaves the Workflow parked on a durable record + with a ``publish()`` that reported success. + """ + backend = MemoryStreamBackend() + key = StreamKey("ns", "wf-1", "first-run-1", "tokens") + await park(backend, key, wait_id=1, gen=4) + # A producer that claimed the generation and then died before signalling. + # Its lease is unexpired, so the claim is not takeable. + await backend.claim_park_generation( + key, 1, 4, claimant="crashed", lease=timedelta(seconds=30) + ) + + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + + offset = await topic.publish("a") + + assert offset is not None + assert len(client.sent) == 1, ( + "publish returned an acknowledged wake that nobody sent; the parked " + "Workflow is stranded on a durable record" + ) + envelope = WakeSignal() + envelope.ParseFromString(client.sent[0].input.payloads[0].data) + assert envelope.park_generation == 4 + assert client.sent[0].request_id == wake_request_id(request(park_generation=4)), ( + "the wake the crashed claimant owed and this one are the same request, " + "so the server collapses them rather than creating a second task" + ) + + +@pytest.mark.asyncio +async def test_a_publish_whose_wake_fails_reports_unacknowledged() -> None: + """Not success. This is the "durable producer" row of the boundary table.""" + backend = MemoryStreamBackend() + client = RecordingClient(fail=True) + topic = make_producer(backend, client).topic("tokens", type=str) + await park(backend, topic.stream_key, wait_id=1, gen=4) + + with pytest.raises(WakeNotAcknowledgedError) as caught: + await topic.publish("a") + + assert caught.value.offset is not None, ( + "the caller has no other way to learn that the append it must not retry " + "did succeed" + ) + assert caught.value.pending + + +@pytest.mark.asyncio +async def test_retrying_the_wake_completes_an_unacknowledged_publish() -> None: + """And appends nothing: only the half that failed is retried.""" + backend = MemoryStreamBackend() + client = RecordingClient(fail=True) + topic = make_producer(backend, client).topic("tokens", type=str) + await park(backend, topic.stream_key, wait_id=1, gen=4) + + with pytest.raises(WakeNotAcknowledgedError) as caught: + await topic.publish("a") + before = await backend.read_after( + topic.stream_key, BEGINNING, max_records=100, block=None + ) + + client.fail = False + assert len(await topic.retry_wake(caught.value.pending)) == 1 + + after = await backend.read_after( + topic.stream_key, BEGINNING, max_records=100, block=None + ) + assert [r.offset for r in after] == [r.offset for r in before], ( + "the retry must complete the wake without re-appending the record" + ) + + +@pytest.mark.asyncio +async def test_the_unacknowledged_state_is_explicit_not_hidden() -> None: + """``wake=False`` says "durable but un-signalled" in the call itself. + + The point is that a caller cannot end up in that state by accident -- it is + reachable only by asking for it, or by a wake that failed loudly. + """ + backend = MemoryStreamBackend() + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + await park(backend, topic.stream_key, wait_id=1, gen=4) + + await topic.publish("a", wake=False) + assert client.sent == [] + + await topic.wake() + assert len(client.sent) == 1 + + +@pytest.mark.asyncio +async def test_a_fence_is_signalled_too() -> None: + """It is the record most likely to find the Workflow parked. + + A fence is what a producer appends when it has nothing more to say, so an + unsignalled one strands the Workflow for its whole idle timeout at exactly + the moment it was waiting to be told. + """ + backend = MemoryStreamBackend() + client = RecordingClient() + topic = make_producer(backend, client).topic("tokens", type=str) + await park(backend, topic.stream_key, wait_id=1, gen=4) + + await topic.finish_writing() + + assert len(client.sent) == 1 + + +# --- deduplication, against a real server ------------------------------------- + + +@workflow.defn +class WaitForeverWorkflow: + """Exists to receive Signals. Never completes on its own.""" + + @workflow.run + async def run(self) -> None: + await asyncio.Future() + + +async def _signalled_events(handle) -> int: # type: ignore[no-untyped-def] + return len( + [ + e + async for e in handle.fetch_history_events() + if e.HasField("workflow_execution_signaled_event_attributes") + ] + ) + + +async def _settled_signal_count(handle, expected: int) -> int: # type: ignore[no-untyped-def] + """Waits for the count to reach ``expected`` and then stay there. + + Counting once races persistence in both directions: too early and a wake + that *was* recorded is missed, too eager and a duplicate that is about to + appear is not. Waiting for the count to reach the expected value and hold + across a further interval distinguishes "the server collapsed them" from + "the second one has not landed yet". + """ + for _ in range(40): + if await _signalled_events(handle) >= expected: + break + await asyncio.sleep(0.1) + await asyncio.sleep(1) + return await _signalled_events(handle) + + +async def _start_waiter(client: Client, task_queue: str): # type: ignore[no-untyped-def] + handle = await client.start_workflow( + WaitForeverWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + description = await handle.describe() + return handle, description.raw_description.workflow_execution_info.first_run_id + + +async def test_two_producers_retrying_one_wake_are_deduplicated_by_the_server( + client: Client, +) -> None: + """ID equality is only half the claim; this is the half that matters. + + Two producers racing to wake one generation must cost one Workflow Task, not + two. Asserting that they derive the same string proves nothing about what + the server does with it -- and the server is the component that has to + collapse them. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker(client, task_queue=task_queue, workflows=[WaitForeverWorkflow]): + handle, first_run = await _start_waiter(client, task_queue) + try: + wake = request( + namespace=client.namespace, + workflow_id=handle.id, + first_execution_run_id=first_run, + park_generation=4, + ) + + # Two producers, one generation, sent independently. + for _ in range(2): + await send_wake_signal(client, wake) + + assert await _settled_signal_count(handle, 1) == 1, ( + "the server recorded both wakes; a producer retrying after an " + "ambiguous failure would then cost a second Workflow Task" + ) + finally: + await handle.terminate() + + +async def test_two_workers_unparked_wakes_are_both_delivered( + client: Client, +) -> None: + """The opposite requirement, and why the counter exists. + + Two Workers shutting down at different times are two separate asks. If their + request IDs collided the server would deduplicate the second away -- turning + a correct retry mechanism into silent loss, which is the failure the sender + identity and counter are there to prevent. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker(client, task_queue=task_queue, workflows=[WaitForeverWorkflow]): + handle, first_run = await _start_waiter(client, task_queue) + try: + for identity in ("worker-a", "worker-b"): + await send_wake_signal( + client, + request( + namespace=client.namespace, + workflow_id=handle.id, + first_execution_run_id=first_run, + park_generation=0, + sender_identity=identity, + wake_counter=1, + ), + ) + + assert await _settled_signal_count(handle, 2) == 2, ( + "one of the two Workers' wakes was deduplicated away; both are " + "separate asks and both must reach the Workflow" + ) + finally: + await handle.terminate() + + +async def test_one_senders_retry_stays_a_single_wake(client: Client) -> None: + """The unparked case still deduplicates a genuine retry. + + The counter distinguishes *attempts*, not sends, so re-sending one attempt + must collapse exactly as a parked wake does -- otherwise the shutdown + sweep's retry would wake the Workflow twice. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker(client, task_queue=task_queue, workflows=[WaitForeverWorkflow]): + handle, first_run = await _start_waiter(client, task_queue) + try: + attempt = request( + namespace=client.namespace, + workflow_id=handle.id, + first_execution_run_id=first_run, + park_generation=0, + sender_identity="worker-a", + wake_counter=7, + ) + await send_wake_signal(client, attempt) + await send_wake_signal(client, attempt) + + assert await _settled_signal_count(handle, 1) == 1 + finally: + await handle.terminate() diff --git a/tests/contrib/external_workflow_streams/test_worker_crash.py b/tests/contrib/external_workflow_streams/test_worker_crash.py new file mode 100644 index 0000000000..5987318e64 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_worker_crash.py @@ -0,0 +1,622 @@ +"""A Worker that dies before its marker commits, and the one that takes over. + +"Marker recording commits the cursor. Reading or delivering a record does not." +The unit tests assert the *consequence* of that inside one manager -- a reset +cursor re-delivers -- but the property is about two Workers and a Workflow Task +that never completed, and nothing in one process can show that. + +So the consuming Worker here runs in a child process and is sent ``SIGKILL`` +while it holds a retained Workflow Task with records consumed on it. No +shutdown, no finalization, no marker. The server times the task out, a second +Worker picks up the retry, and the question is whether it resumes from the last +*committed* cursor -- re-reading everything consumed after it, and nothing +before it. +""" + +from __future__ import annotations + +import asyncio +import sys +import uuid +from datetime import timedelta +from pathlib import Path +from typing import Any + +import pytest + +from temporalio.client import Client, WorkflowUpdateRPCTimeoutOrCancelledError +from temporalio.contrib.external_workflow_streams import ( + ExternalOutputStreamClient, + WorkflowChainKey, +) +from temporalio.contrib.external_workflow_streams._backend import ( + StreamDirection, + StreamKey, +) +from temporalio.contrib.external_workflow_streams._output_backend import ( + OutputStageStatus, +) +from temporalio.contrib.external_workflow_streams._record import BEGINNING +from temporalio.contrib.external_workflow_streams._redis import RedisStreamBackend +from temporalio.worker import Worker +from tests.contrib.external_workflow_streams.conftest import KEY_NAMESPACE, redis_url +from tests.contrib.external_workflow_streams.crash_worker import ( + STREAM_NAME, + CrashConsumeWorkflow, + CrashPublishWorkflow, + CrashUpdatePublishWorkflow, + output_stage_log_key, + read_log_key, + recording_backend, +) +from tests.contrib.external_workflow_streams.test_worker_integration import publish + +TASK_TIMEOUT_SECONDS = 6 +"""Short, because the crash is only visible once the server gives up on the +retained task, and that wait is this test's floor.""" + +FEED_GAP_SECONDS = 0.3 +"""Well under the idle timeout, so the retained task is never released. + +An idle park would complete the task and *commit* the records fed here, which is +exactly the state this test needs not to reach. +""" + +UNCOMMITTED = ["alpha", "beta", "gamma"] +"""Consumed on the retained task the Worker dies holding. Never committed.""" + +LATER = ["delta"] +"""Published to the replacement Worker, so the Run demonstrably continues.""" + +REPO_ROOT = Path(__file__).resolve().parents[3] + + +async def _reads(redis_client: Any, prefix: str, reader: str) -> list[str]: + raw = await redis_client.lrange(read_log_key(prefix, reader), 0, -1) + return [value.decode() if isinstance(value, bytes) else value for value in raw] + + +async def _history(handle: Any) -> list[Any]: + return [e async for e in handle.fetch_history_events()] + + +def _markers(events: list[Any]) -> list[Any]: + return [e for e in events if e.HasField("marker_recorded_event_attributes")] + + +async def _wait_until(predicate: Any, timeout: float, message: str) -> None: + """Polls rather than sleeping a fixed time, so a slow start is not a flake.""" + deadline = asyncio.get_running_loop().time() + timeout + while asyncio.get_running_loop().time() < deadline: + if await predicate(): + return + await asyncio.sleep(0.2) + raise AssertionError(message) + + +async def _spawn_crash_worker( + client: Client, + *, + task_queue: str, + prefix: str, + mode: str, +) -> asyncio.subprocess.Process: + return await asyncio.create_subprocess_exec( + sys.executable, + "-m", + "tests.contrib.external_workflow_streams.crash_worker", + "--address", + client.service_client.config.target_host, + "--namespace", + client.namespace, + "--task-queue", + task_queue, + "--redis-url", + redis_url(), + "--prefix", + prefix, + "--mode", + mode, + cwd=str(REPO_ROOT), + stdout=asyncio.subprocess.DEVNULL, + stderr=asyncio.subprocess.DEVNULL, + ) + + +def _output_marker_tokens(events: list[Any]) -> list[str]: + from temporalio.bridge.proto.external_data import ExternalStreamMarkerData + + tokens: list[str] = [] + for event in events: + if not event.HasField("marker_recorded_event_attributes"): + continue + attributes = event.marker_recorded_event_attributes + if attributes.marker_name != "core_external_stream": + continue + payloads = attributes.details["external_stream"].payloads + if len(payloads) != 1: + continue + marker = ExternalStreamMarkerData() + marker.ParseFromString(payloads[0].data) + if marker.HasField("output"): + tokens.append(marker.output.stage_token) + return tokens + + +async def _output_stage_was_logged(redis_client: Any, prefix: str) -> bool: + return bool(await redis_client.llen(output_stage_log_key(prefix))) + + +@pytest.mark.timeout(240) +async def test_a_crash_before_the_marker_makes_the_next_worker_re_read( + client: Client, redis_worker_id: str +) -> None: + """The durability boundary, with a real dead Worker on one side of it. + + Three things have to hold together, and each fails differently: + + - **Nothing was committed.** The records were consumed on a Workflow Task + that was retained -- so its annotation was still accumulating, unwritten -- + and that task died with its Worker. Asserted against History, because a + marker written before the kill would make the rest of this vacuous. + - **The same offsets are read again.** The replacement Worker's read log has + to contain them. Comparing read logs rather than observing that records + eventually arrived is the difference between proving a re-read and proving + that Redis still had the data. + - **Nothing is lost and nothing is doubled.** The Workflow returns its + records in order, so a record delivered twice or dropped is visible in the + result rather than only in a cursor. + """ + pytest.importorskip("redis.asyncio", reason="redis is not installed") + + prefix = f"{KEY_NAMESPACE}:{redis_worker_id}:{uuid.uuid4().hex}" + publisher = recording_backend(url=redis_url(), prefix=prefix, reader="publisher") + try: + await publisher._client.ping() + except Exception as err: + await publisher.aclose() + pytest.skip(f"Redis is not reachable at {redis_url()}: {err}") + + task_queue = f"tq-{uuid.uuid4()}" + successor = recording_backend(url=redis_url(), prefix=prefix, reader="successor") + child = await asyncio.create_subprocess_exec( + sys.executable, + "-m", + "tests.contrib.external_workflow_streams.crash_worker", + "--address", + client.service_client.config.target_host, + "--namespace", + client.namespace, + "--task-queue", + task_queue, + "--redis-url", + redis_url(), + "--prefix", + prefix, + cwd=str(REPO_ROOT), + stdout=asyncio.subprocess.DEVNULL, + stderr=asyncio.subprocess.DEVNULL, + ) + handle = None + try: + expected = len(UNCOMMITTED) + len(LATER) + handle = await client.start_workflow( + CrashConsumeWorkflow.run, + expected, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + task_timeout=timedelta(seconds=TASK_TIMEOUT_SECONDS), + ) + description = await handle.describe() + key = StreamKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + STREAM_NAME, + ) + + # The child has to be the one holding the task, or the records would + # simply be waiting for the successor and nothing would be read twice. + async def child_took_the_task() -> bool: + return any( + e.HasField("workflow_task_started_event_attributes") + for e in await _history(handle) + ) + + await _wait_until( + child_took_the_task, 60, "the child Worker never picked the task up" + ) + + # One producer session, one continuous sequence: `(session_id, + # sequence)` is the append idempotency key, so restarting the numbering + # would re-use a key with different content, the backend would reject + # it, and the Workflow would appear to hang for reasons far from here. + session = f"crash-{uuid.uuid4()}" + for value in UNCOMMITTED: + await publish(publisher, key, [value], session=session) + await asyncio.sleep(FEED_GAP_SECONDS) + + async def child_read_them() -> bool: + read = await _reads(publisher._client, prefix, "crashed") + return len(read) >= len(UNCOMMITTED) + + await _wait_until( + child_read_them, + 30, + "the child Worker never read the records it is supposed to die " + "holding, so there is nothing for the successor to re-read", + ) + consumed = await publisher.read_after( + key, BEGINNING, max_records=100, block=None + ) + uncommitted_offsets = [str(r.offset) for r in consumed] + + # The kill. Nothing runs: no finalization, no marker, no wake. + child.kill() + await child.wait() + + assert not _markers(await _history(handle)), ( + "a marker was written before the crash, so the records the Worker " + "died holding were committed after all and the re-read below would " + "prove nothing -- the Workflow Task must still have been retained" + ) + + async with Worker( + client, + task_queue=task_queue, + workflows=[CrashConsumeWorkflow], + external_stream_backend=successor, + ): + await _wait_until( + lambda: _read_anything(publisher._client, prefix, "successor"), + 120, + "the replacement Worker never read the stream at all", + ) + await publish(publisher, key, LATER, session=session) + result = await asyncio.wait_for(handle.result(), 120) + + successor_reads = await _reads(publisher._client, prefix, "successor") + assert set(uncommitted_offsets) <= set(successor_reads), ( + "the replacement Worker did not re-read the offsets the dead one " + "had already taken -- an uncommitted consumption was treated as " + f"progress. re-read: {successor_reads}, owed: {uncommitted_offsets}" + ) + assert result == UNCOMMITTED + LATER, ( + "the Run lost or repeated a record across the crash. A marker " + "commits the cursor and only the first record's did, so everything " + f"after it was owed again and exactly once, got {result}" + ) + + assert any( + e.HasField("workflow_task_timed_out_event_attributes") + for e in await _history(handle) + ), ( + "no Workflow Task timed out, so the retained task was handed over " + "rather than dying with its Worker -- which is not the case under " + "test" + ) + finally: + if child.returncode is None: + child.kill() + await child.wait() + if handle is not None: + try: + await handle.terminate() + except Exception: + pass + keys = [k async for k in publisher._client.scan_iter(match=f"{prefix}*")] + if keys: + await publisher._client.delete(*keys) + await publisher.aclose() + await successor.aclose() + + +@pytest.mark.timeout(240) +async def test_real_crash_after_output_stage_aborts_old_token_and_exposes_one_retry( + client: Client, redis_worker_id: str +) -> None: + """SIGKILL after Redis staging leaves no phantom output before retry. + + The child provider blocks only after the pending stage is atomically + durable. Killing that process therefore prevents both the activation report + and graceful reconciliation. A successor retries the Workflow Task with a + fresh token; the reader proves the dead attempt absent from exact-run + History, aborts it, and exposes only the accepted attempt. + """ + pytest.importorskip("redis.asyncio", reason="redis is not installed") + + prefix = f"{KEY_NAMESPACE}:{redis_worker_id}:{uuid.uuid4().hex}" + backend = RedisStreamBackend(url=redis_url(), key_prefix=prefix) + try: + await backend._client.ping() + except Exception as err: + await backend.aclose() + pytest.skip(f"Redis is not reachable at {redis_url()}: {err}") + + task_queue = f"tq-{uuid.uuid4()}" + child = await _spawn_crash_worker( + client, task_queue=task_queue, prefix=prefix, mode="publish" + ) + handle = None + try: + handle = await client.start_workflow( + CrashPublishWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + task_timeout=timedelta(seconds=TASK_TIMEOUT_SECONDS), + ) + await _wait_until( + lambda: _output_stage_was_logged(backend._client, prefix), + 60, + "the child never durably staged Workflow output", + ) + + description = await handle.describe() + chain = WorkflowChainKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + ) + key = chain.stream_key("events", direction=StreamDirection.OUTPUT) + staged_read = await backend.read_output_after( + key, BEGINNING, max_records=100, block=timedelta(0) + ) + assert staged_read.records == () + assert staged_read.pending is not None + abandoned_manifest = staged_read.pending.manifest + abandoned = await backend.output_stage(abandoned_manifest) + assert abandoned is not None + assert abandoned.status is OutputStageStatus.PENDING + assert abandoned.records, "the pending stage is not physically durable" + + before_kill = await _history(handle) + assert _output_marker_tokens(before_kill) == [] + assert not any( + event.HasField("workflow_task_completed_event_attributes") + for event in before_kill + ), "the child reported the producing Workflow Task before it was killed" + + child.kill() + await child.wait() + + async with Worker( + client, + task_queue=task_queue, + workflows=[CrashPublishWorkflow], + external_stream_backend=backend, + ): + assert await asyncio.wait_for(handle.result(), 120) == "done" + output_client = await ExternalOutputStreamClient.connect( + backend=backend, + workflow=chain, + client=client, + ) + visible = [ + item.data + async for item in output_client.topic("events", type=str).subscribe() + ] + + assert visible == ["accepted-after-crash"] + abandoned = await backend.output_stage(abandoned_manifest) + assert abandoned is not None + assert abandoned.status is OutputStageStatus.ABORTED + + events = await _history(handle) + accepted_tokens = _output_marker_tokens(events) + assert len(accepted_tokens) == 1 + assert accepted_tokens[0] != abandoned_manifest.stage_token + assert any( + event.HasField("workflow_task_timed_out_event_attributes") + for event in events + ) + + statuses = await backend._client.hvals(backend._output_status_key(key)) + decoded_statuses = sorted( + value.decode() if isinstance(value, bytes) else value for value in statuses + ) + assert decoded_statuses == ["aborted", "committed"] + finally: + if child.returncode is None: + child.kill() + await child.wait() + if handle is not None: + try: + await handle.terminate() + except Exception: + pass + keys = [key async for key in backend._client.scan_iter(match=f"{prefix}*")] + if keys: + await backend._client.delete(*keys) + await backend.aclose() + + +@pytest.mark.timeout(240) +async def test_disconnected_update_leaves_live_staged_output_undecided_until_timeout( + client: Client, + redis_worker_id: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Cancelling the Update caller cannot decide its staged output. + + This uses a real speculative Update and a child Worker stopped between durable + Redis staging and its server report. A live output reader repeatedly checks + exact-run History after the caller disconnects and must leave the barrier + pending. Only the dead Worker's Workflow Task timeout is a deciding durable + boundary, at which point the reader may abort the abandoned stage. + + Retention loss itself is covered separately with deterministic fake History; + the local test server cannot force expiry of one exact Run on demand. + """ + pytest.importorskip("redis.asyncio", reason="redis is not installed") + + prefix = f"{KEY_NAMESPACE}:{redis_worker_id}:{uuid.uuid4().hex}" + backend = RedisStreamBackend(url=redis_url(), key_prefix=prefix) + try: + await backend._client.ping() + except Exception as err: + await backend.aclose() + pytest.skip(f"Redis is not reachable at {redis_url()}: {err}") + + task_queue = f"tq-{uuid.uuid4()}" + child = await _spawn_crash_worker( + client, task_queue=task_queue, prefix=prefix, mode="update" + ) + handle = None + update_task: asyncio.Task[str] | None = None + read_task: asyncio.Task[Any] | None = None + try: + handle = await client.start_workflow( + CrashUpdatePublishWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + task_timeout=timedelta(seconds=TASK_TIMEOUT_SECONDS), + ) + + async def initial_task_completed() -> bool: + return any( + event.HasField("workflow_task_completed_event_attributes") + for event in await _history(handle) + ) + + await _wait_until( + initial_task_completed, + 60, + "the child never completed the Workflow's initial waiting task", + ) + description = await handle.describe() + chain = WorkflowChainKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + ) + key = chain.stream_key("events", direction=StreamDirection.OUTPUT) + + update_task = asyncio.create_task( + handle.execute_update( + CrashUpdatePublishWorkflow.publish, + "unreported-update-output", + ) + ) + await _wait_until( + lambda: _output_stage_was_logged(backend._client, prefix), + 30, + "the speculative Update handler never durably staged its output", + ) + staged_read = await backend.read_output_after( + key, BEGINNING, max_records=100, block=timedelta(0) + ) + assert staged_read.records == () + assert staged_read.pending is not None + pending = staged_read.pending + manifest = pending.manifest + + output_client = await ExternalOutputStreamClient.connect( + backend=backend, + workflow=chain, + client=client, + ) + topic = output_client.topic("events", type=str) + original_history_decision = topic._history_decision + reconciliation_attempts = 0 + + async def count_live_history_decisions(manifest_to_check: Any) -> Any: + nonlocal reconciliation_attempts + reconciliation_attempts += 1 + return await original_history_decision(manifest_to_check) + + monkeypatch.setattr(topic, "_history_decision", count_live_history_decisions) + + async def reconcile_pending_stage() -> None: + while not await topic._resolve_pending(pending): + await asyncio.sleep(1) + + read_task = asyncio.create_task(reconcile_pending_stage()) + + # This is the caller disconnect. It cancels only its wait for the Update + # result; the server-side speculative Update and its Workflow Task survive. + update_task.cancel() + with pytest.raises(WorkflowUpdateRPCTimeoutOrCancelledError): + await update_task + + async def reader_retried_history() -> bool: + return reconciliation_attempts >= 3 + + await _wait_until( + reader_retried_history, + TASK_TIMEOUT_SECONDS - 1, + "the live reader did not retry ambiguous exact-run History", + ) + undecided = await backend.output_stage(manifest) + assert undecided is not None + assert undecided.status is OutputStageStatus.PENDING + assert not read_task.done() + + events_before_loss = await _history(handle) + relevant = [ + event + for event in events_before_loss + if event.event_id > manifest.history_floor_event_id + ] + # This server delivered the Update speculatively: the handler ran (the + # stage above is durable proof of that), but admission itself cannot + # enter History until the withheld Workflow Task report is accepted. + assert not any( + event.HasField("workflow_execution_update_admitted_event_attributes") + for event in relevant + ) + assert _output_marker_tokens(relevant) == [] + assert not any( + event.HasField(field) + for event in relevant + for field in ( + "workflow_task_completed_event_attributes", + "workflow_task_failed_event_attributes", + "workflow_task_timed_out_event_attributes", + ) + ) + + child.kill() + await child.wait() + + async def abandoned_stage_was_aborted() -> bool: + stage = await backend.output_stage(manifest) + return stage is not None and stage.status is OutputStageStatus.ABORTED + + await _wait_until( + abandoned_stage_was_aborted, + 30, + "the dead Update task timed out but its pending output was not aborted", + ) + await asyncio.wait_for(read_task, 5) + events_after_loss = await _history(handle) + assert any( + event.event_id > manifest.history_floor_event_id + and event.HasField("workflow_task_timed_out_event_attributes") + for event in events_after_loss + ) + assert _output_marker_tokens(events_after_loss) == [] + finally: + if update_task is not None and not update_task.done(): + update_task.cancel() + if read_task is not None and not read_task.done(): + read_task.cancel() + try: + await read_task + except asyncio.CancelledError: + pass + if child.returncode is None: + child.kill() + await child.wait() + if handle is not None: + try: + await handle.terminate() + except Exception: + pass + keys = [key async for key in backend._client.scan_iter(match=f"{prefix}*")] + if keys: + await backend._client.delete(*keys) + await backend.aclose() + + +async def _read_anything(redis_client: Any, prefix: str, reader: str) -> bool: + return bool(await _reads(redis_client, prefix, reader)) diff --git a/tests/contrib/external_workflow_streams/test_worker_handoff.py b/tests/contrib/external_workflow_streams/test_worker_handoff.py new file mode 100644 index 0000000000..3e52796c0e --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_worker_handoff.py @@ -0,0 +1,950 @@ +"""Handing a Run to another Worker: the shutdown sweep, and a failed finalization. + +Two transitions that only exist between Workers, so neither can be shown inside +one: + +- **Shutdown with no open Workflow Task.** Nothing local can create + server-visible work, so the manager sweeps: it asks Core what state the Run is + in through the read-only probe, sends the reserved wake Signal, and waits for + the server to acknowledge it before the watchers go away. The point of the + Signal is that *another* Worker picks the Run up. +- **A finalization that cannot be answered.** "A marker is never written without + its terminal": if the Run's state is gone when Core asks for the terminal, + Python fails the activation, Core writes no marker, and the retry replays from + the previous marker -- losing no record and moving no cursor. +""" + +from __future__ import annotations + +import asyncio +import uuid +from datetime import timedelta +from typing import Any + +import pytest + +from temporalio import workflow +from temporalio.client import Client +from temporalio.contrib.external_workflow_streams._annotation import decode_annotation +from temporalio.contrib.external_workflow_streams._backend import StreamKey +from temporalio.contrib.external_workflow_streams._manager import ( + ReadinessResult, + StreamSubscriptionManager, +) +from temporalio.contrib.external_workflow_streams._record import ( + AFTER, + BEGINNING, + Cursor, +) +from temporalio.worker import Worker +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend +from tests.contrib.external_workflow_streams.test_worker_integration import publish + +with workflow.unsafe.imports_passed_through(): + from temporalio.contrib.external_workflow_streams._api import external_stream + +WAKE_SIGNAL_NAME = "__temporal_external_stream_wake" + +MARKER_DETAILS_KEY = "external_stream" +"""Where Core puts the marker's ``ExternalStreamMarkerData``.""" + +RETENTION_IDLE_TIMEOUT = timedelta(seconds=2) +"""How long a fed Workflow Task stays retained here -- stated, not inherited. + +Case 30 needs a Workflow Task that is *open* at the moment Core asks for the +terminal, and everything the test does between the last record reaching the +Workflow and the shutdown call is bookkeeping read out of local state. The +one-second default makes that bookkeeping share a window with the idle timer, +which is not what this case is about. + +Not larger, because the same timeout also bounds how long the Workflow waits for +its *first* record: a record appended before the subscription's watcher is +running is picked up by the idle timer's park recheck rather than by the +watcher, so raising this raises the test's own floor. Bounded above by the +Workflow Task timeout the test starts the Workflow with in any case -- a task +retained past that does not park, it fails (`wft-lifecycle.md`). +""" + + +# These tests exercise cross-Worker handoff, not the sandbox. Keeping their +# fixture Workflows unsandboxed avoids re-importing this pytest +# assertion-rewritten module while each Worker validates its registrations. +@workflow.defn(sandboxed=False) +class TimerThenStreamWorkflow: + """Consumes one record, takes a timer, then consumes the rest. + + The timer is what puts the Run in the two states this module needs. Its + completion is server-bound, so retention is suppressed: the marker commits + there, and the Run is left cached with **no open Workflow Task** and a live + subscription. Once the timer fires, the Workflow blocks on the stream again + and the next task is retained. + + The marker and the retained task are therefore separated by the *whole* + timer: the marker is written when the timer is **started**, and the retained + task exists only once it has **fired** and its replacement task has reached + this Worker. A test that treats the marker as the signal to start feeding is + reading a clock, not a state -- see the caller. + """ + + @workflow.run + async def run(self, expected: int) -> list[str]: + tokens = external_stream.with_options( + idle_timeout=RETENTION_IDLE_TIMEOUT + ).topic("tokens", type=str) + seen: list[str] = [] + iterator = tokens.subscribe().__aiter__() + seen.append(await iterator.__anext__()) + await asyncio.sleep(0.5) + while len(seen) < expected: + seen.append(await iterator.__anext__()) + return seen + + +@workflow.defn(sandboxed=False) +class LeftWithNoOpenTaskWorkflow: + """Blocks on the stream first, then starts a long timer. The order is the point. + + The ``NoOpenWorkflowTask`` window needs two things at once -- a wait set Core + knows about, and no Workflow Task holding it -- and only this order produces + both: + + 1. **Block before doing anything else.** A completion that carries no command + sends ``WorkflowStreamQuiescent``, and that command is the *only* thing + that registers a wait set with Core. A Workflow that starts a timer on the + same completion it first blocks on never registers one at all: retention + is suppressed, so no quiescent snapshot is ever sent, Core's wait set stays + empty, and a wake Signal then marks nothing ready and creates a Workflow + Task that Core completes with no activation. Such a Run cannot be resumed + by anything -- which is what this case used to try to shut down. + 2. **Start the timer only after a record has woken it.** That command + suppresses retention for *that* completion, so the Workflow Task ends with + the wait set still registered and the Run sits in the window for as long as + the test needs. The timer is also the "unrelated Workflow event" the + shutdown sweep must not wait for. + """ + + @workflow.run + async def run(self, expected: int) -> list[str]: + tokens = external_stream.topic("tokens", type=str) + seen: list[str] = [] + iterator = tokens.subscribe().__aiter__() + seen.append(await iterator.__anext__()) + timer = asyncio.ensure_future(asyncio.sleep(600)) + while len(seen) < expected: + seen.append(await iterator.__anext__()) + timer.cancel() + return seen + + +@pytest.fixture +def backend() -> MemoryStreamBackend: + return MemoryStreamBackend() + + +async def history(handle: Any) -> list[Any]: + return [e async for e in handle.fetch_history_events()] + + +def markers(events: list[Any]) -> list[Any]: + return [e for e in events if e.HasField("marker_recorded_event_attributes")] + + +def wake_signals(events: list[Any]) -> list[Any]: + return [ + e + for e in events + if e.HasField("workflow_execution_signaled_event_attributes") + and e.workflow_execution_signaled_event_attributes.signal_name + == WAKE_SIGNAL_NAME + ] + + +def committed_boundary(marker_event: Any, wait_id: int) -> Cursor: + """Where a marker says one wait's consumption had reached. + + "The cursor did not move" is a statement about the marker; anything derived + from what was published instead would be a statement about timing. + + The annotation's terminal, when it has one. A marker written on a + command-producing completion currently carries no terminal frame at all -- + ``WorkflowStreamProgress`` emits segments only, and ``add_terminal()`` runs + only on the park and finalization paths -- so the end of the last recorded + run is used instead. It names the same position the terminal would. + """ + from temporalio.bridge.proto.external_data import ExternalStreamMarkerData + + details = marker_event.marker_recorded_event_attributes.details + data = ExternalStreamMarkerData() + data.ParseFromString(details[MARKER_DETAILS_KEY].payloads[0].data) + annotation = decode_annotation(data.replay_annotation) + if annotation.terminal is not None: + return annotation.terminal[wait_id] + reached = annotation.header.streams[wait_id].start_cursor + for segment in annotation.segments: + for run in segment.runs: + if run.wait_id == wait_id: + reached = AFTER(run.last_offset) + return reached + + +def marker_bytes(marker_event: Any) -> bytes: + return ( + marker_event.marker_recorded_event_attributes.details[MARKER_DETAILS_KEY] + .payloads[0] + .data + ) + + +def wake_envelope(signal_event: Any) -> Any: + """The wake Signal's own envelope, read the way Core reads it. + + Deliberately not through the ``DataConverter``: the envelope is defined at + the protocol level precisely so Core -- which has no converter -- can read + it. + """ + from temporalio.bridge.proto.external_stream import WakeSignal + + payload = signal_event.workflow_execution_signaled_event_attributes.input.payloads[ + 0 + ] + envelope = WakeSignal() + envelope.ParseFromString(payload.data) + return envelope + + +async def stream_key_for(client: Client, handle: Any) -> StreamKey: + description = await handle.describe() + return StreamKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + "tokens", + ) + + +async def wait_until( + predicate: Any, timeout: float, message: str, interval: float = 0.2 +) -> None: + """Polls rather than sleeping a fixed time, so a slow start is not a flake. + + ``interval`` is worth lowering when the condition is read out of local + process state rather than fetched from the server, and when what happens + *after* it is satisfied is itself time-bounded: the poll gap is then part of + the window the caller has to act in, not merely latency. + """ + deadline = asyncio.get_running_loop().time() + timeout + while asyncio.get_running_loop().time() < deadline: + if await predicate(): + return + await asyncio.sleep(interval) + raise AssertionError(message) + + +async def close_worker(worker: Worker, task: asyncio.Task[Any]) -> None: + """Request orderly shutdown and observe the Worker's terminal task state.""" + if not task.done(): + await asyncio.wait_for(worker.shutdown(), 60) + await asyncio.gather(task, return_exceptions=True) + + +def workflow_worker(worker: Worker) -> Any: + """The private workflow worker these white-box integration tests require.""" + assert worker._workflow_worker is not None + return worker._workflow_worker + + +# --- case 29: shutdown in the NoOpenWorkflowTask window ----------------------- + + +@pytest.mark.timeout(180) +async def test_shutdown_with_no_open_task_hands_the_run_to_another_worker( + client: Client, backend: MemoryStreamBackend +) -> None: + """The window the sweep exists for, with the second Worker that proves it. + + The Run is cached, has a wait set Core knows about, and holds no Workflow + Task. **Nothing else will ever create work for it**: the only other thing + that could produce a Workflow Task here is the Workflow's own 600-second + timer, and waiting for an unrelated event is not a plan. That is why the + sweep exists, and why it runs whether or not a record happens to be waiting + -- the obligation is to hand the Run over, not to announce an append. + + The window is *asserted*, not assumed: the read-only probe is called on the + live Worker before shutdown and must answer ``NoOpenWorkflowTask``. Getting + there needs both of the Workflow's steps, which is what + :class:`LeftWithNoOpenTaskWorkflow` documents. + + Four separate claims, none of which the others imply: + + - no marker is written, because nothing was accumulated to write; + - the Run's state is resolved with the read-only probe rather than a + readiness notification, which would assert a buffered record that the + sweep has no business claiming; + - a **new** wake Signal reaches the Workflow's own History, and it is an + *unparked* one -- ``park_generation = 0``. Counted against the Signals + already there rather than merely found, because an earlier wake in the + same History proves nothing about this one; + - a second Worker takes the resulting Workflow Task, reconstructs the + subscription from the marker, and delivers a record published after the + handover -- without the timer ever firing. + """ + task_queue = f"tq-{uuid.uuid4()}" + session = f"handoff-{uuid.uuid4()}" + worker_a = Worker( + client, + task_queue=task_queue, + workflows=[LeftWithNoOpenTaskWorkflow], + external_stream_backend=backend, + ) + worker_a_task = asyncio.create_task(worker_a.run()) + handle = None + try: + handle = await client.start_workflow( + LeftWithNoOpenTaskWorkflow.run, + 2, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + task_timeout=timedelta(seconds=10), + ) + key = await stream_key_for(client, handle) + + # Published only once the Run's first Workflow Task has ended, which is + # what registers the wait set with Core. A record that arrives before + # that lands on a Run whose wait set is still empty. + await wait_until( + lambda: _has_markers(handle), + 60, + "the Run's first Workflow Task never ended, so its wait set was " + "never registered with Core and no wake could reach it", + ) + await publish(backend, key, ["alpha"], session=session) + await wait_until( + lambda: _timer_started(handle), + 60, + "the record never reached the Workflow, so the completion that " + "leaves the Run in the no-open-task window never happened", + ) + + manager = workflow_worker(worker_a)._external_stream_manager + assert manager is not None, "the manager should exist by now" + real_probe, real_ready = manager._run_status, manager._notify_ready + await wait_until( + lambda: _run_is_registered(manager), + 30, + "the manager holds no Run with subscriptions, so the sweep has " + "nothing to sweep", + ) + run_id = next(iter(manager._runs)) + await wait_until( + lambda: _in_the_window(real_probe, run_id), + 30, + "the Run never reached the no-open-Workflow-Task window, so this " + "case is not testing the transition it names", + ) + + before = markers(await history(handle)) + signals_before = wake_signals(await history(handle)) + + probes: list[str] = [] + readiness: list[str] = [] + + async def counting_probe(probed_run_id: str) -> Any: + probes.append(probed_run_id) + return await real_probe(probed_run_id) + + async def counting_ready( + ready_run_id: str, wait_id: int, generation: int + ) -> Any: + readiness.append(ready_run_id) + return await real_ready(ready_run_id, wait_id, generation) + + manager._run_status = counting_probe + manager._notify_ready = counting_ready + + await asyncio.wait_for(worker_a.shutdown(), 60) + + assert probes, ( + "the sweep never asked Core what state the Run was in; without the " + "read-only probe it cannot know whether a wake is owed at all" + ) + assert not readiness, ( + "the sweep used the readiness call as a probe. Readiness means 'a " + "record is buffered', so using it here asserts something the sweep " + "has no business claiming and manufactures a Workflow Task on the " + "way out" + ) + after = markers(await history(handle)) + assert len(after) == len(before), ( + "a marker was written on the way out of the no-open-task window; " + "nothing was accumulated there, so there was nothing to write" + ) + signals_after = wake_signals(await history(handle)) + assert len(signals_after) > len(signals_before), ( + "the sweep's wake Signal never reached the Workflow's own History, " + f"which still holds the same {len(signals_before)} it did before " + "shutdown. Nothing will ever create a Workflow Task for this Run" + ) + envelope = wake_envelope(signals_after[-1]) + assert envelope.park_generation == 0, ( + "the shutdown wake must be an unparked one -- there is no confirmed " + f"park to name -- got generation {envelope.park_generation}" + ) + + async with Worker( + client, + task_queue=task_queue, + workflows=[LeftWithNoOpenTaskWorkflow], + external_stream_backend=backend, + ): + # Published after the handover, so delivering it proves the second + # Worker rebuilt a live subscription rather than replaying one. + await publish(backend, key, ["beta"], session=session) + result = await asyncio.wait_for(handle.result(), 60) + + assert result == ["alpha", "beta"], ( + "the second Worker did not reconstruct the subscription from the " + f"marker: {result}" + ) + assert not any( + e.HasField("timer_fired_event_attributes") for e in await history(handle) + ), ( + "the Run was resumed by its own 600-second timer rather than by the " + "wake Signal, which is the one thing this case forbids" + ) + finally: + if handle is not None: + try: + await handle.terminate() + except Exception: + pass + await close_worker(worker_a, worker_a_task) + + +async def _has_markers(handle: Any) -> bool: + return bool(markers(await history(handle))) + + +async def _timer_started(handle: Any) -> bool: + return any( + e.HasField("timer_started_event_attributes") for e in await history(handle) + ) + + +async def _run_is_registered(manager: Any) -> bool: + return bool(manager._runs) + + +async def _in_the_window(probe: Any, run_id: str) -> bool: + """Whether Core says this Run has waits registered and no task open. + + Asked through the read-only probe, which is the same call the sweep makes + and the only thing that can tell the three states apart. + """ + status = await probe(run_id) + return getattr(status, "value", status) == "NoOpenWorkflowTask" + + +class _InTheWindow: + """A Run left cached with a live wait set and no open Workflow Task. + + The state both tests below need, and getting there is neither quick nor + obvious: the Workflow must block on the stream *first* so a + ``WorkflowStreamQuiescent`` registers the wait set with Core, park, be woken + by a record, and only then issue a server-bound command so the Workflow Task + ends with the wait set still registered. See + :class:`LeftWithNoOpenTaskWorkflow`. + """ + + def __init__(self, worker: Worker, task: asyncio.Task, handle: Any, key: StreamKey): + self.worker = worker + self.task = task + self.handle = handle + self.key = key + self.manager = workflow_worker(worker)._external_stream_manager + self.run_id = next(iter(self.manager._runs)) + + async def close(self) -> None: + try: + await self.handle.terminate() + except Exception: + pass + await close_worker(self.worker, self.task) + + +async def _leave_a_run_in_the_window( + client: Client, backend: MemoryStreamBackend, session: str +) -> _InTheWindow: + task_queue = f"tq-{uuid.uuid4()}" + worker = Worker( + client, + task_queue=task_queue, + workflows=[LeftWithNoOpenTaskWorkflow], + external_stream_backend=backend, + ) + worker_task = asyncio.create_task(worker.run()) + handle = await client.start_workflow( + LeftWithNoOpenTaskWorkflow.run, + 2, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + task_timeout=timedelta(seconds=10), + ) + key = await stream_key_for(client, handle) + + # The first Workflow Task ends by parking, which is also what registers the + # wait set with Core. A record published before that lands on a Run whose + # wait set is still empty. + await wait_until( + lambda: _has_markers(handle), + 60, + "the Run's first Workflow Task never ended, so it never parked and " + "never registered a wait set with Core", + ) + await wait_until( + lambda: _park_is_installed(backend, key), + 30, + "no park intent was installed, so this Run never confirmed a park and " + "there is no resolved park to reason about", + ) + + await publish(backend, key, ["alpha"], session=session) + await wait_until( + lambda: _timer_started(handle), + 60, + "the record never reached the Workflow, so the completion that leaves " + "the Run in the no-open-task window never happened", + ) + + manager = workflow_worker(worker)._external_stream_manager + assert manager is not None, "the manager should exist by now" + await wait_until( + lambda: _run_is_registered(manager), + 30, + "the manager holds no Run with subscriptions", + ) + return _InTheWindow(worker, worker_task, handle, key) + + +async def _park_is_installed(backend: MemoryStreamBackend, key: StreamKey) -> bool: + return bool(await backend.parked_wait_ids(key)) + + +@pytest.mark.timeout(180) +async def test_a_wake_that_resolves_a_park_removes_its_intent( + client: Client, backend: MemoryStreamBackend +) -> None: + """The invariant, end to end: an intent outlives its park nowhere. + + The park here is confirmed rather than aborted, and it is resolved the way a + live one actually is -- a producer appends, the watcher sends the wake, and + Core clears its ``park_generation``. None of that is visible in the backend, + so if the manager does not take the intent back out on the resolve, nothing + ever will. + + Asserted through ``current_park_generation`` because that is the call the + damage comes through. Every wake path asks it: a producer, to decide what its + Signal names, and this Worker's own shutdown sweep, to decide whether it owes + a parked or an unparked wake. Both get a generation Core has already + discarded, and both wakes are then discarded in turn -- the producer's as + stale, the sweep's by the server, which sees the request ID of the wake that + ended the park. + """ + session = f"resolve-{uuid.uuid4()}" + window = await _leave_a_run_in_the_window(client, backend, session) + try: + assert await backend.parked_wait_ids(window.key) == [], ( + "the park that the wake Signal resolved still has its intent " + "installed in the backend" + ) + for wait_id in (0, 1, 2): + assert await backend.current_park_generation(window.key, wait_id) is None, ( + "a park that is over still reports a generation, so the next " + "wake for this wait will name it instead of being the unparked " + "wake it owes" + ) + finally: + await window.close() + + +@pytest.mark.timeout(180) +async def test_the_shutdown_probe_is_asked_while_core_still_holds_the_run( + client: Client, backend: MemoryStreamBackend +) -> None: + """The sweep's probe has to be asked when its answer still describes something. + + P20 gives the probe four answers and three behaviours: ``WftOpen`` waits for + C15b, ``Parked`` owes nothing, and ``NoOpenWorkflowTask``/``RunNotFound`` + owe the wake. Asked after every activation has been answered, Core has + already dropped the Run and only ``RunNotFound`` is reachable -- so the + sweep still sends its wake and still looks correct, while a parked Run gets + a wake it does not need and a Run holding a Workflow Task gets one that + races Core's own shutdown transition. + + This Run is in the ``NoOpenWorkflowTask`` window, which is asserted on the + live Worker before shutdown, so a sweep asking at the right moment must see + the same thing. The assertion is on what the *sweep* saw, not on what a + probe outside it saw, because the ordering is the only thing under test. + """ + session = f"probe-{uuid.uuid4()}" + window = await _leave_a_run_in_the_window(client, backend, session) + try: + real_probe = window.manager._run_status + await wait_until( + lambda: _in_the_window(real_probe, window.run_id), + 30, + "the Run never reached the no-open-Workflow-Task window, so this " + "case is not testing the ordering it names", + ) + + answers: list[str] = [] + + async def recording_probe(probed_run_id: str) -> Any: + status = await real_probe(probed_run_id) + answers.append(getattr(status, "value", status)) + return status + + window.manager._run_status = recording_probe + + await asyncio.wait_for(window.worker.shutdown(), 60) + + assert answers, "the sweep never probed at all" + assert "RunNotFound" not in answers, ( + "the sweep probed after Core had already dropped the Run, so its " + f"answers were {answers}. Every state P20 distinguishes collapses " + "into RunNotFound there: the wake still goes out, so nothing looks " + "broken, but the Parked and WftOpen branches are unreachable" + ) + assert answers == ["NoOpenWorkflowTask"], ( + "the Run was in the no-open-Workflow-Task window immediately before " + f"shutdown, and the sweep saw {answers}" + ) + finally: + await window.close() + + +# --- case 30: teardown racing finalization ------------------------------------ + + +def _status_value(status: Any) -> Any: + """The probe's answer as a plain string, matched structurally as Core is.""" + return getattr(status, "value", status) + + +def _delivery_boundary(workflow_worker: Any, run_id: str) -> Any: + """Where this Run's deliveries have reached, read from the Run's own state. + + ``blocked_snapshot()`` is the same thing the finalization terminal is built + from, which is what makes it the right thing to wait on: nothing else says + "a Workflow Task is open here and its annotation is not written yet". The + History cannot -- a retained task writes nothing until it ends, which is + exactly the state this case must catch it in. + """ + runtime = workflow_worker._external_stream_runtimes.get(run_id) + return None if runtime is None else runtime.blocked_snapshot() + + +async def _delivery_moved(workflow_worker: Any, run_id: str, previous: Any) -> bool: + """Whether the boundary has moved off ``previous``. + + Movement rather than a value: a cursor is ordered by the provider's own + comparator, so "is it past the record just published" is a question only the + backend can answer, while "did the record land" is a question the boundary + answers by itself. + """ + return _delivery_boundary(workflow_worker, run_id) != previous + + +@pytest.mark.timeout(180) +async def test_a_finalization_that_cannot_be_answered_writes_no_marker( + client: Client, backend: MemoryStreamBackend +) -> None: + """The Run's state vanishes underneath a finalization, and nothing is lost. + + Finalization is answered from the Run's own out-of-sandbox state, and its + only failure mode is that state being gone. The obligation then is not to do + something approximate: **a marker is never written without its terminal**, + because an abandoned Workflow Task commits no cursor and loses no record, + while a truncated annotation is durable and wrong. + + So the Run's entry is removed exactly once, while Core is asking for the + terminal, and the assertions are about what did *not* happen: no marker for + that task, no cursor movement in the marker that already existed, and no + record missing from what the Workflow finally received. The retry replays + from the previous marker and re-consumes everything after it. + """ + task_queue = f"tq-{uuid.uuid4()}" + session = f"finalize-{uuid.uuid4()}" + values = ["alpha", "beta", "gamma", "delta"] + worker_a = Worker( + client, + task_queue=task_queue, + workflows=[TimerThenStreamWorkflow], + external_stream_backend=backend, + ) + worker_a_task = asyncio.create_task(worker_a.run()) + shutdown_task = None + handle = None + try: + handle = await client.start_workflow( + TimerThenStreamWorkflow.run, + len(values), + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + task_timeout=timedelta(seconds=10), + ) + key = await stream_key_for(client, handle) + + # The first record is consumed on a task that also starts a timer, so + # its marker commits: that marker is the "previous" one the retry must + # replay from. + await publish(backend, key, values[:1], session=session) + await wait_until( + lambda: _has_markers(handle), + 60, + "no marker committed the first record, so there is no previous " + "marker for the retry to replay from", + ) + before = markers(await history(handle)) + committed_before = committed_boundary(before[-1], wait_id=1) + marker_before = marker_bytes(before[-1]) + + # Fed from here, and each record is *followed* to the Workflow rather + # than merely published on a schedule. The precondition this case needs + # is a Workflow Task that is open with an annotation still accumulating + # in it, and the marker above says nothing about when that exists: it is + # written when the timer is started, while the retained task appears + # only once the timer has fired and its replacement task has reached + # this Worker. Feeding on a fixed schedule from the marker races the + # timer, and a shutdown that wins the race finds no open Workflow Task + # at all -- Core issues no ``FinalizeExternalStreams``, so the sweep + # runs instead and nothing this case is about ever happens. + # + # The condition waited on is the Run's own delivery boundary, which is + # the snapshot the terminal would be built from. Once it has moved past + # what the previous marker committed, the open task and its unwritten + # annotation are both facts. + workflow_worker_impl = workflow_worker(worker_a) + manager = workflow_worker_impl._external_stream_manager + assert manager is not None, "the manager should exist by now" + await wait_until( + lambda: _run_is_registered(manager), + 30, + "the manager holds no Run with subscriptions, so there is no Run " + "state for a finalization to be answered from", + ) + run_id = next(iter(manager._runs)) + assert _delivery_boundary(workflow_worker_impl, run_id) is not None, ( + "the manager registered a Run the Worker holds no stream runtime " + "for, so the sabotage below would remove nothing" + ) + for value in values[1:3]: + delivered = _delivery_boundary(workflow_worker_impl, run_id) + await publish(backend, key, [value], session=session) + await wait_until( + lambda: _delivery_moved(workflow_worker_impl, run_id, delivered), + 60, + f"{value!r} never reached the Workflow, so no Workflow Task is " + "open here with an unwritten annotation in it and the " + "finalization this case forces would have nothing to fail " + "against", + # The retained task the last delivery leaves behind is what the + # steps after this loop have to run inside, so the poll gap is + # spent out of that window. Short because the condition is a + # dictionary in this process, not a fetch from the server. + interval=0.02, + ) + + assert _status_value(await manager._run_status(run_id)) == "WftOpen", ( + "the records were delivered but the Workflow Task holding them is " + "already gone, so Core has nothing to ask for a terminal on. The " + "task is retained for RETENTION_IDLE_TIMEOUT, which every step " + "since the last delivery is expected to fit inside" + ) + + # The Run's entry disappears underneath the finalization, exactly once. + original = workflow_worker_impl._handle_external_stream_jobs + sabotaged: list[str] = [] + + async def losing_the_run(act: Any, running: Any) -> Any: + if not sabotaged and any( + j.HasField("finalize_external_streams") for j in act.jobs + ): + sabotaged.append(act.run_id) + workflow_worker_impl._external_stream_runtimes.pop(act.run_id, None) + await manager.evict_run(act.run_id) + return await original(act, running) + + workflow_worker_impl._handle_external_stream_jobs = losing_the_run + + # Shutdown is what makes Core ask for the terminal while the task is + # still open. Not awaited: this Worker is deliberately being left unable + # to answer, and the test's subject is what the *server* is left holding. + shutdown_task = asyncio.create_task(worker_a.shutdown()) + + async def the_task_failed() -> bool: + return any( + e.HasField("workflow_task_failed_event_attributes") + for e in await history(handle) + ) + + await wait_until( + the_task_failed, + 90, + "the Workflow Task never failed, so the finalization was answered " + "from somewhere other than the Run state that was removed", + ) + assert sabotaged, "the finalization job never arrived, so nothing was raced" + + events = await history(handle) + failure = next( + e for e in events if e.HasField("workflow_task_failed_event_attributes") + ) + message = failure.workflow_task_failed_event_attributes.failure.message + assert "finalize_external_streams" in message, ( + "the Workflow Task failed for some other reason than the " + f"unanswerable finalization: {message}" + ) + after = markers(events) + assert len(after) == len(before), ( + "Core wrote a marker for a Workflow Task whose terminal was never " + "supplied -- a durable, truncated annotation is exactly what the " + "abandon-and-retry rule exists to prevent" + ) + assert committed_boundary(after[-1], wait_id=1) == committed_before, ( + "the committed cursor moved across a Workflow Task that was " + "abandoned; nothing consumed on it was ever committed" + ) + assert marker_bytes(after[-1]) == marker_before, ( + "the marker that was already in History was rewritten by the " + "abandoned Workflow Task" + ) + + async with Worker( + client, + task_queue=task_queue, + workflows=[TimerThenStreamWorkflow], + external_stream_backend=backend, + ): + await publish(backend, key, values[3:], session=session) + result = await asyncio.wait_for(handle.result(), 90) + + assert result == values, ( + "the retry did not replay from the previous marker: every record " + "consumed on the abandoned task was owed again, exactly once, and " + f"in order, got {result}" + ) + finally: + # Terminated first, then the Worker given a moment to finish shutting + # down: a Worker abandoned mid-shutdown leaves a Core run behind, and + # this suite's later tests are timing-sensitive enough to notice. + if handle is not None: + try: + await handle.terminate() + except Exception: + pass + if shutdown_task is not None: + await asyncio.wait([shutdown_task], timeout=10) + if not shutdown_task.done(): + shutdown_task.cancel() + await asyncio.gather(shutdown_task, return_exceptions=True) + if not worker_a_task.done(): + worker_a_task.cancel() + await asyncio.gather(worker_a_task, return_exceptions=True) + + +# --- the intent the Worker that installed it left behind ----------------------- + + +HANDOFF_RUN_ID = "run-handed-over" + + +def _manager_for(backend: MemoryStreamBackend) -> StreamSubscriptionManager: + """One Worker's manager, wired with only what the park handshake touches. + + No Core and no Signal path: what is under test here is state that lives in + the *backend* and outlives both. + """ + + async def accepted( + run_id: str, # pyright: ignore[reportUnusedParameter] + wait_id: int, # pyright: ignore[reportUnusedParameter] + generation: int, # pyright: ignore[reportUnusedParameter] + ) -> str: + return ReadinessResult.ACCEPTED + + return StreamSubscriptionManager( + backend=backend, + notify_ready=accepted, + watch_block=timedelta(milliseconds=10), + ) + + +async def _nothing_is_parked(backend: MemoryStreamBackend, key: StreamKey) -> bool: + return await backend.parked_wait_ids(key) == [] + + +async def test_a_park_intent_installed_by_a_previous_worker_is_removed( + backend: MemoryStreamBackend, +) -> None: + """The half of the intent invariant that only a hand-off can show. + + ``installed_park_generation`` is a *mirror* of backend state and lives on + the Worker that installed the park. The intent is durable and survives + eviction, a Workflow Task that moved to another Worker, and shutdown; the + mirror survives none of them. So removal keyed on the mirror alone cannot + reach exactly the intents that most need reaching -- the ones whose + installer is gone -- and the Run is then stranded in a way no later resolve + can repair: ``current_park_generation`` keeps answering a generation Core + discarded, every producer wake names it and Core discards it as stale, and + because a parked wake's request ID ignores sender identity the second such + wake is byte-identical to the first and the server deduplicates it away. + + The ordering is the real Worker's: ``ResolveExternalStreamWaits`` is + answered *before* user Workflow code runs, so the new Worker has no + subscription to resolve against at that point and the reconstructed one + arrives afterwards. + """ + key = StreamKey("ns", "wf", "first-run", "tokens") + + worker_a = _manager_for(backend) + worker_a.register(run_id=HANDOFF_RUN_ID, wait_id=1, stream_key=key) + confirmed = not await worker_a.prepare_park(HANDOFF_RUN_ID, 7, {1: BEGINNING}) + assert confirmed, "nothing was appended, so this park must have confirmed" + await worker_a.shutdown() + + assert await backend.parked_wait_ids(key) == [1], ( + "this case is only meaningful if the intent outlived the Worker that " + "installed it -- it is durable backend state, not Worker state" + ) + + worker_b = _manager_for(backend) + try: + await worker_b.resolve_park(HANDOFF_RUN_ID) + worker_b.register( + run_id=HANDOFF_RUN_ID, + wait_id=1, + stream_key=key, + ) + + await wait_until( + lambda: _nothing_is_parked(backend, key), + 5, + "the intent of a park no Worker is sitting in is still installed, " + "so the next producer wake names a dead generation instead of the " + "unparked wake it owes", + interval=0.01, + ) + assert await backend.current_park_generation(key, 1) is None, ( + "a park that ended with the Worker that confirmed it still reports " + "a generation" + ) + + # What the producer actually asks, in the order it asks it: the record + # lands, and then the wake is chosen from what the backend says is + # parked. Nothing is, so it is the unparked wake -- which Core always + # accepts as a recheck request -- rather than generation 7, which it + # discards as stale. + await publish(backend, key, ["beta"], session=f"handoff-{uuid.uuid4()}") + assert await backend.parked_wait_ids(key) == [] + finally: + await worker_b.shutdown() diff --git a/tests/contrib/external_workflow_streams/test_worker_integration.py b/tests/contrib/external_workflow_streams/test_worker_integration.py new file mode 100644 index 0000000000..cf797719f7 --- /dev/null +++ b/tests/contrib/external_workflow_streams/test_worker_integration.py @@ -0,0 +1,1644 @@ +"""P10a/P10b/P11/P19 — the wiring, exercised through a real Worker. + +Everything below runs against a live server with a real Workflow Task loop. The +unit tests elsewhere prove each piece in isolation; these prove the pieces are +actually connected to each other, which is the only thing isolation cannot show. + +PYTEST_DONT_REWRITE: sandboxed fixture Workflows re-import this module, so pytest's +injected imports would make sandbox validation depend on pytest's import locks. +""" + +# pyright: reportMissingParameterType=false +from __future__ import annotations + +import asyncio +import threading +import time +import uuid +from collections.abc import Sequence +from datetime import timedelta +from pathlib import Path + +import pytest + +import temporalio.api.common.v1 +import temporalio.api.enums.v1 +import temporalio.converter +from temporalio import workflow +from temporalio.client import Client +from temporalio.contrib.external_workflow_streams._backend import StreamKey +from temporalio.contrib.external_workflow_streams._errors import ( + METRIC_DECODE, + METRIC_INTEGRITY, + METRIC_STORAGE, +) +from temporalio.contrib.external_workflow_streams._record import ( + Offset, + RecordKind, + StreamRecord, +) +from temporalio.contrib.external_workflow_streams._wake import ( + WakeRequest, + send_wake_signal, + wake_request_id, +) +from temporalio.runtime import MetricBuffer, Runtime, TelemetryConfig +from temporalio.service import RPCError, RPCStatusCode +from temporalio.testing import WorkflowEnvironment +from temporalio.worker import Worker +from tests.contrib.external_workflow_streams.memory_backend import MemoryStreamBackend + +with workflow.unsafe.imports_passed_through(): + from temporalio.contrib.external_workflow_streams._api import external_stream + + +@workflow.defn +class ConsumeTokensWorkflow: + """Consumes a fixed number of records, then returns them.""" + + def __init__(self) -> None: + self._seen: list[str] = [] + + @workflow.run + async def run(self, expected: int) -> list[str]: + tokens = external_stream.with_options(idle_timeout=timedelta(seconds=30)).topic( + "tokens", type=str + ) + + async for token in tokens.subscribe(): + self._seen.append(token) + if len(self._seen) >= expected: + break + return self._seen + + @workflow.query + def seen(self) -> list[str]: + return self._seen + + +@workflow.defn +class CountTokensWorkflow: + """Consumes records and returns only how many, never the values. + + Deliberately distinct from :py:class:`ConsumeTokensWorkflow`: a Workflow that + returns its records puts them in its own result, and History would then + contain them for a reason that has nothing to do with streams. + """ + + @workflow.run + async def run(self, expected: int) -> int: + tokens = external_stream.with_options(idle_timeout=timedelta(seconds=30)).topic( + "tokens", type=str + ) + + seen = 0 + async for _ in tokens.subscribe(): + seen += 1 + if seen >= expected: + break + return seen + + +@workflow.defn +class SubscribeOnlyWorkflow: + """Subscribes and blocks forever, so the task is retained.""" + + @workflow.run + async def run(self) -> None: + tokens = external_stream.topic("tokens", type=str) + async for _ in tokens.subscribe(): + pass + + +@workflow.defn +class TimerSuppressedSubscriptionWorkflow: + """Subscribes and waits, with a long timer stopping the task being retained. + + Deliberately distinct from :py:class:`SubscribeOnlyWorkflow`: a Workflow + blocked on the stream *alone* leaves a retained Workflow Task, which is a + different shutdown transition. The timer makes every task complete, so the + Run sits cached with a live subscription and no open Workflow Task -- the + state that receives no eviction activation at shutdown at all. + """ + + @workflow.run + async def run(self) -> None: + tokens = external_stream.topic("tokens", type=str) + timer = asyncio.ensure_future(asyncio.sleep(600)) + try: + async for _ in tokens.subscribe(): + pass + finally: + timer.cancel() + + +@pytest.fixture +def backend() -> MemoryStreamBackend: + return MemoryStreamBackend() + + +async def publish( + backend: MemoryStreamBackend, + key: StreamKey, + values: list[str], + *, + session: str = "producer", +) -> None: + """Appends records the way a producer would, encoded for the topic's type. + + The sequence continues across calls within one session, because + ``(session_id, sequence)`` is the idempotency key: restarting it would + re-use a key with different content, which the contract rejects outright. + """ + import temporalio.converter + from temporalio.contrib.external_workflow_streams._codec import StreamPayloadCodec + + codec = StreamPayloadCodec(temporalio.converter.DataConverter.default, str) + start = _sequences.get((id(backend), session), 0) + for i, value in enumerate(values, start=start): + await backend.append( + key, + StreamRecord(RecordKind.DATA, await codec.encode(value), session, i), + ) + _sequences[(id(backend), session)] = start + len(values) + + +#: Per-producer-session sequence, so successive publishes do not collide. +_sequences: dict[tuple[int, str], int] = {} + + +async def test_a_workflow_consumes_records_it_never_read_itself( + client: Client, backend: MemoryStreamBackend +) -> None: + """The vertical slice: subscribe, retain, notify, resolve, drain, deliver. + + Nothing in the Workflow ever touches the backend. The Worker's manager reads + it on its own loop, buffers, and tells Core; Core turns that into an + activation; and the Workflow's drain is a buffer pop that cannot block. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[ConsumeTokensWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + ConsumeTokensWorkflow.run, + 3, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + description = await handle.describe() + key = StreamKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + "tokens", + ) + + # Let the Workflow reach its subscription before anything is published, + # so this exercises the wake path rather than a buffer that was already + # full when the first activation ran. + await asyncio.sleep(1) + await publish(backend, key, ["alpha", "beta", "gamma"]) + + assert await asyncio.wait_for(handle.result(), 30) == [ + "alpha", + "beta", + "gamma", + ] + + +async def test_no_stream_payload_reaches_history( + client: Client, backend: MemoryStreamBackend +) -> None: + """The entire point of the feature, asserted against real History. + + History cost scales with consumption batches, not with item count -- so a + record's bytes must appear nowhere in it, and the marker that *is* there + must be bounded metadata rather than the payloads. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[CountTokensWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + CountTokensWorkflow.run, + 2, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + description = await handle.describe() + key = StreamKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + "tokens", + ) + await asyncio.sleep(1) + await publish(backend, key, ["secret-payload-one", "secret-payload-two"]) + assert await asyncio.wait_for(handle.result(), 30) == 2 + + events = [e async for e in handle.fetch_history_events()] + raw = b"".join(e.SerializeToString() for e in events) + + assert b"secret-payload-one" not in raw, ( + "a stream payload reached Temporal History, which is the one thing " + "this feature exists to prevent" + ) + assert b"secret-payload-two" not in raw + + # And History did not grow an event per record either: cost scales with + # consumption batches, not with item count. + markers = [e for e in events if e.HasField("marker_recorded_event_attributes")] + assert len(markers) <= 2, ( + f"expected marker cost to be per Workflow Task, got {len(markers)} " + "markers for 2 records" + ) + + +async def test_a_blocked_workflow_retains_its_task_rather_than_completing_it( + client: Client, backend: MemoryStreamBackend +) -> None: + """Retention is what lets the next record arrive on the *same* task. + + Without it every record would cost a Workflow Task round trip, which is the + cost model the feature exists to avoid. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[SubscribeOnlyWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + SubscribeOnlyWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + await asyncio.sleep(2) + + try: + history = [e async for e in handle.fetch_history_events()] + completed = [ + e + for e in history + if e.HasField("workflow_task_completed_event_attributes") + ] + # Exactly one task has completed -- the one that ran up to the + # subscription. The one holding the subscription open has not, and + # will not until a record arrives or the idle timeout fires. + assert len(completed) <= 1, ( + f"expected the subscription's task to be retained, got " + f"{len(completed)} completed tasks" + ) + finally: + await handle.terminate() + + +async def test_a_workflow_without_a_configured_backend_says_so( + client: Client, +) -> None: + """The failure names the Worker option rather than an attribute error.""" + task_queue = f"tq-{uuid.uuid4()}" + async with Worker(client, task_queue=task_queue, workflows=[SubscribeOnlyWorkflow]): + handle = await client.start_workflow( + SubscribeOnlyWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + await asyncio.sleep(2) + try: + history = [e async for e in handle.fetch_history_events()] + failures = [ + e + for e in history + if e.HasField("workflow_task_failed_event_attributes") + ] + assert failures, "expected the Workflow Task to fail" + message = failures[0].workflow_task_failed_event_attributes.failure.message + assert "external_stream_backend" in message, ( + f"the failure should name the Worker option, got: {message}" + ) + finally: + await handle.terminate() + + +@workflow.defn +class TimerThenConsumeWorkflow: + """Consumes one record, sleeps, then consumes another. + + The sleep is the point: a completion carrying a timer is server-bound, so + retention is suppressed and the Workflow Task ends. The second record + therefore arrives with **no open Workflow Task**, which is the window only + the wake Signal covers. + """ + + @workflow.run + async def run(self) -> list[str]: + tokens = external_stream.with_options(idle_timeout=timedelta(seconds=30)).topic( + "tokens", type=str + ) + seen: list[str] = [] + subscription = tokens.subscribe() + iterator = subscription.__aiter__() + + seen.append(await iterator.__anext__()) + # A real timer: the completion that carries it cannot ask for retention. + await asyncio.sleep(2) + seen.append(await iterator.__anext__()) + return seen + + +async def test_an_append_with_no_open_task_wakes_the_workflow( + client: Client, backend: MemoryStreamBackend +) -> None: + """The wake Signal path, end to end, with nothing else able to explain it. + + The second record is published while the Workflow is sleeping, so no + Workflow Task is open to accept local readiness. Nothing else will create + one -- the sleep's own timer fires on its own schedule and the Workflow + would then find the record only by luck of timing. The watcher observing + `NoOpenWorkflowTask` and sending the Signal is what makes the delivery + reliable rather than incidental. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[TimerThenConsumeWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + TimerThenConsumeWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + description = await handle.describe() + key = StreamKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + "tokens", + ) + await asyncio.sleep(1) + await publish(backend, key, ["first"]) + + # Let the Workflow consume it and settle into its sleep, so the second + # append lands squarely in the no-open-task window. + await asyncio.sleep(1) + await publish(backend, key, ["second"]) + + assert await asyncio.wait_for(handle.result(), 60) == ["first", "second"] + + events = [e async for e in handle.fetch_history_events()] + signals = [ + e + for e in events + if e.HasField("workflow_execution_signaled_event_attributes") + and e.workflow_execution_signaled_event_attributes.signal_name + == "__temporal_external_stream_wake" + ] + assert signals, ( + "the record was delivered without a wake Signal, so this passed by " + "timing rather than by the mechanism under test" + ) + + +async def test_every_marker_a_run_writes_is_a_complete_annotation( + client: Client, backend: MemoryStreamBackend +) -> None: + """Each marker carries its own header and ends with its own terminal. + + A Run that writes two markers is where this can go wrong, and + :py:class:`TimerThenConsumeWorkflow` writes exactly two: the completion that + carries its timer ends one Workflow Task, and the completion that returns + ends the next. Core writes and clears the accumulated annotation on both, so + an accumulator that carried its header across the first of them would begin + the second annotation at whatever frame came next -- decoded as a schema + version, since that is what leads an annotation. + + Decoding every marker independently is the assertion: the annotation is + opaque to Core, so nothing between here and History would notice one that + cannot be read back, and the failure would surface only on the replay that + needed it. + """ + from temporalio.bridge.proto.external_data import ExternalStreamMarkerData + from temporalio.contrib.external_workflow_streams._annotation import ( + decode_annotation, + ) + + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[TimerThenConsumeWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + TimerThenConsumeWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + description = await handle.describe() + key = StreamKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + "tokens", + ) + await asyncio.sleep(1) + await publish(backend, key, ["first"]) + await asyncio.sleep(1) + await publish(backend, key, ["second"]) + + assert await asyncio.wait_for(handle.result(), 60) == ["first", "second"] + + events = [e async for e in handle.fetch_history_events()] + markers = [ + e + for e in events + if e.HasField("marker_recorded_event_attributes") + and e.marker_recorded_event_attributes.marker_name == "core_external_stream" + ] + assert len(markers) >= 2, ( + "this Run should write one marker per Workflow Task it ended, and " + f"it ended at least two; got {len(markers)}" + ) + # Decoded first, all of them: a marker that lost its header fails here, + # and it is the *second* one that loses it, so a per-marker assertion + # interleaved with the decoding would stop on the first marker's own + # shortcoming and never reach the one under test. + annotations = [] + for index, event in enumerate(markers): + data = ExternalStreamMarkerData() + data.ParseFromString( + event.marker_recorded_event_attributes.details["external_stream"] + .payloads[0] + .data + ) + annotations.append(decode_annotation(data.replay_annotation)) + + for index, annotation in enumerate(annotations): + assert annotation.header.streams, ( + f"marker {index} records no stream in its header, so replay of " + "it has nothing to start from" + ) + assert annotation.terminal is not None, ( + f"marker {index} has no terminal, so nothing in it says where " + "the Workflow Task's deliveries stopped (ADR-008)" + ) + + +@workflow.defn +class TimerThenFirstRecordWorkflow: + """Starts a timer and blocks on the stream in the **same** activation. + + Deliberately distinct from :py:class:`TimerThenConsumeWorkflow`, which + consumes a record *before* its timer: that one is quiescent once before the + server-bound command ever appears, so its wait set is already registered + with Core and only has to survive. Here the very first block rides a + completion that carries a timer, so nothing has registered the wait yet -- + which is the only difference between the two, and the whole case. + """ + + @workflow.run + async def run(self) -> list[str]: + tokens = external_stream.with_options(idle_timeout=timedelta(seconds=30)).topic( + "tokens", type=str + ) + iterator = tokens.subscribe().__aiter__() + # Long enough that it cannot be what resumes this Workflow. If a record + # is ever returned, the stream delivered it. + timer = asyncio.ensure_future(asyncio.sleep(600)) + try: + return [await iterator.__anext__()] + finally: + timer.cancel() + + +async def test_a_first_block_that_rides_a_server_bound_command_is_still_wakeable( + client: Client, backend: MemoryStreamBackend +) -> None: + """The wait set has to be registered even by a completion that is reported. + + A completion carrying a timer must not ask for retention -- the server has + to be told about the timer -- but the subscriptions it leaves behind are + still active, and the wake Signal that covers that window can only resume a + Run whose waits Core knows about. Without the registration every append + produces a wake Signal, every wake Signal produces a Workflow Task with no + jobs in it, and the Workflow is unresumable for the life of the timer. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[TimerThenFirstRecordWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + TimerThenFirstRecordWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + description = await handle.describe() + key = StreamKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + "tokens", + ) + + # Published after the Workflow has settled, so the record arrives with + # no open Workflow Task and the wake Signal is the only way in. + await asyncio.sleep(1) + await publish(backend, key, ["first"]) + + try: + # Short on purpose: the timer is 600 seconds, so nothing but the + # stream can finish this Workflow and a longer wait would only make + # a known gap slower to report. + assert await asyncio.wait_for(handle.result(), 15) == ["first"] + finally: + # The Workflow reaching its own end is the success case here, so + # terminating unconditionally turns a pass into a NOT_FOUND error -- + # which is how the same mistake presented in the rollover tests. + try: + await handle.terminate() + except RPCError as err: + if err.status is not RPCStatusCode.NOT_FOUND: + raise + + +def _stream_markers(events: list) -> list: # type: ignore[type-arg] + return [ + e + for e in events + if e.HasField("marker_recorded_event_attributes") + and e.marker_recorded_event_attributes.marker_name == "core_external_stream" + ] + + +async def _wait_for_markers( + handle, count: int, message: str, timeout: float = 30 +) -> None: # type: ignore[no-untyped-def] + """Polls until History holds ``count`` stream markers. + + Polled rather than slept: a marker is written when the Workflow Task ends, + and waiting a fixed time instead would make every assertion after it a + statement about how fast this machine is. + """ + deadline = asyncio.get_running_loop().time() + timeout + while asyncio.get_running_loop().time() < deadline: + if ( + len(_stream_markers([e async for e in handle.fetch_history_events()])) + >= count + ): + return + await asyncio.sleep(0.3) + raise AssertionError(message) + + +class ReplayFailureTrace: + """Correlates the manager events that distinguish issues 2 and 7.""" + + def __init__(self) -> None: + self._started = time.monotonic() + self.events: list[dict[str, object]] = [] + self.managers: list[object] = [] + + def record(self, event: str, **details: object) -> None: + entry: dict[str, object] = { + "at_seconds": round(time.monotonic() - self._started, 6), + "event": event, + } + entry.update(details) + self.events.append(entry) + + def record_explicit_wake(self, request: WakeRequest) -> None: + self.record( + "explicit_wake", + request_id=wake_request_id(request), + workflow_id=request.workflow_id, + run_id=request.first_execution_run_id, + stream_name=request.stream_name, + wait_id=request.wait_id, + park_generation=request.park_generation, + sender_identity=request.sender_identity, + wake_counter=request.wake_counter, + ) + + def record_explicit_wake_acknowledged(self, request: WakeRequest) -> None: + self.record("explicit_wake_acknowledged", request_id=wake_request_id(request)) + + def manager_wake_bursts_after_explicit_ack(self) -> list[int]: + """Counts wakes between each Core-accepted task completion.""" + acknowledged = next( + ( + index + for index, event in enumerate(self.events) + if event["event"] == "explicit_wake_acknowledged" + ), + None, + ) + if acknowledged is None: + raise AssertionError("the explicit wake was never acknowledged") + bursts: list[int] = [] + current: set[object] = set() + for event in self.events[acknowledged + 1 :]: + if event["event"] == "manager_wake_started": + current.add(event["request_id"]) + elif event["event"] == "workflow_task_completion_accepted": + bursts.append(len(current)) + current = set() + bursts.append(len(current)) + return bursts + + +@pytest.fixture +def replay_failure_trace(monkeypatch: pytest.MonkeyPatch) -> ReplayFailureTrace: + """Captures diagnostics without changing readiness or wake behavior.""" + import temporalio.contrib.external_workflow_streams._wake as wake_module + from temporalio.contrib.external_workflow_streams._manager import ( + StreamSubscriptionManager, + ) + + trace = ReplayFailureTrace() + original_init = StreamSubscriptionManager.__init__ + original_notify = StreamSubscriptionManager._notify_ready_with_retries + original_evict = StreamSubscriptionManager.evict_run + original_started = StreamSubscriptionManager.note_workflow_task_started + original_completed = StreamSubscriptionManager.note_workflow_task_completed + original_send_wake = wake_module.send_wake_signal + + def traced_init(self, *args, **kwargs): # type: ignore[no-untyped-def] + original_init(self, *args, **kwargs) + trace.managers.append(self) + trace.record("manager_created", sender_identity=self.wake_sender_identity) + + async def traced_notify(self, subscription): # type: ignore[no-untyped-def] + generation = subscription.current_wait_generation() + result = await original_notify(self, subscription) + trace.record( + "readiness", + run_id=subscription.run_id, + wait_id=subscription.wait_id, + generation=generation, + result=result, + buffered=subscription.buffered, + wakes_owed=subscription.wakes_owed, + ) + return result + + async def traced_evict(self, run_id: str) -> None: # type: ignore[no-untyped-def] + trace.record("evict_started", run_id=run_id) + await original_evict(self, run_id) + trace.record("evict_finished", run_id=run_id) + + def traced_completed(self, run_id: str, *, terminal: bool) -> None: # type: ignore[no-untyped-def] + trace.record( + "workflow_task_completion_accepted", + run_id=run_id, + terminal=terminal, + ) + original_completed(self, run_id, terminal=terminal) + + def traced_started(self, run_id: str) -> None: # type: ignore[no-untyped-def] + trace.record("workflow_task_started", run_id=run_id) + original_started(self, run_id) + + async def traced_send_wake( + client: Client, + request: WakeRequest, + *, + producer_session_id: str = "", + ) -> str: + request_id = wake_request_id(request) + trace.record( + "manager_wake_started", + request_id=request_id, + workflow_id=request.workflow_id, + run_id=request.first_execution_run_id, + stream_name=request.stream_name, + wait_id=request.wait_id, + park_generation=request.park_generation, + sender_identity=request.sender_identity, + wake_counter=request.wake_counter, + ) + try: + result = await original_send_wake( + client, request, producer_session_id=producer_session_id + ) + except BaseException as err: + trace.record( + "manager_wake_failed", + request_id=request_id, + error=repr(err), + ) + raise + trace.record("manager_wake_acknowledged", request_id=request_id) + return result + + monkeypatch.setattr(StreamSubscriptionManager, "__init__", traced_init) + monkeypatch.setattr( + StreamSubscriptionManager, "_notify_ready_with_retries", traced_notify + ) + monkeypatch.setattr(StreamSubscriptionManager, "evict_run", traced_evict) + monkeypatch.setattr( + StreamSubscriptionManager, + "note_workflow_task_started", + traced_started, + ) + monkeypatch.setattr( + StreamSubscriptionManager, + "note_workflow_task_completed", + traced_completed, + ) + monkeypatch.setattr(wake_module, "send_wake_signal", traced_send_wake) + return trace + + +async def _replay_failure_diagnosis( + handle, # type: ignore[no-untyped-def] + backend: MemoryStreamBackend, + key: StreamKey, + trace: ReplayFailureTrace, + artifact_dir: Path, + phase: str, +) -> str: + """Preserves the evidence a timeout cleanup used to destroy.""" + lines = [f"diagnostic phase: {phase}"] + try: + history = await handle.fetch_history() + artifact = artifact_dir / f"{handle.id}-{phase}-history.json" + artifact.write_text(history.to_json()) + lines.append(f"complete History: {artifact}") + workflow_task_timeline = [] + for event in history.events: + attributes = event.WhichOneof("attributes") or "?" + if attributes.startswith("workflow_task_"): + item: dict[str, object] = { + "event_id": event.event_id, + "attributes": attributes, + } + if event.HasField("workflow_task_failed_event_attributes"): + failed = event.workflow_task_failed_event_attributes + item["cause"] = failed.cause + item["message"] = failed.failure.message[:300] + workflow_task_timeline.append(item) + lines.append(f"Workflow Task timeline: {workflow_task_timeline}") + except Exception as err: # noqa: BLE001 -- diagnostics must not mask timeout + lines.append(f"History unavailable: {err!r}") + try: + description = await handle.describe() + lines.append(f"workflow status: {description.status}") + except Exception as err: # noqa: BLE001 + lines.append(f"describe unavailable: {err!r}") + try: + lines.append( + "stream state: " + f"records={backend.all_records(key)} " + f"parked_wait_ids={await backend.parked_wait_ids(key)}" + ) + except Exception as err: # noqa: BLE001 + lines.append(f"stream state unavailable: {err!r}") + for manager in trace.managers: + subscriptions = getattr(manager, "_runs", {}) + for run_id, waits in subscriptions.items(): + for wait_id, subscription in waits.items(): + lines.append( + f"subscription {run_id}/{wait_id}: " + f"buffered={subscription.buffered} " + f"committed={subscription.committed_cursor} " + f"delivery={subscription.delivery_cursor} " + f"prefetch={subscription.prefetch_cursor} " + f"generation={subscription.current_wait_generation()} " + f"wakes_owed={subscription.wakes_owed} " + f"cancelled={subscription._cancelled} " + f"watcher_done=" + f"{None if subscription._watcher is None else subscription._watcher.done()}" + ) + lines.append(f"correlated manager trace: {trace.events}") + return "\n".join(lines) + + +async def _await_replayed_result( + handle, # type: ignore[no-untyped-def] + expected: list[str], + backend: MemoryStreamBackend, + key: StreamKey, + trace: ReplayFailureTrace, + artifact_dir: Path, +) -> None: + """Waits for the result and retains both timeout and late-observation state.""" + try: + result = await asyncio.wait_for(handle.result(), 30) + except asyncio.TimeoutError: + at_timeout = await _replay_failure_diagnosis( + handle, backend, key, trace, artifact_dir, "at-timeout" + ) + try: + late_result = await asyncio.wait_for(handle.result(), 10) + except asyncio.TimeoutError: + after_observation = await _replay_failure_diagnosis( + handle, backend, key, trace, artifact_dir, "after-observation" + ) + raise AssertionError( + "the replayed Run did not complete within 30 seconds and remained " + "non-terminal for the 10-second observation window\n\n" + f"{at_timeout}\n\n{after_observation}" + ) from None + raise AssertionError( + "the replayed Run missed the 30-second bound but completed during " + f"the observation window with {late_result!r}\n\n{at_timeout}" + ) from None + assert result == expected + + +@workflow.defn +class ParkedAcrossEvictionWorkflow: + """Consumes records with nothing else to do, so every task ends in a park. + + No timer, no other command: each Workflow Task is retained until the idle + timeout parks it, which is what puts a marker in History and leaves the Run + parked rather than merely cached. Both are needed here -- the marker is what + makes the Run *replay* after eviction, and the park is what makes the wake + Signal the only thing that can bring a Workflow Task back. + """ + + @workflow.run + async def run(self, expected: int) -> list[str]: + tokens = external_stream.topic("tokens", type=str) + seen: list[str] = [] + async for token in tokens.subscribe(): + seen.append(token) + if len(seen) >= expected: + break + return seen + + +@workflow.defn +class FillerWorkflow: + """Occupies the one cache slot, which is what evicts the Run under test.""" + + @workflow.run + async def run(self) -> str: + return "done" + + +async def test_a_replayed_run_re_registers_its_wait_set( + client: Client, + backend: MemoryStreamBackend, + replay_failure_trace: ReplayFailureTrace, + tmp_path: Path, +) -> None: + """A Run that comes back through replay has to end replay registered. + + Core's wait set is per-Worker runtime state: it is not in History and it + does not survive eviction, so a replayed Run rebuilds it or has none. The + only thing that builds it is ``WorkflowStreamQuiescent``, and a completion + that reports no progress because it is replaying must still report the + snapshot -- what is already in History is the *annotation*, not the + registration. + + The Run is deliberately left **parked** before it is evicted, so nothing + else can explain the resume: there is no timer to fire, no watcher left + alive on this Worker, and no open Workflow Task. A wake Signal creates the + replacement task, and everything after that depends on the replayed Run + knowing what it is subscribed to. Without it every later readiness is + answered as though the Run had no subscriptions at all, each wake produces a + Workflow Task with no activation in it, and the record sits in the stream + while the Workflow waits. + """ + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[ParkedAcrossEvictionWorkflow, FillerWorkflow], + external_stream_backend=backend, + # One slot, so running anything else evicts the Run under test. Eviction + # is what discards the wait set, the watchers, and the buffers, leaving + # replay to rebuild all three. + max_cached_workflows=1, + max_concurrent_workflow_tasks=2, + ): + handle = await client.start_workflow( + ParkedAcrossEvictionWorkflow.run, + 2, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + description = await handle.describe() + first_run_id = description.raw_description.workflow_execution_info.first_run_id + key = StreamKey(client.namespace, handle.id, first_run_id, "tokens") + + # Each wait is for a *park*, not for a duration: the record is published + # only once the Workflow has actually parked, so it is delivered by a + # wake rather than into a Workflow Task that happened to still be open. + await _wait_for_markers(handle, 1, "the Run never parked at all") + + # Consumed live, so the Run has a marker to replay from and a cursor + # that must be honoured -- a replay that re-delivered "alpha" would show + # up in the result rather than in a timeout. + await publish(backend, key, ["alpha"]) + await _wait_for_markers( + handle, + 2, + "the Run never parked again after consuming its first record, so it " + "was never in the state this case evicts from", + ) + + # Evicts the Run: its wait set, watchers, and buffers all go with it. + await client.execute_workflow( + FillerWorkflow.run, id=f"filler-{uuid.uuid4()}", task_queue=task_queue + ) + + # Appended with nothing on this Worker watching for it, then woken the + # way a durable producer wakes a parked Run. The Signal is the only + # thing that creates a Workflow Task here, and it carries no records -- + # so if the replayed Run does not know its own subscription, the + # Workflow Task it creates has nothing in it. + await publish(backend, key, ["beta"]) + request = WakeRequest( + namespace=client.namespace, + workflow_id=handle.id, + first_execution_run_id=first_run_id, + stream_name="tokens", + wait_id=1, + park_generation=0, + sender_identity=f"test-{uuid.uuid4()}", + wake_counter=1, + ) + replay_failure_trace.record_explicit_wake(request) + await send_wake_signal(client, request) + replay_failure_trace.record_explicit_wake_acknowledged(request) + + try: + await _await_replayed_result( + handle, + ["alpha", "beta"], + backend, + key, + replay_failure_trace, + tmp_path, + ) + bursts = replay_failure_trace.manager_wake_bursts_after_explicit_ack() + assert max(bursts, default=0) <= 1, ( + "readiness produced multiple wakes before Core accepted the " + f"task completion that one wake caused: {replay_failure_trace.events}" + ) + finally: + try: + await handle.terminate() + except Exception: + pass + + +@workflow.defn +class FloodedCountWorkflow: + """Consumes far more records than one activation is allowed to deliver.""" + + def __init__(self) -> None: + self._seen = 0 + + @workflow.run + async def run(self, expected: int) -> int: + tokens = external_stream.with_options(idle_timeout=timedelta(seconds=30)).topic( + "tokens", type=str + ) + async for _ in tokens.subscribe(): + self._seen += 1 + if self._seen >= expected: + break + return self._seen + + @workflow.query + def seen(self) -> int: + return self._seen + + +async def test_a_flood_larger_than_one_activations_budget_is_delivered_in_full( + client: Client, backend: MemoryStreamBackend +) -> None: + """The whole budget mechanism, end to end, on a real Workflow Task loop. + + Three things have to be connected for this to finish, and each of them hangs + the Workflow permanently on its own: + + - the per-activation budget, or the first ``activate()`` never returns and + the Workflow Task dies on the 2-second deadlock timeout -- on every retry, + because the records are still there; + - the reset at the start of each activation, or delivery stops for good once + the first budget is spent; + - the readiness re-arm, or the records the budget left buffered are never + announced again, because the watcher moved its prefetch cursor past them + when it buffered them. + + So a plain "did it finish" assertion is not a weak one here: nothing about + this Workflow completes if any part of the mechanism is missing. + """ + from temporalio.contrib.external_workflow_streams._api import ( + MAX_RECORDS_PER_ACTIVATION, + ) + + total = 3 * MAX_RECORDS_PER_ACTIVATION + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + client, + task_queue=task_queue, + workflows=[FloodedCountWorkflow], + external_stream_backend=backend, + ): + handle = await client.start_workflow( + FloodedCountWorkflow.run, + total, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + description = await handle.describe() + key = StreamKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + "tokens", + ) + await asyncio.sleep(1) + # One session, one continuous sequence: `(session_id, sequence)` is the + # idempotency key, so restarting the numbering would re-use a key with + # different content and the backend would reject the append. + await publish(backend, key, [f"t{i}" for i in range(total)]) + + assert await asyncio.wait_for(handle.result(), 60) == total + + +async def test_a_clean_shutdown_sweeps_and_tears_down_the_manager( + client: Client, backend: MemoryStreamBackend +) -> None: + """P20's sweep has to run on the path users actually take. + + A normal ``Worker.shutdown()`` completes with no worker task raising, so + nothing is replaced by ``drain_poll_queue()``. Wiring the sweep there alone + means the whole mechanism is dead on the ordinary path: the manager never + enters its shutting-down state, every Run stays registered, watchers go on + polling, and the buffers and backend connections go out with the process -- + while the wake Signal that would hand the Run to another Worker is never + even considered. + + The Workflow here keeps a long timer running alongside its subscription, so + every Workflow Task completes and the Run sits *cached, subscribed, and + holding no Workflow Task* when shutdown begins. That is the state the sweep + exists for, and the one that gets no eviction activation of its own. + """ + task_queue = f"tq-{uuid.uuid4()}" + worker = Worker( + client, + task_queue=task_queue, + workflows=[TimerSuppressedSubscriptionWorkflow], + external_stream_backend=backend, + ) + worker_task = asyncio.create_task(worker.run()) + handle = None + try: + handle = await client.start_workflow( + TimerSuppressedSubscriptionWorkflow.run, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + # Wait for the subscription to actually reach the manager rather than + # sleeping: with no Run registered there is nothing to sweep and the + # assertions below would hold for the wrong reason. + deadline = asyncio.get_running_loop().time() + 30 + while asyncio.get_running_loop().time() < deadline: + manager = worker._workflow_worker._external_stream_manager # type: ignore[union-attr] + if manager is not None and manager.runs_with_subscriptions(): + break + await asyncio.sleep(0.2) + manager = worker._workflow_worker._external_stream_manager # type: ignore[union-attr] + assert manager is not None and manager.runs_with_subscriptions(), ( + "the Workflow never registered a subscription, so there is nothing " + "for a shutdown sweep to find" + ) + + await asyncio.wait_for(worker.shutdown(), 60) + + assert manager._shutting_down, ( + "the manager was never told the Worker is shutting down, so no Run " + "was probed and no owed wake was sent" + ) + assert manager.runs_with_subscriptions() == [], ( + "subscriptions survived shutdown: their watchers, buffers, and " + "backend connections leak with the process" + ) + finally: + if handle is not None: + try: + await handle.terminate() + except Exception: + pass + if not worker_task.done(): + worker_task.cancel() + await asyncio.gather(worker_task, return_exceptions=True) + + +# --- P19 / ADR-011: decoding is Worker-side work, not Workflow-thread work --- + +#: What the probe codec saw each time it decoded a stream record: the thread it +#: ran on, and the module the running event loop's class comes from. Both are +#: recorded because either one alone is arguable -- a Workflow's deterministic +#: loop is identified by its class, and the executor thread by its name. +_decode_sites: list[tuple[str, str]] = [] + +#: Marks the one payload this probe reacts to, so ordinary activation payloads +#: -- arguments, results, headers -- are passed through untouched and cannot be +#: mistaken for a stream record's decode. +_STREAM_SENTINEL = b"external-stream-decode-probe" + +#: Longer than the Worker's 2-second deadlock timeout. A codec is user code and +#: may legitimately take this long: it can be fetching an external payload or +#: talking to a KMS. Awaited on the Worker's loop it costs nothing but latency; +#: awaited on the Workflow thread it is either a deadlocked Workflow Task or a +#: Workflow command synthesized out of a codec's internals. +_SLOW_DECODE = timedelta(seconds=2.5) + + +class ProbeCodec(temporalio.converter.PayloadCodec): + """A pass-through codec that reports where the stream record was decoded. + + Only the record carrying :py:data:`_STREAM_SENTINEL` is treated specially, + so the Worker's own ``decode_activation`` work -- which is *supposed* to run + on the Worker's loop -- adds no observations of its own. + """ + + async def encode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + return list(payloads) + + def __init__(self, delay: float | None = None) -> None: + self.delay = _SLOW_DECODE.total_seconds() if delay is None else delay + + async def decode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + for payload in payloads: + if _STREAM_SENTINEL not in payload.data: + continue + loop = asyncio.get_running_loop() + _decode_sites.append( + (threading.current_thread().name, type(loop).__module__) + ) + # Real asynchronous work, of the length a codec is allowed to + # take: an external-payload fetch or a KMS round trip is I/O, and + # the Worker's loop is where the Worker awaits every other + # payload's codec. + await asyncio.sleep(self.delay) + return list(payloads) + + +async def test_a_slow_codec_decodes_off_the_workflow_thread( + client: Client, backend: MemoryStreamBackend +) -> None: + """A stream record's payload codec must not run inside ``activate()``. + + ``DataConverter.decode`` is three things: an external-payload retrieval, a + user ``PayloadCodec``, and a payload *converter*. The first two are + arbitrary asynchronous work -- a network fetch, a KMS round trip -- and the + Worker awaits them for every ordinary activation payload *before* handing + the activation to the Workflow executor, precisely so that no Workflow Task + can be failed by them. Only the third, a synchronous conversion, belongs on + the Workflow thread. + + A stream record is a payload like any other, so the same split has to hold + for it: by the time a record reaches ``activate()`` its codec has already + run on the Worker's loop, and the Workflow thread does nothing but convert + already-prepared bytes into a value. + + Running the codec on the Workflow thread instead is wrong three ways at + once, and the assertions below name them: the awaited work happens inside a + deterministic event loop that cannot perform I/O, it can exceed the + 2-second deadlock timeout and fail the Workflow Task for a perfectly + healthy codec, and anything the codec awaits is synthesized as Workflow + commands. + """ + _decode_sites.clear() + config = client.config() + config["data_converter"] = temporalio.converter.DataConverter( + payload_codec=ProbeCodec() + ) + codec_client = Client(**config) + + task_queue = f"tq-{uuid.uuid4()}" + async with Worker( + codec_client, + task_queue=task_queue, + workflows=[CountTokensWorkflow], + external_stream_backend=backend, + ): + handle = await codec_client.start_workflow( + CountTokensWorkflow.run, + 1, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + description = await handle.describe() + key = StreamKey( + client.namespace, + handle.id, + description.raw_description.workflow_execution_info.first_run_id, + "tokens", + ) + + # Published after the Workflow is already blocked on its subscription, + # so the record arrives through the live wake path rather than sitting + # in a buffer the first activation happens to find. + await asyncio.sleep(1) + from temporalio.contrib.external_workflow_streams._codec import ( + StreamPayloadCodec, + ) + + producer_codec = StreamPayloadCodec(codec_client.data_converter, str) + await backend.append( + key, + StreamRecord( + RecordKind.DATA, + await producer_codec.encode(_STREAM_SENTINEL.decode()), + "producer", + 0, + ), + ) + + # A Workflow Task failed by a deadlock is *retried*, and every retry + # re-decodes the same record and deadlocks again, so a Workflow that + # never returns is the visible form of that failure. + assert await asyncio.wait_for(handle.result(), 30) == 1 + + assert _decode_sites, ( + "the probe codec never saw the stream record, so this test proved " + "nothing about where decoding runs" + ) + for thread_name, loop_module in _decode_sites: + assert not thread_name.startswith("temporal_workflow_"), ( + f"the stream record's codec ran on the Workflow executor thread " + f"({thread_name}): arbitrary async work inside activate(), under " + "the 2-second deadlock timeout" + ) + assert loop_module != "temporalio.worker._workflow_instance", ( + "the stream record's codec was awaited on the Workflow's " + "deterministic event loop, which performs no I/O and turns every " + "await into a Workflow command" + ) + + # The codec's own awaits must not have become Workflow commands. A timer in + # History here is not a slow Workflow -- it is a Workflow whose History + # depends on what a codec did internally, and replay reproduces it only for + # as long as the codec behaves identically. + events = [e async for e in handle.fetch_history_events()] + assert not [e for e in events if e.HasField("timer_started_event_attributes")], ( + "the codec's await was turned into a Workflow timer command" + ) + + +async def test_a_replayed_record_is_prepared_off_the_workflow_thread( + env: WorkflowEnvironment, + replay_failure_trace: ReplayFailureTrace, + tmp_path: Path, +) -> None: + """Replay delivers down the same drain, so it prepares the same way. + + A replayed record never passes a watcher: its bytes are read and validated + when the replay job is prepared, before the Workflow thread runs at all. If + preparation happened only in the watcher, replay would hand raw producer + bytes to a Workflow whose converter has a codec -- and the choice would be + between running that codec inside ``activate()`` and yielding whatever the + undecoded bytes convert to. Preparing both paths in the same place is what + makes replay indistinguishable from live delivery, which is the property the + whole buffer design rests on. + """ + _decode_sites.clear() + backend = MemoryStreamBackend() + client = await env.connect_client( + data_converter=temporalio.converter.DataConverter(payload_codec=ProbeCodec(0)) + ) + task_queue = f"tq-{uuid.uuid4()}" + handle = None + async with Worker( + client, + task_queue=task_queue, + workflows=[ParkedAcrossEvictionWorkflow, FillerWorkflow], + external_stream_backend=backend, + max_cached_workflows=1, + max_concurrent_workflow_tasks=2, + ): + try: + handle = await client.start_workflow( + ParkedAcrossEvictionWorkflow.run, + 2, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + description = await handle.describe() + first_run_id = ( + description.raw_description.workflow_execution_info.first_run_id + ) + key = StreamKey(client.namespace, handle.id, first_run_id, "tokens") + + await _wait_for_markers(handle, 1, "the Run never parked at all") + await publish( + backend, + key, + [_STREAM_SENTINEL.decode()], + session=f"producer-{uuid.uuid4()}", + ) + await _wait_for_markers( + handle, + 2, + "the Run never committed a marker, so nothing would be replayed", + ) + + # Evicts the Run, so the next Workflow Task replays the marker -- + # and re-delivers the recorded record through the replay path. + await client.execute_workflow( + FillerWorkflow.run, id=f"filler-{uuid.uuid4()}", task_queue=task_queue + ) + await publish(backend, key, ["beta"], session=f"producer-{uuid.uuid4()}") + request = WakeRequest( + namespace=client.namespace, + workflow_id=handle.id, + first_execution_run_id=first_run_id, + stream_name="tokens", + wait_id=1, + park_generation=0, + sender_identity=f"test-{uuid.uuid4()}", + wake_counter=1, + ) + replay_failure_trace.record_explicit_wake(request) + await send_wake_signal(client, request) + replay_failure_trace.record_explicit_wake_acknowledged(request) + + await _await_replayed_result( + handle, + [_STREAM_SENTINEL.decode(), "beta"], + backend, + key, + replay_failure_trace, + tmp_path, + ) + bursts = replay_failure_trace.manager_wake_bursts_after_explicit_ack() + assert max(bursts, default=0) <= 1, ( + "replayed readiness produced multiple wakes before Core accepted " + f"the task completion one wake caused: {replay_failure_trace.events}" + ) + finally: + if handle is not None: + try: + await handle.terminate() + except Exception: + pass + + assert _decode_sites, "the probe codec never saw the replayed record" + for thread_name, loop_module in _decode_sites: + assert not thread_name.startswith("temporal_workflow_"), ( + f"a replayed record's codec ran on the Workflow executor thread " + f"({thread_name})" + ) + assert loop_module != "temporalio.worker._workflow_instance", ( + "a replayed record's codec was awaited on the Workflow's " + "deterministic event loop" + ) + + +# --- P18: the failure taxonomy, connected end to end ------------------------- + + +class TaxonomyCodec(temporalio.converter.PayloadCodec): + """A pass-through codec that can be made to reject one known value. + + Stands in for the ordinary way row three happens: a consumer whose + converter or codec stops matching the producer's, with the stream itself + untouched. Scoped to the one sentinel value so the Workflow's own arguments + and result -- which travel through this same codec -- are unaffected. + """ + + def __init__(self) -> None: + self.fail = False + + async def encode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + return list(payloads) + + async def decode( + self, payloads: Sequence[temporalio.api.common.v1.Payload] + ) -> list[temporalio.api.common.v1.Payload]: + if self.fail and any(_TAXONOMY_SENTINEL in p.data for p in payloads): + raise RuntimeError("this codec cannot read the producer's payloads") + return list(payloads) + + +class UnreachableBackend(MemoryStreamBackend): + """A backend whose recorded-range read can be made to fail. + + Only ``read_range`` -- the replay read -- so the live path that wrote the + marker is unaffected and the failure lands where row one describes it: + reading back a range the marker already committed. + """ + + def __init__(self) -> None: + super().__init__() + self.fail_reads = False + + async def read_range( + self, key: StreamKey, first: Offset, last: Offset + ) -> list[StreamRecord]: + if self.fail_reads: + raise ConnectionError("the backend is unreachable") + return await super().read_range(key, first, last) + + +_TAXONOMY_SENTINEL = b"taxonomy-sentinel" + +#: The three counters, so each case can assert that the other two stayed silent. +_TAXONOMY_METRICS = ( + METRIC_STORAGE, + METRIC_INTEGRITY, + METRIC_DECODE, +) + + +def _failure_types(failure) -> list[str]: # type: ignore[no-untyped-def] + """Every application failure type in a failure's cause chain.""" + types = [] + while True: + types.append(failure.application_failure_info.type) + if not failure.HasField("cause"): + return types + failure = failure.cause + + +async def _await_failed_task(handle, message: str, timeout: float = 30): # type: ignore[no-untyped-def] + """The first ``WorkflowTaskFailed`` event, or an assertion naming why not.""" + deadline = asyncio.get_running_loop().time() + timeout + while asyncio.get_running_loop().time() < deadline: + for event in [e async for e in handle.fetch_history_events()]: + if event.HasField("workflow_task_failed_event_attributes"): + return event.workflow_task_failed_event_attributes + await asyncio.sleep(0.3) + raise AssertionError(message) + + +@pytest.mark.parametrize( + ["row", "error_type", "metric"], + [ + ("storage", "StreamStorageError", METRIC_STORAGE), + ("integrity", "StreamIntegrityError", METRIC_INTEGRITY), + ("decode", "StreamDecodeError", METRIC_DECODE), + ], +) +async def test_each_failure_row_is_reported_as_its_own( + env: WorkflowEnvironment, row: str, error_type: str, metric: str +) -> None: + """Rows one to three, each reaching an operator as itself. + + All three are Workflow Task failures, and the server retries a Workflow Task + failure whatever caused it -- so the retry says nothing about which one + happened. What distinguishes them is exactly two things, and this asserts + both: the completion carries the **external-storage failure cause**, which + separates all three from an ordinary Workflow bug, and **one counter** + increments while the other two stay silent, which is what lets an alert on + integrity loss mean integrity loss. + + They are reached the way an operator would meet them, through a marker that + was written live and then replayed against a backend that has since changed: + + - the recorded range cannot be read at all -- transient, clears itself; + - the recorded range reads back short -- the record was trimmed or expired, + and an operator has to repair the backend or terminate the Run; + - the range reads back exactly as recorded and the consumer's codec cannot + decode it -- the stream is undamaged and the configuration is wrong. + + The third is the one the taxonomy is most easily collapsed on, because it + is the only one whose failure surfaces from *inside* ``activate()``: the + record is prepared on the Worker's loop, its failure travels with it, and + the delivery that would have yielded its value raises it. + """ + backend = UnreachableBackend() + codec = TaxonomyCodec() + buffer = MetricBuffer(10000) + runtime = Runtime(telemetry=TelemetryConfig(metrics=buffer)) + client = await env.connect_client( + runtime=runtime, + data_converter=temporalio.converter.DataConverter(payload_codec=codec), + ) + + task_queue = f"tq-{uuid.uuid4()}" + handle = None + async with Worker( + client, + task_queue=task_queue, + workflows=[ParkedAcrossEvictionWorkflow, FillerWorkflow], + external_stream_backend=backend, + # One slot, so the filler Workflow evicts the Run under test and the + # next Workflow Task has to replay the marker. + max_cached_workflows=1, + max_concurrent_workflow_tasks=2, + ): + try: + handle = await client.start_workflow( + ParkedAcrossEvictionWorkflow.run, + 2, + id=f"wf-{uuid.uuid4()}", + task_queue=task_queue, + ) + description = await handle.describe() + first_run_id = ( + description.raw_description.workflow_execution_info.first_run_id + ) + key = StreamKey(client.namespace, handle.id, first_run_id, "tokens") + + await _wait_for_markers(handle, 1, "the Run never parked at all") + # Consumed and committed live, so there is a recorded range to + # replay -- and the codec is asked for this value twice, once here + # while it still works and once on replay. + await publish( + backend, + key, + [_TAXONOMY_SENTINEL.decode()], + session=f"producer-{uuid.uuid4()}", + ) + await _wait_for_markers( + handle, + 2, + "the Run never committed a marker for the record it consumed, " + "so there is no recorded range for replay to read", + ) + + # Evicts the Run: buffers, watchers, and wait set all go with it, + # and the next Workflow Task replays the marker. + await client.execute_workflow( + FillerWorkflow.run, id=f"filler-{uuid.uuid4()}", task_queue=task_queue + ) + + if row == "storage": + backend.fail_reads = True + elif row == "integrity": + data = [r for r in backend.all_records(key) if not r.is_control] + assert data, "nothing was ever appended, so nothing can be lost" + assert data[0].offset is not None + await backend.delete_for_test(key, data[0].offset) + else: + codec.fail = True + + # The wake is what creates the Workflow Task that replays. + await send_wake_signal( + client, + WakeRequest( + namespace=client.namespace, + workflow_id=handle.id, + first_execution_run_id=first_run_id, + stream_name="tokens", + wait_id=1, + park_generation=0, + sender_identity=f"test-{uuid.uuid4()}", + wake_counter=1, + ), + ) + + failed = await _await_failed_task( + handle, + f"the {row} damage produced no Workflow Task failure at all, so " + "the Run either never replayed or silently delivered something", + ) + finally: + if handle is not None: + try: + await handle.terminate() + except Exception: + pass + + assert error_type in _failure_types(failed.failure), ( + f"the {row} failure reached the server as " + f"{_failure_types(failed.failure)} rather than as {error_type}, so an " + "operator cannot tell which of the three rows happened" + ) + assert ( + failed.cause + == temporalio.api.enums.v1.WorkflowTaskFailedCause.WORKFLOW_TASK_FAILED_CAUSE_EXTERNAL_STORAGE_FAILURE + ), ( + "the Workflow Task failed without the external-storage cause, so it is " + "indistinguishable from a bug in the Workflow's own code" + ) + + counted = { + update.metric.name + for update in buffer.retrieve_updates() + if update.metric.name in _TAXONOMY_METRICS + } + assert counted == {metric}, ( + f"the {row} failure incremented {sorted(counted)} rather than only " + f"{metric}; an alert on one row must not be diluted by another" + ) diff --git a/uv.lock b/uv.lock index cee51de68d..5184b131c3 100644 --- a/uv.lock +++ b/uv.lock @@ -4026,6 +4026,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/97/1b/295bf2fa3e740131778065e5ffa2c481f0e7210182d408e9a2c244ff5b0c/readme_renderer-45.0-py3-none-any.whl", hash = "sha256:3385ed220117104a2bceb4a9dac8c5fdf6d1f96890d7ea2a9c7174fd5c84091f", size = 14134, upload-time = "2026-06-09T21:05:15.85Z" }, ] +[[package]] +name = "redis" +version = "8.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "async-timeout", marker = "python_full_version < '3.11.3'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a8/99/604f0b666d4c616d891cf77ebb9db6bb21601344c051aebf1b72b9ff915f/redis-8.1.0.tar.gz", hash = "sha256:6e1a19beef9225c83efd689c7e6b7da2d5215b1f42cd13b7fc3714d0a09c7b25", size = 5254356, upload-time = "2026-07-30T08:51:00.269Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/66/9d/c5731f6e3608663d4d3656fd8d3aecee8b509c3082818f5a13eae925baea/redis-8.1.0-py3-none-any.whl", hash = "sha256:a4fe1aac3d3b3cc791d4b3d5931c5a956045dc951ee74d1c913ee3ac4d2ee9fb", size = 560618, upload-time = "2026-07-30T08:50:58.497Z" }, +] + [[package]] name = "referencing" version = "0.37.0" @@ -4802,6 +4814,7 @@ dev = [ { name = "pytest-rerunfailures" }, { name = "pytest-timeout" }, { name = "pytest-xdist" }, + { name = "redis" }, { name = "ruff" }, { name = "setuptools" }, { name = "strands-agents" }, @@ -4886,6 +4899,7 @@ dev = [ { name = "pytest-rerunfailures", specifier = ">=16.1" }, { name = "pytest-timeout", specifier = "~=2.2" }, { name = "pytest-xdist", specifier = ">=3.6,<4" }, + { name = "redis", specifier = ">=5,<9" }, { name = "ruff", specifier = ">=0.15.12,<0.16" }, { name = "setuptools", specifier = "<82" }, { name = "strands-agents", specifier = ">=1.51.0" },