124 lines
5.2 KiB
Python
124 lines
5.2 KiB
Python
"""POSIX shared-memory ring writer compatible with the C++ IPC ring."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import mmap
|
|
import os
|
|
|
|
from python_app.models.run_config_schema import RingEndpointModel
|
|
from python_app.models.run_config_validation import validate_ring_endpoint
|
|
from python_app.orchestration.shm import atomics
|
|
from python_app.orchestration.shm.ring_layout import HEADER, MAGIC, SLOT_HEADER_SIZE, VERSION, ring_path
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ShmRingWriter:
|
|
"""The sole producer owns the ring and replaces incompatible stale segments.
|
|
|
|
Replacement uses unlink plus a new inode, preserving existing readers' mappings.
|
|
Compatible segments retain their unread payloads and sequence counters.
|
|
"""
|
|
|
|
def __init__(self, ring_name: str, capacity: int, slot_size_bytes: int) -> None:
|
|
self._path = ring_path(ring_name)
|
|
validate_ring_endpoint(RingEndpointModel(ring_name, capacity, slot_size_bytes))
|
|
self._ring_name = ring_name
|
|
self._capacity = capacity
|
|
self._slot_size_bytes = slot_size_bytes
|
|
self._mapped_size = HEADER.size + capacity * (SLOT_HEADER_SIZE + slot_size_bytes)
|
|
self._mmap = self._open_owned()
|
|
|
|
def _open_owned(self) -> mmap.mmap:
|
|
"""Reuse a compatible segment, otherwise replace it without truncating it."""
|
|
try:
|
|
with self._path.open("r+b", buffering=0) as file:
|
|
if os.fstat(file.fileno()).st_size == self._mapped_size:
|
|
mapping = mmap.mmap(file.fileno(), self._mapped_size)
|
|
try:
|
|
if self._header_matches(mapping):
|
|
return mapping
|
|
except BaseException:
|
|
mapping.close()
|
|
raise
|
|
mapping.close()
|
|
logger.warning("Ring %s: stale segment, recreating", self._ring_name)
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
self._path.unlink(missing_ok=True)
|
|
fd = os.open(self._path, os.O_RDWR | os.O_CREAT | os.O_EXCL, 0o660)
|
|
with os.fdopen(fd, "r+b", buffering=0) as file:
|
|
file.truncate(self._mapped_size)
|
|
mapping = mmap.mmap(file.fileno(), self._mapped_size)
|
|
try:
|
|
# New slots are already zero-filled. Publish the version last so readers
|
|
# cannot accept a partially initialized header.
|
|
mapping[:HEADER.size] = HEADER.pack(MAGIC, 0, self._capacity, self._slot_size_bytes, 0, 0, 0, 0)
|
|
atomics.store(mapping, 8, VERSION, 4)
|
|
except BaseException:
|
|
mapping.close()
|
|
raise
|
|
return mapping
|
|
|
|
def _header_matches(self, mapping: mmap.mmap) -> bool:
|
|
return (
|
|
atomics.load(mapping, 8, 4) == VERSION
|
|
and mapping[:8] == MAGIC
|
|
and atomics.load(mapping, 12, 4) == self._capacity
|
|
and atomics.load(mapping, 16, 4) == self._slot_size_bytes
|
|
)
|
|
|
|
def close(self) -> None:
|
|
"""Close the mapping and its owned file descriptor; safe to repeat."""
|
|
self._mmap.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:
|
|
# 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.
|
|
atomics.advance(self._mmap, 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
|
|
|
|
# Invalidate the previous publication before changing its bytes. A lapped
|
|
# consumer must not accept the old sequence during this overwrite.
|
|
self._write_u64(slot_offset + 8, 0)
|
|
atomics.fence()
|
|
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 slot_size_bytes(self) -> int:
|
|
"""Return maximum payload size per slot."""
|
|
return self._slot_size_bytes
|
|
|
|
def _read_u64(self, offset: int) -> int:
|
|
return atomics.load(self._mmap, offset, 8)
|
|
|
|
def _write_u32(self, offset: int, value: int) -> None:
|
|
atomics.store(self._mmap, offset, value, 4)
|
|
|
|
def _write_u64(self, offset: int, value: int) -> None:
|
|
atomics.store(self._mmap, offset, value, 8)
|