161 lines
6.2 KiB
Python
161 lines
6.2 KiB
Python
"""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()
|
|
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 _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("<I", self._mmap, offset)[0]
|
|
|
|
def _read_u64(self, offset: int) -> int:
|
|
return struct.unpack_from("<Q", self._mmap, offset)[0]
|
|
|
|
def _write_u32(self, offset: int, value: int) -> None:
|
|
struct.pack_into("<I", self._mmap, offset, int(value))
|
|
|
|
def _write_u64(self, offset: int, value: int) -> None:
|
|
struct.pack_into("<Q", self._mmap, offset, int(value))
|