"""POSIX shared-memory ring writer compatible with the C++ IPC ring.""" from __future__ import annotations import mmap import os from pathlib import Path import struct from typing import Final _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. self._file.truncate(self._mapped_size) created = True self._mmap = mmap.mmap(self._file.fileno(), self._mapped_size) if created: self._initialize_header() 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(): self._mmap.close() self._file.close() self._unlink_if_present() 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) self._file.truncate(self._mapped_size) self._mmap = mmap.mmap(self._file.fileno(), self._mapped_size) self._initialize_header() 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() 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: self._write_u64(32, read_seq + 1) dropped = self._read_u64(40) self._write_u64(40, dropped + 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 _validate_header(self) -> None: magic = self._mmap[:8] version = self._read_u32(8) capacity = self._read_u32(12) slot_size_bytes = self._read_u32(16) if magic != _MAGIC: raise RuntimeError(f"Shared memory ring magic mismatch for {self._ring_name}") if version != _VERSION: raise RuntimeError(f"Shared memory ring version mismatch for {self._ring_name}") if capacity != self._capacity or slot_size_bytes != self._slot_size_bytes: raise RuntimeError(f"Shared memory ring geometry mismatch for {self._ring_name}") def _header_matches(self) -> bool: """Return whether the existing segment's header matches this ring's geometry. Non-throwing counterpart of `_validate_header` used by the owner to decide whether a same-sized pre-existing segment can be reused or must be 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("