radar_system/python_app/orchestration/shm/ring_reader.py

"""POSIX shared-memory ring reader implementation."""

from __future__ import annotations

import logging
import mmap
import time

from python_app.models.dataset_model import ResultCollection, SweepCollection
from python_app.orchestration.shm import atomics
from python_app.orchestration.shm.decoder import (
    PREPROC_MAGIC,
    RAW_MAGIC,
    decode_result_collection,
    decode_trace_collection,
)
from python_app.orchestration.shm.ring_layout import HEADER, MAGIC, SLOT_HEADER_SIZE, VERSION, ring_path

logger = logging.getLogger(__name__)


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`."""
        self.name = ring_name
        self._path = ring_path(ring_name)
        self._wait_for_ring_file(timeout_s=open_timeout_s, poll_s=open_poll_s)

        with self._path.open("r+b", buffering=0) as file:
            self._mmap = mmap.mmap(file.fileno(), 0)
        try:
            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
        logger.debug(
            "Opened SHM ring reader %s (capacity=%d, slot_size=%d bytes)",
            self.name,
            self.capacity,
            self.slot_size_bytes,
        )

    def close(self) -> None:
        """Close the mapping and its owned file descriptor; safe to repeat."""
        self._mmap.close()
        logger.debug("Closed SHM ring reader %s", self.name)

    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.
            logger.debug("Ring %s: slot lapped before read, resyncing to write_seq=%d", self.name, write_seq)
            atomics.advance(self._mmap, 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:
            logger.debug(
                "Ring %s: payload_size %d exceeds slot %d, resyncing",
                self.name,
                payload_size,
                self.slot_size_bytes,
            )
            atomics.advance(self._mmap, 32, write_seq)
            return None

        payload_offset = slot_offset + SLOT_HEADER_SIZE
        payload = bytes(self._mmap[payload_offset : payload_offset + payload_size])
        atomics.fence()

        # 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:
            logger.debug("Ring %s: slot overwritten mid-copy, discarding torn payload", self.name)
            atomics.advance(self._mmap, 32, write_seq)
            return None

        atomics.advance(self._mmap, 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 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)
        atomics.advance(self._mmap, 32, write_seq)
        logger.debug("Ring %s: dropped %d unread payload(s)", self.name, dropped)
        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:
                capacity = self.capacity
                slot_size = self.slot_size_bytes
                expected_size = HEADER.size + capacity * (SLOT_HEADER_SIZE + slot_size)
                if capacity > 0 and slot_size > 0 and slot_size % 8 == 0 and expected_size == len(self._mmap):
                    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.name}: "
                        f"got=0x{got_magic}, expected=0x{expected_magic}"
                    )
                if version != VERSION:
                    raise RuntimeError(
                        f"Shared memory ring version mismatch for {self.name}: "
                        f"got={version}, expected={VERSION}"
                    )
                raise RuntimeError(f"Shared memory ring geometry mismatch for {self.name}")

            time.sleep(poll_s)

    def _wait_for_ring_file(self, timeout_s: float, poll_s: float) -> None:
        """Wait for the creator to allocate the ring before mapping its header."""
        deadline = time.monotonic() + timeout_s
        while True:
            try:
                if self._path.stat().st_size >= HEADER.size:
                    return
            except FileNotFoundError:
                pass
            if time.monotonic() >= deadline:
                raise FileNotFoundError(f"Shared memory ring does not exist: {self.name}")
            time.sleep(poll_s)

    def _read_u32(self, offset: int) -> int:
        """Read little-endian u32 at mmap offset."""
        return atomics.load(self._mmap, offset, 4)

    def _read_u64(self, offset: int) -> int:
        """Read little-endian u64 at mmap offset."""
        return atomics.load(self._mmap, offset, 8)