test_checkpoint_monotonicity.py

"""Is the executions row's "immutable" workflow version actually immutable?

``SqlCheckpointStore._update`` refuses four things by name, and the method
carries a comment saying exactly why a refusal has to be written into the
statement rather than decided in Python:

    Re-stating the bound as the UPDATE's own predicate is what makes the guard
    hold, because the database evaluates it against the row as it is at write
    time rather than as this worker last saw it.

That sentence is applied to one of the four. ``schedule_sequence`` is carried
into the UPDATE's ``WHERE``; ``workflow_id``, ``workflow_version`` and
``workflow_snapshot`` are compared against the row object read a moment
earlier and appear nowhere in the write. So are the two decisions taken from
that same stale read -- *is the version still unset, is the snapshot still
unset* -- whose answers are then written unconditionally.

Under PostgreSQL READ COMMITTED that difference is the whole story. Every
statement takes its own snapshot, so two workers reading the row a millisecond
apart both see ``workflow_version IS NULL``, both decide to fill it in, and
both write. Nothing refuses either of them, and the row ends up carrying a
version that one of the two workers never agreed to -- while the store's own
error message calls that column immutable.

The sibling suite ``tests/infrastructure/test_checkpoint_concurrency.py``
cannot see this. It races the *sequence*, which is the one bound that is in
the predicate, so it passes against a store that leaves the other three
outside it.

Two workers on one execution is not a hypothetical here: the queue is
at-least-once (``docs/failure-semantics.md`` section 9), a job that outlives
its visibility window is delivered again, and the second worker recompiles a
workflow that may have been edited since. That is precisely how two different
``workflow_version`` values arrive at one ``execution_id``.

Timing is not left to the scheduler. Every writer is parked on a real row lock
taken by a connection of the harness's own, so each of them provably finishes
its read before any of them can write, and they are then released together to
queue for the row in whatever order the database picks. What each *client was
told* is recorded as a history and the assertions read that, never a single
call's return value.
"""

from __future__ import annotations

import asyncio
from datetime import UTC, datetime
from typing import TYPE_CHECKING, NamedTuple
from uuid import uuid4

import pytest
from sqlalchemy import select, text, update

from flowforge.domain.checkpoint import (
    CHECKPOINT_SCHEMA_VERSION,
    ExecutionCheckpoint,
    deserialize_checkpoint,
)
from flowforge.domain.execution import ExecutionState, ExecutionStatus, NodeExecution, NodeStatus
from flowforge.domain.ids import ExecutionId, NodeId, WorkflowId
from flowforge.infrastructure.db import Base, create_engine, create_session_factory
from flowforge.infrastructure.models import ExecutionRow
from flowforge.infrastructure.persistence import SqlCheckpointStore, SqlExecutionRepository
from jepsen.conftest import requires_postgres
from jepsen.history import History

if TYPE_CHECKING:
    from collections.abc import AsyncIterator, Sequence

    from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker

pytestmark = [pytest.mark.integration, requires_postgres()]

_EXECUTION = ExecutionId(uuid4())
_WORKFLOW = WorkflowId(uuid4())
_STARTED_AT = datetime(2026, 1, 1, tzinfo=UTC)

#: The row's ordering mark, and the mark every racing writer carries. Holding
#: it constant keeps the one bound that *is* in the predicate out of the way,
#: so a refusal can only have come from the bounds under test.
_SEQUENCE = 5

#: One workflow version per worker, all different. Duplicate delivery of an
#: edited workflow is what produces this: same execution, recompiled twice.
_VERSIONS = (1, 2, 3, 4)

#: Backends parked on a lock in this database. The harness's own polling
#: connection is never one of them, because it is running rather than waiting.
_BLOCKED_WRITERS = text(
    "SELECT count(*) FROM pg_stat_activity "
    "WHERE datname = current_database() "
    "AND wait_event_type = 'Lock' "
    "AND query ILIKE '%executions%'"
)


class Stored(NamedTuple):
    """What the row says, and what the envelope inside it says.

    Both are read because a resume reads both: the column is what the guard
    claims to protect, and the envelope is what a restarted worker actually
    restores from. A store that let them disagree would restore an execution
    against a definition the row says it is not running.
    """

    workflow_version: int | None
    schedule_sequence: int
    envelope_version: int
    envelope_sequence: int


@pytest.fixture
async def postgres_engine(postgres_dsn: str) -> AsyncIterator[AsyncEngine]:
    """A real PostgreSQL engine, with this suite's tables created around it.

    Every session gets its own connection, which is the point: the sibling
    suites route every session through one in-memory SQLite connection, and a
    defect that needs two transactions reading the same row cannot occur at
    all when there is only ever one.
    """
    engine = create_engine(postgres_dsn)
    async with engine.begin() as connection:
        await connection.run_sync(Base.metadata.create_all)
    yield engine
    async with engine.begin() as connection:
        await connection.run_sync(Base.metadata.drop_all)
    await engine.dispose()


