from inspect import signature
from multiprocessing.shared_memory import SharedMemory
from struct import calcsize, pack, unpack
from types import TracebackType
from typing import Optional, Self, Any
from math import ceil
from posix_ipc import Semaphore, O_CREX
from edri.config.setting import CHUNK_SIZE
__all__ = ["SharedMemoryPipe"]
_SHM_SUPPORTS_TRACK = "track" in signature(SharedMemory).parameters
[docs]
class SharedMemoryPipe:
"""
A bounded, single-producer/single-consumer byte pipe implemented on top of
POSIX shared memory + POSIX semaphores.
The shared memory region is divided into `max_slots` fixed-size slots.
Each slot contains a small header followed by up to `slot_size - HEADER_SIZE`
bytes of payload.
Synchronization:
- `slots` semaphore counts free slots (writer acquires before writing).
- `items` semaphore counts filled slots (reader acquires before reading).
Termination:
- EOF is signaled by writing a slot whose header has `eof=1`.
- `read()` returns None on EOF; the iterator stops.
Notes:
- Intended for local IPC between processes on the same host.
- The writer end owns creation of SHM and semaphores; the reader end
attaches by name.
"""
HEADER_FMT = "I B 3x" # uint32 size, uint8 eof, 3-byte pad
HEADER_SIZE = calcsize(HEADER_FMT)
_ALIGNMENT = 64 # bytes – tune here if you like
@staticmethod
def _aligned(n: int, alignment: int = _ALIGNMENT) -> int:
"""Round `n` up to the next multiple of `alignment` bytes."""
return ((n + alignment - 1) // alignment) * alignment
@classmethod
def _compute_slot_size(cls,
*,
total_size: int | None,
max_slots: int,
min_slot_size: int) -> int:
"""
Choose an efficient slot size.
If `total_size` is provided, this derives a slot size large enough that
`total_size` bytes can be transported in at most `max_slots` slots.
The computed size:
- reserves space for the per-slot header,
- is rounded up to `_ALIGNMENT` bytes (cache friendliness),
- is never smaller than `min_slot_size`.
If `total_size` is None, returns `min_slot_size` unchanged.
"""
if total_size is None:
return min_slot_size
# Bytes of **payload** we must be able to carry per slot
payload_per_slot = ceil(total_size / max_slots)
# Add header, align, and honour the caller’s minimum
candidate = cls._aligned(payload_per_slot + cls.HEADER_SIZE)
return max(min_slot_size, candidate)
def __init__(self,
name: Optional[str] = None,
*,
max_slots: int = 64,
slot_size: int = CHUNK_SIZE,
total_size: int | None = None,
_closed: bool = False,):
"""
Create a shared-memory-backed pipe endpoint.
This class has two modes:
**Writer mode** (default):
- `name` is None
- creates a new shared memory segment sized `slot_size * max_slots`
- creates two POSIX semaphores:
- `<shm.name>_items` (initial 0): number of filled slots
- `<shm.name>_slots` (initial `max_slots`): number of free slots
**Reader mode**:
- `name` is the shared memory name created by the writer
- attaches to the existing shared memory segment and semaphores
Parameters
----------
name:
Shared memory name to attach to. If None, a new shared memory region is
created and this instance becomes the writer end.
max_slots:
Number of fixed-size slots in the ring buffer. Higher values increase
buffering and reduce contention, at the cost of more shared memory.
slot_size:
Minimum slot size in bytes (including header). Each slot holds a
`HEADER_SIZE` header plus up to `slot_size - HEADER_SIZE` payload bytes.
If `total_size` is provided, this acts as a lower bound only.
total_size:
Optional expected total number of payload bytes to be sent. If provided,
the constructor may increase `slot_size` (aligned to `_ALIGNMENT`) so
that the full payload can be transmitted in at most `max_slots` chunks,
reducing chunk count and semaphore traffic.
_closed:
Internal flag used for pickling / state restoration. When True, no OS
resources are created or attached.
Attributes set
--------------
name:
The shared memory name (generated in writer mode, provided in reader mode).
is_writer:
True if this endpoint created the shared memory (writer mode).
slot_size:
Final slot size actually used (may be larger than the `slot_size` argument
when `total_size` is provided).
max_slots:
Number of slots in the ring buffer.
local_index:
Slot cursor for this endpoint (write index for writer, read index for reader).
Notes
-----
- On Python versions where `SharedMemory` supports the `track` parameter,
this implementation disables resource tracking via `track=False`.
- This implementation assumes a single writer and a single reader, each with
its own `local_index`.
"""
self.name = name
self.max_slots = max_slots
self.total_size = total_size
self.slot_size = self._compute_slot_size(total_size=total_size,
max_slots=max_slots,
min_slot_size=slot_size)
self.is_writer = self.name is None
self._closed = _closed
self.local_index = 0
self.shm: SharedMemory | None = None
shm_kwargs: dict[str, Any] = {}
if _SHM_SUPPORTS_TRACK:
shm_kwargs["track"] = False
if self._closed:
return
if self.is_writer:
shm_kwargs.update({"create": True, "size": self.slot_size * self.max_slots})
self.shm = SharedMemory(**shm_kwargs)
self.name = self.shm.name
self._items_sem_name = self.shm.name + "_items"
self._slots_sem_name = self.shm.name + "_slots"
self.items = Semaphore(self._items_sem_name, flags=O_CREX, initial_value=0)
self.slots = Semaphore(self._slots_sem_name, flags=O_CREX, initial_value=self.max_slots)
else:
assert self.name is not None # reader mode always has a name
shm_kwargs["name"] = self.name
self._items_sem_name = self.name + "_items"
self._slots_sem_name = self.name + "_slots"
self.shm = SharedMemory(**shm_kwargs)
self.items = Semaphore(self._items_sem_name)
self.slots = Semaphore(self._slots_sem_name)
@property
def _buf(self) -> memoryview:
"""The shared memory buffer; raises if the endpoint is closed/detached."""
assert self.shm is not None and self.shm.buf is not None
return self.shm.buf
def _write_slot(self, data: bytes, eof: bool) -> None:
"""
Write a single slot worth of data into the ring buffer.
This method:
1) waits for a free slot (`slots.acquire()`),
2) writes the slot header (payload length + EOF flag),
3) writes the payload bytes,
4) advances the local ring index,
5) signals availability to the reader (`items.release()`).
Parameters
----------
data:
Payload bytes to store in the slot. Must be <= max payload size
(`slot_size - HEADER_SIZE`).
eof:
Whether this slot marks end-of-stream. When `eof` is True, the
reader will return None from `read()` after consuming this slot.
Notes
-----
This is an internal primitive; `write()` handles chunking automatically.
"""
self.slots.acquire()
slot_index = self.local_index % self.max_slots
offset = slot_index * self.slot_size
header = pack(self.HEADER_FMT, len(data), int(eof))
buf = self._buf
buf[offset: offset + self.HEADER_SIZE] = header
buf[offset + self.HEADER_SIZE: offset + self.HEADER_SIZE + len(data)] = data
self.local_index += 1
self.items.release()
[docs]
def write(self, data: bytes, /, *, close: bool = False) -> None:
"""
Write bytes into the pipe, splitting into slot-sized chunks as needed.
Parameters
----------
data:
The bytes to send.
close:
If True, send end-of-stream after the data. All of `data` is written
first, then a dedicated empty EOF slot is emitted (via `close()`), so
the reader observes every byte and *then* stops. Reliable even when
`data` is empty.
Notes
-----
- Chunk size is `slot_size - HEADER_SIZE`.
- EOF is always its own empty slot; data slots never carry the EOF flag.
(Marking a data-carrying slot as EOF would make `read()` drop that
slot's payload, since it returns None on EOF.)
"""
max_payload = self.slot_size - self.HEADER_SIZE
total_chunks = (len(data) + max_payload - 1) // max_payload
for i in range(total_chunks):
chunk = data[i * max_payload: (i + 1) * max_payload]
self._write_slot(chunk, eof=False)
if close:
self.close()
[docs]
def read(self) -> Optional[bytes]:
"""
Read the next chunk from the pipe.
Returns
-------
bytes | None
The next payload chunk, or None if an EOF slot was encountered.
Behavior
--------
This method:
1) waits for an available item (`items.acquire()`),
2) reads the header to determine payload length and EOF flag,
3) copies out the payload bytes,
4) advances the local ring index,
5) frees the slot (`slots.release()`).
Notes
-----
Returning None is the EOF sentinel used by `__next__` to stop iteration.
"""
self.items.acquire()
slot_index = self.local_index % self.max_slots
offset = slot_index * self.slot_size
buf = self._buf
header = buf[offset: offset + self.HEADER_SIZE]
size, eof = unpack(self.HEADER_FMT, header)
data = bytes(buf[offset + self.HEADER_SIZE
: offset + self.HEADER_SIZE + size])
self.local_index += 1
self.slots.release()
return None if eof else data
[docs]
def reader(self) -> "SharedMemoryPipe":
"""
Create a reader endpoint attached to the same underlying shared memory.
Returns
-------
SharedMemoryPipe
A new instance configured as the reader end of this pipe.
Raises
------
RuntimeError
If called on an instance that is already a reader.
Notes
-----
The returned reader shares the same SHM name and semaphore names,
but maintains its own local read index.
"""
if not self.is_writer:
raise RuntimeError("Only writer can spawn reader")
return SharedMemoryPipe(name=self.name,
max_slots=self.max_slots,
slot_size=self.slot_size,
total_size=self.total_size)
[docs]
def close(self) -> None:
"""
Close this endpoint.
Writer semantics:
- Enqueues an explicit EOF slot (empty payload with eof=1),
then closes the shared memory handle.
Reader semantics:
- Closes the shared memory handle and unlinks the shared memory
segment (removing it from the system namespace).
Notes
-----
- This method is idempotent.
- Depending on your lifecycle expectations, you may prefer that the
writer is responsible for unlinking SHM instead of the reader.
"""
if self._closed:
return
self._closed = True
assert self.shm is not None
if self.is_writer:
# send EOF downstream
self._write_slot(b"", eof=True)
self.shm.close()
else:
self.shm.close()
self.shm.unlink()
self.shm = None # guard against accidental reuse
def __enter__(self) -> Self:
return self
def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None,
exc_tb: TracebackType | None) -> None:
self.close()
def __iter__(self) -> Self:
"""
Return an iterator over incoming chunks (reader side only).
Raises
------
RuntimeError
If called on the writer end.
"""
if self.is_writer:
raise RuntimeError("Cannot iterate over the writer end of the pipe")
return self
def __next__(self) -> bytes:
"""
Yield the next chunk from the pipe.
On EOF, this closes the reader endpoint and raises StopIteration.
"""
if self._closed:
raise StopIteration
chunk = self.read()
if chunk is None:
self.close()
raise StopIteration
return chunk
def __getstate__(self) -> dict[str, Any]:
"""
Support pickling of the reader endpoint.
The writer cannot be pickled because it owns system resources and is
responsible for emitting EOF and coordinating lifetime.
Returns
-------
dict
Minimal state needed to reattach to the shared memory region and
semaphores in another process.
Raises
------
RuntimeError
If invoked on the writer end.
"""
if self.is_writer:
raise RuntimeError("Writer cannot be pickled")
return {
"name": self.name,
"max_slots": self.max_slots,
"slot_size": self.slot_size,
"total_size": self.total_size,
"_closed": self._closed,
}
def __setstate__(self, state: dict[str, Any]) -> None:
"""
Restore a pickled reader endpoint by reattaching to SHM/semaphores.
Parameters
----------
state:
The dictionary produced by `__getstate__`.
"""
SharedMemoryPipe.__init__(self, **state)