#!/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())