547 lines
19 KiB
Python
Executable File
547 lines
19 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Pure control protocol for the independent TG3 data recorder.
|
|
|
|
The ROS node is deliberately kept as a thin transport adapter. This module
|
|
owns JSON validation, idempotence, session/capture matching, acknowledgement
|
|
state, and the bridge-heartbeat watchdog, and can therefore be tested without
|
|
ROS installed.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import math
|
|
import time
|
|
from collections import OrderedDict
|
|
from collections.abc import Callable, Mapping, Sequence
|
|
from pathlib import Path
|
|
from typing import Any, Protocol
|
|
|
|
|
|
PROTOCOL_VERSION = 1
|
|
ACTIVE_MANAGER_STATES = {
|
|
"start_pending",
|
|
"starting",
|
|
"recording",
|
|
"stop_pending",
|
|
"finalizing",
|
|
}
|
|
|
|
|
|
class RecorderBackend(Protocol):
|
|
"""The non-blocking subset exposed by ``DataRecorderManager``."""
|
|
|
|
def request_start(
|
|
self,
|
|
session_id: str,
|
|
reason: str = "button",
|
|
*,
|
|
episode_id: str | None = None,
|
|
) -> str | None: ...
|
|
|
|
def request_stop(self, reason: str = "operator_button") -> bool: ...
|
|
|
|
def status(self) -> dict[str, Any]: ...
|
|
|
|
def shutdown(
|
|
self, *, wait: bool = True, timeout: float | None = None
|
|
) -> bool: ...
|
|
|
|
|
|
class ProtocolError(ValueError):
|
|
"""A control message is malformed or invalid for the current session."""
|
|
|
|
def __init__(self, code: str, message: str) -> None:
|
|
super().__init__(message)
|
|
self.code = code
|
|
|
|
|
|
def _nonempty_text(value: Any, name: str, *, maximum: int = 128) -> str:
|
|
if (
|
|
not isinstance(value, str)
|
|
or not value
|
|
or value.strip() != value
|
|
or len(value) > maximum
|
|
or any(ord(character) < 0x20 for character in value)
|
|
):
|
|
raise ProtocolError(
|
|
f"invalid_{name}",
|
|
f"{name} must be non-empty text up to {maximum} characters",
|
|
)
|
|
return value
|
|
|
|
|
|
def _capture_id(value: Any) -> str:
|
|
capture_id = _nonempty_text(value, "capture_id")
|
|
if (
|
|
not capture_id[0].isalnum()
|
|
or Path(capture_id).name != capture_id
|
|
or capture_id in (".", "..")
|
|
or any(
|
|
character
|
|
not in "-_.0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
|
for character in capture_id
|
|
)
|
|
):
|
|
raise ProtocolError(
|
|
"invalid_capture_id",
|
|
"capture_id must be one safe path component beginning with a letter or digit",
|
|
)
|
|
return capture_id
|
|
|
|
|
|
def _event_sequence(value: Any) -> int:
|
|
if type(value) is not int or value < 0 or value > (2**63 - 1):
|
|
raise ProtocolError(
|
|
"invalid_event_seq",
|
|
"event_seq must be an integer from 0 through 2^63-1",
|
|
)
|
|
return value
|
|
|
|
|
|
def _optional_unix_time(value: Any) -> float | None:
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
|
raise ProtocolError(
|
|
"invalid_sent_unix_s", "sent_unix_s must be a finite number"
|
|
)
|
|
result = float(value)
|
|
if not math.isfinite(result) or result < 0.0:
|
|
raise ProtocolError(
|
|
"invalid_sent_unix_s", "sent_unix_s must be a finite number"
|
|
)
|
|
return result
|
|
|
|
|
|
def parse_control_message(text: str) -> dict[str, Any]:
|
|
"""Parse one strict but forward-compatible recorder control object."""
|
|
|
|
if not isinstance(text, str) or not text or len(text.encode("utf-8")) > 16384:
|
|
raise ProtocolError(
|
|
"invalid_json", "control payload must be 1 through 16384 UTF-8 bytes"
|
|
)
|
|
try:
|
|
raw = json.loads(text)
|
|
except (json.JSONDecodeError, UnicodeError) as error:
|
|
raise ProtocolError("invalid_json", f"invalid JSON: {error}") from error
|
|
if not isinstance(raw, Mapping):
|
|
raise ProtocolError("invalid_json_object", "control JSON must be an object")
|
|
if raw.get("version") != PROTOCOL_VERSION:
|
|
raise ProtocolError(
|
|
"unsupported_version", f"version must equal {PROTOCOL_VERSION}"
|
|
)
|
|
|
|
command = raw.get("command")
|
|
if command not in ("start", "stop", "heartbeat"):
|
|
raise ProtocolError(
|
|
"invalid_command", "command must be start, stop, or heartbeat"
|
|
)
|
|
event_seq = _event_sequence(raw.get("event_seq"))
|
|
request_id = _nonempty_text(raw.get("request_id"), "request_id")
|
|
session_id = _nonempty_text(
|
|
raw.get("teleop_session_id"), "teleop_session_id"
|
|
)
|
|
capture_id = _capture_id(raw.get("capture_id"))
|
|
sent_unix_s = _optional_unix_time(raw.get("sent_unix_s"))
|
|
|
|
reason_value = raw.get("reason", "bridge_heartbeat" if command == "heartbeat" else "button")
|
|
reason = _nonempty_text(reason_value, "reason", maximum=512)
|
|
return {
|
|
"version": PROTOCOL_VERSION,
|
|
"command": command,
|
|
"event_seq": event_seq,
|
|
"request_id": request_id,
|
|
"teleop_session_id": session_id,
|
|
"capture_id": capture_id,
|
|
"reason": reason,
|
|
"sent_unix_s": sent_unix_s,
|
|
}
|
|
|
|
|
|
def topics_without_publishers(
|
|
required_topics: Sequence[str],
|
|
publisher_lookup: Callable[[str], Sequence[Any]],
|
|
) -> tuple[str, ...]:
|
|
"""Return required topics that currently have no live publisher endpoint.
|
|
|
|
ROS graph discovery can retain a topic name solely because this process is
|
|
subscribed to it. Looking up publisher endpoints is therefore the
|
|
meaningful start preflight; topic-name presence alone is insufficient.
|
|
Exceptions deliberately propagate so the caller can report graph-query
|
|
failure separately from an ordinary missing publisher.
|
|
"""
|
|
|
|
missing: list[str] = []
|
|
for topic in required_topics:
|
|
if not publisher_lookup(topic):
|
|
missing.append(topic)
|
|
return tuple(missing)
|
|
|
|
|
|
class RecorderControlProtocol:
|
|
"""Idempotent bridge-to-recorder protocol and heartbeat watchdog."""
|
|
|
|
def __init__(
|
|
self,
|
|
recorder: RecorderBackend,
|
|
*,
|
|
heartbeat_timeout_seconds: float = 3.0,
|
|
monotonic: Callable[[], float] = time.monotonic,
|
|
unix_time: Callable[[], float] = time.time,
|
|
start_preflight: Callable[[], tuple[bool, str]] | None = None,
|
|
request_cache_size: int = 512,
|
|
session_cache_size: int = 64,
|
|
) -> None:
|
|
if (
|
|
not math.isfinite(heartbeat_timeout_seconds)
|
|
or heartbeat_timeout_seconds <= 0.0
|
|
):
|
|
raise ValueError("heartbeat timeout must be positive and finite")
|
|
if request_cache_size <= 0 or session_cache_size <= 0:
|
|
raise ValueError("protocol cache sizes must be positive")
|
|
self.recorder = recorder
|
|
self.heartbeat_timeout_seconds = float(heartbeat_timeout_seconds)
|
|
self._monotonic = monotonic
|
|
self._unix_time = unix_time
|
|
self._start_preflight = start_preflight
|
|
self._request_cache_size = request_cache_size
|
|
self._session_cache_size = session_cache_size
|
|
self._request_cache: OrderedDict[str, tuple[str, bool, str, int]] = (
|
|
OrderedDict()
|
|
)
|
|
self._last_sequence_by_session: OrderedDict[str, int] = OrderedDict()
|
|
self._capture_id: str | None = None
|
|
self._session_id: str | None = None
|
|
self._last_heartbeat_monotonic: float | None = None
|
|
self._watchdog_stop_requested = False
|
|
self._last_protocol_error = ""
|
|
self._ack_event_seq: int | None = None
|
|
self._ack_request_id: str | None = None
|
|
self._ack_accepted: bool | None = None
|
|
self._ack_code = "initialized"
|
|
self._accepted_count = 0
|
|
self._rejected_count = 0
|
|
self._duplicate_count = 0
|
|
self._watchdog_stop_count = 0
|
|
|
|
def handle_json(self, text: str, *, now: float | None = None) -> dict[str, Any]:
|
|
"""Handle one payload without blocking on recorder I/O."""
|
|
|
|
timestamp = self._monotonic() if now is None else float(now)
|
|
try:
|
|
event = parse_control_message(text)
|
|
except ProtocolError as error:
|
|
self._ack_event_seq = None
|
|
self._ack_request_id = None
|
|
self._ack_accepted = False
|
|
self._ack_code = error.code
|
|
self._last_protocol_error = str(error)
|
|
self._rejected_count += 1
|
|
return self.status(now=timestamp)
|
|
|
|
fingerprint = json.dumps(event, sort_keys=True, separators=(",", ":"))
|
|
request_id = event["request_id"]
|
|
cached = self._request_cache.get(request_id)
|
|
if cached is not None:
|
|
cached_fingerprint, accepted, code, event_seq = cached
|
|
if cached_fingerprint != fingerprint:
|
|
return self._reject(
|
|
event,
|
|
"request_id_reused",
|
|
"request_id was already used for a different payload",
|
|
fingerprint=fingerprint,
|
|
cache=False,
|
|
now=timestamp,
|
|
)
|
|
self._request_cache.move_to_end(request_id)
|
|
self._duplicate_count += 1
|
|
self._set_ack(event_seq, request_id, accepted, code)
|
|
return self.status(now=timestamp)
|
|
|
|
session_id = event["teleop_session_id"]
|
|
last_sequence = self._last_sequence_by_session.get(session_id)
|
|
if last_sequence is not None and event["event_seq"] <= last_sequence:
|
|
return self._reject(
|
|
event,
|
|
"stale_event_seq",
|
|
f"event_seq must be greater than the previous value {last_sequence}",
|
|
fingerprint=fingerprint,
|
|
now=timestamp,
|
|
)
|
|
|
|
self._remember_sequence(session_id, event["event_seq"])
|
|
self._last_protocol_error = ""
|
|
try:
|
|
accepted, code = self._execute(event, timestamp)
|
|
message = (
|
|
""
|
|
if accepted
|
|
else (self._last_protocol_error or code)
|
|
)
|
|
except Exception as error:
|
|
accepted = False
|
|
code = "recorder_exception"
|
|
message = f"{type(error).__name__}: {error}"
|
|
|
|
if accepted:
|
|
self._accepted_count += 1
|
|
self._last_protocol_error = ""
|
|
else:
|
|
self._rejected_count += 1
|
|
self._last_protocol_error = message
|
|
self._remember_request(
|
|
request_id,
|
|
fingerprint,
|
|
accepted,
|
|
code,
|
|
event["event_seq"],
|
|
)
|
|
self._set_ack(
|
|
event["event_seq"],
|
|
request_id,
|
|
accepted,
|
|
code,
|
|
)
|
|
return self.status(now=timestamp)
|
|
|
|
def poll(self, *, now: float | None = None) -> dict[str, Any]:
|
|
"""Advance the watchdog and return the current status snapshot."""
|
|
|
|
timestamp = self._monotonic() if now is None else float(now)
|
|
recorder_status = self._safe_recorder_status()
|
|
manager_state = str(recorder_status.get("state", "unknown"))
|
|
if (
|
|
manager_state in ("start_pending", "starting", "recording")
|
|
and self._last_heartbeat_monotonic is not None
|
|
and not self._watchdog_stop_requested
|
|
and timestamp - self._last_heartbeat_monotonic
|
|
> self.heartbeat_timeout_seconds
|
|
):
|
|
try:
|
|
accepted = self.recorder.request_stop("bridge_heartbeat_timeout")
|
|
except Exception as error:
|
|
self._last_protocol_error = (
|
|
f"watchdog stop failed: {type(error).__name__}: {error}"
|
|
)
|
|
else:
|
|
if accepted:
|
|
self._watchdog_stop_requested = True
|
|
self._watchdog_stop_count += 1
|
|
return self._status_from_snapshot(recorder_status, timestamp)
|
|
|
|
def status(self, *, now: float | None = None) -> dict[str, Any]:
|
|
timestamp = self._monotonic() if now is None else float(now)
|
|
return self._status_from_snapshot(self._safe_recorder_status(), timestamp)
|
|
|
|
def shutdown(self, *, timeout: float | None = None) -> bool:
|
|
"""Request clean recorder finalization when the supervisor exits."""
|
|
|
|
return self.recorder.shutdown(wait=True, timeout=timeout)
|
|
|
|
def _execute(self, event: Mapping[str, Any], now: float) -> tuple[bool, str]:
|
|
command = event["command"]
|
|
if command == "start":
|
|
return self._start(event, now)
|
|
if command == "stop":
|
|
return self._stop(event)
|
|
return self._heartbeat(event, now)
|
|
|
|
def _start(self, event: Mapping[str, Any], now: float) -> tuple[bool, str]:
|
|
recorder_status = self._safe_recorder_status()
|
|
manager_state = str(recorder_status.get("state", "unknown"))
|
|
if manager_state in ACTIVE_MANAGER_STATES:
|
|
if self._matches_context(event):
|
|
return True, "already_active"
|
|
return False, "recorder_busy"
|
|
if manager_state not in ("idle",):
|
|
return False, "recorder_unavailable"
|
|
if self._start_preflight is not None:
|
|
allowed, detail = self._start_preflight()
|
|
if not allowed:
|
|
self._last_protocol_error = detail
|
|
return False, "preflight_failed"
|
|
|
|
actual_capture_id = self.recorder.request_start(
|
|
event["teleop_session_id"],
|
|
event["reason"],
|
|
episode_id=event["capture_id"],
|
|
)
|
|
if actual_capture_id is None:
|
|
return False, "recorder_busy"
|
|
if actual_capture_id != event["capture_id"]:
|
|
# A mismatched directory would make bridge samples and manifests
|
|
# impossible to correlate, so immediately fail safe by stopping.
|
|
self.recorder.request_stop("capture_id_mismatch")
|
|
return False, "capture_id_mismatch"
|
|
self._capture_id = actual_capture_id
|
|
self._session_id = event["teleop_session_id"]
|
|
self._last_heartbeat_monotonic = now
|
|
self._watchdog_stop_requested = False
|
|
return True, "start_accepted"
|
|
|
|
def _stop(self, event: Mapping[str, Any]) -> tuple[bool, str]:
|
|
recorder_status = self._safe_recorder_status()
|
|
manager_state = str(recorder_status.get("state", "unknown"))
|
|
if manager_state not in ACTIVE_MANAGER_STATES:
|
|
if self._matches_context(event) or self._capture_id is None:
|
|
return True, "already_stopped"
|
|
return False, "capture_not_active"
|
|
if not self._matches_context(event):
|
|
return False, "capture_mismatch"
|
|
accepted = self.recorder.request_stop(event["reason"])
|
|
if not accepted:
|
|
return False, "stop_rejected"
|
|
self._watchdog_stop_requested = True
|
|
return True, "stop_accepted"
|
|
|
|
def _heartbeat(
|
|
self, event: Mapping[str, Any], now: float
|
|
) -> tuple[bool, str]:
|
|
if not self._matches_context(event):
|
|
return False, "capture_mismatch"
|
|
recorder_status = self._safe_recorder_status()
|
|
manager_state = str(recorder_status.get("state", "unknown"))
|
|
if manager_state not in ACTIVE_MANAGER_STATES:
|
|
return True, "already_stopped"
|
|
self._last_heartbeat_monotonic = now
|
|
return True, "heartbeat_accepted"
|
|
|
|
def _matches_context(self, event: Mapping[str, Any]) -> bool:
|
|
return (
|
|
self._capture_id == event["capture_id"]
|
|
and self._session_id == event["teleop_session_id"]
|
|
)
|
|
|
|
def _reject(
|
|
self,
|
|
event: Mapping[str, Any],
|
|
code: str,
|
|
message: str,
|
|
*,
|
|
fingerprint: str,
|
|
cache: bool = True,
|
|
now: float,
|
|
) -> dict[str, Any]:
|
|
self._rejected_count += 1
|
|
self._last_protocol_error = message
|
|
self._set_ack(
|
|
event["event_seq"], event["request_id"], False, code
|
|
)
|
|
if cache:
|
|
self._remember_request(
|
|
event["request_id"],
|
|
fingerprint,
|
|
False,
|
|
code,
|
|
event["event_seq"],
|
|
)
|
|
return self.status(now=now)
|
|
|
|
def _remember_sequence(self, session_id: str, sequence: int) -> None:
|
|
self._last_sequence_by_session[session_id] = sequence
|
|
self._last_sequence_by_session.move_to_end(session_id)
|
|
while len(self._last_sequence_by_session) > self._session_cache_size:
|
|
self._last_sequence_by_session.popitem(last=False)
|
|
|
|
def _remember_request(
|
|
self,
|
|
request_id: str,
|
|
fingerprint: str,
|
|
accepted: bool,
|
|
code: str,
|
|
event_seq: int,
|
|
) -> None:
|
|
self._request_cache[request_id] = (
|
|
fingerprint,
|
|
accepted,
|
|
code,
|
|
event_seq,
|
|
)
|
|
self._request_cache.move_to_end(request_id)
|
|
while len(self._request_cache) > self._request_cache_size:
|
|
self._request_cache.popitem(last=False)
|
|
|
|
def _set_ack(
|
|
self,
|
|
event_seq: int,
|
|
request_id: str,
|
|
accepted: bool,
|
|
code: str,
|
|
) -> None:
|
|
self._ack_event_seq = event_seq
|
|
self._ack_request_id = request_id
|
|
self._ack_accepted = accepted
|
|
self._ack_code = code
|
|
|
|
def _safe_recorder_status(self) -> dict[str, Any]:
|
|
try:
|
|
status = self.recorder.status()
|
|
except Exception as error:
|
|
self._last_protocol_error = (
|
|
f"recorder status failed: {type(error).__name__}: {error}"
|
|
)
|
|
return {
|
|
"state": "unknown",
|
|
"recording": False,
|
|
"last_error": self._last_protocol_error,
|
|
}
|
|
return dict(status)
|
|
|
|
def _status_from_snapshot(
|
|
self, recorder_status: Mapping[str, Any], now: float
|
|
) -> dict[str, Any]:
|
|
manager_state = str(recorder_status.get("state", "unknown"))
|
|
if manager_state in ("start_pending", "starting"):
|
|
state = "starting"
|
|
elif manager_state == "recording":
|
|
state = "recording"
|
|
elif manager_state in ("stop_pending", "finalizing"):
|
|
state = "stopping"
|
|
elif manager_state == "idle":
|
|
last_result = recorder_status.get("last_result")
|
|
state = last_result if last_result in ("ready", "failed") else "idle"
|
|
else:
|
|
state = "failed"
|
|
|
|
heartbeat_age: float | None = None
|
|
if self._last_heartbeat_monotonic is not None:
|
|
heartbeat_age = max(0.0, now - self._last_heartbeat_monotonic)
|
|
manager_error = recorder_status.get("last_error")
|
|
last_error = (
|
|
str(manager_error)
|
|
if isinstance(manager_error, str) and manager_error
|
|
else self._last_protocol_error
|
|
)
|
|
return {
|
|
"version": PROTOCOL_VERSION,
|
|
"state": state,
|
|
"capture_id": self._capture_id,
|
|
"teleop_session_id": self._session_id,
|
|
"ack_event_seq": self._ack_event_seq,
|
|
"ack_request_id": self._ack_request_id,
|
|
"ack_accepted": self._ack_accepted,
|
|
"ack_code": self._ack_code,
|
|
"last_error": last_error,
|
|
"stop_reason": recorder_status.get("stop_reason"),
|
|
"heartbeat_age_s": heartbeat_age,
|
|
"heartbeat_timeout_s": self.heartbeat_timeout_seconds,
|
|
"updated_unix_s": float(self._unix_time()),
|
|
"statistics": {
|
|
"accepted": self._accepted_count,
|
|
"rejected": self._rejected_count,
|
|
"duplicates": self._duplicate_count,
|
|
"watchdog_stops": self._watchdog_stop_count,
|
|
},
|
|
"recorder": dict(recorder_status),
|
|
}
|
|
|
|
|
|
__all__ = [
|
|
"PROTOCOL_VERSION",
|
|
"ProtocolError",
|
|
"RecorderControlProtocol",
|
|
"parse_control_message",
|
|
"topics_without_publishers",
|
|
]
|