"""POSIX shared-memory ring writer compatible with the C++ IPC ring.""" from __future__ import annotations import logging import mmap import os from pathlib import Path import struct from typing import Final logger = logging.getLogger(__name__) _HEADER_SIZE: Final[int] = 64 _SLOT_HEADER_SIZE: Final[int] = 16 _MAGIC: Final[bytes] = b"RDRRING2" _VERSION: Final[int] = 1 class ShmRingWriter: """Write binary payloads into the shared-memory ring used by C++ workers. The writer is the sole *owner* of the rings it opens: there is exactly one producer per ring (the acquisition producer for the raw/raw_tap rings). On a geometry mismatch with a pre-existing segment (e.g. a stale ring left by a prior run with a different sweep config), the owner unlinks and recreates the segment from scratch rather than truncating in place or diverging silently — mirroring the clean-shm-on-restart contract on the C++/deploy side (#13). A non-owner must never recreate a ring; readers and C++ consumers only ever attach to an existing one. """ def __init__(self, ring_name: str, capacity: int, slot_size_bytes: int) -> None: """Open or create a POSIX SHM ring by name (as the ring owner).""" if not ring_name.startswith("/"): raise ValueError("ring_name must start with '/'") if capacity <= 0: raise ValueError("capacity must be > 0") if slot_size_bytes <= 0: raise ValueError("slot_size_bytes must be > 0") self._ring_name = ring_name self._capacity = int(capacity) self._slot_size_bytes = int(slot_size_bytes) self._mapped_size = _HEADER_SIZE + self._capacity * (_SLOT_HEADER_SIZE + self._slot_size_bytes) self._path = Path("/dev/shm") / ring_name[1:] self._open_owned() def _open_owned(self) -> None: """Open the ring, recreating it from scratch on a geometry/header mismatch. As the single owner of this ring we may safely discard a stale segment: a size or header mismatch means the existing segment belongs to an earlier, incompatible run, so we unlink it and create a fresh one instead of mapping an inconsistent layout. """ created = not self._path.exists() fd = os.open(self._path, os.O_RDWR | os.O_CREAT, 0o660) self._file = os.fdopen(fd, "r+b", buffering=0) if created or self._path.stat().st_size != self._mapped_size: # Wrong-sized stale segment: drop it entirely and recreate, so the file # and any future mapping agree on geometry instead of being truncated # under a producer/consumer that still expects the old layout. if not created: logger.info("Ring %s: stale segment with wrong size, recreating", self._ring_name) self._file.truncate(self._mapped_size) created = True self._mmap = mmap.mmap(self._file.fileno(), self._mapped_size) if created: self._initialize_header() logger.debug( "Created SHM ring writer %s (capacity=%d, slot_size=%d bytes)", self._ring_name, self._capacity, self._slot_size_bytes, ) return # Size matched but the header geometry/magic does not: the owner recreates # rather than diverge. Unlink and reopen as a brand-new ring. if not self._header_matches(): logger.warning("Ring %s: header/geometry mismatch on existing segment, recreating", self._ring_name) self._mmap.close() self._file.close() self._unlink_if_present() fd = os.open(self._path, os.O_RDWR | os.O_CREAT, 0o660) self._file = os.fdopen(fd, "r+b", buffering=0) self._file.truncate(self._mapped_size) self._mmap = mmap.mmap(self._file.fileno(), self._mapped_size) self._initialize_header() else: logger.debug("Reusing existing SHM ring writer %s", self._ring_name) def _unlink_if_present(self) -> None: """Remove the backing /dev/shm file if it exists (owner-only operation).""" try: self._path.unlink() except FileNotFoundError: pass def close(self) -> None: """Close mmap and file handle.""" self._mmap.close() self._file.close() logger.debug("Closed SHM ring writer %s", self._ring_name) def push(self, payload: bytes) -> bool: """Push one payload with overwrite-oldest semantics on overflow.""" if len(payload) > self._slot_size_bytes: return False write_seq = self._read_u64(24) read_seq = self._read_u64(32) if max(0, write_seq - read_seq) >= self._capacity: # Ring full: the consumer is not keeping up, so the oldest unread slot is # overwritten (overwrite-oldest). Logged at DEBUG to avoid flooding the # log when a backlog persists across many pushes. logger.debug("Ring %s: full, overwriting oldest unread slot", self._ring_name) # Advance the consumer cursor past the slot we are about to overwrite, but # re-read it first and move it only forward: a concurrent reader may have # already advanced it, and clobbering that backward would re-deliver an # already-consumed slot as a duplicate. self._write_u64(32, max(self._read_u64(32), read_seq + 1)) self._write_u64(40, self._read_u64(40) + 1) index = write_seq % self._capacity slot_offset = _HEADER_SIZE + index * (_SLOT_HEADER_SIZE + self._slot_size_bytes) payload_offset = slot_offset + _SLOT_HEADER_SIZE self._write_u32(slot_offset, len(payload)) self._write_u32(slot_offset + 4, 0) self._mmap[payload_offset : payload_offset + len(payload)] = payload self._write_u64(slot_offset + 8, write_seq + 1) self._write_u64(24, write_seq + 1) return True @property def ring_name(self) -> str: """Return the POSIX SHM ring name.""" return self._ring_name @property def slot_size_bytes(self) -> int: """Return maximum payload size per slot.""" return self._slot_size_bytes def _initialize_header(self) -> None: self._mmap[:] = b"\x00" * self._mapped_size self._mmap[:8] = _MAGIC self._write_u32(8, _VERSION) self._write_u32(12, self._capacity) self._write_u32(16, self._slot_size_bytes) self._write_u32(20, 0) self._write_u64(24, 0) self._write_u64(32, 0) self._write_u64(40, 0) def _header_matches(self) -> bool: """Return whether the existing segment's header matches this ring's geometry. Used by the owner to decide whether a same-sized pre-existing segment can be reused as-is or must be unlinked and recreated. """ return ( self._mmap[:8] == _MAGIC and self._read_u32(8) == _VERSION and self._read_u32(12) == self._capacity and self._read_u32(16) == self._slot_size_bytes ) def _read_u32(self, offset: int) -> int: return struct.unpack_from(" int: return struct.unpack_from(" None: struct.pack_into(" None: struct.pack_into("