def _checkpoint(
    *, workflow_version: int, schedule_sequence: int = _SEQUENCE
) -> ExecutionCheckpoint:
    """A checkpoint for the shared execution, carrying one workflow version."""
    node = NodeExecution(NodeId("a")).start(1).finish(NodeStatus.SUCCEEDED)
    state = ExecutionState(
        execution_id=_EXECUTION,
        workflow_id=_WORKFLOW,
        workflow_version=workflow_version,
        workflow_snapshot={"name": "historical", "version": workflow_version},
        schedule_sequence=schedule_sequence,
        status=ExecutionStatus.RUNNING,
        started_at=_STARTED_AT,
        nodes={node.node_id: node},
    )
    return ExecutionCheckpoint(
        schema_version=CHECKPOINT_SCHEMA_VERSION,
        workflow_version=workflow_version,
        state=state,
    )


async def _seed_versionless_row(factory: async_sessionmaker[AsyncSession]) -> None:
    """Create the row the way the scheduler does, before a version is known.

    ``ExecutionState.workflow_version`` is documented as ``None`` for
    "pre-resolution pending jobs", and the execution repository writes that
    ``None`` straight through. So the column the checkpoint store calls
    immutable begins life unset, and whichever checkpoint arrives first is
    what fills it in. That is the state this suite attacks, and it is reached
    through the production writer rather than by hand-crafting a row.
    """
    await SqlExecutionRepository(factory).save(
        ExecutionState(
            execution_id=_EXECUTION,
            workflow_id=_WORKFLOW,
            status=ExecutionStatus.RUNNING,
            started_at=_STARTED_AT,
            schedule_sequence=_SEQUENCE,
        )
    )


async def _stored(factory: async_sessionmaker[AsyncSession]) -> Stored:
    """Read back the durable row and the checkpoint envelope it holds."""
    async with factory() as session:
        row = await session.get(ExecutionRow, _EXECUTION)
        assert row is not None
        assert row.checkpoint is not None
        envelope = deserialize_checkpoint(row.checkpoint)
        return Stored(
            workflow_version=row.workflow_version,
            schedule_sequence=int(row.schedule_sequence),
            envelope_version=envelope.workflow_version,
            envelope_sequence=envelope.state.schedule_sequence,
        )


async def _await_blocked_writers(engine: AsyncEngine, count: int, within: float = 20.0) -> None:
    """Block until ``count`` backends are queued behind the row lock.

    Waiting on the condition rather than on a duration is what makes the
    interleave a fact instead of a hope: once a writer is parked on the lock
    it has already finished its read, so releasing the lock at this point
    guarantees every writer decided against the same row.
    """
    deadline = asyncio.get_running_loop().time() + within
    waiting = 0
    while asyncio.get_running_loop().time() < deadline:
        async with engine.connect() as connection:
            waiting = int((await connection.execute(_BLOCKED_WRITERS)).scalar_one())
        if waiting >= count:
            return
        await asyncio.sleep(0.05)
    raise TimeoutError(f"only {waiting} of {count} writers reached the row lock within {within}s")


async def _checkpoint_once(
    store: SqlCheckpointStore, history: History, process: int, checkpoint: ExecutionCheckpoint
) -> str | None:
    """Save one checkpoint, recording what this client was told, and why.

    A ``ValueError`` is the store's own refusal, decided and raised before it
    commits, so the client knows the write did not land: that is ``fail``, and
    the message is returned so a test can check that the refusal names the
    bound it was actually about. Any other exception leaves the operation
    ``info`` and propagates, because a client that lost its connection
    mid-statement never learned whether the row moved.
    """
    version = checkpoint.workflow_version
    with history.attempt(process, "checkpoint", version) as op:
        try:
            await store.save(checkpoint)
        except ValueError as refusal:
            op.fail(version)
            return str(refusal)
        op.ok(version)
    return None


async def _checkpoint_together(
    engine: AsyncEngine,
    factory: async_sessionmaker[AsyncSession],
    checkpoints: Sequence[ExecutionCheckpoint],
) -> tuple[History, list[str]]:
    """Hold every writer at the row lock, then release them all at once.

    The lock is taken by a connection belonging to the harness and changes
    nothing; it exists only to stop any writer committing while another is
    still reading. Every worker therefore reads the row in the same state, and
    then queues for it -- which is exactly the shape of two workers arriving
    on one execution from a redelivered job.
    """
    history = History()
    store = SqlCheckpointStore(factory)
    async with engine.connect() as blocker:
        await blocker.execute(
            select(ExecutionRow.execution_id)
            .where(ExecutionRow.execution_id == _EXECUTION)
            .with_for_update()
        )
        workers = [
            asyncio.create_task(_checkpoint_once(store, history, process, checkpoint))
            for process, checkpoint in enumerate(checkpoints)
        ]
        await _await_blocked_writers(engine, len(workers))
        await blocker.rollback()
        refusals = [refusal for refusal in await asyncio.gather(*workers) if refusal is not None]
    return history, refusals


