216 lines
8.5 KiB
Python
216 lines
8.5 KiB
Python
"""POSIX shared-memory ring reader implementation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import mmap
|
|
from pathlib import Path
|
|
import struct
|
|
import time
|
|
from typing import Final
|
|
|
|
from python_app.models.dataset_model import ResultCollection, SweepCollection
|
|
from python_app.orchestration.shm.decoder import (
|
|
PREPROC_MAGIC,
|
|
RAW_MAGIC,
|
|
decode_result_collection,
|
|
decode_trace_collection,
|
|
)
|
|
|
|
_HEADER_SIZE: Final[int] = 64
|
|
_SLOT_HEADER_SIZE: Final[int] = 16
|
|
_MAGIC: Final[bytes] = b"RDRRING2"
|
|
_VERSION: Final[int] = 1
|
|
|
|
|
|
class ShmRingReader:
|
|
"""Read binary payloads from a lock-free ring in `/dev/shm`."""
|
|
|
|
def __init__(self, ring_name: str, open_timeout_s: float = 2.0, open_poll_s: float = 0.01) -> None:
|
|
"""Open and validate ring by name, for example `/radar_results`."""
|
|
if not ring_name.startswith("/"):
|
|
raise ValueError("ring_name must start with '/'")
|
|
|
|
self._ring_name = ring_name
|
|
self._path = Path("/dev/shm") / ring_name[1:]
|
|
self._wait_for_ring_file(timeout_s=open_timeout_s, poll_s=open_poll_s)
|
|
|
|
self._file = self._path.open("r+b", buffering=0)
|
|
self._mmap: mmap.mmap | None = None
|
|
try:
|
|
self._mmap = mmap.mmap(self._file.fileno(), 0)
|
|
self._validate_header_with_wait(timeout_s=1.0, poll_s=0.002)
|
|
except BaseException:
|
|
# A fail-fast open (absent/incompatible ring) must not leak the fd/mapping.
|
|
self.close()
|
|
raise
|
|
|
|
def close(self) -> None:
|
|
"""Close mmap and file handle."""
|
|
if self._mmap is not None:
|
|
self._mmap.close()
|
|
self._mmap = None
|
|
self._file.close()
|
|
|
|
def pop_payload(self) -> bytes | None:
|
|
"""Read next payload from ring, or `None` if no unread payload exists."""
|
|
write_seq = self._read_u64(24)
|
|
read_seq = self._read_u64(32)
|
|
if read_seq >= write_seq:
|
|
return None
|
|
|
|
index = read_seq % self.capacity
|
|
slot_offset = _HEADER_SIZE + index * (_SLOT_HEADER_SIZE + self.slot_size_bytes)
|
|
|
|
# Seqlock read mirroring the C++ pop: a slot is valid for read_seq R only if
|
|
# its sequence equals R+1 and is unchanged across the payload copy (i.e. the
|
|
# producer did not overwrite this slot mid-copy). Sequence and payload_size
|
|
# are read first; the slot is only accepted after the re-read confirms both.
|
|
sequence = self._read_u64(slot_offset + 8)
|
|
if sequence != read_seq + 1:
|
|
# Producer overwrote this slot before we read it. Resync to latest.
|
|
self._write_u64(32, write_seq)
|
|
return None
|
|
|
|
payload_size = self._read_u32(slot_offset)
|
|
# Bound payload_size against the slot before slicing so a torn/garbage size
|
|
# can never read out of the slot region; resync and skip on violation.
|
|
if payload_size > self.slot_size_bytes:
|
|
self._write_u64(32, write_seq)
|
|
return None
|
|
|
|
payload_offset = slot_offset + _SLOT_HEADER_SIZE
|
|
payload = bytes(self._mmap[payload_offset : payload_offset + payload_size])
|
|
|
|
# Re-read the slot sequence after the copy; if it changed, the producer
|
|
# overwrote this slot mid-copy and the payload is torn — discard and resync.
|
|
if self._read_u64(slot_offset + 8) != read_seq + 1:
|
|
self._write_u64(32, write_seq)
|
|
return None
|
|
|
|
self._write_u64(32, read_seq + 1)
|
|
return payload
|
|
|
|
def pop_raw_collection(self) -> SweepCollection | None:
|
|
"""Read next raw collection from ring."""
|
|
payload = self.pop_payload()
|
|
if payload is None:
|
|
return None
|
|
return decode_trace_collection(payload, RAW_MAGIC)
|
|
|
|
def pop_preprocessed_collection(self) -> SweepCollection | None:
|
|
"""Read next preprocessed collection from ring."""
|
|
payload = self.pop_payload()
|
|
if payload is None:
|
|
return None
|
|
return decode_trace_collection(payload, PREPROC_MAGIC)
|
|
|
|
def pop_result_collection(self) -> ResultCollection | None:
|
|
"""Read next processed result collection from ring."""
|
|
payload = self.pop_payload()
|
|
if payload is None:
|
|
return None
|
|
return decode_result_collection(payload)
|
|
|
|
def peek_latest_payload(self) -> bytes | None:
|
|
"""Return the most recently published payload WITHOUT consuming it.
|
|
|
|
Reads the newest slot through the seqlock and never advances `read_seq`, so a
|
|
viewer can peek the freshest frame while the ring's real consumer keeps its own
|
|
cursor — the two coexist without stealing each other's payloads. Inherently
|
|
latest-wins: always the freshest published frame, or `None` when nothing has
|
|
been published yet or the slot is being overwritten at this instant.
|
|
"""
|
|
write_seq = self._read_u64(24)
|
|
if write_seq == 0:
|
|
return None
|
|
|
|
latest_seq = write_seq - 1
|
|
index = latest_seq % self.capacity
|
|
slot_offset = _HEADER_SIZE + index * (_SLOT_HEADER_SIZE + self.slot_size_bytes)
|
|
|
|
# Seqlock read with NO cursor advance: accept the slot only if its sequence
|
|
# equals the published value both before and after the copy (i.e. the producer
|
|
# did not lap this slot mid-read). Never touch read_seq, so the real consumer
|
|
# is undisturbed.
|
|
if self._read_u64(slot_offset + 8) != latest_seq + 1:
|
|
return None
|
|
payload_size = self._read_u32(slot_offset)
|
|
if payload_size > self.slot_size_bytes:
|
|
return None
|
|
payload_offset = slot_offset + _SLOT_HEADER_SIZE
|
|
payload = bytes(self._mmap[payload_offset : payload_offset + payload_size])
|
|
if self._read_u64(slot_offset + 8) != latest_seq + 1:
|
|
return None
|
|
return payload
|
|
|
|
def peek_latest_result_collection(self) -> ResultCollection | None:
|
|
"""Return the most recently published result collection without consuming it."""
|
|
payload = self.peek_latest_payload()
|
|
if payload is None:
|
|
return None
|
|
return decode_result_collection(payload)
|
|
|
|
def drop_all(self) -> int:
|
|
"""Mark all unread slots as consumed and return number of dropped payloads."""
|
|
write_seq = self._read_u64(24)
|
|
read_seq = self._read_u64(32)
|
|
if read_seq >= write_seq:
|
|
return 0
|
|
dropped = int(write_seq - read_seq)
|
|
self._write_u64(32, write_seq)
|
|
return dropped
|
|
|
|
@property
|
|
def capacity(self) -> int:
|
|
"""Number of slots in ring."""
|
|
return self._read_u32(12)
|
|
|
|
@property
|
|
def slot_size_bytes(self) -> int:
|
|
"""Maximum payload size per slot."""
|
|
return self._read_u32(16)
|
|
|
|
def _validate_header_with_wait(self, timeout_s: float, poll_s: float) -> None:
|
|
"""Wait for ring header to contain expected magic and version."""
|
|
deadline = time.monotonic() + timeout_s
|
|
while True:
|
|
magic = self._mmap[:8]
|
|
version = self._read_u32(8)
|
|
if magic == _MAGIC and version == _VERSION:
|
|
return
|
|
|
|
if time.monotonic() >= deadline:
|
|
if magic != _MAGIC:
|
|
got_magic = bytes(magic).hex()
|
|
expected_magic = _MAGIC.hex()
|
|
raise RuntimeError(
|
|
f"Shared memory ring magic mismatch for {self._ring_name}: "
|
|
f"got=0x{got_magic}, expected=0x{expected_magic}"
|
|
)
|
|
raise RuntimeError(
|
|
f"Shared memory ring version mismatch for {self._ring_name}: "
|
|
f"got={version}, expected={_VERSION}"
|
|
)
|
|
|
|
time.sleep(poll_s)
|
|
|
|
def _wait_for_ring_file(self, timeout_s: float, poll_s: float) -> None:
|
|
"""Wait until ring file appears in `/dev/shm`."""
|
|
deadline = time.monotonic() + timeout_s
|
|
while not self._path.exists():
|
|
if time.monotonic() >= deadline:
|
|
raise FileNotFoundError(f"Shared memory ring does not exist: {self._ring_name}")
|
|
time.sleep(poll_s)
|
|
|
|
def _read_u32(self, offset: int) -> int:
|
|
"""Read little-endian u32 at mmap offset."""
|
|
return struct.unpack_from("<I", self._mmap, offset)[0]
|
|
|
|
def _read_u64(self, offset: int) -> int:
|
|
"""Read little-endian u64 at mmap offset."""
|
|
return struct.unpack_from("<Q", self._mmap, offset)[0]
|
|
|
|
def _write_u64(self, offset: int, value: int) -> None:
|
|
"""Write little-endian u64 at mmap offset."""
|
|
struct.pack_into("<Q", self._mmap, offset, value)
|