mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-10 06:23:43 +08:00
firestar5683
1bf86e10c3
Provide offline log decoding and torque-surface data contracts for fleet learning.
104 lines
3.8 KiB
Python
104 lines
3.8 KiB
Python
"""Bounded, all-or-nothing decode of an already admitted local log segment."""
|
|
|
|
import bz2
|
|
import io
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
import zstandard as zstd
|
|
|
|
from openpilot.cereal import log
|
|
|
|
|
|
MAX_COMPRESSED_BYTES = 32 * 1024 * 1024
|
|
MAX_EXPANDED_BYTES = 64 * 1024 * 1024
|
|
MAX_MESSAGES = 250_000
|
|
EVENT_TRAVERSAL_WORDS = 1_000_000
|
|
CHUNK_BYTES = 64 * 1024
|
|
|
|
|
|
class LogDecodeError(ValueError):
|
|
"""The whole segment is unavailable; callers must discard every event."""
|
|
|
|
|
|
def _not_cancelled(cancelled: Callable[[], bool]) -> None:
|
|
try:
|
|
if cancelled():
|
|
raise LogDecodeError("segment decode cancelled")
|
|
except LogDecodeError:
|
|
raise
|
|
except Exception as error:
|
|
raise LogDecodeError("segment cancellation check unavailable") from error
|
|
|
|
|
|
def _bounded_read(reader: Any, cancelled: Callable[[], bool]) -> bytes:
|
|
chunks: list[bytes] = []
|
|
total = 0
|
|
while True:
|
|
_not_cancelled(cancelled)
|
|
chunk = reader.read(min(CHUNK_BYTES, MAX_EXPANDED_BYTES + 1 - total))
|
|
if not chunk:
|
|
break
|
|
total += len(chunk)
|
|
if total > MAX_EXPANDED_BYTES:
|
|
raise LogDecodeError("expanded segment exceeds 64 MiB")
|
|
chunks.append(chunk)
|
|
return b"".join(chunks)
|
|
|
|
|
|
def _expand(data: bytes, codec: str, cancelled: Callable[[], bool]) -> bytes:
|
|
try:
|
|
if codec == "zst":
|
|
with zstd.ZstdDecompressor(max_window_size=MAX_EXPANDED_BYTES).stream_reader(io.BytesIO(data)) as reader:
|
|
expanded = _bounded_read(reader, cancelled)
|
|
# stream_reader in pinned zstandard 0.25 returns short data for a
|
|
# truncated frame. Validate the complete frame after proving its output
|
|
# fits our bound, and reject a concatenated/trailing second frame.
|
|
decoder = zstd.ZstdDecompressor(max_window_size=MAX_EXPANDED_BYTES).decompressobj()
|
|
verified = decoder.decompress(data)
|
|
if not decoder.eof or decoder.unused_data or verified != expanded:
|
|
raise LogDecodeError("incomplete or trailing zstd frame")
|
|
elif codec == "bz2":
|
|
with bz2.BZ2File(io.BytesIO(data), "rb") as reader:
|
|
expanded = _bounded_read(reader, cancelled)
|
|
decoder = bz2.BZ2Decompressor()
|
|
verified = decoder.decompress(data)
|
|
if not decoder.eof or decoder.unused_data or verified != expanded:
|
|
raise LogDecodeError("incomplete or trailing bzip2 stream")
|
|
else:
|
|
raise LogDecodeError("unsupported segment compression")
|
|
_not_cancelled(cancelled)
|
|
return expanded
|
|
except LogDecodeError:
|
|
raise
|
|
except (OSError, EOFError, ValueError, zstd.ZstdError) as error:
|
|
raise LogDecodeError("corrupt or incomplete compressed segment") from error
|
|
|
|
|
|
def decode_segment(data: bytes, codec: str, *, cancelled: Callable[[], bool] = lambda: False,
|
|
retain: frozenset[str] | None = None) -> tuple[Any, ...]:
|
|
"""Return all complete Event readers, or reject the entire segment."""
|
|
if type(data) is not bytes or len(data) > MAX_COMPRESSED_BYTES or type(codec) is not str or \
|
|
codec not in ("zst", "bz2"):
|
|
raise LogDecodeError("unsupported or oversized compressed segment")
|
|
_not_cancelled(cancelled)
|
|
expanded = _expand(data, codec, cancelled)
|
|
events: list[Any] = []
|
|
try:
|
|
for count, event in enumerate(log.Event.read_multiple_bytes(expanded, traversal_limit_in_words=EVENT_TRAVERSAL_WORDS)):
|
|
_not_cancelled(cancelled)
|
|
if count >= MAX_MESSAGES:
|
|
raise LogDecodeError("segment message count exceeded")
|
|
# Force envelope and union decoding before a reader can escape; a later
|
|
# malformed record invalidates the entire segment, never a prefix.
|
|
int(event.logMonoTime)
|
|
kind = event.which()
|
|
if retain is None or kind in retain:
|
|
events.append(event)
|
|
except LogDecodeError:
|
|
raise
|
|
except Exception as error:
|
|
raise LogDecodeError("corrupt Cereal Event segment") from error
|
|
_not_cancelled(cancelled)
|
|
return tuple(events)
|