archive-clients/e2e/scenarios/protocol_ordering.py

#!/usr/bin/env python3
"""Exercise duplicate and reordered job events across a real WebSocket."""

from __future__ import annotations

import asyncio
import json
import time
import urllib.request
from pathlib import Path
from uuid import uuid4

from websockets.asyncio.client import connect

from archive_clients.protocol import decode, encode, new_envelope
from archive_control.v1 import client_pb2, common_pb2, control_pb2, job_pb2


CONTROL_HTTP = "http://127.0.0.1:18081/test/v1"
CONTROL_WS = "ws://127.0.0.1:8765/archive_control"
CLIENT_ID = "protocol-probe"


def http_json(method: str, path: str, value: object | None = None):
    encoded = (
        json.dumps(value, separators=(",", ":")).encode()
        if value is not None
        else None
    )
    request = urllib.request.Request(
        f"{CONTROL_HTTP}{path}",
        data=encoded,
        method=method,
        headers={"Content-Type": "application/json"} if encoded else {},
    )
    with urllib.request.urlopen(request, timeout=15) as response:
        return json.load(response)


def registration():
    envelope = new_envelope()
    request = envelope.register_request
    request.protocol_version.CopyFrom(envelope.protocol_version)
    request.client.client_id = CLIENT_ID
    request.client.display_name = "E2E protocol ordering probe"
    request.client.role = common_pb2.CLIENT_ROLE_CACHE
    request.connection_instance_id = str(uuid4())
    request.shared_token = Path(
        "/run/secrets/archive_control_token"
    ).read_text(encoding="utf-8").strip()
    request.capabilities.max_envelope_bytes = 1024 * 1024
    return envelope


def acknowledgement(command_id: str):
    envelope = new_envelope()
    envelope.command_ack.command_id = command_id
    envelope.command_ack.status = (
        control_pb2.COMMAND_ACK_STATUS_ACCEPTED
    )
    return envelope


def heartbeat_ack(sequence: int):
    envelope = new_envelope()
    envelope.heartbeat_ack.sequence = sequence
    return envelope


def job_event(
    job_id: str,
    *,
    event_id: str,
    sequence: int,
    revision: int,
    event_type: int,
    state: int,
):
    envelope = new_envelope()
    event = envelope.job_event
    event.event_id = event_id
    event.job_id = job_id
    event.sequence = sequence
    event.job_revision = revision
    event.type = event_type
    event.state = state
    event.occurred_at.GetCurrentTime()
    if state == job_pb2.JOB_STATE_FAILED:
        event.error.code = common_pb2.ERROR_CODE_PRECONDITION_FAILED
        event.error.message = "protocol ordering probe terminal event"
    return envelope


async def receive_payload(websocket, expected: set[str], timeout: float = 15):
    deadline = time.monotonic() + timeout
    while True:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            raise RuntimeError(f"timed out waiting for {sorted(expected)}")
        envelope = decode(await asyncio.wait_for(websocket.recv(), remaining))
        payload = envelope.WhichOneof("payload")
        if payload == "heartbeat":
            await websocket.send(
                encode(heartbeat_ack(envelope.heartbeat.sequence))
            )
            continue
        if payload in expected:
            return envelope


async def main() -> None:
    async with connect(
        CONTROL_WS,
        ping_interval=None,
        proxy=None,
        max_size=1024 * 1024,
    ) as websocket:
        await websocket.send(encode(registration()))
        response = await receive_payload(websocket, {"register_response"})
        if (
            response.register_response.status
            != client_pb2.REGISTRATION_STATUS_ACCEPTED
        ):
            raise RuntimeError("protocol probe registration was rejected")

        lease = await asyncio.to_thread(
            http_json,
            "POST",
            "/protocol/order-probe",
            {"client_id": CLIENT_ID},
        )
        if not lease.get("delivered"):
            raise RuntimeError("protocol probe command was not delivered")
        assigned = await receive_payload(websocket, {"command"})
        if assigned.command.WhichOneof("payload") != "assign_job":
            raise RuntimeError("protocol probe received the wrong command")
        if assigned.command.command_id != lease["command_id"]:
            raise RuntimeError("protocol probe received the wrong command ID")
        await websocket.send(
            encode(acknowledgement(assigned.command.command_id))
        )

        job_id = lease["job_id"]
        first = job_event(
            job_id,
            event_id=str(uuid4()),
            sequence=1,
            revision=0,
            event_type=control_pb2.JOB_EVENT_TYPE_ASSIGNED,
            state=job_pb2.JOB_STATE_PREPARING,
        )
        second = job_event(
            job_id,
            event_id=str(uuid4()),
            sequence=2,
            revision=1,
            event_type=control_pb2.JOB_EVENT_TYPE_FAILED,
            state=job_pb2.JOB_STATE_FAILED,
        )

        second_frame = encode(second)
        await websocket.send(second_frame)
        observed: set[str] = set()
        while observed != {"snapshot", "error"}:
            envelope = await receive_payload(
                websocket, {"command", "protocol_error"}
            )
            if envelope.WhichOneof("payload") == "protocol_error":
                error = envelope.protocol_error
                if error.offending_message_id != second.message_id:
                    raise RuntimeError("gap error identified the wrong envelope")
                if (
                    "expected event sequence 1, received 2"
                    not in error.error.message
                ):
                    raise RuntimeError("gap error did not explain the ordering fault")
                observed.add("error")
                continue
            if (
                envelope.command.WhichOneof("payload")
                != "request_job_snapshot"
                or list(envelope.command.request_job_snapshot.job_ids)
                != [job_id]
            ):
                raise RuntimeError("gap did not request the expected snapshot")
            await websocket.send(
                encode(acknowledgement(envelope.command.command_id))
            )
            observed.add("snapshot")

        first_frame = encode(first)
        await websocket.send(first_frame)
        await websocket.send(first_frame)
        await websocket.send(second_frame)

        deadline = time.monotonic() + 15
        last = None
        while time.monotonic() < deadline:
            last = await asyncio.to_thread(
                http_json, "GET", f"/jobs/{job_id}"
            )
            if (
                last["state"] == "JOB_STATE_FAILED"
                and int(last["last_event_sequence"]) == 2
                and int(last["revision"]) == 1
            ):
                print(json.dumps({
                    "duplicate_sequence": 1,
                    "gap_sequence": 2,
                    "job_id": job_id,
                    "recovered_terminal_state": last["state"],
                    "snapshot_requested": True,
                }, sort_keys=True))
                return
            await asyncio.sleep(0.1)
        raise RuntimeError(f"protocol ordering probe did not converge: {last}")


if __name__ == "__main__":
    asyncio.run(main())