async def test_one_writer_is_refused_a_second_workflow_version(
    postgres_engine: AsyncEngine,
) -> None:
    """The guard as written, with nothing racing it.

    This is the control that makes the next test a concurrency failure rather
    than a missing feature: the store does refuse a conflicting version, and
    it refuses it by name.
    """
    factory = create_session_factory(postgres_engine)
    await _seed_versionless_row(factory)
    store = SqlCheckpointStore(factory)

    await store.save(_checkpoint(workflow_version=1))
    with pytest.raises(ValueError, match="workflow version is immutable"):
        await store.save(_checkpoint(workflow_version=2))

    assert await _stored(factory) == Stored(1, _SEQUENCE, 1, _SEQUENCE)


async def test_a_writer_that_queued_for_the_row_re_reads_the_sequence_bound(
    postgres_engine: AsyncEngine,
) -> None:
    """The bound that *is* in the predicate holds, and here is why.

    A writer whose UPDATE queues behind an uncommitted one does not get to
    apply its statement against the row it read. PostgreSQL re-fetches the
    newly committed version and re-evaluates the ``WHERE`` clause against it
    (``EvaluatePlanQual``), so a predicate naming the bound is checked against
    the row as it is at write time.

    This is the mechanism the whole file turns on, so it is demonstrated
    rather than assumed: it is what makes ``schedule_sequence`` safe, and its
    absence is what makes the other three bounds unsafe.
    """
    factory = create_session_factory(postgres_engine)
    await _seed_versionless_row(factory)
    store = SqlCheckpointStore(factory)
    await store.save(_checkpoint(workflow_version=1))

    async with postgres_engine.begin() as blocker:
        await blocker.execute(
            update(ExecutionRow)
            .where(ExecutionRow.execution_id == _EXECUTION)
            .values(schedule_sequence=_SEQUENCE + 4)
        )
        stale = asyncio.create_task(
            store.save(_checkpoint(workflow_version=1, schedule_sequence=_SEQUENCE + 1))
        )
        await _await_blocked_writers(postgres_engine, 1)

    with pytest.raises(ValueError, match="cannot move backwards"):
        await stale

    assert await _stored(factory) == Stored(1, _SEQUENCE + 4, 1, _SEQUENCE), (
        "a writer that read the row before an overtaking commit was allowed to "
        "apply its statement anyway; the ordering mark went backwards"
    )


async def test_concurrent_checkpoints_cannot_disagree_about_the_workflow_version(
    postgres_engine: AsyncEngine,
) -> None:
    """No client is told its checkpoint landed while the row says otherwise.

    Every worker here reads the same versionless row, and every one of them
    carries a different ``workflow_version``. Only one of those versions can
    survive, so the store owes every other worker a refusal: telling a client
    the write succeeded, and then storing a different version than the one it
    sent, is the store contradicting its own error message.

    The property is stated over the recorded history rather than over any one
    call, because the defect is not that a call returns the wrong thing -- each
    call, on its own, returns exactly what a single-writer store would.
    """
    factory = create_session_factory(postgres_engine)
    await _seed_versionless_row(factory)

    history, refusals = await _checkpoint_together(
        postgres_engine,
        factory,
        [_checkpoint(workflow_version=version) for version in _VERSIONS],
    )
    stored = await _stored(factory)

    accepted = {op.value for op in history.completed("checkpoint", "ok")}
    assert accepted, f"every writer was refused, so nothing was proved\n{history.render()}"
    assert accepted == {stored.workflow_version}, (
        f"{len(accepted)} workers were told their checkpoint was stored, carrying "
        f"workflow versions {sorted(accepted)}, but the row holds "
        f"{stored.workflow_version}: a column the store calls immutable was "
        f"written by writers no guard ever refused\n{history.render()}"
    )
    assert stored.envelope_version == stored.workflow_version, (
        f"the row says workflow version {stored.workflow_version} and the checkpoint "
        f"envelope inside it says {stored.envelope_version}; a resume would restore "
        f"against a definition the row denies\n{history.render()}"
    )
    assert all("workflow version is immutable" in refusal for refusal in refusals), (
        "a version conflict was reported as something else, so a caller cannot "
        f"tell which bound it broke: {refusals}"
    )


async def test_duplicate_delivery_of_one_workflow_version_is_never_refused(
    postgres_engine: AsyncEngine,
) -> None:
    """The ordinary redelivery case keeps working, and keeps its frontier.

    This is the case the queue actually produces: the same job delivered
    twice, so both workers carry the same version at the same mark. Neither
    contradicts the other, so neither may be refused -- a store that answered
    "you lost the race" here would be throwing away a checkpoint for no reason,
    and losing a checkpoint is the failure the sequence guard exists to
    prevent in the first place.

    It is here to keep a fix for the test above from being paid for with data
    loss.
    """
    factory = create_session_factory(postgres_engine)
    await _seed_versionless_row(factory)

    history, refusals = await _checkpoint_together(
        postgres_engine,
        factory,
        [_checkpoint(workflow_version=1) for _ in _VERSIONS],
    )

    assert not refusals, (
        f"duplicate delivery of one workflow version was refused: {refusals}\n{history.render()}"
    )
    assert len(history.completed("checkpoint", "ok")) == len(_VERSIONS)
    assert await _stored(factory) == Stored(1, _SEQUENCE, 1, _SEQUENCE)