# tractor: structured concurrent "actors".
# Copyright 2018-eternity Tyler Goodlet.
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see .
'''
IPC Reliable RingBuffer implementation
'''
from __future__ import annotations
import struct
from typing import (
TypeVar,
ContextManager,
AsyncContextManager
)
from contextlib import (
contextmanager as cm,
asynccontextmanager as acm
)
from multiprocessing.shared_memory import SharedMemory
import trio
from msgspec import (
Struct,
to_builtins
)
from msgspec.msgpack import (
Encoder,
Decoder,
)
from tractor.log import get_logger
from tractor._exceptions import (
InternalError
)
from tractor.ipc._mp_bs import disable_mantracker
from tractor.linux._fdshare import (
share_fds,
unshare_fds,
request_fds_from
)
from tractor.linux.eventfd import (
open_eventfd,
EFDReadCancelled,
EventFD
)
from tractor._state import current_actor
log = get_logger(__name__)
disable_mantracker()
_DEFAULT_RB_SIZE = 10 * 1024
class RBToken(Struct, frozen=True):
'''
RingBuffer token contains necesary info to open resources of a ringbuf,
even in the case that ringbuf was not allocated by current actor.
'''
owner: str # if owner != `current_actor().name` we must use fdshare
shm_name: str
write_eventfd: int # used to signal writer ptr advance
wrap_eventfd: int # used to signal reader ready after wrap around
eof_eventfd: int # used to signal writer closed
buf_size: int # size in bytes of underlying shared memory buffer
def as_msg(self):
return to_builtins(self)
@classmethod
def from_msg(cls, msg: dict) -> RBToken:
if isinstance(msg, RBToken):
return msg
return RBToken(**msg)
@property
def fds(self) -> tuple[int, int, int]:
return (
self.write_eventfd,
self.wrap_eventfd,
self.eof_eventfd
)
def alloc_ringbuf(
shm_name: str,
buf_size: int = _DEFAULT_RB_SIZE,
) -> tuple[SharedMemory, RBToken]:
'''
Allocate OS resources for a ringbuf.
'''
shm = SharedMemory(
name=shm_name,
size=buf_size,
create=True
)
token = RBToken(
owner=current_actor().name,
shm_name=shm_name,
write_eventfd=open_eventfd(),
wrap_eventfd=open_eventfd(),
eof_eventfd=open_eventfd(),
buf_size=buf_size
)
# register fds for sharing
share_fds(
shm_name,
token.fds,
)
return shm, token
@cm
def open_ringbuf_sync(
shm_name: str,
buf_size: int = _DEFAULT_RB_SIZE,
) -> ContextManager[RBToken]:
'''
Handle resources for a ringbuf (shm, eventfd), yield `RBToken` to
be used with `attach_to_ringbuf_sender` and `attach_to_ringbuf_receiver`,
post yield maybe unshare fds and unlink shared memory
'''
shm: SharedMemory | None = None
token: RBToken | None = None
try:
shm, token = alloc_ringbuf(shm_name, buf_size=buf_size)
yield token
finally:
if token:
unshare_fds(shm_name)
if shm:
shm.unlink()
@acm
async def open_ringbuf(
shm_name: str,
buf_size: int = _DEFAULT_RB_SIZE,
) -> AsyncContextManager[RBToken]:
'''
Helper to use `open_ringbuf_sync` inside an async with block.
'''
with open_ringbuf_sync(
shm_name,
buf_size=buf_size
) as token:
yield token
@cm
def open_ringbufs_sync(
shm_names: list[str],
buf_sizes: int | list[str] = _DEFAULT_RB_SIZE,
) -> ContextManager[tuple[RBToken]]:
'''
Handle resources for multiple ringbufs at once.
'''
# maybe convert single int into list
if isinstance(buf_sizes, int):
buf_size = [buf_sizes] * len(shm_names)
# ensure len(shm_names) == len(buf_sizes)
if (
isinstance(buf_sizes, list)
and
len(buf_sizes) != len(shm_names)
):
raise ValueError(
'Expected buf_size list to be same length as shm_names'
)
# allocate resources
rings: list[tuple[SharedMemory, RBToken]] = [
alloc_ringbuf(shm_name, buf_size=buf_size)
for shm_name, buf_size in zip(shm_names, buf_size)
]
try:
yield tuple([token for _, token in rings])
finally:
# attempt fd unshare and shm unlink for each
for shm, token in rings:
try:
unshare_fds(token.shm_name)
except RuntimeError:
log.exception(f'while unsharing fds of {token}')
shm.unlink()
@acm
async def open_ringbufs(
shm_names: list[str],
buf_sizes: int | list[str] = _DEFAULT_RB_SIZE,
) -> AsyncContextManager[tuple[RBToken]]:
'''
Helper to use `open_ringbufs_sync` inside an async with block.
'''
with open_ringbufs_sync(
shm_names,
buf_sizes=buf_sizes
) as tokens:
yield tokens
@cm
def open_ringbuf_pair_sync(
shm_name: str,
buf_size: int = _DEFAULT_RB_SIZE
) -> ContextManager[tuple(RBToken, RBToken)]:
'''
Handle resources for a ringbuf pair to be used for
bidirectional messaging.
'''
with open_ringbufs_sync(
[
f'{shm_name}.send',
f'{shm_name}.recv'
],
buf_sizes=buf_size
) as tokens:
yield tokens
@acm
async def open_ringbuf_pair(
shm_name: str,
buf_size: int = _DEFAULT_RB_SIZE
) -> AsyncContextManager[tuple[RBToken, RBToken]]:
'''
Helper to use `open_ringbuf_pair_sync` inside an async with block.
'''
with open_ringbuf_pair_sync(
shm_name,
buf_size=buf_size
) as tokens:
yield tokens
Buffer = bytes | bytearray | memoryview
'''
IPC Reliable Ring Buffer
`eventfd(2)` is used for wrap around sync, to signal writes to
the reader and end of stream.
In order to guarantee full messages are received, all bytes
sent by `RingBufferSendChannel` are preceded with a 4 byte header
which decodes into a uint32 indicating the actual size of the
next full payload.
'''
PayloadT = TypeVar('PayloadT')
class RingBufferSendChannel(trio.abc.SendChannel[PayloadT]):
'''
Ring Buffer sender side implementation
Do not use directly! manage with `attach_to_ringbuf_sender`
after having opened a ringbuf context with `open_ringbuf`.
Optional batch mode:
If `batch_size` > 1 messages wont get sent immediately but will be
stored until `batch_size` messages are pending, then it will send
them all at once.
`batch_size` can be changed dynamically but always call, `flush()`
right before.
'''
def __init__(
self,
token: RBToken,
batch_size: int = 1,
cleanup: bool = False,
encoder: Encoder | None = None
):
self._token = RBToken.from_msg(token)
self.batch_size = batch_size
# ringbuf os resources
self._shm: SharedMemory | None = None
self._write_event = EventFD(self._token.write_eventfd, 'w')
self._wrap_event = EventFD(self._token.wrap_eventfd, 'r')
self._eof_event = EventFD(self._token.eof_eventfd, 'w')
# current write pointer
self._ptr: int = 0
# when `batch_size` > 1 store messages on `self._batch` and write them
# all, once `len(self._batch) == `batch_size`
self._batch: list[bytes] = []
# close shm & fds on exit?
self._cleanup: bool = cleanup
self._enc: Encoder | None = encoder
# have we closed this ringbuf?
# set to `False` on `.open()`
self._is_closed: bool = True
# ensure no concurrent `.send_all()` calls
self._send_all_lock = trio.StrictFIFOLock()
# ensure no concurrent `.send()` calls
self._send_lock = trio.StrictFIFOLock()
# ensure no concurrent `.flush()` calls
self._flush_lock = trio.StrictFIFOLock()
@property
def closed(self) -> bool:
return self._is_closed
@property
def name(self) -> str:
if not self._shm:
raise ValueError('shared memory not initialized yet!')
return self._shm.name
@property
def size(self) -> int:
return self._token.buf_size
@property
def ptr(self) -> int:
return self._ptr
@property
def write_fd(self) -> int:
return self._write_event.fd
@property
def wrap_fd(self) -> int:
return self._wrap_event.fd
@property
def pending_msgs(self) -> int:
return len(self._batch)
@property
def must_flush(self) -> bool:
return self.pending_msgs >= self.batch_size
async def _wait_wrap(self):
await self._wrap_event.read()
async def send_all(self, data: Buffer):
if self.closed:
raise trio.ClosedResourceError
if self._send_all_lock.locked():
raise trio.BusyResourceError
async with self._send_all_lock:
# while data is larger than the remaining buf
target_ptr = self.ptr + len(data)
while target_ptr > self.size:
# write all bytes that fit
remaining = self.size - self.ptr
self._shm.buf[self.ptr:] = data[:remaining]
# signal write and wait for reader wrap around
self._write_event.write(remaining)
await self._wait_wrap()
# wrap around and trim already written bytes
self._ptr = 0
data = data[remaining:]
target_ptr = self._ptr + len(data)
# remaining data fits on buffer
self._shm.buf[self.ptr:target_ptr] = data
self._write_event.write(len(data))
self._ptr = target_ptr
async def wait_send_all_might_not_block(self):
return
async def flush(
self,
new_batch_size: int | None = None
) -> None:
if self.closed:
raise trio.ClosedResourceError
async with self._flush_lock:
for msg in self._batch:
await self.send_all(msg)
self._batch = []
if new_batch_size:
self.batch_size = new_batch_size
async def send(self, value: PayloadT) -> None:
if self.closed:
raise trio.ClosedResourceError
if self._send_lock.locked():
raise trio.BusyResourceError
raw_value: bytes = (
value
if isinstance(value, bytes)
else
self._enc.encode(value)
)
async with self._send_lock:
msg: bytes = struct.pack(" 0:
await self.flush()
await self.send_all(msg)
return
self._batch.append(msg)
if self.must_flush:
await self.flush()
def open(self):
try:
self._shm = SharedMemory(
name=self._token.shm_name,
size=self._token.buf_size,
create=False
)
self._write_event.open()
self._wrap_event.open()
self._eof_event.open()
self._is_closed = False
except Exception as e:
e.add_note(f'while opening sender for {self._token.as_msg()}')
raise e
def _close(self):
self._eof_event.write(
self._ptr if self._ptr > 0 else self.size
)
if self._cleanup:
self._write_event.close()
self._wrap_event.close()
self._eof_event.close()
self._shm.close()
self._is_closed = True
async def aclose(self):
if self.closed:
return
self._close()
async def __aenter__(self):
self.open()
return self
class RingBufferReceiveChannel(trio.abc.ReceiveChannel[PayloadT]):
'''
Ring Buffer receiver side implementation
Do not use directly! manage with `attach_to_ringbuf_receiver`
after having opened a ringbuf context with `open_ringbuf`.
'''
def __init__(
self,
token: RBToken,
cleanup: bool = True,
decoder: Decoder | None = None
):
self._token = RBToken.from_msg(token)
# ringbuf os resources
self._shm: SharedMemory | None = None
self._write_event = EventFD(self._token.write_eventfd, 'w')
self._wrap_event = EventFD(self._token.wrap_eventfd, 'r')
self._eof_event = EventFD(self._token.eof_eventfd, 'r')
# current read ptr
self._ptr: int = 0
# current write_ptr (max bytes we can read from buf)
self._write_ptr: int = 0
# end ptr is used when EOF is signaled, it will contain maximun
# readable position on buf
self._end_ptr: int = -1
# close shm & fds on exit?
self._cleanup: bool = cleanup
# have we closed this ringbuf?
# set to `False` on `.open()`
self._is_closed: bool = True
self._dec: Decoder | None = decoder
# ensure no concurrent `.receive_some()` calls
self._receive_some_lock = trio.StrictFIFOLock()
# ensure no concurrent `.receive_exactly()` calls
self._receive_exactly_lock = trio.StrictFIFOLock()
# ensure no concurrent `.receive()` calls
self._receive_lock = trio.StrictFIFOLock()
@property
def closed(self) -> bool:
return self._is_closed
@property
def name(self) -> str:
if not self._shm:
raise ValueError('shared memory not initialized yet!')
return self._shm.name
@property
def size(self) -> int:
return self._token.buf_size
@property
def ptr(self) -> int:
return self._ptr
@property
def write_fd(self) -> int:
return self._write_event.fd
@property
def wrap_fd(self) -> int:
return self._wrap_event.fd
@property
def eof_was_signaled(self) -> bool:
return self._end_ptr != -1
async def _eof_monitor_task(self):
'''
Long running EOF event monitor, automatically run in bg by
`attach_to_ringbuf_receiver` context manager, if EOF event
is set its value will be the end pointer (highest valid
index to be read from buf, after setting the `self._end_ptr`
we close the write event which should cancel any blocked
`self._write_event.read()`s on it.
'''
try:
self._end_ptr = await self._eof_event.read()
except EFDReadCancelled:
...
except trio.Cancelled:
...
finally:
# closing write_event should trigger `EFDReadCancelled`
# on any pending read
self._write_event.close()
def receive_nowait(self, max_bytes: int = _DEFAULT_RB_SIZE) -> bytes:
'''
Try to receive any bytes we can without blocking or raise
`trio.WouldBlock`.
Returns b'' when no more bytes can be read (EOF signaled & read all).
'''
if max_bytes < 1:
raise ValueError("max_bytes must be >= 1")
# in case `end_ptr` is set that means eof was signaled.
# it will be >= `write_ptr`, use it for delta calc
highest_ptr = max(self._write_ptr, self._end_ptr)
delta = highest_ptr - self._ptr
# no more bytes to read
if delta == 0:
# if `end_ptr` is set that means we read all bytes before EOF
if self.eof_was_signaled:
return b''
# signal the need to wait on `write_event`
raise trio.WouldBlock
# dont overflow caller
delta = min(delta, max_bytes)
target_ptr = self._ptr + delta
# fetch next segment and advance ptr
segment = bytes(self._shm.buf[self._ptr:target_ptr])
self._ptr = target_ptr
if self._ptr == self.size:
# reached the end, signal wrap around
self._ptr = 0
self._write_ptr = 0
self._wrap_event.write(1)
return segment
async def receive_some(self, max_bytes: int = _DEFAULT_RB_SIZE) -> bytes:
'''
Receive up to `max_bytes`, if no `max_bytes` is provided
a reasonable default is used.
Can return < max_bytes.
'''
if self.closed:
raise trio.ClosedResourceError
if self._receive_some_lock.locked():
raise trio.BusyResourceError
async with self._receive_some_lock:
try:
# attempt direct read
return self.receive_nowait(max_bytes=max_bytes)
except trio.WouldBlock as e:
# we have read all we can, see if new data is available
if not self.eof_was_signaled:
# if we havent been signaled about EOF yet
try:
# wait next write and advance `write_ptr`
delta = await self._write_event.read()
self._write_ptr += delta
# yield lock and re-enter
except (
EFDReadCancelled, # read was cancelled with cscope
trio.Cancelled, # read got cancelled from outside
trio.BrokenResourceError # OSError EBADF happened while reading
):
# while waiting for new data `self._write_event` was closed
try:
# if eof was signaled receive no wait will not raise
# trio.WouldBlock and will push remaining until EOF
return self.receive_nowait(max_bytes=max_bytes)
except trio.WouldBlock:
# eof was not signaled but `self._wrap_event` is closed
# this means send side closed without EOF signal
return b''
else:
# shouldnt happen because receive_nowait does not raise
# trio.WouldBlock when `end_ptr` is set
raise InternalError(
'self._end_ptr is set but receive_nowait raised trio.WouldBlock'
) from e
return await self.receive_some(max_bytes=max_bytes)
async def receive_exactly(self, num_bytes: int) -> bytes:
'''
Fetch bytes until we read exactly `num_bytes` or EOC.
'''
if self.closed:
raise trio.ClosedResourceError
if self._receive_exactly_lock.locked():
raise trio.BusyResourceError
async with self._receive_exactly_lock:
payload = b''
while len(payload) < num_bytes:
remaining = num_bytes - len(payload)
new_bytes = await self.receive_some(
max_bytes=remaining
)
if new_bytes == b'':
break
payload += new_bytes
if payload == b'':
raise trio.EndOfChannel
return payload
async def receive(self, raw: bool = False) -> PayloadT:
'''
Receive a complete payload or raise EOC
'''
if self.closed:
raise trio.ClosedResourceError
if self._receive_lock.locked():
raise trio.BusyResourceError
async with self._receive_lock:
header: bytes = await self.receive_exactly(4)
size: int
size, = struct.unpack(" tuple[bytes, PayloadT]:
if not self._dec:
raise RuntimeError('iter_raw_pair requires decoder')
while True:
try:
raw = await self.receive(raw=True)
yield raw, self._dec.decode(raw)
except trio.EndOfChannel:
break
def open(self):
try:
self._shm = SharedMemory(
name=self._token.shm_name,
size=self._token.buf_size,
create=False
)
self._write_event.open()
self._wrap_event.open()
self._eof_event.open()
self._is_closed = False
except Exception as e:
e.add_note(f'while opening receiver for {self._token.as_msg()}')
raise e
def close(self):
if self._cleanup:
self._write_event.close()
self._wrap_event.close()
self._eof_event.close()
self._shm.close()
self._is_closed = True
async def aclose(self):
if self.closed:
return
self.close()
async def __aenter__(self):
self.open()
return self
async def _maybe_obtain_shared_resources(token: RBToken):
token = RBToken.from_msg(token)
# maybe token wasn't allocated by current actor
if token.owner != current_actor().name:
# use fdshare module to retrieve a copy of the FDs
fds = await request_fds_from(
token.owner,
token.shm_name
)
write, wrap, eof = fds
# rebuild token using FDs copies
token = RBToken(
owner=token.owner,
shm_name=token.shm_name,
write_eventfd=write,
wrap_eventfd=wrap,
eof_eventfd=eof,
buf_size=token.buf_size
)
return token
@acm
async def attach_to_ringbuf_receiver(
token: RBToken,
cleanup: bool = True,
decoder: Decoder | None = None
) -> AsyncContextManager[RingBufferReceiveChannel]:
'''
Attach a RingBufferReceiveChannel from a previously opened
RBToken.
Requires tractor runtime to be up in order to support opening a ringbuf
originally allocated by a different actor.
Launches `receiver._eof_monitor_task` in a `trio.Nursery`.
'''
token = await _maybe_obtain_shared_resources(token)
async with (
trio.open_nursery(strict_exception_groups=False) as n,
RingBufferReceiveChannel(
token,
cleanup=cleanup,
decoder=decoder
) as receiver
):
n.start_soon(receiver._eof_monitor_task)
yield receiver
@acm
async def attach_to_ringbuf_sender(
token: RBToken,
batch_size: int = 1,
cleanup: bool = True,
encoder: Encoder | None = None
) -> AsyncContextManager[RingBufferSendChannel]:
'''
Attach a RingBufferSendChannel from a previously opened
RBToken.
Requires tractor runtime to be up in order to support opening a ringbuf
originally allocated by a different actor.
'''
token = await _maybe_obtain_shared_resources(token)
async with RingBufferSendChannel(
token,
batch_size=batch_size,
cleanup=cleanup,
encoder=encoder
) as sender:
yield sender
class RingBufferChannel(trio.abc.Channel[bytes]):
'''
Combine `RingBufferSendChannel` and `RingBufferReceiveChannel`
in order to expose the bidirectional `trio.abc.Channel` API.
'''
def __init__(
self,
sender: RingBufferSendChannel,
receiver: RingBufferReceiveChannel
):
self._sender = sender
self._receiver = receiver
@property
def batch_size(self) -> int:
return self._sender.batch_size
@batch_size.setter
def batch_size(self, value: int) -> None:
self._sender.batch_size = value
@property
def pending_msgs(self) -> int:
return self._sender.pending_msgs
async def send_all(self, value: bytes) -> None:
await self._sender.send_all(value)
async def wait_send_all_might_not_block(self):
await self._sender.wait_send_all_might_not_block()
async def flush(
self,
new_batch_size: int | None = None
) -> None:
await self._sender.flush(new_batch_size=new_batch_size)
async def send(self, value: bytes) -> None:
await self._sender.send(value)
async def send_eof(self) -> None:
await self._sender.send_eof()
def receive_nowait(self, max_bytes: int = _DEFAULT_RB_SIZE) -> bytes:
return self._receiver.receive_nowait(max_bytes=max_bytes)
async def receive_some(self, max_bytes: int = _DEFAULT_RB_SIZE) -> bytes:
return await self._receiver.receive_some(max_bytes=max_bytes)
async def receive_exactly(self, num_bytes: int) -> bytes:
return await self._receiver.receive_exactly(num_bytes)
async def receive(self) -> bytes:
return await self._receiver.receive()
async def aclose(self):
await self._receiver.aclose()
await self._sender.aclose()
@acm
async def attach_to_ringbuf_channel(
token_in: RBToken,
token_out: RBToken,
batch_size: int = 1,
cleanup_in: bool = True,
cleanup_out: bool = True,
encoder: Encoder | None = None,
decoder: Decoder | None = None
) -> AsyncContextManager[trio.StapledStream]:
'''
Attach to two previously opened `RBToken`s and return a `RingBufferChannel`
'''
async with (
attach_to_ringbuf_receiver(
token_in,
cleanup=cleanup_in,
decoder=decoder
) as receiver,
attach_to_ringbuf_sender(
token_out,
batch_size=batch_size,
cleanup=cleanup_out,
encoder=encoder
) as sender,
):
yield RingBufferChannel(sender, receiver)