From eeb9a7d61ba6865d9045e315d98f08024fa8a93e Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Wed, 12 Mar 2025 16:13:40 -0300 Subject: [PATCH 01/19] IPC ring bug impl with async read --- default.nix | 18 +++ pyproject.toml | 1 + tests/test_shm.py | 80 +++++++++++ tractor/_shm.py | 299 +++++++++++++++++++++++++++++++++++++++++- tractor/_spawn.py | 12 +- tractor/_supervise.py | 4 + uv.lock | 30 +++++ 7 files changed, 438 insertions(+), 6 deletions(-) create mode 100644 default.nix diff --git a/default.nix b/default.nix new file mode 100644 index 00000000..5a936971 --- /dev/null +++ b/default.nix @@ -0,0 +1,18 @@ +{ pkgs ? import {} }: +let + nativeBuildInputs = with pkgs; [ + stdenv.cc.cc.lib + uv + ]; + +in +pkgs.mkShell { + inherit nativeBuildInputs; + + LD_LIBRARY_PATH = pkgs.lib.makeLibraryPath nativeBuildInputs; + + shellHook = '' + set -e + uv venv .venv --python=3.12 + ''; +} diff --git a/pyproject.toml b/pyproject.toml index b3e9e100..fd67bff2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ dependencies = [ "pdbp>=1.6,<2", # windows only (from `pdbp`) # typed IPC msging "msgspec>=0.19.0", + "cffi>=1.17.1", ] # ------ project ------ diff --git a/tests/test_shm.py b/tests/test_shm.py index 2b7a382f..db0b1818 100644 --- a/tests/test_shm.py +++ b/tests/test_shm.py @@ -2,7 +2,10 @@ Shared mem primitives and APIs. """ +import time import uuid +import string +import random # import numpy import pytest @@ -11,6 +14,7 @@ import tractor from tractor._shm import ( open_shm_list, attach_shm_list, + EventFD, open_ringbuffer_sender, open_ringbuffer_receiver, ) @@ -165,3 +169,79 @@ def test_parent_writer_child_reader( await portal.cancel_actor() trio.run(main) + + +def random_string(size=256): + return ''.join(random.choice(string.ascii_lowercase) for i in range(size)) + + +async def child_read_shm( + msg_amount: int, + key: str, + write_event_fd: int, + wrap_event_fd: int, + max_bytes: int, +) -> None: + log = tractor.log.get_console_log(level='info') + recvd_msgs = 0 + start_ts = time.time() + async with open_ringbuffer_receiver( + write_event_fd, + wrap_event_fd, + key, + max_bytes=max_bytes + ) as receiver: + while recvd_msgs < msg_amount: + msg = await receiver.receive_some() + msgs = bytes(msg).split(b'\n') + first = msgs[0] + last = msgs[-2] + log.info((receiver.ptr - len(msg), receiver.ptr, first[:10], last[:10])) + recvd_msgs += len(msgs) + + end_ts = time.time() + elapsed = end_ts - start_ts + elapsed_ms = int(elapsed * 1000) + log.info(f'elapsed ms: {elapsed_ms}') + log.info(f'msg/sec: {int(msg_amount / elapsed):,}') + log.info(f'bytes/sec: {int(max_bytes / elapsed):,}') + +def test_ring_buff(): + log = tractor.log.get_console_log(level='info') + msg_amount = 100_000 + log.info(f'generating {msg_amount} messages...') + msgs = [ + f'[{i:08}]: {random_string()}\n'.encode('utf-8') + for i in range(msg_amount) + ] + buf_size = sum((len(m) for m in msgs)) + log.info(f'done! buffer size: {buf_size}') + async def main(): + with ( + EventFD(initval=0) as write_event, + EventFD(initval=0) as wrap_event, + ): + async with ( + tractor.open_nursery() as an, + open_ringbuffer_sender( + write_event.fd, + wrap_event.fd, + max_bytes=buf_size + ) as sender + ): + await an.run_in_actor( + child_read_shm, + msg_amount=msg_amount, + key=sender.key, + write_event_fd=write_event.fd, + wrap_event_fd=wrap_event.fd, + max_bytes=buf_size, + proc_kwargs={ + 'pass_fds': (write_event.fd, wrap_event.fd) + } + ) + for msg in msgs: + await sender.send_all(msg) + + + trio.run(main) diff --git a/tractor/_shm.py b/tractor/_shm.py index f8295105..7c177bc5 100644 --- a/tractor/_shm.py +++ b/tractor/_shm.py @@ -25,6 +25,7 @@ considered optional within the context of this runtime-library. from __future__ import annotations from sys import byteorder import time +import platform from typing import Optional from multiprocessing import shared_memory as shm from multiprocessing.shared_memory import ( @@ -32,7 +33,7 @@ from multiprocessing.shared_memory import ( ShareableList, ) -from msgspec import Struct +from msgspec import Struct, to_builtins import tractor from .log import get_logger @@ -142,7 +143,7 @@ class NDToken(Struct, frozen=True): ).descr def as_msg(self): - return self.to_dict() + return to_builtins(self) @classmethod def from_msg(cls, msg: dict) -> NDToken: @@ -831,3 +832,297 @@ def attach_shm_list( name=key, readonly=readonly, ) + + +if platform.system() == 'Linux': + import os + import errno + import string + import random + from contextlib import asynccontextmanager as acm + + import cffi + import trio + + ffi = cffi.FFI() + + # Declare the C functions and types we plan to use. + # - eventfd: for creating the event file descriptor + # - write: for writing to the file descriptor + # - read: for reading from the file descriptor + # - close: for closing the file descriptor + ffi.cdef( + ''' + int eventfd(unsigned int initval, int flags); + + ssize_t write(int fd, const void *buf, size_t count); + ssize_t read(int fd, void *buf, size_t count); + + int close(int fd); + ''' + ) + + # Open the default dynamic library (essentially 'libc' in most cases) + C = ffi.dlopen(None) + + # Constants from , if needed. + EFD_SEMAPHORE = 1 << 0 # 0x1 + EFD_CLOEXEC = 1 << 1 # 0x2 + EFD_NONBLOCK = 1 << 2 # 0x4 + + + def open_eventfd(initval: int = 0, flags: int = 0) -> int: + ''' + Open an eventfd with the given initial value and flags. + Returns the file descriptor on success, otherwise raises OSError. + ''' + fd = C.eventfd(initval, flags) + if fd < 0: + raise OSError(errno.errorcode[ffi.errno], 'eventfd failed') + return fd + + def write_eventfd(fd: int, value: int) -> int: + ''' + Write a 64-bit integer (uint64_t) to the eventfd's counter. + ''' + # Create a uint64_t* in C, store `value` + data_ptr = ffi.new('uint64_t *', value) + + # Call write(fd, data_ptr, 8) + # We expect to write exactly 8 bytes (sizeof(uint64_t)) + ret = C.write(fd, data_ptr, 8) + if ret < 0: + raise OSError(errno.errorcode[ffi.errno], 'write to eventfd failed') + return ret + + def read_eventfd(fd: int) -> int: + ''' + Read a 64-bit integer (uint64_t) from the eventfd, returning the value. + Reading resets the counter to 0 (unless using EFD_SEMAPHORE). + ''' + # Allocate an 8-byte buffer in C for reading + buf = ffi.new('char[]', 8) + + ret = C.read(fd, buf, 8) + if ret < 0: + raise OSError(errno.errorcode[ffi.errno], 'read from eventfd failed') + # Convert the 8 bytes we read into a Python integer + data_bytes = ffi.unpack(buf, 8) # returns a Python bytes object of length 8 + value = int.from_bytes(data_bytes, byteorder='little', signed=False) + return value + + def close_eventfd(fd: int) -> int: + ''' + Close the eventfd. + ''' + ret = C.close(fd) + if ret < 0: + raise OSError(errno.errorcode[ffi.errno], 'close failed') + + + class EventFD: + + def __init__( + self, + initval: int = 0, + flags: int = 0, + fd: int | None = None, + omode: str = 'r' + ): + self._initval: int = initval + self._flags: int = flags + self._fd: int | None = fd + self._omode: str = omode + self._fobj = None + + @property + def fd(self) -> int | None: + return self._fd + + def write(self, value: int) -> int: + return write_eventfd(self._fd, value) + + async def read(self) -> int: + return await trio.to_thread.run_sync(read_eventfd, self._fd) + + def open(self): + if not self._fd: + self._fd = open_eventfd( + initval=self._initval, flags=self._flags) + + else: + self._fobj = os.fdopen(self._fd, self._omode) + + def close(self): + if self._fobj: + self._fobj.close() + return + + if self._fd: + close_eventfd(self._fd) + + def __enter__(self): + self.open() + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.close() + + + class RingBuffSender(trio.abc.SendStream): + + def __init__( + self, + shm: SharedMemory, + write_event: EventFD, + wrap_event: EventFD, + start_ptr: int = 0 + ): + self._shm: SharedMemory = shm + self._write_event = write_event + self._wrap_event = wrap_event + self._ptr = start_ptr + + @property + def key(self) -> str: + return self._shm.name + + @property + def size(self) -> int: + return self._shm.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 + + async def send_all(self, data: bytes | bytearray | memoryview): + target_ptr = self.ptr + len(data) + if target_ptr > self.size: + remaining = self.size - self.ptr + self._shm.buf[self.ptr:] = data[:remaining] + self._write_event.write(remaining) + await self._wrap_event.read() + self._ptr = 0 + data = data[remaining:] + target_ptr = self._ptr + len(data) + + 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): + ... + + async def aclose(self): + ... + + async def __aenter__(self): + self._write_event.open() + self._wrap_event.open() + return self + + async def __aexit__(self, exc_type, exc_value, traceback): + await self.aclose() + + + class RingBuffReceiver(trio.abc.ReceiveStream): + + def __init__( + self, + shm: SharedMemory, + write_event: EventFD, + wrap_event: EventFD, + start_ptr: int = 0 + ): + self._shm: SharedMemory = shm + self._write_event = write_event + self._wrap_event = wrap_event + self._ptr = start_ptr + + @property + def key(self) -> str: + return self._shm.name + + @property + def size(self) -> int: + return self._shm.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 + + async def receive_some(self, max_bytes: int | None = None) -> bytes: + delta = await self._write_event.read() + next_ptr = self._ptr + delta + segment = bytes(self._shm.buf[self._ptr:next_ptr]) + self._ptr = next_ptr + if self.ptr == self.size: + self._ptr = 0 + self._wrap_event.write(1) + return segment + + async def aclose(self): + ... + + async def __aenter__(self): + self._write_event.open() + self._wrap_event.open() + return self + + async def __aexit__(self, exc_type, exc_value, traceback): + await self.aclose() + + @acm + async def open_ringbuffer_sender( + write_event_fd: int, + wrap_event_fd: int, + key: str | None = None, + max_bytes: int = 10 * 1024, + start_ptr: int = 0, + ) -> RingBuffSender: + if not key: + key: str = ''.join(random.choice(string.ascii_lowercase) for i in range(32)) + + shm = SharedMemory( + name=key, + size=max_bytes, + create=True + ) + async with RingBuffSender( + shm, EventFD(fd=write_event_fd, omode='w'), EventFD(fd=wrap_event_fd), start_ptr=start_ptr + ) as s: + yield s + + @acm + async def open_ringbuffer_receiver( + write_event_fd: int, + wrap_event_fd: int, + key: str, + max_bytes: int = 10 * 1024, + start_ptr: int = 0, + ) -> RingBuffSender: + shm = SharedMemory( + name=key, + size=max_bytes, + create=False + ) + async with RingBuffReceiver( + shm, EventFD(fd=write_event_fd), EventFD(fd=wrap_event_fd, omode='w'), start_ptr=start_ptr + ) as r: + yield r diff --git a/tractor/_spawn.py b/tractor/_spawn.py index 3159508d..dc2429d9 100644 --- a/tractor/_spawn.py +++ b/tractor/_spawn.py @@ -399,7 +399,8 @@ async def new_proc( *, infect_asyncio: bool = False, - task_status: TaskStatus[Portal] = trio.TASK_STATUS_IGNORED + task_status: TaskStatus[Portal] = trio.TASK_STATUS_IGNORED, + proc_kwargs: dict[str, any] = {} ) -> None: @@ -419,6 +420,7 @@ async def new_proc( _runtime_vars, # run time vars infect_asyncio=infect_asyncio, task_status=task_status, + proc_kwargs=proc_kwargs ) @@ -434,7 +436,8 @@ async def trio_proc( _runtime_vars: dict[str, Any], # serialized and sent to _child *, infect_asyncio: bool = False, - task_status: TaskStatus[Portal] = trio.TASK_STATUS_IGNORED + task_status: TaskStatus[Portal] = trio.TASK_STATUS_IGNORED, + proc_kwargs: dict[str, any] = {} ) -> None: ''' @@ -475,7 +478,7 @@ async def trio_proc( proc: trio.Process|None = None try: try: - proc: trio.Process = await trio.lowlevel.open_process(spawn_cmd) + proc: trio.Process = await trio.lowlevel.open_process(spawn_cmd, **proc_kwargs) log.runtime( 'Started new child\n' f'|_{proc}\n' @@ -640,7 +643,8 @@ async def mp_proc( _runtime_vars: dict[str, Any], # serialized and sent to _child *, infect_asyncio: bool = False, - task_status: TaskStatus[Portal] = trio.TASK_STATUS_IGNORED + task_status: TaskStatus[Portal] = trio.TASK_STATUS_IGNORED, + proc_kwargs: dict[str, any] = {} ) -> None: diff --git a/tractor/_supervise.py b/tractor/_supervise.py index bc6bc983..052a5f4c 100644 --- a/tractor/_supervise.py +++ b/tractor/_supervise.py @@ -141,6 +141,7 @@ class ActorNursery: # a `._ria_nursery` since the dependent APIs have been # removed! nursery: trio.Nursery|None = None, + proc_kwargs: dict[str, any] = {} ) -> Portal: ''' @@ -204,6 +205,7 @@ class ActorNursery: parent_addr, _rtv, # run time vars infect_asyncio=infect_asyncio, + proc_kwargs=proc_kwargs ) ) @@ -227,6 +229,7 @@ class ActorNursery: enable_modules: list[str] | None = None, loglevel: str | None = None, # set log level per subactor infect_asyncio: bool = False, + proc_kwargs: dict[str, any] = {}, **kwargs, # explicit args to ``fn`` @@ -257,6 +260,7 @@ class ActorNursery: # use the run_in_actor nursery nursery=self._ria_nursery, infect_asyncio=infect_asyncio, + proc_kwargs=proc_kwargs ) # XXX: don't allow stream funcs diff --git a/uv.lock b/uv.lock index e1c409f5..76b22243 100644 --- a/uv.lock +++ b/uv.lock @@ -20,10 +20,38 @@ dependencies = [ ] sdist = { url = "https://files.pythonhosted.org/packages/fc/97/c783634659c2920c3fc70419e3af40972dbaf758daa229a7d6ea6135c90d/cffi-1.17.1.tar.gz", hash = "sha256:1c39c6016c32bc48dd54561950ebd6836e1670f2ae46128f67cf49e789c52824", size = 516621 } wheels = [ + { url = "https://files.pythonhosted.org/packages/6b/f4/927e3a8899e52a27fa57a48607ff7dc91a9ebe97399b357b85a0c7892e00/cffi-1.17.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:a45e3c6913c5b87b3ff120dcdc03f6131fa0065027d0ed7ee6190736a74cd401", size = 182264 }, + { url = "https://files.pythonhosted.org/packages/6c/f5/6c3a8efe5f503175aaddcbea6ad0d2c96dad6f5abb205750d1b3df44ef29/cffi-1.17.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:30c5e0cb5ae493c04c8b42916e52ca38079f1b235c2f8ae5f4527b963c401caf", size = 178651 }, + { url = "https://files.pythonhosted.org/packages/94/dd/a3f0118e688d1b1a57553da23b16bdade96d2f9bcda4d32e7d2838047ff7/cffi-1.17.1-cp311-cp311-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f75c7ab1f9e4aca5414ed4d8e5c0e303a34f4421f8a0d47a4d019ceff0ab6af4", size = 445259 }, + { url = "https://files.pythonhosted.org/packages/2e/ea/70ce63780f096e16ce8588efe039d3c4f91deb1dc01e9c73a287939c79a6/cffi-1.17.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a1ed2dd2972641495a3ec98445e09766f077aee98a1c896dcb4ad0d303628e41", size = 469200 }, + { url = "https://files.pythonhosted.org/packages/1c/a0/a4fa9f4f781bda074c3ddd57a572b060fa0df7655d2a4247bbe277200146/cffi-1.17.1-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:46bf43160c1a35f7ec506d254e5c890f3c03648a4dbac12d624e4490a7046cd1", size = 477235 }, + { url = "https://files.pythonhosted.org/packages/62/12/ce8710b5b8affbcdd5c6e367217c242524ad17a02fe5beec3ee339f69f85/cffi-1.17.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a24ed04c8ffd54b0729c07cee15a81d964e6fee0e3d4d342a27b020d22959dc6", size = 459721 }, + { url = "https://files.pythonhosted.org/packages/ff/6b/d45873c5e0242196f042d555526f92aa9e0c32355a1be1ff8c27f077fd37/cffi-1.17.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:610faea79c43e44c71e1ec53a554553fa22321b65fae24889706c0a84d4ad86d", size = 467242 }, + { url = "https://files.pythonhosted.org/packages/1a/52/d9a0e523a572fbccf2955f5abe883cfa8bcc570d7faeee06336fbd50c9fc/cffi-1.17.1-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:a9b15d491f3ad5d692e11f6b71f7857e7835eb677955c00cc0aefcd0669adaf6", size = 477999 }, + { url = "https://files.pythonhosted.org/packages/44/74/f2a2460684a1a2d00ca799ad880d54652841a780c4c97b87754f660c7603/cffi-1.17.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:de2ea4b5833625383e464549fec1bc395c1bdeeb5f25c4a3a82b5a8c756ec22f", size = 454242 }, + { url = "https://files.pythonhosted.org/packages/f8/4a/34599cac7dfcd888ff54e801afe06a19c17787dfd94495ab0c8d35fe99fb/cffi-1.17.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:fc48c783f9c87e60831201f2cce7f3b2e4846bf4d8728eabe54d60700b318a0b", size = 478604 }, { url = "https://files.pythonhosted.org/packages/34/33/e1b8a1ba29025adbdcda5fb3a36f94c03d771c1b7b12f726ff7fef2ebe36/cffi-1.17.1-cp311-cp311-win32.whl", hash = "sha256:85a950a4ac9c359340d5963966e3e0a94a676bd6245a4b55bc43949eee26a655", size = 171727 }, { url = "https://files.pythonhosted.org/packages/3d/97/50228be003bb2802627d28ec0627837ac0bf35c90cf769812056f235b2d1/cffi-1.17.1-cp311-cp311-win_amd64.whl", hash = "sha256:caaf0640ef5f5517f49bc275eca1406b0ffa6aa184892812030f04c2abf589a0", size = 181400 }, + { url = "https://files.pythonhosted.org/packages/5a/84/e94227139ee5fb4d600a7a4927f322e1d4aea6fdc50bd3fca8493caba23f/cffi-1.17.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:805b4371bf7197c329fcb3ead37e710d1bca9da5d583f5073b799d5c5bd1eee4", size = 183178 }, + { url = "https://files.pythonhosted.org/packages/da/ee/fb72c2b48656111c4ef27f0f91da355e130a923473bf5ee75c5643d00cca/cffi-1.17.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:733e99bc2df47476e3848417c5a4540522f234dfd4ef3ab7fafdf555b082ec0c", size = 178840 }, + { url = "https://files.pythonhosted.org/packages/cc/b6/db007700f67d151abadf508cbfd6a1884f57eab90b1bb985c4c8c02b0f28/cffi-1.17.1-cp312-cp312-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1257bdabf294dceb59f5e70c64a3e2f462c30c7ad68092d01bbbfb1c16b1ba36", size = 454803 }, + { url = "https://files.pythonhosted.org/packages/1a/df/f8d151540d8c200eb1c6fba8cd0dfd40904f1b0682ea705c36e6c2e97ab3/cffi-1.17.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:da95af8214998d77a98cc14e3a3bd00aa191526343078b530ceb0bd710fb48a5", size = 478850 }, + { url = "https://files.pythonhosted.org/packages/28/c0/b31116332a547fd2677ae5b78a2ef662dfc8023d67f41b2a83f7c2aa78b1/cffi-1.17.1-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d63afe322132c194cf832bfec0dc69a99fb9bb6bbd550f161a49e9e855cc78ff", size = 485729 }, + { url = "https://files.pythonhosted.org/packages/91/2b/9a1ddfa5c7f13cab007a2c9cc295b70fbbda7cb10a286aa6810338e60ea1/cffi-1.17.1-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f79fc4fc25f1c8698ff97788206bb3c2598949bfe0fef03d299eb1b5356ada99", size = 471256 }, + { url = "https://files.pythonhosted.org/packages/b2/d5/da47df7004cb17e4955df6a43d14b3b4ae77737dff8bf7f8f333196717bf/cffi-1.17.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b62ce867176a75d03a665bad002af8e6d54644fad99a3c70905c543130e39d93", size = 479424 }, + { url = "https://files.pythonhosted.org/packages/0b/ac/2a28bcf513e93a219c8a4e8e125534f4f6db03e3179ba1c45e949b76212c/cffi-1.17.1-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:386c8bf53c502fff58903061338ce4f4950cbdcb23e2902d86c0f722b786bbe3", size = 484568 }, + { url = "https://files.pythonhosted.org/packages/d4/38/ca8a4f639065f14ae0f1d9751e70447a261f1a30fa7547a828ae08142465/cffi-1.17.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:4ceb10419a9adf4460ea14cfd6bc43d08701f0835e979bf821052f1805850fe8", size = 488736 }, { url = "https://files.pythonhosted.org/packages/86/c5/28b2d6f799ec0bdecf44dced2ec5ed43e0eb63097b0f58c293583b406582/cffi-1.17.1-cp312-cp312-win32.whl", hash = "sha256:a08d7e755f8ed21095a310a693525137cfe756ce62d066e53f502a83dc550f65", size = 172448 }, { url = "https://files.pythonhosted.org/packages/50/b9/db34c4755a7bd1cb2d1603ac3863f22bcecbd1ba29e5ee841a4bc510b294/cffi-1.17.1-cp312-cp312-win_amd64.whl", hash = "sha256:51392eae71afec0d0c8fb1a53b204dbb3bcabcb3c9b807eedf3e1e6ccf2de903", size = 181976 }, + { url = "https://files.pythonhosted.org/packages/8d/f8/dd6c246b148639254dad4d6803eb6a54e8c85c6e11ec9df2cffa87571dbe/cffi-1.17.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:f3a2b4222ce6b60e2e8b337bb9596923045681d71e5a082783484d845390938e", size = 182989 }, + { url = "https://files.pythonhosted.org/packages/8b/f1/672d303ddf17c24fc83afd712316fda78dc6fce1cd53011b839483e1ecc8/cffi-1.17.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:0984a4925a435b1da406122d4d7968dd861c1385afe3b45ba82b750f229811e2", size = 178802 }, + { url = "https://files.pythonhosted.org/packages/0e/2d/eab2e858a91fdff70533cab61dcff4a1f55ec60425832ddfdc9cd36bc8af/cffi-1.17.1-cp313-cp313-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d01b12eeeb4427d3110de311e1774046ad344f5b1a7403101878976ecd7a10f3", size = 454792 }, + { url = "https://files.pythonhosted.org/packages/75/b2/fbaec7c4455c604e29388d55599b99ebcc250a60050610fadde58932b7ee/cffi-1.17.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:706510fe141c86a69c8ddc029c7910003a17353970cff3b904ff0686a5927683", size = 478893 }, + { url = "https://files.pythonhosted.org/packages/4f/b7/6e4a2162178bf1935c336d4da8a9352cccab4d3a5d7914065490f08c0690/cffi-1.17.1-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:de55b766c7aa2e2a3092c51e0483d700341182f08e67c63630d5b6f200bb28e5", size = 485810 }, + { url = "https://files.pythonhosted.org/packages/c7/8a/1d0e4a9c26e54746dc08c2c6c037889124d4f59dffd853a659fa545f1b40/cffi-1.17.1-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c59d6e989d07460165cc5ad3c61f9fd8f1b4796eacbd81cee78957842b834af4", size = 471200 }, + { url = "https://files.pythonhosted.org/packages/26/9f/1aab65a6c0db35f43c4d1b4f580e8df53914310afc10ae0397d29d697af4/cffi-1.17.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dd398dbc6773384a17fe0d3e7eeb8d1a21c2200473ee6806bb5e6a8e62bb73dd", size = 479447 }, + { url = "https://files.pythonhosted.org/packages/5f/e4/fb8b3dd8dc0e98edf1135ff067ae070bb32ef9d509d6cb0f538cd6f7483f/cffi-1.17.1-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:3edc8d958eb099c634dace3c7e16560ae474aa3803a5df240542b305d14e14ed", size = 484358 }, + { url = "https://files.pythonhosted.org/packages/f1/47/d7145bf2dc04684935d57d67dff9d6d795b2ba2796806bb109864be3a151/cffi-1.17.1-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:72e72408cad3d5419375fc87d289076ee319835bdfa2caad331e377589aebba9", size = 488469 }, { url = "https://files.pythonhosted.org/packages/bf/ee/f94057fa6426481d663b88637a9a10e859e492c73d0384514a17d78ee205/cffi-1.17.1-cp313-cp313-win32.whl", hash = "sha256:e03eab0a8677fa80d646b5ddece1cbeaf556c313dcfac435ba11f107ba117b5d", size = 172475 }, { url = "https://files.pythonhosted.org/packages/7c/fc/6a8cb64e5f0324877d503c854da15d76c1e50eb722e320b15345c4d0c6de/cffi-1.17.1-cp313-cp313-win_amd64.whl", hash = "sha256:f6a16c31041f09ead72d69f583767292f750d24913dadacf5756b966aacb3f1a", size = 182009 }, ] @@ -321,6 +349,7 @@ name = "tractor" version = "0.1.0a6.dev0" source = { editable = "." } dependencies = [ + { name = "cffi" }, { name = "colorlog" }, { name = "msgspec" }, { name = "pdbp" }, @@ -342,6 +371,7 @@ dev = [ [package.metadata] requires-dist = [ + { name = "cffi", specifier = ">=1.17.1" }, { name = "colorlog", specifier = ">=6.8.2,<7" }, { name = "msgspec", specifier = ">=0.19.0" }, { name = "pdbp", specifier = ">=1.6,<2" }, -- 2.34.1 From 5afe0a02643a803e24a5e5ecaeeae4d87f62957c Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Thu, 13 Mar 2025 20:17:04 -0300 Subject: [PATCH 02/19] General improvements EventFD class now expects the fd to already be init with open_eventfd RingBuff Sender and Receiver fully manage SharedMemory and EventFD lifecycles, no aditional ctx mngrs needed Separate ring buf tests into its own test bed Add parametrization to test and cancellation Add docstrings Add simple testing data gen module .samples --- tests/test_ringbuf.py | 212 ++++++++++++++++++++++++++++++++++++ tests/test_shm.py | 80 -------------- tractor/_shm.py | 195 +++++++++++++++++++-------------- tractor/_testing/samples.py | 35 ++++++ 4 files changed, 360 insertions(+), 162 deletions(-) create mode 100644 tests/test_ringbuf.py create mode 100644 tractor/_testing/samples.py diff --git a/tests/test_ringbuf.py b/tests/test_ringbuf.py new file mode 100644 index 00000000..b81ea5f9 --- /dev/null +++ b/tests/test_ringbuf.py @@ -0,0 +1,212 @@ +import time + +import trio +import pytest +import tractor +from tractor._shm import ( + EFD_NONBLOCK, + open_eventfd, + RingBuffSender, + RingBuffReceiver +) +from tractor._testing.samples import generate_sample_messages + + +@tractor.context +async def child_read_shm( + ctx: tractor.Context, + msg_amount: int, + shm_key: str, + write_eventfd: int, + wrap_eventfd: int, + buf_size: int, + total_bytes: int, + flags: int = 0, +) -> None: + recvd_bytes = 0 + await ctx.started() + start_ts = time.time() + async with RingBuffReceiver( + shm_key, + write_eventfd, + wrap_eventfd, + buf_size=buf_size, + flags=flags + ) as receiver: + while recvd_bytes < total_bytes: + msg = await receiver.receive_some() + recvd_bytes += len(msg) + + # make sure we dont hold any memoryviews + # before the ctx manager aclose() + msg = None + + end_ts = time.time() + elapsed = end_ts - start_ts + elapsed_ms = int(elapsed * 1000) + + print(f'\n\telapsed ms: {elapsed_ms}') + print(f'\tmsg/sec: {int(msg_amount / elapsed):,}') + print(f'\tbytes/sec: {int(recvd_bytes / elapsed):,}') + + +@tractor.context +async def child_write_shm( + ctx: tractor.Context, + msg_amount: int, + rand_min: int, + rand_max: int, + shm_key: str, + write_eventfd: int, + wrap_eventfd: int, + buf_size: int, +) -> None: + msgs, total_bytes = generate_sample_messages( + msg_amount, + rand_min=rand_min, + rand_max=rand_max, + ) + await ctx.started(total_bytes) + async with RingBuffSender( + shm_key, + write_eventfd, + wrap_eventfd, + buf_size=buf_size + ) as sender: + for msg in msgs: + await sender.send_all(msg) + + +@pytest.mark.parametrize( + 'msg_amount,rand_min,rand_max,buf_size', + [ + # simple case, fixed payloads, large buffer + (100_000, 0, 0, 10 * 1024), + + # guaranteed wrap around on every write + (100, 10 * 1024, 20 * 1024, 10 * 1024), + + # large payload size, but large buffer + (10_000, 256 * 1024, 512 * 1024, 10 * 1024 * 1024) + ], + ids=[ + 'fixed_payloads_large_buffer', + 'wrap_around_every_write', + 'large_payloads_large_buffer', + ] +) +def test_ring_buff( + msg_amount: int, + rand_min: int, + rand_max: int, + buf_size: int +): + write_eventfd = open_eventfd() + wrap_eventfd = open_eventfd() + + proc_kwargs = { + 'pass_fds': (write_eventfd, wrap_eventfd) + } + + shm_key = 'test_ring_buff' + + common_kwargs = { + 'msg_amount': msg_amount, + 'shm_key': shm_key, + 'write_eventfd': write_eventfd, + 'wrap_eventfd': wrap_eventfd, + 'buf_size': buf_size + } + + async def main(): + async with tractor.open_nursery() as an: + send_p = await an.start_actor( + 'ring_sender', + enable_modules=[__name__], + proc_kwargs=proc_kwargs + ) + recv_p = await an.start_actor( + 'ring_receiver', + enable_modules=[__name__], + proc_kwargs=proc_kwargs + ) + async with ( + send_p.open_context( + child_write_shm, + rand_min=rand_min, + rand_max=rand_max, + **common_kwargs + ) as (sctx, total_bytes), + recv_p.open_context( + child_read_shm, + **common_kwargs, + total_bytes=total_bytes, + ) as (sctx, _sent), + ): + await recv_p.result() + + await send_p.cancel_actor() + await recv_p.cancel_actor() + + + trio.run(main) + + +@tractor.context +async def child_blocked_receiver( + ctx: tractor.Context, + shm_key: str, + write_eventfd: int, + wrap_eventfd: int, + flags: int = 0 +): + async with RingBuffReceiver( + shm_key, + write_eventfd, + wrap_eventfd, + flags=flags + ) as receiver: + await ctx.started() + await receiver.receive_some() + + +def test_ring_reader_cancel(): + flags = EFD_NONBLOCK + write_eventfd = open_eventfd(flags=flags) + wrap_eventfd = open_eventfd() + + proc_kwargs = { + 'pass_fds': (write_eventfd, wrap_eventfd) + } + + shm_key = 'test_ring_cancel' + + async def main(): + async with ( + tractor.open_nursery() as an, + RingBuffSender( + shm_key, + write_eventfd, + wrap_eventfd, + ) as _sender, + ): + recv_p = await an.start_actor( + 'ring_blocked_receiver', + enable_modules=[__name__], + proc_kwargs=proc_kwargs + ) + async with ( + recv_p.open_context( + child_blocked_receiver, + write_eventfd=write_eventfd, + wrap_eventfd=wrap_eventfd, + shm_key=shm_key, + flags=flags + ) as (sctx, _sent), + ): + await trio.sleep(1) + await an.cancel() + + + with pytest.raises(tractor._exceptions.ContextCancelled): + trio.run(main) diff --git a/tests/test_shm.py b/tests/test_shm.py index db0b1818..2b7a382f 100644 --- a/tests/test_shm.py +++ b/tests/test_shm.py @@ -2,10 +2,7 @@ Shared mem primitives and APIs. """ -import time import uuid -import string -import random # import numpy import pytest @@ -14,7 +11,6 @@ import tractor from tractor._shm import ( open_shm_list, attach_shm_list, - EventFD, open_ringbuffer_sender, open_ringbuffer_receiver, ) @@ -169,79 +165,3 @@ def test_parent_writer_child_reader( await portal.cancel_actor() trio.run(main) - - -def random_string(size=256): - return ''.join(random.choice(string.ascii_lowercase) for i in range(size)) - - -async def child_read_shm( - msg_amount: int, - key: str, - write_event_fd: int, - wrap_event_fd: int, - max_bytes: int, -) -> None: - log = tractor.log.get_console_log(level='info') - recvd_msgs = 0 - start_ts = time.time() - async with open_ringbuffer_receiver( - write_event_fd, - wrap_event_fd, - key, - max_bytes=max_bytes - ) as receiver: - while recvd_msgs < msg_amount: - msg = await receiver.receive_some() - msgs = bytes(msg).split(b'\n') - first = msgs[0] - last = msgs[-2] - log.info((receiver.ptr - len(msg), receiver.ptr, first[:10], last[:10])) - recvd_msgs += len(msgs) - - end_ts = time.time() - elapsed = end_ts - start_ts - elapsed_ms = int(elapsed * 1000) - log.info(f'elapsed ms: {elapsed_ms}') - log.info(f'msg/sec: {int(msg_amount / elapsed):,}') - log.info(f'bytes/sec: {int(max_bytes / elapsed):,}') - -def test_ring_buff(): - log = tractor.log.get_console_log(level='info') - msg_amount = 100_000 - log.info(f'generating {msg_amount} messages...') - msgs = [ - f'[{i:08}]: {random_string()}\n'.encode('utf-8') - for i in range(msg_amount) - ] - buf_size = sum((len(m) for m in msgs)) - log.info(f'done! buffer size: {buf_size}') - async def main(): - with ( - EventFD(initval=0) as write_event, - EventFD(initval=0) as wrap_event, - ): - async with ( - tractor.open_nursery() as an, - open_ringbuffer_sender( - write_event.fd, - wrap_event.fd, - max_bytes=buf_size - ) as sender - ): - await an.run_in_actor( - child_read_shm, - msg_amount=msg_amount, - key=sender.key, - write_event_fd=write_event.fd, - wrap_event_fd=wrap_event.fd, - max_bytes=buf_size, - proc_kwargs={ - 'pass_fds': (write_event.fd, wrap_event.fd) - } - ) - for msg in msgs: - await sender.send_all(msg) - - - trio.run(main) diff --git a/tractor/_shm.py b/tractor/_shm.py index 7c177bc5..5038e77a 100644 --- a/tractor/_shm.py +++ b/tractor/_shm.py @@ -837,8 +837,6 @@ def attach_shm_list( if platform.system() == 'Linux': import os import errno - import string - import random from contextlib import asynccontextmanager as acm import cffi @@ -862,19 +860,21 @@ if platform.system() == 'Linux': ''' ) + # Open the default dynamic library (essentially 'libc' in most cases) C = ffi.dlopen(None) # Constants from , if needed. - EFD_SEMAPHORE = 1 << 0 # 0x1 - EFD_CLOEXEC = 1 << 1 # 0x2 - EFD_NONBLOCK = 1 << 2 # 0x4 + EFD_SEMAPHORE = 1 + EFD_CLOEXEC = 0o2000000 + EFD_NONBLOCK = 0o4000 def open_eventfd(initval: int = 0, flags: int = 0) -> int: ''' Open an eventfd with the given initial value and flags. Returns the file descriptor on success, otherwise raises OSError. + ''' fd = C.eventfd(initval, flags) if fd < 0: @@ -884,6 +884,7 @@ if platform.system() == 'Linux': def write_eventfd(fd: int, value: int) -> int: ''' Write a 64-bit integer (uint64_t) to the eventfd's counter. + ''' # Create a uint64_t* in C, store `value` data_ptr = ffi.new('uint64_t *', value) @@ -899,6 +900,7 @@ if platform.system() == 'Linux': ''' Read a 64-bit integer (uint64_t) from the eventfd, returning the value. Reading resets the counter to 0 (unless using EFD_SEMAPHORE). + ''' # Allocate an 8-byte buffer in C for reading buf = ffi.new('char[]', 8) @@ -914,6 +916,7 @@ if platform.system() == 'Linux': def close_eventfd(fd: int) -> int: ''' Close the eventfd. + ''' ret = C.close(fd) if ret < 0: @@ -921,17 +924,19 @@ if platform.system() == 'Linux': class EventFD: + ''' + Use a previously opened eventfd(2), meant to be used in + sub-actors after root actor opens the eventfds then passes + them through pass_fds + + ''' def __init__( self, - initval: int = 0, - flags: int = 0, - fd: int | None = None, - omode: str = 'r' + fd: int, + omode: str ): - self._initval: int = initval - self._flags: int = flags - self._fd: int | None = fd + self._fd: int = fd self._omode: str = omode self._fobj = None @@ -943,23 +948,15 @@ if platform.system() == 'Linux': return write_eventfd(self._fd, value) async def read(self) -> int: + #TODO: how to handle signals? return await trio.to_thread.run_sync(read_eventfd, self._fd) def open(self): - if not self._fd: - self._fd = open_eventfd( - initval=self._initval, flags=self._flags) - - else: - self._fobj = os.fdopen(self._fd, self._omode) + self._fobj = os.fdopen(self._fd, self._omode) def close(self): if self._fobj: self._fobj.close() - return - - if self._fd: - close_eventfd(self._fd) def __enter__(self): self.open() @@ -970,18 +967,34 @@ if platform.system() == 'Linux': class RingBuffSender(trio.abc.SendStream): + ''' + IPC Reliable Ring Buffer sender side implementation + + `eventfd(2)` is used for wrap around sync, and also to signal + writes to the reader. + + TODO: if blocked on wrap around event wait it will not respond + to signals, fix soon TM + ''' def __init__( self, - shm: SharedMemory, - write_event: EventFD, - wrap_event: EventFD, - start_ptr: int = 0 + shm_key: str, + write_eventfd: int, + wrap_eventfd: int, + start_ptr: int = 0, + buf_size: int = 10 * 1024, + clean_shm_on_exit: bool = True ): - self._shm: SharedMemory = shm - self._write_event = write_event - self._wrap_event = wrap_event + self._shm = SharedMemory( + name=shm_key, + size=buf_size, + create=True + ) + self._write_event = EventFD(write_eventfd, 'w') + self._wrap_event = EventFD(wrap_eventfd, 'r') self._ptr = start_ptr + self.clean_shm_on_exit = clean_shm_on_exit @property def key(self) -> str: @@ -1004,25 +1017,37 @@ if platform.system() == 'Linux': return self._wrap_event.fd async def send_all(self, data: bytes | bytearray | memoryview): + # while data is larger than the remaining buf target_ptr = self.ptr + len(data) - if target_ptr > self.size: + 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._wrap_event.read() + + # 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): - ... + raise NotImplementedError async def aclose(self): - ... + self._write_event.close() + self._wrap_event.close() + if self.clean_shm_on_exit: + self._shm.unlink() + + else: + self._shm.close() async def __aenter__(self): self._write_event.open() @@ -1034,18 +1059,37 @@ if platform.system() == 'Linux': class RingBuffReceiver(trio.abc.ReceiveStream): + ''' + IPC Reliable Ring Buffer receiver side implementation + + `eventfd(2)` is used for wrap around sync, and also to signal + writes to the reader. + + Unless eventfd(2) object is opened with EFD_NONBLOCK flag, + calls to `receive_some` will block the signal handling, + on the main thread, for now solution is using polling, + working on a way to unblock GIL during read(2) to allow + signal processing on the main thread. + ''' def __init__( self, - shm: SharedMemory, - write_event: EventFD, - wrap_event: EventFD, - start_ptr: int = 0 + shm_key: str, + write_eventfd: int, + wrap_eventfd: int, + start_ptr: int = 0, + buf_size: int = 10 * 1024, + flags: int = 0 ): - self._shm: SharedMemory = shm - self._write_event = write_event - self._wrap_event = wrap_event + self._shm = SharedMemory( + name=shm_key, + size=buf_size, + create=False + ) + self._write_event = EventFD(write_eventfd, 'w') + self._wrap_event = EventFD(wrap_eventfd, 'r') self._ptr = start_ptr + self._flags = flags @property def key(self) -> str: @@ -1067,18 +1111,44 @@ if platform.system() == 'Linux': def wrap_fd(self) -> int: return self._wrap_event.fd - async def receive_some(self, max_bytes: int | None = None) -> bytes: - delta = await self._write_event.read() + async def receive_some( + self, + max_bytes: int | None = None, + nb_timeout: float = 0.1 + ) -> memoryview: + # if non blocking eventfd enabled, do polling + # until next write, this allows signal handling + if self._flags | EFD_NONBLOCK: + delta = None + while delta is None: + try: + delta = await self._write_event.read() + + except OSError as e: + if e.errno == 'EAGAIN': + continue + + raise e + + else: + delta = await self._write_event.read() + + # fetch next segment and advance ptr next_ptr = self._ptr + delta - segment = bytes(self._shm.buf[self._ptr:next_ptr]) + segment = self._shm.buf[self._ptr:next_ptr] self._ptr = next_ptr + if self.ptr == self.size: + # reached the end, signal wrap around self._ptr = 0 self._wrap_event.write(1) + return segment async def aclose(self): - ... + self._write_event.close() + self._wrap_event.close() + self._shm.close() async def __aenter__(self): self._write_event.open() @@ -1087,42 +1157,3 @@ if platform.system() == 'Linux': async def __aexit__(self, exc_type, exc_value, traceback): await self.aclose() - - @acm - async def open_ringbuffer_sender( - write_event_fd: int, - wrap_event_fd: int, - key: str | None = None, - max_bytes: int = 10 * 1024, - start_ptr: int = 0, - ) -> RingBuffSender: - if not key: - key: str = ''.join(random.choice(string.ascii_lowercase) for i in range(32)) - - shm = SharedMemory( - name=key, - size=max_bytes, - create=True - ) - async with RingBuffSender( - shm, EventFD(fd=write_event_fd, omode='w'), EventFD(fd=wrap_event_fd), start_ptr=start_ptr - ) as s: - yield s - - @acm - async def open_ringbuffer_receiver( - write_event_fd: int, - wrap_event_fd: int, - key: str, - max_bytes: int = 10 * 1024, - start_ptr: int = 0, - ) -> RingBuffSender: - shm = SharedMemory( - name=key, - size=max_bytes, - create=False - ) - async with RingBuffReceiver( - shm, EventFD(fd=write_event_fd), EventFD(fd=wrap_event_fd, omode='w'), start_ptr=start_ptr - ) as r: - yield r diff --git a/tractor/_testing/samples.py b/tractor/_testing/samples.py new file mode 100644 index 00000000..a87a22c4 --- /dev/null +++ b/tractor/_testing/samples.py @@ -0,0 +1,35 @@ +import os +import random + + +def generate_sample_messages( + amount: int, + rand_min: int = 0, + rand_max: int = 0, + silent: bool = False +) -> tuple[list[bytes], int]: + + msgs = [] + size = 0 + + if not silent: + print(f'\ngenerating {amount} messages...') + + for i in range(amount): + msg = f'[{i:08}]'.encode('utf-8') + + if rand_max > 0: + msg += os.urandom( + random.randint(rand_min, rand_max)) + + size += len(msg) + + msgs.append(msg) + + if not silent and i and i % 10_000 == 0: + print(f'{i} generated') + + if not silent: + print(f'done, {size:,} bytes in total') + + return msgs, size -- 2.34.1 From 7b8b9d6805f65dd70414ee642a49324ccd542fbc Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Thu, 13 Mar 2025 20:41:30 -0300 Subject: [PATCH 03/19] move tractor._ipc.py into tractor.ipc._chan.py --- tractor/__init__.py | 2 +- tractor/_context.py | 4 ++-- tractor/_discovery.py | 2 +- tractor/_exceptions.py | 2 +- tractor/_portal.py | 2 +- tractor/_root.py | 2 +- tractor/_rpc.py | 2 +- tractor/_runtime.py | 2 +- tractor/_streaming.py | 2 +- tractor/devx/_debug.py | 2 +- tractor/ipc/__init__.py | 11 +++++++++++ tractor/{_ipc.py => ipc/_chan.py} | 0 tractor/log.py | 4 ++-- 13 files changed, 24 insertions(+), 13 deletions(-) create mode 100644 tractor/ipc/__init__.py rename tractor/{_ipc.py => ipc/_chan.py} (100%) diff --git a/tractor/__init__.py b/tractor/__init__.py index 0c011a22..6fac747f 100644 --- a/tractor/__init__.py +++ b/tractor/__init__.py @@ -64,7 +64,7 @@ from ._root import ( run_daemon as run_daemon, open_root_actor as open_root_actor, ) -from ._ipc import Channel as Channel +from .ipc import Channel as Channel from ._portal import Portal as Portal from ._runtime import Actor as Actor # from . import hilevel as hilevel diff --git a/tractor/_context.py b/tractor/_context.py index 201e920a..d93d7759 100644 --- a/tractor/_context.py +++ b/tractor/_context.py @@ -89,7 +89,7 @@ from .msg import ( pretty_struct, _ops as msgops, ) -from ._ipc import ( +from .ipc import ( Channel, ) from ._streaming import ( @@ -105,7 +105,7 @@ from ._state import ( if TYPE_CHECKING: from ._portal import Portal from ._runtime import Actor - from ._ipc import MsgTransport + from .ipc import MsgTransport from .devx._frame_stack import ( CallerInfo, ) diff --git a/tractor/_discovery.py b/tractor/_discovery.py index a681c63b..1c3cbff0 100644 --- a/tractor/_discovery.py +++ b/tractor/_discovery.py @@ -29,7 +29,7 @@ from contextlib import asynccontextmanager as acm from tractor.log import get_logger from .trionics import gather_contexts -from ._ipc import _connect_chan, Channel +from .ipc import _connect_chan, Channel from ._portal import ( Portal, open_portal, diff --git a/tractor/_exceptions.py b/tractor/_exceptions.py index f9e18e18..8442ecfd 100644 --- a/tractor/_exceptions.py +++ b/tractor/_exceptions.py @@ -65,7 +65,7 @@ if TYPE_CHECKING: from ._context import Context from .log import StackLevelAdapter from ._stream import MsgStream - from ._ipc import Channel + from .ipc import Channel log = get_logger('tractor') diff --git a/tractor/_portal.py b/tractor/_portal.py index cee10c47..c8a781a7 100644 --- a/tractor/_portal.py +++ b/tractor/_portal.py @@ -43,7 +43,7 @@ from .trionics import maybe_open_nursery from ._state import ( current_actor, ) -from ._ipc import Channel +from .ipc import Channel from .log import get_logger from .msg import ( # Error, diff --git a/tractor/_root.py b/tractor/_root.py index 2a9beaa3..35639c15 100644 --- a/tractor/_root.py +++ b/tractor/_root.py @@ -43,7 +43,7 @@ from .devx import _debug from . import _spawn from . import _state from . import log -from ._ipc import _connect_chan +from .ipc import _connect_chan from ._exceptions import is_multi_cancelled diff --git a/tractor/_rpc.py b/tractor/_rpc.py index c5daed9e..6dfecd38 100644 --- a/tractor/_rpc.py +++ b/tractor/_rpc.py @@ -42,7 +42,7 @@ from trio import ( TaskStatus, ) -from ._ipc import Channel +from .ipc import Channel from ._context import ( Context, ) diff --git a/tractor/_runtime.py b/tractor/_runtime.py index 890a690a..2c8dbbd9 100644 --- a/tractor/_runtime.py +++ b/tractor/_runtime.py @@ -73,7 +73,7 @@ from tractor.msg import ( pretty_struct, types as msgtypes, ) -from ._ipc import Channel +from .ipc import Channel from ._context import ( mk_context, Context, diff --git a/tractor/_streaming.py b/tractor/_streaming.py index 2ff2d41c..21e59214 100644 --- a/tractor/_streaming.py +++ b/tractor/_streaming.py @@ -56,7 +56,7 @@ from tractor.msg import ( if TYPE_CHECKING: from ._runtime import Actor from ._context import Context - from ._ipc import Channel + from .ipc import Channel log = get_logger(__name__) diff --git a/tractor/devx/_debug.py b/tractor/devx/_debug.py index c6ca1d89..b95640dc 100644 --- a/tractor/devx/_debug.py +++ b/tractor/devx/_debug.py @@ -91,7 +91,7 @@ from tractor._state import ( if TYPE_CHECKING: from trio.lowlevel import Task from threading import Thread - from tractor._ipc import Channel + from tractor.ipc import Channel from tractor._runtime import ( Actor, ) diff --git a/tractor/ipc/__init__.py b/tractor/ipc/__init__.py new file mode 100644 index 00000000..d056acaa --- /dev/null +++ b/tractor/ipc/__init__.py @@ -0,0 +1,11 @@ +from ._chan import ( + _connect_chan, + MsgTransport, + Channel +) + +__all__ = [ + '_connect_chan', + 'MsgTransport', + 'Channel' +] diff --git a/tractor/_ipc.py b/tractor/ipc/_chan.py similarity index 100% rename from tractor/_ipc.py rename to tractor/ipc/_chan.py diff --git a/tractor/log.py b/tractor/log.py index 74e0321b..48b5cbd4 100644 --- a/tractor/log.py +++ b/tractor/log.py @@ -92,7 +92,7 @@ class StackLevelAdapter(LoggerAdapter): ) -> None: ''' IPC transport level msg IO; generally anything below - `._ipc.Channel` and friends. + `.ipc.Channel` and friends. ''' return self.log(5, msg) @@ -285,7 +285,7 @@ def get_logger( # NOTE: for handling for modules that use ``get_logger(__name__)`` # we make the following stylistic choice: # - always avoid duplicate project-package token - # in msg output: i.e. tractor.tractor _ipc.py in header + # in msg output: i.e. tractor.tractor.ipc._chan.py in header # looks ridiculous XD # - never show the leaf module name in the {name} part # since in python the {filename} is always this same -- 2.34.1 From 6b4d08d03064ce6c385106a17ba1f75d6712eadb Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Thu, 13 Mar 2025 20:59:14 -0300 Subject: [PATCH 04/19] Move tractor._shm to tractor.ipc._shm --- tests/test_ringbuf.py | 2 +- tests/test_shm.py | 2 +- tractor/ipc/__init__.py | 2 +- tractor/{ => ipc}/_shm.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) rename tractor/{ => ipc}/_shm.py (99%) diff --git a/tests/test_ringbuf.py b/tests/test_ringbuf.py index b81ea5f9..e4011768 100644 --- a/tests/test_ringbuf.py +++ b/tests/test_ringbuf.py @@ -3,7 +3,7 @@ import time import trio import pytest import tractor -from tractor._shm import ( +from tractor.ipc._shm import ( EFD_NONBLOCK, open_eventfd, RingBuffSender, diff --git a/tests/test_shm.py b/tests/test_shm.py index 2b7a382f..ddeb67aa 100644 --- a/tests/test_shm.py +++ b/tests/test_shm.py @@ -8,7 +8,7 @@ import uuid import pytest import trio import tractor -from tractor._shm import ( +from tractor.ipc._shm import ( open_shm_list, attach_shm_list, ) diff --git a/tractor/ipc/__init__.py b/tractor/ipc/__init__.py index d056acaa..2a401bf6 100644 --- a/tractor/ipc/__init__.py +++ b/tractor/ipc/__init__.py @@ -7,5 +7,5 @@ from ._chan import ( __all__ = [ '_connect_chan', 'MsgTransport', - 'Channel' + 'Channel', ] diff --git a/tractor/_shm.py b/tractor/ipc/_shm.py similarity index 99% rename from tractor/_shm.py rename to tractor/ipc/_shm.py index 5038e77a..752f81ff 100644 --- a/tractor/_shm.py +++ b/tractor/ipc/_shm.py @@ -36,7 +36,7 @@ from multiprocessing.shared_memory import ( from msgspec import Struct, to_builtins import tractor -from .log import get_logger +from tractor.log import get_logger _USE_POSIX = getattr(shm, '_USE_POSIX', False) -- 2.34.1 From 1a83626f26cc6012b7c36f632223e9d6df39a860 Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Thu, 13 Mar 2025 21:10:23 -0300 Subject: [PATCH 05/19] Move linux specifics from tractor.ipc._shm into tractor.ipc._linux --- tests/test_ringbuf.py | 2 +- tractor/ipc/__init__.py | 26 +++- tractor/ipc/_linux.py | 324 +++++++++++++++++++++++++++++++++++++++ tractor/ipc/_shm.py | 326 ---------------------------------------- 4 files changed, 343 insertions(+), 335 deletions(-) create mode 100644 tractor/ipc/_linux.py diff --git a/tests/test_ringbuf.py b/tests/test_ringbuf.py index e4011768..1c4e88f9 100644 --- a/tests/test_ringbuf.py +++ b/tests/test_ringbuf.py @@ -3,7 +3,7 @@ import time import trio import pytest import tractor -from tractor.ipc._shm import ( +from tractor.ipc import ( EFD_NONBLOCK, open_eventfd, RingBuffSender, diff --git a/tractor/ipc/__init__.py b/tractor/ipc/__init__.py index 2a401bf6..1b548b65 100644 --- a/tractor/ipc/__init__.py +++ b/tractor/ipc/__init__.py @@ -1,11 +1,21 @@ +import platform + from ._chan import ( - _connect_chan, - MsgTransport, - Channel + _connect_chan as _connect_chan, + MsgTransport as MsgTransport, + Channel as Channel ) -__all__ = [ - '_connect_chan', - 'MsgTransport', - 'Channel', -] +if platform.system() == 'Linux': + from ._linux import ( + EFD_SEMAPHORE as EFD_SEMAPHORE, + EFD_CLOEXEC as EFD_CLOEXEC, + EFD_NONBLOCK as EFD_NONBLOCK, + open_eventfd as open_eventfd, + write_eventfd as write_eventfd, + read_eventfd as read_eventfd, + close_eventfd as close_eventfd, + EventFD as EventFD, + RingBuffSender as RingBuffSender, + RingBuffReceiver as RingBuffReceiver + ) diff --git a/tractor/ipc/_linux.py b/tractor/ipc/_linux.py new file mode 100644 index 00000000..2a9eabc1 --- /dev/null +++ b/tractor/ipc/_linux.py @@ -0,0 +1,324 @@ + +import os +import errno + +from multiprocessing.shared_memory import SharedMemory + +import cffi +import trio + +ffi = cffi.FFI() + +# Declare the C functions and types we plan to use. +# - eventfd: for creating the event file descriptor +# - write: for writing to the file descriptor +# - read: for reading from the file descriptor +# - close: for closing the file descriptor +ffi.cdef( + ''' + int eventfd(unsigned int initval, int flags); + + ssize_t write(int fd, const void *buf, size_t count); + ssize_t read(int fd, void *buf, size_t count); + + int close(int fd); + ''' +) + + +# Open the default dynamic library (essentially 'libc' in most cases) +C = ffi.dlopen(None) + +# Constants from , if needed. +EFD_SEMAPHORE = 1 +EFD_CLOEXEC = 0o2000000 +EFD_NONBLOCK = 0o4000 + + +def open_eventfd(initval: int = 0, flags: int = 0) -> int: + ''' + Open an eventfd with the given initial value and flags. + Returns the file descriptor on success, otherwise raises OSError. + + ''' + fd = C.eventfd(initval, flags) + if fd < 0: + raise OSError(errno.errorcode[ffi.errno], 'eventfd failed') + return fd + +def write_eventfd(fd: int, value: int) -> int: + ''' + Write a 64-bit integer (uint64_t) to the eventfd's counter. + + ''' + # Create a uint64_t* in C, store `value` + data_ptr = ffi.new('uint64_t *', value) + + # Call write(fd, data_ptr, 8) + # We expect to write exactly 8 bytes (sizeof(uint64_t)) + ret = C.write(fd, data_ptr, 8) + if ret < 0: + raise OSError(errno.errorcode[ffi.errno], 'write to eventfd failed') + return ret + +def read_eventfd(fd: int) -> int: + ''' + Read a 64-bit integer (uint64_t) from the eventfd, returning the value. + Reading resets the counter to 0 (unless using EFD_SEMAPHORE). + + ''' + # Allocate an 8-byte buffer in C for reading + buf = ffi.new('char[]', 8) + + ret = C.read(fd, buf, 8) + if ret < 0: + raise OSError(errno.errorcode[ffi.errno], 'read from eventfd failed') + # Convert the 8 bytes we read into a Python integer + data_bytes = ffi.unpack(buf, 8) # returns a Python bytes object of length 8 + value = int.from_bytes(data_bytes, byteorder='little', signed=False) + return value + +def close_eventfd(fd: int) -> int: + ''' + Close the eventfd. + + ''' + ret = C.close(fd) + if ret < 0: + raise OSError(errno.errorcode[ffi.errno], 'close failed') + + +class EventFD: + ''' + Use a previously opened eventfd(2), meant to be used in + sub-actors after root actor opens the eventfds then passes + them through pass_fds + + ''' + + def __init__( + self, + fd: int, + omode: str + ): + self._fd: int = fd + self._omode: str = omode + self._fobj = None + + @property + def fd(self) -> int | None: + return self._fd + + def write(self, value: int) -> int: + return write_eventfd(self._fd, value) + + async def read(self) -> int: + #TODO: how to handle signals? + return await trio.to_thread.run_sync(read_eventfd, self._fd) + + def open(self): + self._fobj = os.fdopen(self._fd, self._omode) + + def close(self): + if self._fobj: + self._fobj.close() + + def __enter__(self): + self.open() + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.close() + + +class RingBuffSender(trio.abc.SendStream): + ''' + IPC Reliable Ring Buffer sender side implementation + + `eventfd(2)` is used for wrap around sync, and also to signal + writes to the reader. + + TODO: if blocked on wrap around event wait it will not respond + to signals, fix soon TM + ''' + + def __init__( + self, + shm_key: str, + write_eventfd: int, + wrap_eventfd: int, + start_ptr: int = 0, + buf_size: int = 10 * 1024, + clean_shm_on_exit: bool = True + ): + self._shm = SharedMemory( + name=shm_key, + size=buf_size, + create=True + ) + self._write_event = EventFD(write_eventfd, 'w') + self._wrap_event = EventFD(wrap_eventfd, 'r') + self._ptr = start_ptr + self.clean_shm_on_exit = clean_shm_on_exit + + @property + def key(self) -> str: + return self._shm.name + + @property + def size(self) -> int: + return self._shm.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 + + async def send_all(self, data: bytes | bytearray | memoryview): + # 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._wrap_event.read() + + # 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): + raise NotImplementedError + + async def aclose(self): + self._write_event.close() + self._wrap_event.close() + if self.clean_shm_on_exit: + self._shm.unlink() + + else: + self._shm.close() + + async def __aenter__(self): + self._write_event.open() + self._wrap_event.open() + return self + + async def __aexit__(self, exc_type, exc_value, traceback): + await self.aclose() + + +class RingBuffReceiver(trio.abc.ReceiveStream): + ''' + IPC Reliable Ring Buffer receiver side implementation + + `eventfd(2)` is used for wrap around sync, and also to signal + writes to the reader. + + Unless eventfd(2) object is opened with EFD_NONBLOCK flag, + calls to `receive_some` will block the signal handling, + on the main thread, for now solution is using polling, + working on a way to unblock GIL during read(2) to allow + signal processing on the main thread. + ''' + + def __init__( + self, + shm_key: str, + write_eventfd: int, + wrap_eventfd: int, + start_ptr: int = 0, + buf_size: int = 10 * 1024, + flags: int = 0 + ): + self._shm = SharedMemory( + name=shm_key, + size=buf_size, + create=False + ) + self._write_event = EventFD(write_eventfd, 'w') + self._wrap_event = EventFD(wrap_eventfd, 'r') + self._ptr = start_ptr + self._flags = flags + + @property + def key(self) -> str: + return self._shm.name + + @property + def size(self) -> int: + return self._shm.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 + + async def receive_some( + self, + max_bytes: int | None = None, + nb_timeout: float = 0.1 + ) -> memoryview: + # if non blocking eventfd enabled, do polling + # until next write, this allows signal handling + if self._flags | EFD_NONBLOCK: + delta = None + while delta is None: + try: + delta = await self._write_event.read() + + except OSError as e: + if e.errno == 'EAGAIN': + continue + + raise e + + else: + delta = await self._write_event.read() + + # fetch next segment and advance ptr + next_ptr = self._ptr + delta + segment = self._shm.buf[self._ptr:next_ptr] + self._ptr = next_ptr + + if self.ptr == self.size: + # reached the end, signal wrap around + self._ptr = 0 + self._wrap_event.write(1) + + return segment + + async def aclose(self): + self._write_event.close() + self._wrap_event.close() + self._shm.close() + + async def __aenter__(self): + self._write_event.open() + self._wrap_event.open() + return self + + async def __aexit__(self, exc_type, exc_value, traceback): + await self.aclose() diff --git a/tractor/ipc/_shm.py b/tractor/ipc/_shm.py index 752f81ff..0ee8bf23 100644 --- a/tractor/ipc/_shm.py +++ b/tractor/ipc/_shm.py @@ -25,7 +25,6 @@ considered optional within the context of this runtime-library. from __future__ import annotations from sys import byteorder import time -import platform from typing import Optional from multiprocessing import shared_memory as shm from multiprocessing.shared_memory import ( @@ -832,328 +831,3 @@ def attach_shm_list( name=key, readonly=readonly, ) - - -if platform.system() == 'Linux': - import os - import errno - from contextlib import asynccontextmanager as acm - - import cffi - import trio - - ffi = cffi.FFI() - - # Declare the C functions and types we plan to use. - # - eventfd: for creating the event file descriptor - # - write: for writing to the file descriptor - # - read: for reading from the file descriptor - # - close: for closing the file descriptor - ffi.cdef( - ''' - int eventfd(unsigned int initval, int flags); - - ssize_t write(int fd, const void *buf, size_t count); - ssize_t read(int fd, void *buf, size_t count); - - int close(int fd); - ''' - ) - - - # Open the default dynamic library (essentially 'libc' in most cases) - C = ffi.dlopen(None) - - # Constants from , if needed. - EFD_SEMAPHORE = 1 - EFD_CLOEXEC = 0o2000000 - EFD_NONBLOCK = 0o4000 - - - def open_eventfd(initval: int = 0, flags: int = 0) -> int: - ''' - Open an eventfd with the given initial value and flags. - Returns the file descriptor on success, otherwise raises OSError. - - ''' - fd = C.eventfd(initval, flags) - if fd < 0: - raise OSError(errno.errorcode[ffi.errno], 'eventfd failed') - return fd - - def write_eventfd(fd: int, value: int) -> int: - ''' - Write a 64-bit integer (uint64_t) to the eventfd's counter. - - ''' - # Create a uint64_t* in C, store `value` - data_ptr = ffi.new('uint64_t *', value) - - # Call write(fd, data_ptr, 8) - # We expect to write exactly 8 bytes (sizeof(uint64_t)) - ret = C.write(fd, data_ptr, 8) - if ret < 0: - raise OSError(errno.errorcode[ffi.errno], 'write to eventfd failed') - return ret - - def read_eventfd(fd: int) -> int: - ''' - Read a 64-bit integer (uint64_t) from the eventfd, returning the value. - Reading resets the counter to 0 (unless using EFD_SEMAPHORE). - - ''' - # Allocate an 8-byte buffer in C for reading - buf = ffi.new('char[]', 8) - - ret = C.read(fd, buf, 8) - if ret < 0: - raise OSError(errno.errorcode[ffi.errno], 'read from eventfd failed') - # Convert the 8 bytes we read into a Python integer - data_bytes = ffi.unpack(buf, 8) # returns a Python bytes object of length 8 - value = int.from_bytes(data_bytes, byteorder='little', signed=False) - return value - - def close_eventfd(fd: int) -> int: - ''' - Close the eventfd. - - ''' - ret = C.close(fd) - if ret < 0: - raise OSError(errno.errorcode[ffi.errno], 'close failed') - - - class EventFD: - ''' - Use a previously opened eventfd(2), meant to be used in - sub-actors after root actor opens the eventfds then passes - them through pass_fds - - ''' - - def __init__( - self, - fd: int, - omode: str - ): - self._fd: int = fd - self._omode: str = omode - self._fobj = None - - @property - def fd(self) -> int | None: - return self._fd - - def write(self, value: int) -> int: - return write_eventfd(self._fd, value) - - async def read(self) -> int: - #TODO: how to handle signals? - return await trio.to_thread.run_sync(read_eventfd, self._fd) - - def open(self): - self._fobj = os.fdopen(self._fd, self._omode) - - def close(self): - if self._fobj: - self._fobj.close() - - def __enter__(self): - self.open() - return self - - def __exit__(self, exc_type, exc_value, traceback): - self.close() - - - class RingBuffSender(trio.abc.SendStream): - ''' - IPC Reliable Ring Buffer sender side implementation - - `eventfd(2)` is used for wrap around sync, and also to signal - writes to the reader. - - TODO: if blocked on wrap around event wait it will not respond - to signals, fix soon TM - ''' - - def __init__( - self, - shm_key: str, - write_eventfd: int, - wrap_eventfd: int, - start_ptr: int = 0, - buf_size: int = 10 * 1024, - clean_shm_on_exit: bool = True - ): - self._shm = SharedMemory( - name=shm_key, - size=buf_size, - create=True - ) - self._write_event = EventFD(write_eventfd, 'w') - self._wrap_event = EventFD(wrap_eventfd, 'r') - self._ptr = start_ptr - self.clean_shm_on_exit = clean_shm_on_exit - - @property - def key(self) -> str: - return self._shm.name - - @property - def size(self) -> int: - return self._shm.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 - - async def send_all(self, data: bytes | bytearray | memoryview): - # 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._wrap_event.read() - - # 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): - raise NotImplementedError - - async def aclose(self): - self._write_event.close() - self._wrap_event.close() - if self.clean_shm_on_exit: - self._shm.unlink() - - else: - self._shm.close() - - async def __aenter__(self): - self._write_event.open() - self._wrap_event.open() - return self - - async def __aexit__(self, exc_type, exc_value, traceback): - await self.aclose() - - - class RingBuffReceiver(trio.abc.ReceiveStream): - ''' - IPC Reliable Ring Buffer receiver side implementation - - `eventfd(2)` is used for wrap around sync, and also to signal - writes to the reader. - - Unless eventfd(2) object is opened with EFD_NONBLOCK flag, - calls to `receive_some` will block the signal handling, - on the main thread, for now solution is using polling, - working on a way to unblock GIL during read(2) to allow - signal processing on the main thread. - ''' - - def __init__( - self, - shm_key: str, - write_eventfd: int, - wrap_eventfd: int, - start_ptr: int = 0, - buf_size: int = 10 * 1024, - flags: int = 0 - ): - self._shm = SharedMemory( - name=shm_key, - size=buf_size, - create=False - ) - self._write_event = EventFD(write_eventfd, 'w') - self._wrap_event = EventFD(wrap_eventfd, 'r') - self._ptr = start_ptr - self._flags = flags - - @property - def key(self) -> str: - return self._shm.name - - @property - def size(self) -> int: - return self._shm.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 - - async def receive_some( - self, - max_bytes: int | None = None, - nb_timeout: float = 0.1 - ) -> memoryview: - # if non blocking eventfd enabled, do polling - # until next write, this allows signal handling - if self._flags | EFD_NONBLOCK: - delta = None - while delta is None: - try: - delta = await self._write_event.read() - - except OSError as e: - if e.errno == 'EAGAIN': - continue - - raise e - - else: - delta = await self._write_event.read() - - # fetch next segment and advance ptr - next_ptr = self._ptr + delta - segment = self._shm.buf[self._ptr:next_ptr] - self._ptr = next_ptr - - if self.ptr == self.size: - # reached the end, signal wrap around - self._ptr = 0 - self._wrap_event.write(1) - - return segment - - async def aclose(self): - self._write_event.close() - self._wrap_event.close() - self._shm.close() - - async def __aenter__(self): - self._write_event.open() - self._wrap_event.open() - return self - - async def __aexit__(self, exc_type, exc_value, traceback): - await self.aclose() -- 2.34.1 From 8de9ab291ec2fc4fcd45e33f12f8c3629509d038 Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Thu, 13 Mar 2025 21:15:16 -0300 Subject: [PATCH 06/19] Move RingBuffSender|Receiver to its own tractor.ipc._ringbuf module --- tractor/ipc/__init__.py | 3 + tractor/ipc/_linux.py | 195 -------------------------------------- tractor/ipc/_ringbuf.py | 201 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 204 insertions(+), 195 deletions(-) create mode 100644 tractor/ipc/_ringbuf.py diff --git a/tractor/ipc/__init__.py b/tractor/ipc/__init__.py index 1b548b65..c0b58951 100644 --- a/tractor/ipc/__init__.py +++ b/tractor/ipc/__init__.py @@ -16,6 +16,9 @@ if platform.system() == 'Linux': read_eventfd as read_eventfd, close_eventfd as close_eventfd, EventFD as EventFD, + ) + + from ._ringbuf import ( RingBuffSender as RingBuffSender, RingBuffReceiver as RingBuffReceiver ) diff --git a/tractor/ipc/_linux.py b/tractor/ipc/_linux.py index 2a9eabc1..53feac6a 100644 --- a/tractor/ipc/_linux.py +++ b/tractor/ipc/_linux.py @@ -2,8 +2,6 @@ import os import errno -from multiprocessing.shared_memory import SharedMemory - import cffi import trio @@ -129,196 +127,3 @@ class EventFD: def __exit__(self, exc_type, exc_value, traceback): self.close() - - -class RingBuffSender(trio.abc.SendStream): - ''' - IPC Reliable Ring Buffer sender side implementation - - `eventfd(2)` is used for wrap around sync, and also to signal - writes to the reader. - - TODO: if blocked on wrap around event wait it will not respond - to signals, fix soon TM - ''' - - def __init__( - self, - shm_key: str, - write_eventfd: int, - wrap_eventfd: int, - start_ptr: int = 0, - buf_size: int = 10 * 1024, - clean_shm_on_exit: bool = True - ): - self._shm = SharedMemory( - name=shm_key, - size=buf_size, - create=True - ) - self._write_event = EventFD(write_eventfd, 'w') - self._wrap_event = EventFD(wrap_eventfd, 'r') - self._ptr = start_ptr - self.clean_shm_on_exit = clean_shm_on_exit - - @property - def key(self) -> str: - return self._shm.name - - @property - def size(self) -> int: - return self._shm.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 - - async def send_all(self, data: bytes | bytearray | memoryview): - # 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._wrap_event.read() - - # 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): - raise NotImplementedError - - async def aclose(self): - self._write_event.close() - self._wrap_event.close() - if self.clean_shm_on_exit: - self._shm.unlink() - - else: - self._shm.close() - - async def __aenter__(self): - self._write_event.open() - self._wrap_event.open() - return self - - async def __aexit__(self, exc_type, exc_value, traceback): - await self.aclose() - - -class RingBuffReceiver(trio.abc.ReceiveStream): - ''' - IPC Reliable Ring Buffer receiver side implementation - - `eventfd(2)` is used for wrap around sync, and also to signal - writes to the reader. - - Unless eventfd(2) object is opened with EFD_NONBLOCK flag, - calls to `receive_some` will block the signal handling, - on the main thread, for now solution is using polling, - working on a way to unblock GIL during read(2) to allow - signal processing on the main thread. - ''' - - def __init__( - self, - shm_key: str, - write_eventfd: int, - wrap_eventfd: int, - start_ptr: int = 0, - buf_size: int = 10 * 1024, - flags: int = 0 - ): - self._shm = SharedMemory( - name=shm_key, - size=buf_size, - create=False - ) - self._write_event = EventFD(write_eventfd, 'w') - self._wrap_event = EventFD(wrap_eventfd, 'r') - self._ptr = start_ptr - self._flags = flags - - @property - def key(self) -> str: - return self._shm.name - - @property - def size(self) -> int: - return self._shm.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 - - async def receive_some( - self, - max_bytes: int | None = None, - nb_timeout: float = 0.1 - ) -> memoryview: - # if non blocking eventfd enabled, do polling - # until next write, this allows signal handling - if self._flags | EFD_NONBLOCK: - delta = None - while delta is None: - try: - delta = await self._write_event.read() - - except OSError as e: - if e.errno == 'EAGAIN': - continue - - raise e - - else: - delta = await self._write_event.read() - - # fetch next segment and advance ptr - next_ptr = self._ptr + delta - segment = self._shm.buf[self._ptr:next_ptr] - self._ptr = next_ptr - - if self.ptr == self.size: - # reached the end, signal wrap around - self._ptr = 0 - self._wrap_event.write(1) - - return segment - - async def aclose(self): - self._write_event.close() - self._wrap_event.close() - self._shm.close() - - async def __aenter__(self): - self._write_event.open() - self._wrap_event.open() - return self - - async def __aexit__(self, exc_type, exc_value, traceback): - await self.aclose() diff --git a/tractor/ipc/_ringbuf.py b/tractor/ipc/_ringbuf.py new file mode 100644 index 00000000..0895381f --- /dev/null +++ b/tractor/ipc/_ringbuf.py @@ -0,0 +1,201 @@ +from multiprocessing.shared_memory import SharedMemory + +import trio + +from ._linux import ( + EFD_NONBLOCK, + EventFD +) + + +class RingBuffSender(trio.abc.SendStream): + ''' + IPC Reliable Ring Buffer sender side implementation + + `eventfd(2)` is used for wrap around sync, and also to signal + writes to the reader. + + TODO: if blocked on wrap around event wait it will not respond + to signals, fix soon TM + ''' + + def __init__( + self, + shm_key: str, + write_eventfd: int, + wrap_eventfd: int, + start_ptr: int = 0, + buf_size: int = 10 * 1024, + clean_shm_on_exit: bool = True + ): + self._shm = SharedMemory( + name=shm_key, + size=buf_size, + create=True + ) + self._write_event = EventFD(write_eventfd, 'w') + self._wrap_event = EventFD(wrap_eventfd, 'r') + self._ptr = start_ptr + self.clean_shm_on_exit = clean_shm_on_exit + + @property + def key(self) -> str: + return self._shm.name + + @property + def size(self) -> int: + return self._shm.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 + + async def send_all(self, data: bytes | bytearray | memoryview): + # 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._wrap_event.read() + + # 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): + raise NotImplementedError + + async def aclose(self): + self._write_event.close() + self._wrap_event.close() + if self.clean_shm_on_exit: + self._shm.unlink() + + else: + self._shm.close() + + async def __aenter__(self): + self._write_event.open() + self._wrap_event.open() + return self + + async def __aexit__(self, exc_type, exc_value, traceback): + await self.aclose() + + +class RingBuffReceiver(trio.abc.ReceiveStream): + ''' + IPC Reliable Ring Buffer receiver side implementation + + `eventfd(2)` is used for wrap around sync, and also to signal + writes to the reader. + + Unless eventfd(2) object is opened with EFD_NONBLOCK flag, + calls to `receive_some` will block the signal handling, + on the main thread, for now solution is using polling, + working on a way to unblock GIL during read(2) to allow + signal processing on the main thread. + ''' + + def __init__( + self, + shm_key: str, + write_eventfd: int, + wrap_eventfd: int, + start_ptr: int = 0, + buf_size: int = 10 * 1024, + flags: int = 0 + ): + self._shm = SharedMemory( + name=shm_key, + size=buf_size, + create=False + ) + self._write_event = EventFD(write_eventfd, 'w') + self._wrap_event = EventFD(wrap_eventfd, 'r') + self._ptr = start_ptr + self._flags = flags + + @property + def key(self) -> str: + return self._shm.name + + @property + def size(self) -> int: + return self._shm.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 + + async def receive_some( + self, + max_bytes: int | None = None, + nb_timeout: float = 0.1 + ) -> memoryview: + # if non blocking eventfd enabled, do polling + # until next write, this allows signal handling + if self._flags | EFD_NONBLOCK: + delta = None + while delta is None: + try: + delta = await self._write_event.read() + + except OSError as e: + if e.errno == 'EAGAIN': + continue + + raise e + + else: + delta = await self._write_event.read() + + # fetch next segment and advance ptr + next_ptr = self._ptr + delta + segment = self._shm.buf[self._ptr:next_ptr] + self._ptr = next_ptr + + if self.ptr == self.size: + # reached the end, signal wrap around + self._ptr = 0 + self._wrap_event.write(1) + + return segment + + async def aclose(self): + self._write_event.close() + self._wrap_event.close() + self._shm.close() + + async def __aenter__(self): + self._write_event.open() + self._wrap_event.open() + return self + + async def __aexit__(self, exc_type, exc_value, traceback): + await self.aclose() -- 2.34.1 From 9980bb2bd0400feca3e6fb30a5f7828e707a4476 Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Thu, 13 Mar 2025 21:25:50 -0300 Subject: [PATCH 07/19] Add module headers and fix spacing on tractor._ipc._linux --- tractor/ipc/__init__.py | 17 +++++++++++++++++ tractor/ipc/_linux.py | 22 ++++++++++++++++++++++ tractor/ipc/_ringbuf.py | 19 +++++++++++++++++++ 3 files changed, 58 insertions(+) diff --git a/tractor/ipc/__init__.py b/tractor/ipc/__init__.py index c0b58951..59fc1e16 100644 --- a/tractor/ipc/__init__.py +++ b/tractor/ipc/__init__.py @@ -1,3 +1,20 @@ +# 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 . + + import platform from ._chan import ( diff --git a/tractor/ipc/_linux.py b/tractor/ipc/_linux.py index 53feac6a..12b00260 100644 --- a/tractor/ipc/_linux.py +++ b/tractor/ipc/_linux.py @@ -1,4 +1,22 @@ +# 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 . +''' +Linux specifics, for now we are only exposing EventFD + +''' import os import errno @@ -27,6 +45,7 @@ ffi.cdef( # Open the default dynamic library (essentially 'libc' in most cases) C = ffi.dlopen(None) + # Constants from , if needed. EFD_SEMAPHORE = 1 EFD_CLOEXEC = 0o2000000 @@ -44,6 +63,7 @@ def open_eventfd(initval: int = 0, flags: int = 0) -> int: raise OSError(errno.errorcode[ffi.errno], 'eventfd failed') return fd + def write_eventfd(fd: int, value: int) -> int: ''' Write a 64-bit integer (uint64_t) to the eventfd's counter. @@ -59,6 +79,7 @@ def write_eventfd(fd: int, value: int) -> int: raise OSError(errno.errorcode[ffi.errno], 'write to eventfd failed') return ret + def read_eventfd(fd: int) -> int: ''' Read a 64-bit integer (uint64_t) from the eventfd, returning the value. @@ -76,6 +97,7 @@ def read_eventfd(fd: int) -> int: value = int.from_bytes(data_bytes, byteorder='little', signed=False) return value + def close_eventfd(fd: int) -> int: ''' Close the eventfd. diff --git a/tractor/ipc/_ringbuf.py b/tractor/ipc/_ringbuf.py index 0895381f..50a9eff1 100644 --- a/tractor/ipc/_ringbuf.py +++ b/tractor/ipc/_ringbuf.py @@ -1,3 +1,22 @@ +# 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 multiprocessing.shared_memory import SharedMemory import trio -- 2.34.1 From f799e9ac51e7123e59aac9ce5a4225cb36e61774 Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Thu, 13 Mar 2025 22:43:02 -0300 Subject: [PATCH 08/19] Handle cancelation on EventFD.read --- tests/test_ringbuf.py | 7 ++----- tractor/ipc/_linux.py | 6 ++++-- tractor/ipc/_ringbuf.py | 6 +++--- 3 files changed, 9 insertions(+), 10 deletions(-) diff --git a/tests/test_ringbuf.py b/tests/test_ringbuf.py index 1c4e88f9..9e457b2a 100644 --- a/tests/test_ringbuf.py +++ b/tests/test_ringbuf.py @@ -4,7 +4,6 @@ import trio import pytest import tractor from tractor.ipc import ( - EFD_NONBLOCK, open_eventfd, RingBuffSender, RingBuffReceiver @@ -95,7 +94,7 @@ async def child_write_shm( 'large_payloads_large_buffer', ] ) -def test_ring_buff( +def test_ringbuf( msg_amount: int, rand_min: int, rand_max: int, @@ -171,8 +170,7 @@ async def child_blocked_receiver( def test_ring_reader_cancel(): - flags = EFD_NONBLOCK - write_eventfd = open_eventfd(flags=flags) + write_eventfd = open_eventfd() wrap_eventfd = open_eventfd() proc_kwargs = { @@ -201,7 +199,6 @@ def test_ring_reader_cancel(): write_eventfd=write_eventfd, wrap_eventfd=wrap_eventfd, shm_key=shm_key, - flags=flags ) as (sctx, _sent), ): await trio.sleep(1) diff --git a/tractor/ipc/_linux.py b/tractor/ipc/_linux.py index 12b00260..88d80d1c 100644 --- a/tractor/ipc/_linux.py +++ b/tractor/ipc/_linux.py @@ -133,8 +133,10 @@ class EventFD: return write_eventfd(self._fd, value) async def read(self) -> int: - #TODO: how to handle signals? - return await trio.to_thread.run_sync(read_eventfd, self._fd) + return await trio.to_thread.run_sync( + read_eventfd, self._fd, + abandon_on_cancel=True + ) def open(self): self._fobj = os.fdopen(self._fd, self._omode) diff --git a/tractor/ipc/_ringbuf.py b/tractor/ipc/_ringbuf.py index 50a9eff1..77a30ab8 100644 --- a/tractor/ipc/_ringbuf.py +++ b/tractor/ipc/_ringbuf.py @@ -45,7 +45,7 @@ class RingBuffSender(trio.abc.SendStream): wrap_eventfd: int, start_ptr: int = 0, buf_size: int = 10 * 1024, - clean_shm_on_exit: bool = True + unlink_on_exit: bool = True ): self._shm = SharedMemory( name=shm_key, @@ -55,7 +55,7 @@ class RingBuffSender(trio.abc.SendStream): self._write_event = EventFD(write_eventfd, 'w') self._wrap_event = EventFD(wrap_eventfd, 'r') self._ptr = start_ptr - self.clean_shm_on_exit = clean_shm_on_exit + self.unlink_on_exit = unlink_on_exit @property def key(self) -> str: @@ -104,7 +104,7 @@ class RingBuffSender(trio.abc.SendStream): async def aclose(self): self._write_event.close() self._wrap_event.close() - if self.clean_shm_on_exit: + if self.unlink_on_exit: self._shm.unlink() else: -- 2.34.1 From 6ac6fd56c0d62bd1cae3a1994af0d19dca40e8e5 Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Thu, 13 Mar 2025 22:47:45 -0300 Subject: [PATCH 09/19] Address some of fomo\'s comments --- tractor/ipc/_ringbuf.py | 6 ------ tractor/ipc/_shm.py | 5 ++++- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/tractor/ipc/_ringbuf.py b/tractor/ipc/_ringbuf.py index 77a30ab8..0a4f3819 100644 --- a/tractor/ipc/_ringbuf.py +++ b/tractor/ipc/_ringbuf.py @@ -115,9 +115,6 @@ class RingBuffSender(trio.abc.SendStream): self._wrap_event.open() return self - async def __aexit__(self, exc_type, exc_value, traceback): - await self.aclose() - class RingBuffReceiver(trio.abc.ReceiveStream): ''' @@ -215,6 +212,3 @@ class RingBuffReceiver(trio.abc.ReceiveStream): self._write_event.open() self._wrap_event.open() return self - - async def __aexit__(self, exc_type, exc_value, traceback): - await self.aclose() diff --git a/tractor/ipc/_shm.py b/tractor/ipc/_shm.py index 0ee8bf23..c101af30 100644 --- a/tractor/ipc/_shm.py +++ b/tractor/ipc/_shm.py @@ -32,7 +32,10 @@ from multiprocessing.shared_memory import ( ShareableList, ) -from msgspec import Struct, to_builtins +from msgspec import ( + Struct, + to_builtins +) import tractor from tractor.log import get_logger -- 2.34.1 From 59c8c7bfe36daf75d5e1dfa8c6af23c4fd7e1b0f Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Thu, 13 Mar 2025 23:12:20 -0300 Subject: [PATCH 10/19] Make ring buf api use pickle-able RBToken --- tests/test_ringbuf.py | 165 ++++++++++++++++------------------------ tractor/ipc/__init__.py | 4 +- tractor/ipc/_ringbuf.py | 90 +++++++++++++++------- 3 files changed, 130 insertions(+), 129 deletions(-) diff --git a/tests/test_ringbuf.py b/tests/test_ringbuf.py index 9e457b2a..64fb37e9 100644 --- a/tests/test_ringbuf.py +++ b/tests/test_ringbuf.py @@ -4,7 +4,8 @@ import trio import pytest import tractor from tractor.ipc import ( - open_eventfd, + open_ringbuf, + RBToken, RingBuffSender, RingBuffReceiver ) @@ -15,22 +16,16 @@ from tractor._testing.samples import generate_sample_messages async def child_read_shm( ctx: tractor.Context, msg_amount: int, - shm_key: str, - write_eventfd: int, - wrap_eventfd: int, + token: RBToken, buf_size: int, total_bytes: int, - flags: int = 0, ) -> None: recvd_bytes = 0 await ctx.started() start_ts = time.time() async with RingBuffReceiver( - shm_key, - write_eventfd, - wrap_eventfd, + token, buf_size=buf_size, - flags=flags ) as receiver: while recvd_bytes < total_bytes: msg = await receiver.receive_some() @@ -55,9 +50,7 @@ async def child_write_shm( msg_amount: int, rand_min: int, rand_max: int, - shm_key: str, - write_eventfd: int, - wrap_eventfd: int, + token: RBToken, buf_size: int, ) -> None: msgs, total_bytes = generate_sample_messages( @@ -67,9 +60,7 @@ async def child_write_shm( ) await ctx.started(total_bytes) async with RingBuffSender( - shm_key, - write_eventfd, - wrap_eventfd, + token, buf_size=buf_size ) as sender: for msg in msgs: @@ -100,52 +91,48 @@ def test_ringbuf( rand_max: int, buf_size: int ): - write_eventfd = open_eventfd() - wrap_eventfd = open_eventfd() - - proc_kwargs = { - 'pass_fds': (write_eventfd, wrap_eventfd) - } - - shm_key = 'test_ring_buff' - - common_kwargs = { - 'msg_amount': msg_amount, - 'shm_key': shm_key, - 'write_eventfd': write_eventfd, - 'wrap_eventfd': wrap_eventfd, - 'buf_size': buf_size - } - async def main(): - async with tractor.open_nursery() as an: - send_p = await an.start_actor( - 'ring_sender', - enable_modules=[__name__], - proc_kwargs=proc_kwargs - ) - recv_p = await an.start_actor( - 'ring_receiver', - enable_modules=[__name__], - proc_kwargs=proc_kwargs - ) - async with ( - send_p.open_context( - child_write_shm, - rand_min=rand_min, - rand_max=rand_max, - **common_kwargs - ) as (sctx, total_bytes), - recv_p.open_context( - child_read_shm, - **common_kwargs, - total_bytes=total_bytes, - ) as (sctx, _sent), - ): - await recv_p.result() + with open_ringbuf( + 'test_ringbuf', + buf_size=buf_size + ) as token: + proc_kwargs = { + 'pass_fds': (token.write_eventfd, token.wrap_eventfd) + } - await send_p.cancel_actor() - await recv_p.cancel_actor() + common_kwargs = { + 'msg_amount': msg_amount, + 'token': token, + 'buf_size': buf_size + } + async with tractor.open_nursery() as an: + send_p = await an.start_actor( + 'ring_sender', + enable_modules=[__name__], + proc_kwargs=proc_kwargs + ) + recv_p = await an.start_actor( + 'ring_receiver', + enable_modules=[__name__], + proc_kwargs=proc_kwargs + ) + async with ( + send_p.open_context( + child_write_shm, + rand_min=rand_min, + rand_max=rand_max, + **common_kwargs + ) as (sctx, total_bytes), + recv_p.open_context( + child_read_shm, + **common_kwargs, + total_bytes=total_bytes, + ) as (sctx, _sent), + ): + await recv_p.result() + + await send_p.cancel_actor() + await recv_p.cancel_actor() trio.run(main) @@ -154,55 +141,35 @@ def test_ringbuf( @tractor.context async def child_blocked_receiver( ctx: tractor.Context, - shm_key: str, - write_eventfd: int, - wrap_eventfd: int, - flags: int = 0 + token: RBToken ): - async with RingBuffReceiver( - shm_key, - write_eventfd, - wrap_eventfd, - flags=flags - ) as receiver: + async with RingBuffReceiver(token) as receiver: await ctx.started() await receiver.receive_some() def test_ring_reader_cancel(): - write_eventfd = open_eventfd() - wrap_eventfd = open_eventfd() - - proc_kwargs = { - 'pass_fds': (write_eventfd, wrap_eventfd) - } - - shm_key = 'test_ring_cancel' - async def main(): - async with ( - tractor.open_nursery() as an, - RingBuffSender( - shm_key, - write_eventfd, - wrap_eventfd, - ) as _sender, - ): - recv_p = await an.start_actor( - 'ring_blocked_receiver', - enable_modules=[__name__], - proc_kwargs=proc_kwargs - ) + with open_ringbuf('test_ring_cancel') as token: async with ( - recv_p.open_context( - child_blocked_receiver, - write_eventfd=write_eventfd, - wrap_eventfd=wrap_eventfd, - shm_key=shm_key, - ) as (sctx, _sent), + tractor.open_nursery() as an, + RingBuffSender(token) as _sender, ): - await trio.sleep(1) - await an.cancel() + recv_p = await an.start_actor( + 'ring_blocked_receiver', + enable_modules=[__name__], + proc_kwargs={ + 'pass_fds': (token.write_eventfd, token.wrap_eventfd) + } + ) + async with ( + recv_p.open_context( + child_blocked_receiver, + token=token + ) as (sctx, _sent), + ): + await trio.sleep(1) + await an.cancel() with pytest.raises(tractor._exceptions.ContextCancelled): diff --git a/tractor/ipc/__init__.py b/tractor/ipc/__init__.py index 59fc1e16..ec6217a1 100644 --- a/tractor/ipc/__init__.py +++ b/tractor/ipc/__init__.py @@ -36,6 +36,8 @@ if platform.system() == 'Linux': ) from ._ringbuf import ( + RBToken as RBToken, RingBuffSender as RingBuffSender, - RingBuffReceiver as RingBuffReceiver + RingBuffReceiver as RingBuffReceiver, + open_ringbuf ) diff --git a/tractor/ipc/_ringbuf.py b/tractor/ipc/_ringbuf.py index 0a4f3819..c590e8e2 100644 --- a/tractor/ipc/_ringbuf.py +++ b/tractor/ipc/_ringbuf.py @@ -17,16 +17,65 @@ IPC Reliable RingBuffer implementation ''' +from __future__ import annotations +from contextlib import contextmanager as cm from multiprocessing.shared_memory import SharedMemory import trio +from msgspec import ( + Struct, + to_builtins +) from ._linux import ( EFD_NONBLOCK, + open_eventfd, EventFD ) +class RBToken(Struct, frozen=True): + ''' + RingBuffer token contains necesary info to open the two + eventfds and the shared memory + + ''' + shm_name: str + write_eventfd: int + wrap_eventfd: int + + 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) + + +@cm +def open_ringbuf( + shm_name: str, + buf_size: int = 10 * 1024, + write_efd_flags: int = 0, + wrap_efd_flags: int = 0 +) -> RBToken: + shm = SharedMemory( + name=shm_name, + size=buf_size, + create=True + ) + token = RBToken( + shm_name=shm_name, + write_eventfd=open_eventfd(flags=write_efd_flags), + wrap_eventfd=open_eventfd(flags=wrap_efd_flags) + ) + yield token + shm.close() + + class RingBuffSender(trio.abc.SendStream): ''' IPC Reliable Ring Buffer sender side implementation @@ -34,28 +83,22 @@ class RingBuffSender(trio.abc.SendStream): `eventfd(2)` is used for wrap around sync, and also to signal writes to the reader. - TODO: if blocked on wrap around event wait it will not respond - to signals, fix soon TM ''' - def __init__( self, - shm_key: str, - write_eventfd: int, - wrap_eventfd: int, + token: RBToken, start_ptr: int = 0, buf_size: int = 10 * 1024, - unlink_on_exit: bool = True ): + token = RBToken.from_msg(token) self._shm = SharedMemory( - name=shm_key, + name=token.shm_name, size=buf_size, - create=True + create=False ) - self._write_event = EventFD(write_eventfd, 'w') - self._wrap_event = EventFD(wrap_eventfd, 'r') + self._write_event = EventFD(token.write_eventfd, 'w') + self._wrap_event = EventFD(token.wrap_eventfd, 'r') self._ptr = start_ptr - self.unlink_on_exit = unlink_on_exit @property def key(self) -> str: @@ -104,11 +147,7 @@ class RingBuffSender(trio.abc.SendStream): async def aclose(self): self._write_event.close() self._wrap_event.close() - if self.unlink_on_exit: - self._shm.unlink() - - else: - self._shm.close() + self._shm.close() async def __aenter__(self): self._write_event.open() @@ -123,29 +162,22 @@ class RingBuffReceiver(trio.abc.ReceiveStream): `eventfd(2)` is used for wrap around sync, and also to signal writes to the reader. - Unless eventfd(2) object is opened with EFD_NONBLOCK flag, - calls to `receive_some` will block the signal handling, - on the main thread, for now solution is using polling, - working on a way to unblock GIL during read(2) to allow - signal processing on the main thread. ''' - def __init__( self, - shm_key: str, - write_eventfd: int, - wrap_eventfd: int, + token: RBToken, start_ptr: int = 0, buf_size: int = 10 * 1024, flags: int = 0 ): + token = RBToken.from_msg(token) self._shm = SharedMemory( - name=shm_key, + name=token.shm_name, size=buf_size, create=False ) - self._write_event = EventFD(write_eventfd, 'w') - self._wrap_event = EventFD(wrap_eventfd, 'r') + self._write_event = EventFD(token.write_eventfd, 'w') + self._wrap_event = EventFD(token.wrap_eventfd, 'r') self._ptr = start_ptr self._flags = flags -- 2.34.1 From 6b155849b7f852b610a15eb993db7e4c075d9f83 Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Fri, 14 Mar 2025 00:25:10 -0300 Subject: [PATCH 11/19] Add buf_size to RBToken and add sender cancel test, move disable_mantracker to its own _mp_bs module --- tests/test_ringbuf.py | 53 +++++++++++++++++++++++++++++++---------- tractor/ipc/_mp_bs.py | 45 ++++++++++++++++++++++++++++++++++ tractor/ipc/_ringbuf.py | 29 +++++++++++++--------- tractor/ipc/_shm.py | 29 +--------------------- 4 files changed, 105 insertions(+), 51 deletions(-) create mode 100644 tractor/ipc/_mp_bs.py diff --git a/tests/test_ringbuf.py b/tests/test_ringbuf.py index 64fb37e9..28af7b83 100644 --- a/tests/test_ringbuf.py +++ b/tests/test_ringbuf.py @@ -17,16 +17,12 @@ async def child_read_shm( ctx: tractor.Context, msg_amount: int, token: RBToken, - buf_size: int, total_bytes: int, ) -> None: recvd_bytes = 0 await ctx.started() start_ts = time.time() - async with RingBuffReceiver( - token, - buf_size=buf_size, - ) as receiver: + async with RingBuffReceiver(token) as receiver: while recvd_bytes < total_bytes: msg = await receiver.receive_some() recvd_bytes += len(msg) @@ -51,7 +47,6 @@ async def child_write_shm( rand_min: int, rand_max: int, token: RBToken, - buf_size: int, ) -> None: msgs, total_bytes = generate_sample_messages( msg_amount, @@ -59,10 +54,7 @@ async def child_write_shm( rand_max=rand_max, ) await ctx.started(total_bytes) - async with RingBuffSender( - token, - buf_size=buf_size - ) as sender: + async with RingBuffSender(token) as sender: for msg in msgs: await sender.send_all(msg) @@ -103,7 +95,6 @@ def test_ringbuf( common_kwargs = { 'msg_amount': msg_amount, 'token': token, - 'buf_size': buf_size } async with tractor.open_nursery() as an: send_p = await an.start_actor( @@ -150,7 +141,7 @@ async def child_blocked_receiver( def test_ring_reader_cancel(): async def main(): - with open_ringbuf('test_ring_cancel') as token: + with open_ringbuf('test_ring_cancel_reader') as token: async with ( tractor.open_nursery() as an, RingBuffSender(token) as _sender, @@ -174,3 +165,41 @@ def test_ring_reader_cancel(): with pytest.raises(tractor._exceptions.ContextCancelled): trio.run(main) + + +@tractor.context +async def child_blocked_sender( + ctx: tractor.Context, + token: RBToken +): + async with RingBuffSender(token) as sender: + await ctx.started() + await sender.send_all(b'this will wrap') + + +def test_ring_sender_cancel(): + async def main(): + with open_ringbuf( + 'test_ring_cancel_sender', + buf_size=1 + ) as token: + async with tractor.open_nursery() as an: + recv_p = await an.start_actor( + 'ring_blocked_sender', + enable_modules=[__name__], + proc_kwargs={ + 'pass_fds': (token.write_eventfd, token.wrap_eventfd) + } + ) + async with ( + recv_p.open_context( + child_blocked_sender, + token=token + ) as (sctx, _sent), + ): + await trio.sleep(1) + await an.cancel() + + + with pytest.raises(tractor._exceptions.ContextCancelled): + trio.run(main) diff --git a/tractor/ipc/_mp_bs.py b/tractor/ipc/_mp_bs.py new file mode 100644 index 00000000..e51aa9ae --- /dev/null +++ b/tractor/ipc/_mp_bs.py @@ -0,0 +1,45 @@ +# 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 . +''' +Utils to tame mp non-SC madeness + +''' +def disable_mantracker(): + ''' + Disable all ``multiprocessing``` "resource tracking" machinery since + it's an absolute multi-threaded mess of non-SC madness. + + ''' + from multiprocessing import resource_tracker as mantracker + + # Tell the "resource tracker" thing to fuck off. + class ManTracker(mantracker.ResourceTracker): + def register(self, name, rtype): + pass + + def unregister(self, name, rtype): + pass + + def ensure_running(self): + pass + + # "know your land and know your prey" + # https://www.dailymotion.com/video/x6ozzco + mantracker._resource_tracker = ManTracker() + mantracker.register = mantracker._resource_tracker.register + mantracker.ensure_running = mantracker._resource_tracker.ensure_running + mantracker.unregister = mantracker._resource_tracker.unregister + mantracker.getfd = mantracker._resource_tracker.getfd diff --git a/tractor/ipc/_ringbuf.py b/tractor/ipc/_ringbuf.py index c590e8e2..6337eea1 100644 --- a/tractor/ipc/_ringbuf.py +++ b/tractor/ipc/_ringbuf.py @@ -32,6 +32,10 @@ from ._linux import ( open_eventfd, EventFD ) +from ._mp_bs import disable_mantracker + + +disable_mantracker() class RBToken(Struct, frozen=True): @@ -43,6 +47,7 @@ class RBToken(Struct, frozen=True): shm_name: str write_eventfd: int wrap_eventfd: int + buf_size: int def as_msg(self): return to_builtins(self) @@ -67,13 +72,17 @@ def open_ringbuf( size=buf_size, create=True ) - token = RBToken( - shm_name=shm_name, - write_eventfd=open_eventfd(flags=write_efd_flags), - wrap_eventfd=open_eventfd(flags=wrap_efd_flags) - ) - yield token - shm.close() + try: + token = RBToken( + shm_name=shm_name, + write_eventfd=open_eventfd(flags=write_efd_flags), + wrap_eventfd=open_eventfd(flags=wrap_efd_flags), + buf_size=buf_size + ) + yield token + + finally: + shm.unlink() class RingBuffSender(trio.abc.SendStream): @@ -88,12 +97,11 @@ class RingBuffSender(trio.abc.SendStream): self, token: RBToken, start_ptr: int = 0, - buf_size: int = 10 * 1024, ): token = RBToken.from_msg(token) self._shm = SharedMemory( name=token.shm_name, - size=buf_size, + size=token.buf_size, create=False ) self._write_event = EventFD(token.write_eventfd, 'w') @@ -167,13 +175,12 @@ class RingBuffReceiver(trio.abc.ReceiveStream): self, token: RBToken, start_ptr: int = 0, - buf_size: int = 10 * 1024, flags: int = 0 ): token = RBToken.from_msg(token) self._shm = SharedMemory( name=token.shm_name, - size=buf_size, + size=token.buf_size, create=False ) self._write_event = EventFD(token.write_eventfd, 'w') diff --git a/tractor/ipc/_shm.py b/tractor/ipc/_shm.py index c101af30..9868ac73 100644 --- a/tractor/ipc/_shm.py +++ b/tractor/ipc/_shm.py @@ -38,6 +38,7 @@ from msgspec import ( ) import tractor +from tractor.ipc._mp_bs import disable_mantracker from tractor.log import get_logger @@ -57,34 +58,6 @@ except ImportError: log = get_logger(__name__) -def disable_mantracker(): - ''' - Disable all ``multiprocessing``` "resource tracking" machinery since - it's an absolute multi-threaded mess of non-SC madness. - - ''' - from multiprocessing import resource_tracker as mantracker - - # Tell the "resource tracker" thing to fuck off. - class ManTracker(mantracker.ResourceTracker): - def register(self, name, rtype): - pass - - def unregister(self, name, rtype): - pass - - def ensure_running(self): - pass - - # "know your land and know your prey" - # https://www.dailymotion.com/video/x6ozzco - mantracker._resource_tracker = ManTracker() - mantracker.register = mantracker._resource_tracker.register - mantracker.ensure_running = mantracker._resource_tracker.ensure_running - mantracker.unregister = mantracker._resource_tracker.unregister - mantracker.getfd = mantracker._resource_tracker.getfd - - disable_mantracker() -- 2.34.1 From 9b2161506fc1b15e79072f4253e72637e71bacc5 Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Sun, 16 Mar 2025 14:14:32 -0300 Subject: [PATCH 12/19] Break out transport protocol and tcp specifics into their own submodules under tractor.ipc --- tractor/ipc/__init__.py | 11 +- tractor/ipc/_chan.py | 434 +------------------------------------- tractor/ipc/_tcp.py | 406 +++++++++++++++++++++++++++++++++++ tractor/ipc/_transport.py | 74 +++++++ 4 files changed, 498 insertions(+), 427 deletions(-) create mode 100644 tractor/ipc/_tcp.py create mode 100644 tractor/ipc/_transport.py diff --git a/tractor/ipc/__init__.py b/tractor/ipc/__init__.py index ec6217a1..cd16a139 100644 --- a/tractor/ipc/__init__.py +++ b/tractor/ipc/__init__.py @@ -17,9 +17,16 @@ import platform +from ._transport import MsgTransport as MsgTransport + +from ._tcp import ( + get_stream_addrs as get_stream_addrs, + MsgpackTCPStream as MsgpackTCPStream +) + from ._chan import ( _connect_chan as _connect_chan, - MsgTransport as MsgTransport, + get_msg_transport as get_msg_transport, Channel as Channel ) @@ -39,5 +46,5 @@ if platform.system() == 'Linux': RBToken as RBToken, RingBuffSender as RingBuffSender, RingBuffReceiver as RingBuffReceiver, - open_ringbuf + open_ringbuf as open_ringbuf ) diff --git a/tractor/ipc/_chan.py b/tractor/ipc/_chan.py index 83186147..1b6ba29f 100644 --- a/tractor/ipc/_chan.py +++ b/tractor/ipc/_chan.py @@ -19,455 +19,39 @@ Inter-process comms abstractions """ from __future__ import annotations -from collections.abc import ( - AsyncGenerator, - AsyncIterator, -) +from collections.abc import AsyncGenerator from contextlib import ( asynccontextmanager as acm, contextmanager as cm, ) import platform from pprint import pformat -import struct import typing from typing import ( Any, - Callable, - runtime_checkable, - Protocol, - Type, - TypeVar, + Type ) -import msgspec -from tricycle import BufferedReceiveStream import trio +from tractor.ipc._transport import MsgTransport +from tractor.ipc._tcp import ( + MsgpackTCPStream, + get_stream_addrs +) from tractor.log import get_logger from tractor._exceptions import ( MsgTypeError, pack_from_raise, - TransportClosed, - _mk_send_mte, - _mk_recv_mte, -) -from tractor.msg import ( - _ctxvar_MsgCodec, - # _codec, XXX see `self._codec` sanity/debug checks - MsgCodec, - types as msgtypes, - pretty_struct, ) +from tractor.msg import MsgCodec + log = get_logger(__name__) _is_windows = platform.system() == 'Windows' -def get_stream_addrs( - stream: trio.SocketStream -) -> tuple[ - tuple[str, int], # local - tuple[str, int], # remote -]: - ''' - Return the `trio` streaming transport prot's socket-addrs for - both the local and remote sides as a pair. - - ''' - # rn, should both be IP sockets - lsockname = stream.socket.getsockname() - rsockname = stream.socket.getpeername() - return ( - tuple(lsockname[:2]), - tuple(rsockname[:2]), - ) - - -# from tractor.msg.types import MsgType -# ?TODO? this should be our `Union[*msgtypes.__spec__]` alias now right..? -# => BLEH, except can't bc prots must inherit typevar or param-spec -# vars.. -MsgType = TypeVar('MsgType') - - -# TODO: break up this mod into a subpkg so we can start adding new -# backends and move this type stuff into a dedicated file.. Bo -# -@runtime_checkable -class MsgTransport(Protocol[MsgType]): -# -# ^-TODO-^ consider using a generic def and indexing with our -# eventual msg definition/types? -# - https://docs.python.org/3/library/typing.html#typing.Protocol - - stream: trio.SocketStream - drained: list[MsgType] - - def __init__(self, stream: trio.SocketStream) -> None: - ... - - # XXX: should this instead be called `.sendall()`? - async def send(self, msg: MsgType) -> None: - ... - - async def recv(self) -> MsgType: - ... - - def __aiter__(self) -> MsgType: - ... - - def connected(self) -> bool: - ... - - # defining this sync otherwise it causes a mypy error because it - # can't figure out it's a generator i guess?..? - def drain(self) -> AsyncIterator[dict]: - ... - - @property - def laddr(self) -> tuple[str, int]: - ... - - @property - def raddr(self) -> tuple[str, int]: - ... - - -# TODO: typing oddity.. not sure why we have to inherit here, but it -# seems to be an issue with `get_msg_transport()` returning -# a `Type[Protocol]`; probably should make a `mypy` issue? -class MsgpackTCPStream(MsgTransport): - ''' - A ``trio.SocketStream`` delivering ``msgpack`` formatted data - using the ``msgspec`` codec lib. - - ''' - layer_key: int = 4 - name_key: str = 'tcp' - - # TODO: better naming for this? - # -[ ] check how libp2p does naming for such things? - codec_key: str = 'msgpack' - - def __init__( - self, - stream: trio.SocketStream, - prefix_size: int = 4, - - # XXX optionally provided codec pair for `msgspec`: - # https://jcristharif.com/msgspec/extending.html#mapping-to-from-native-types - # - # TODO: define this as a `Codec` struct which can be - # overriden dynamically by the application/runtime? - codec: tuple[ - Callable[[Any], Any]|None, # coder - Callable[[type, Any], Any]|None, # decoder - ]|None = None, - - ) -> None: - - self.stream = stream - assert self.stream.socket - - # should both be IP sockets - self._laddr, self._raddr = get_stream_addrs(stream) - - # create read loop instance - self._aiter_pkts = self._iter_packets() - self._send_lock = trio.StrictFIFOLock() - - # public i guess? - self.drained: list[dict] = [] - - self.recv_stream = BufferedReceiveStream( - transport_stream=stream - ) - self.prefix_size = prefix_size - - # allow for custom IPC msg interchange format - # dynamic override Bo - self._task = trio.lowlevel.current_task() - - # XXX for ctxvar debug only! - # self._codec: MsgCodec = ( - # codec - # or - # _codec._ctxvar_MsgCodec.get() - # ) - - async def _iter_packets(self) -> AsyncGenerator[dict, None]: - ''' - Yield `bytes`-blob decoded packets from the underlying TCP - stream using the current task's `MsgCodec`. - - This is a streaming routine implemented as an async generator - func (which was the original design, but could be changed?) - and is allocated by a `.__call__()` inside `.__init__()` where - it is assigned to the `._aiter_pkts` attr. - - ''' - decodes_failed: int = 0 - - while True: - try: - header: bytes = await self.recv_stream.receive_exactly(4) - except ( - ValueError, - ConnectionResetError, - - # not sure entirely why we need this but without it we - # seem to be getting racy failures here on - # arbiter/registry name subs.. - trio.BrokenResourceError, - - ) as trans_err: - - loglevel = 'transport' - match trans_err: - # case ( - # ConnectionResetError() - # ): - # loglevel = 'transport' - - # peer actor (graceful??) TCP EOF but `tricycle` - # seems to raise a 0-bytes-read? - case ValueError() if ( - 'unclean EOF' in trans_err.args[0] - ): - pass - - # peer actor (task) prolly shutdown quickly due - # to cancellation - case trio.BrokenResourceError() if ( - 'Connection reset by peer' in trans_err.args[0] - ): - pass - - # unless the disconnect condition falls under "a - # normal operation breakage" we usualy console warn - # about it. - case _: - loglevel: str = 'warning' - - - raise TransportClosed( - message=( - f'IPC transport already closed by peer\n' - f'x]> {type(trans_err)}\n' - f' |_{self}\n' - ), - loglevel=loglevel, - ) from trans_err - - # XXX definitely can happen if transport is closed - # manually by another `trio.lowlevel.Task` in the - # same actor; we use this in some simulated fault - # testing for ex, but generally should never happen - # under normal operation! - # - # NOTE: as such we always re-raise this error from the - # RPC msg loop! - except trio.ClosedResourceError as closure_err: - raise TransportClosed( - message=( - f'IPC transport already manually closed locally?\n' - f'x]> {type(closure_err)} \n' - f' |_{self}\n' - ), - loglevel='error', - raise_on_report=( - closure_err.args[0] == 'another task closed this fd' - or - closure_err.args[0] in ['another task closed this fd'] - ), - ) from closure_err - - # graceful TCP EOF disconnect - if header == b'': - raise TransportClosed( - message=( - f'IPC transport already gracefully closed\n' - f']>\n' - f' |_{self}\n' - ), - loglevel='transport', - # cause=??? # handy or no? - ) - - size: int - size, = struct.unpack(" None: - ''' - Send a msgpack encoded py-object-blob-as-msg over TCP. - - If `strict_types == True` then a `MsgTypeError` will be raised on any - invalid msg type - - ''' - __tracebackhide__: bool = hide_tb - - # XXX see `trio._sync.AsyncContextManagerMixin` for details - # on the `.acquire()`/`.release()` sequencing.. - async with self._send_lock: - - # NOTE: lookup the `trio.Task.context`'s var for - # the current `MsgCodec`. - codec: MsgCodec = _ctxvar_MsgCodec.get() - - # XXX for ctxvar debug only! - # if self._codec.pld_spec != codec.pld_spec: - # self._codec = codec - # log.runtime( - # f'Using new codec in {self}.send()\n' - # f'codec: {self._codec}\n\n' - # f'msg: {msg}\n' - # ) - - if type(msg) not in msgtypes.__msg_types__: - if strict_types: - raise _mk_send_mte( - msg, - codec=codec, - ) - else: - log.warning( - 'Sending non-`Msg`-spec msg?\n\n' - f'{msg}\n' - ) - - try: - bytes_data: bytes = codec.encode(msg) - except TypeError as _err: - typerr = _err - msgtyperr: MsgTypeError = _mk_send_mte( - msg, - codec=codec, - message=( - f'IPC-msg-spec violation in\n\n' - f'{pretty_struct.Struct.pformat(msg)}' - ), - src_type_error=typerr, - ) - raise msgtyperr from typerr - - # supposedly the fastest says, - # https://stackoverflow.com/a/54027962 - size: bytes = struct.pack(" - # except BaseException as _err: - # err = _err - # if not isinstance(err, MsgTypeError): - # __tracebackhide__: bool = False - # raise - - @property - def laddr(self) -> tuple[str, int]: - return self._laddr - - @property - def raddr(self) -> tuple[str, int]: - return self._raddr - - async def recv(self) -> Any: - return await self._aiter_pkts.asend(None) - - async def drain(self) -> AsyncIterator[dict]: - ''' - Drain the stream's remaining messages sent from - the far end until the connection is closed by - the peer. - - ''' - try: - async for msg in self._iter_packets(): - self.drained.append(msg) - except TransportClosed: - for msg in self.drained: - yield msg - - def __aiter__(self): - return self._aiter_pkts - - def connected(self) -> bool: - return self.stream.socket.fileno() != -1 - - def get_msg_transport( key: tuple[str, str], diff --git a/tractor/ipc/_tcp.py b/tractor/ipc/_tcp.py new file mode 100644 index 00000000..03185f82 --- /dev/null +++ b/tractor/ipc/_tcp.py @@ -0,0 +1,406 @@ +# 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 . +''' +TCP implementation of tractor.ipc._transport.MsgTransport protocol + +''' +from __future__ import annotations +from collections.abc import ( + AsyncGenerator, + AsyncIterator, +) +import struct +from typing import ( + Any, + Callable, + Type, +) + +import msgspec +from tricycle import BufferedReceiveStream +import trio + +from tractor.log import get_logger +from tractor._exceptions import ( + MsgTypeError, + TransportClosed, + _mk_send_mte, + _mk_recv_mte, +) +from tractor.msg import ( + _ctxvar_MsgCodec, + # _codec, XXX see `self._codec` sanity/debug checks + MsgCodec, + types as msgtypes, + pretty_struct, +) +from tractor.ipc import MsgTransport + + +log = get_logger(__name__) + + +def get_stream_addrs( + stream: trio.SocketStream +) -> tuple[ + tuple[str, int], # local + tuple[str, int], # remote +]: + ''' + Return the `trio` streaming transport prot's socket-addrs for + both the local and remote sides as a pair. + + ''' + # rn, should both be IP sockets + lsockname = stream.socket.getsockname() + rsockname = stream.socket.getpeername() + return ( + tuple(lsockname[:2]), + tuple(rsockname[:2]), + ) + + +# TODO: typing oddity.. not sure why we have to inherit here, but it +# seems to be an issue with `get_msg_transport()` returning +# a `Type[Protocol]`; probably should make a `mypy` issue? +class MsgpackTCPStream(MsgTransport): + ''' + A ``trio.SocketStream`` delivering ``msgpack`` formatted data + using the ``msgspec`` codec lib. + + ''' + layer_key: int = 4 + name_key: str = 'tcp' + + # TODO: better naming for this? + # -[ ] check how libp2p does naming for such things? + codec_key: str = 'msgpack' + + def __init__( + self, + stream: trio.SocketStream, + prefix_size: int = 4, + + # XXX optionally provided codec pair for `msgspec`: + # https://jcristharif.com/msgspec/extending.html#mapping-to-from-native-types + # + # TODO: define this as a `Codec` struct which can be + # overriden dynamically by the application/runtime? + codec: tuple[ + Callable[[Any], Any]|None, # coder + Callable[[type, Any], Any]|None, # decoder + ]|None = None, + + ) -> None: + + self.stream = stream + assert self.stream.socket + + # should both be IP sockets + self._laddr, self._raddr = get_stream_addrs(stream) + + # create read loop instance + self._aiter_pkts = self._iter_packets() + self._send_lock = trio.StrictFIFOLock() + + # public i guess? + self.drained: list[dict] = [] + + self.recv_stream = BufferedReceiveStream( + transport_stream=stream + ) + self.prefix_size = prefix_size + + # allow for custom IPC msg interchange format + # dynamic override Bo + self._task = trio.lowlevel.current_task() + + # XXX for ctxvar debug only! + # self._codec: MsgCodec = ( + # codec + # or + # _codec._ctxvar_MsgCodec.get() + # ) + + async def _iter_packets(self) -> AsyncGenerator[dict, None]: + ''' + Yield `bytes`-blob decoded packets from the underlying TCP + stream using the current task's `MsgCodec`. + + This is a streaming routine implemented as an async generator + func (which was the original design, but could be changed?) + and is allocated by a `.__call__()` inside `.__init__()` where + it is assigned to the `._aiter_pkts` attr. + + ''' + decodes_failed: int = 0 + + while True: + try: + header: bytes = await self.recv_stream.receive_exactly(4) + except ( + ValueError, + ConnectionResetError, + + # not sure entirely why we need this but without it we + # seem to be getting racy failures here on + # arbiter/registry name subs.. + trio.BrokenResourceError, + + ) as trans_err: + + loglevel = 'transport' + match trans_err: + # case ( + # ConnectionResetError() + # ): + # loglevel = 'transport' + + # peer actor (graceful??) TCP EOF but `tricycle` + # seems to raise a 0-bytes-read? + case ValueError() if ( + 'unclean EOF' in trans_err.args[0] + ): + pass + + # peer actor (task) prolly shutdown quickly due + # to cancellation + case trio.BrokenResourceError() if ( + 'Connection reset by peer' in trans_err.args[0] + ): + pass + + # unless the disconnect condition falls under "a + # normal operation breakage" we usualy console warn + # about it. + case _: + loglevel: str = 'warning' + + + raise TransportClosed( + message=( + f'IPC transport already closed by peer\n' + f'x)> {type(trans_err)}\n' + f' |_{self}\n' + ), + loglevel=loglevel, + ) from trans_err + + # XXX definitely can happen if transport is closed + # manually by another `trio.lowlevel.Task` in the + # same actor; we use this in some simulated fault + # testing for ex, but generally should never happen + # under normal operation! + # + # NOTE: as such we always re-raise this error from the + # RPC msg loop! + except trio.ClosedResourceError as closure_err: + raise TransportClosed( + message=( + f'IPC transport already manually closed locally?\n' + f'x)> {type(closure_err)} \n' + f' |_{self}\n' + ), + loglevel='error', + raise_on_report=( + closure_err.args[0] == 'another task closed this fd' + or + closure_err.args[0] in ['another task closed this fd'] + ), + ) from closure_err + + # graceful TCP EOF disconnect + if header == b'': + raise TransportClosed( + message=( + f'IPC transport already gracefully closed\n' + f')>\n' + f'|_{self}\n' + ), + loglevel='transport', + # cause=??? # handy or no? + ) + + size: int + size, = struct.unpack(" None: + ''' + Send a msgpack encoded py-object-blob-as-msg over TCP. + + If `strict_types == True` then a `MsgTypeError` will be raised on any + invalid msg type + + ''' + __tracebackhide__: bool = hide_tb + + # XXX see `trio._sync.AsyncContextManagerMixin` for details + # on the `.acquire()`/`.release()` sequencing.. + async with self._send_lock: + + # NOTE: lookup the `trio.Task.context`'s var for + # the current `MsgCodec`. + codec: MsgCodec = _ctxvar_MsgCodec.get() + + # XXX for ctxvar debug only! + # if self._codec.pld_spec != codec.pld_spec: + # self._codec = codec + # log.runtime( + # f'Using new codec in {self}.send()\n' + # f'codec: {self._codec}\n\n' + # f'msg: {msg}\n' + # ) + + if type(msg) not in msgtypes.__msg_types__: + if strict_types: + raise _mk_send_mte( + msg, + codec=codec, + ) + else: + log.warning( + 'Sending non-`Msg`-spec msg?\n\n' + f'{msg}\n' + ) + + try: + bytes_data: bytes = codec.encode(msg) + except TypeError as _err: + typerr = _err + msgtyperr: MsgTypeError = _mk_send_mte( + msg, + codec=codec, + message=( + f'IPC-msg-spec violation in\n\n' + f'{pretty_struct.Struct.pformat(msg)}' + ), + src_type_error=typerr, + ) + raise msgtyperr from typerr + + # supposedly the fastest says, + # https://stackoverflow.com/a/54027962 + size: bytes = struct.pack(" + # except BaseException as _err: + # err = _err + # if not isinstance(err, MsgTypeError): + # __tracebackhide__: bool = False + # raise + + @property + def laddr(self) -> tuple[str, int]: + return self._laddr + + @property + def raddr(self) -> tuple[str, int]: + return self._raddr + + async def recv(self) -> Any: + return await self._aiter_pkts.asend(None) + + async def drain(self) -> AsyncIterator[dict]: + ''' + Drain the stream's remaining messages sent from + the far end until the connection is closed by + the peer. + + ''' + try: + async for msg in self._iter_packets(): + self.drained.append(msg) + except TransportClosed: + for msg in self.drained: + yield msg + + def __aiter__(self): + return self._aiter_pkts + + def connected(self) -> bool: + return self.stream.socket.fileno() != -1 diff --git a/tractor/ipc/_transport.py b/tractor/ipc/_transport.py new file mode 100644 index 00000000..24e03a90 --- /dev/null +++ b/tractor/ipc/_transport.py @@ -0,0 +1,74 @@ +# 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 . +''' +typing.Protocol based generic msg API, implement this class to add backends for +tractor.ipc.Channel + +''' +import trio +from typing import ( + runtime_checkable, + Protocol, + TypeVar, +) +from collections.abc import AsyncIterator + + +# from tractor.msg.types import MsgType +# ?TODO? this should be our `Union[*msgtypes.__spec__]` alias now right..? +# => BLEH, except can't bc prots must inherit typevar or param-spec +# vars.. +MsgType = TypeVar('MsgType') + + +@runtime_checkable +class MsgTransport(Protocol[MsgType]): +# +# ^-TODO-^ consider using a generic def and indexing with our +# eventual msg definition/types? +# - https://docs.python.org/3/library/typing.html#typing.Protocol + + stream: trio.SocketStream + drained: list[MsgType] + + def __init__(self, stream: trio.SocketStream) -> None: + ... + + # XXX: should this instead be called `.sendall()`? + async def send(self, msg: MsgType) -> None: + ... + + async def recv(self) -> MsgType: + ... + + def __aiter__(self) -> MsgType: + ... + + def connected(self) -> bool: + ... + + # defining this sync otherwise it causes a mypy error because it + # can't figure out it's a generator i guess?..? + def drain(self) -> AsyncIterator[dict]: + ... + + @property + def laddr(self) -> tuple[str, int]: + ... + + @property + def raddr(self) -> tuple[str, int]: + ... -- 2.34.1 From ba353bf46f64ea5e3126d59c083d2fad6ecf429c Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Sun, 16 Mar 2025 17:50:13 -0300 Subject: [PATCH 13/19] Better encapsulate RingBuff ctx managment methods and support non ipc usage Add trio.StrictFIFOLock on sender.send_all Support max_bytes argument on receive_some, keep track of write_ptr on receiver Add max_bytes receive test test_ringbuf_max_bytes Add docstrings to all ringbuf tests Remove EFD_NONBLOCK support, not necesary anymore since we can use abandon_on_cancel=True on trio.to_thread.run_sync Close eventfd's after usage on open_ringbuf --- tests/test_ringbuf.py | 54 ++++++++++++ tractor/ipc/_ringbuf.py | 180 ++++++++++++++++++++++------------------ 2 files changed, 153 insertions(+), 81 deletions(-) diff --git a/tests/test_ringbuf.py b/tests/test_ringbuf.py index 28af7b83..52cf0836 100644 --- a/tests/test_ringbuf.py +++ b/tests/test_ringbuf.py @@ -58,6 +58,8 @@ async def child_write_shm( for msg in msgs: await sender.send_all(msg) + print('writer exit') + @pytest.mark.parametrize( 'msg_amount,rand_min,rand_max,buf_size', @@ -83,6 +85,15 @@ def test_ringbuf( rand_max: int, buf_size: int ): + ''' + - Open a new ring buf on root actor + - Create a sender subactor and generate {msg_amount} messages + optionally with a random amount of bytes at the end of each, + return total_bytes on `ctx.started`, then send all messages + - Create a receiver subactor and receive until total_bytes are + read, print simple perf stats. + + ''' async def main(): with open_ringbuf( 'test_ringbuf', @@ -140,6 +151,11 @@ async def child_blocked_receiver( def test_ring_reader_cancel(): + ''' + Test that a receiver blocked on eventfd(2) read responds to + cancellation. + + ''' async def main(): with open_ringbuf('test_ring_cancel_reader') as token: async with ( @@ -178,6 +194,11 @@ async def child_blocked_sender( def test_ring_sender_cancel(): + ''' + Test that a sender blocked on eventfd(2) read responds to + cancellation. + + ''' async def main(): with open_ringbuf( 'test_ring_cancel_sender', @@ -203,3 +224,36 @@ def test_ring_sender_cancel(): with pytest.raises(tractor._exceptions.ContextCancelled): trio.run(main) + + +def test_ringbuf_max_bytes(): + ''' + Test that RingBuffReceiver.receive_some's max_bytes optional + argument works correctly, send a msg of size 100, then + force receive of messages with max_bytes == 1, wait until + 100 of these messages are received, then compare join of + msgs with original message + + ''' + msg = b''.join(str(i % 10).encode() for i in range(100)) + msgs = [] + + async def main(): + with open_ringbuf( + 'test_ringbuf_max_bytes', + buf_size=10 + ) as token: + async with ( + trio.open_nursery() as n, + RingBuffSender(token, is_ipc=False) as sender, + RingBuffReceiver(token, is_ipc=False) as receiver + ): + n.start_soon(sender.send_all, msg) + while len(msgs) < len(msg): + msg_part = await receiver.receive_some(max_bytes=1) + msg_part = bytes(msg_part) + assert len(msg_part) == 1 + msgs.append(msg_part) + + trio.run(main) + assert msg == b''.join(msgs) diff --git a/tractor/ipc/_ringbuf.py b/tractor/ipc/_ringbuf.py index 6337eea1..304454ed 100644 --- a/tractor/ipc/_ringbuf.py +++ b/tractor/ipc/_ringbuf.py @@ -28,11 +28,15 @@ from msgspec import ( ) from ._linux import ( - EFD_NONBLOCK, open_eventfd, + close_eventfd, EventFD ) from ._mp_bs import disable_mantracker +from tractor.log import get_logger + + +log = get_logger(__name__) disable_mantracker() @@ -64,8 +68,6 @@ class RBToken(Struct, frozen=True): def open_ringbuf( shm_name: str, buf_size: int = 10 * 1024, - write_efd_flags: int = 0, - wrap_efd_flags: int = 0 ) -> RBToken: shm = SharedMemory( name=shm_name, @@ -75,16 +77,21 @@ def open_ringbuf( try: token = RBToken( shm_name=shm_name, - write_eventfd=open_eventfd(flags=write_efd_flags), - wrap_eventfd=open_eventfd(flags=wrap_efd_flags), + write_eventfd=open_eventfd(), + wrap_eventfd=open_eventfd(), buf_size=buf_size ) yield token + close_eventfd(token.write_eventfd) + close_eventfd(token.wrap_eventfd) finally: shm.unlink() +Buffer = bytes | bytearray | memoryview + + class RingBuffSender(trio.abc.SendStream): ''' IPC Reliable Ring Buffer sender side implementation @@ -97,24 +104,26 @@ class RingBuffSender(trio.abc.SendStream): self, token: RBToken, start_ptr: int = 0, + is_ipc: bool = True ): - token = RBToken.from_msg(token) - self._shm = SharedMemory( - name=token.shm_name, - size=token.buf_size, - create=False - ) - self._write_event = EventFD(token.write_eventfd, 'w') - self._wrap_event = EventFD(token.wrap_eventfd, 'r') + self._token = RBToken.from_msg(token) + self._shm: SharedMemory | None = None + self._write_event = EventFD(self._token.write_eventfd, 'w') + self._wrap_event = EventFD(self._token.wrap_eventfd, 'r') self._ptr = start_ptr + self._is_ipc = is_ipc + self._send_lock = trio.StrictFIFOLock() + @property - def key(self) -> str: + 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._shm.size + return self._token.buf_size @property def ptr(self) -> int: @@ -128,38 +137,48 @@ class RingBuffSender(trio.abc.SendStream): def wrap_fd(self) -> int: return self._wrap_event.fd - async def send_all(self, data: bytes | bytearray | memoryview): - # 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._wrap_event.read() + async def send_all(self, data: Buffer): + async with self._send_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._wrap_event.read() - # wrap around and trim already written bytes - self._ptr = 0 - data = data[remaining:] - target_ptr = self._ptr + len(data) + # 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 + # 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): raise NotImplementedError - async def aclose(self): - self._write_event.close() - self._wrap_event.close() - self._shm.close() - - async def __aenter__(self): + def open(self): + self._shm = SharedMemory( + name=self._token.shm_name, + size=self._token.buf_size, + create=False + ) self._write_event.open() self._wrap_event.open() + + async def aclose(self): + if self._is_ipc: + self._write_event.close() + self._wrap_event.close() + self._shm.close() + + async def __aenter__(self): + self.open() return self @@ -175,26 +194,25 @@ class RingBuffReceiver(trio.abc.ReceiveStream): self, token: RBToken, start_ptr: int = 0, - flags: int = 0 + is_ipc: bool = True ): - token = RBToken.from_msg(token) - self._shm = SharedMemory( - name=token.shm_name, - size=token.buf_size, - create=False - ) - self._write_event = EventFD(token.write_eventfd, 'w') - self._wrap_event = EventFD(token.wrap_eventfd, 'r') + self._token = RBToken.from_msg(token) + self._shm: SharedMemory | None = None + self._write_event = EventFD(self._token.write_eventfd, 'w') + self._wrap_event = EventFD(self._token.wrap_eventfd, 'r') self._ptr = start_ptr - self._flags = flags + self._write_ptr = start_ptr + self._is_ipc = is_ipc @property - def key(self) -> str: + 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._shm.size + return self._token.buf_size @property def ptr(self) -> int: @@ -208,46 +226,46 @@ class RingBuffReceiver(trio.abc.ReceiveStream): def wrap_fd(self) -> int: return self._wrap_event.fd - async def receive_some( - self, - max_bytes: int | None = None, - nb_timeout: float = 0.1 - ) -> memoryview: - # if non blocking eventfd enabled, do polling - # until next write, this allows signal handling - if self._flags | EFD_NONBLOCK: - delta = None - while delta is None: - try: - delta = await self._write_event.read() - - except OSError as e: - if e.errno == 'EAGAIN': - continue - - raise e - - else: + async def receive_some(self, max_bytes: int | None = None) -> memoryview: + delta = self._write_ptr - self._ptr + if delta == 0: delta = await self._write_event.read() + self._write_ptr += delta + + if isinstance(max_bytes, int): + if max_bytes == 0: + raise ValueError('if set, max_bytes must be > 0') + delta = min(delta, max_bytes) + + target_ptr = self._ptr + delta # fetch next segment and advance ptr - next_ptr = self._ptr + delta - segment = self._shm.buf[self._ptr:next_ptr] - self._ptr = next_ptr + segment = self._shm.buf[self._ptr:target_ptr] + self._ptr = target_ptr - if self.ptr == self.size: + 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 aclose(self): - self._write_event.close() - self._wrap_event.close() - self._shm.close() - - async def __aenter__(self): + def open(self): + self._shm = SharedMemory( + name=self._token.shm_name, + size=self._token.buf_size, + create=False + ) self._write_event.open() self._wrap_event.open() + + async def aclose(self): + if self._is_ipc: + self._write_event.close() + self._wrap_event.close() + self._shm.close() + + async def __aenter__(self): + self.open() return self -- 2.34.1 From be818a720ab3978f6b48ebde48a58f09fbfe92e7 Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Sun, 16 Mar 2025 23:57:26 -0300 Subject: [PATCH 14/19] Switch `tractor.ipc.MsgTransport.stream` type to `trio.abc.Stream` Add EOF signaling mechanism Support proper `receive_some` end of stream semantics Add StapledStream non-ipc test Create MsgpackRBStream similar to MsgpackTCPStream for buffered whole-msg reads Add EventFD.read cancellation on EventFD.close mechanism using cancel scope Add test for eventfd cancellation Improve and add docstrings --- tests/test_eventfd.py | 32 +++ tests/test_ringbuf.py | 197 ++++++++++++--- tractor/_testing/samples.py | 4 + tractor/ipc/__init__.py | 9 +- tractor/ipc/_linux.py | 36 ++- tractor/ipc/_ringbuf.py | 463 +++++++++++++++++++++++++++++++++--- tractor/ipc/_tcp.py | 1 - tractor/ipc/_transport.py | 4 +- 8 files changed, 671 insertions(+), 75 deletions(-) create mode 100644 tests/test_eventfd.py diff --git a/tests/test_eventfd.py b/tests/test_eventfd.py new file mode 100644 index 00000000..3d757169 --- /dev/null +++ b/tests/test_eventfd.py @@ -0,0 +1,32 @@ +import trio +import pytest +from tractor.ipc import ( + open_eventfd, + EFDReadCancelled, + EventFD +) + + +def test_eventfd_read_cancellation(): + ''' + Ensure EventFD.read raises EFDReadCancelled if EventFD.close() + is called. + + ''' + fd = open_eventfd() + + async def _read(event: EventFD): + with pytest.raises(EFDReadCancelled): + await event.read() + + async def main(): + async with trio.open_nursery() as n: + with ( + EventFD(fd, 'w') as event, + trio.fail_after(3) + ): + n.start_soon(_read, event) + await trio.sleep(0.2) + event.close() + + trio.run(main) diff --git a/tests/test_ringbuf.py b/tests/test_ringbuf.py index 52cf0836..3aa32cf8 100644 --- a/tests/test_ringbuf.py +++ b/tests/test_ringbuf.py @@ -5,11 +5,16 @@ import pytest import tractor from tractor.ipc import ( open_ringbuf, + attach_to_ringbuf_receiver, + attach_to_ringbuf_sender, + attach_to_ringbuf_pair, + attach_to_ringbuf_stream, RBToken, - RingBuffSender, - RingBuffReceiver ) -from tractor._testing.samples import generate_sample_messages +from tractor._testing.samples import ( + generate_single_byte_msgs, + generate_sample_messages +) @tractor.context @@ -17,20 +22,14 @@ async def child_read_shm( ctx: tractor.Context, msg_amount: int, token: RBToken, - total_bytes: int, ) -> None: recvd_bytes = 0 await ctx.started() start_ts = time.time() - async with RingBuffReceiver(token) as receiver: - while recvd_bytes < total_bytes: - msg = await receiver.receive_some() + async with attach_to_ringbuf_receiver(token) as receiver: + async for msg in receiver: recvd_bytes += len(msg) - # make sure we dont hold any memoryviews - # before the ctx manager aclose() - msg = None - end_ts = time.time() elapsed = end_ts - start_ts elapsed_ms = int(elapsed * 1000) @@ -38,6 +37,7 @@ async def child_read_shm( print(f'\n\telapsed ms: {elapsed_ms}') print(f'\tmsg/sec: {int(msg_amount / elapsed):,}') print(f'\tbytes/sec: {int(recvd_bytes / elapsed):,}') + print(f'\treceived bytes: {recvd_bytes}') @tractor.context @@ -54,7 +54,7 @@ async def child_write_shm( rand_max=rand_max, ) await ctx.started(total_bytes) - async with RingBuffSender(token) as sender: + async with attach_to_ringbuf_sender(token, cleanup=False) as sender: for msg in msgs: await sender.send_all(msg) @@ -99,14 +99,8 @@ def test_ringbuf( 'test_ringbuf', buf_size=buf_size ) as token: - proc_kwargs = { - 'pass_fds': (token.write_eventfd, token.wrap_eventfd) - } + proc_kwargs = {'pass_fds': token.fds} - common_kwargs = { - 'msg_amount': msg_amount, - 'token': token, - } async with tractor.open_nursery() as an: send_p = await an.start_actor( 'ring_sender', @@ -121,14 +115,15 @@ def test_ringbuf( async with ( send_p.open_context( child_write_shm, + token=token, + msg_amount=msg_amount, rand_min=rand_min, rand_max=rand_max, - **common_kwargs ) as (sctx, total_bytes), recv_p.open_context( child_read_shm, - **common_kwargs, - total_bytes=total_bytes, + token=token, + msg_amount=msg_amount ) as (sctx, _sent), ): await recv_p.result() @@ -145,7 +140,7 @@ async def child_blocked_receiver( ctx: tractor.Context, token: RBToken ): - async with RingBuffReceiver(token) as receiver: + async with attach_to_ringbuf_receiver(token) as receiver: await ctx.started() await receiver.receive_some() @@ -160,13 +155,13 @@ def test_ring_reader_cancel(): with open_ringbuf('test_ring_cancel_reader') as token: async with ( tractor.open_nursery() as an, - RingBuffSender(token) as _sender, + attach_to_ringbuf_sender(token) as _sender, ): recv_p = await an.start_actor( 'ring_blocked_receiver', enable_modules=[__name__], proc_kwargs={ - 'pass_fds': (token.write_eventfd, token.wrap_eventfd) + 'pass_fds': token.fds } ) async with ( @@ -188,7 +183,7 @@ async def child_blocked_sender( ctx: tractor.Context, token: RBToken ): - async with RingBuffSender(token) as sender: + async with attach_to_ringbuf_sender(token) as sender: await ctx.started() await sender.send_all(b'this will wrap') @@ -209,7 +204,7 @@ def test_ring_sender_cancel(): 'ring_blocked_sender', enable_modules=[__name__], proc_kwargs={ - 'pass_fds': (token.write_eventfd, token.wrap_eventfd) + 'pass_fds': token.fds } ) async with ( @@ -235,7 +230,7 @@ def test_ringbuf_max_bytes(): msgs with original message ''' - msg = b''.join(str(i % 10).encode() for i in range(100)) + msg = generate_single_byte_msgs(100) msgs = [] async def main(): @@ -245,15 +240,153 @@ def test_ringbuf_max_bytes(): ) as token: async with ( trio.open_nursery() as n, - RingBuffSender(token, is_ipc=False) as sender, - RingBuffReceiver(token, is_ipc=False) as receiver + attach_to_ringbuf_sender(token, cleanup=False) as sender, + attach_to_ringbuf_receiver(token, cleanup=False) as receiver ): - n.start_soon(sender.send_all, msg) + async def _send_and_close(): + await sender.send_all(msg) + await sender.aclose() + + n.start_soon(_send_and_close) while len(msgs) < len(msg): msg_part = await receiver.receive_some(max_bytes=1) - msg_part = bytes(msg_part) assert len(msg_part) == 1 msgs.append(msg_part) trio.run(main) assert msg == b''.join(msgs) + + +def test_stapled_ringbuf(): + ''' + Open two ringbufs and give tokens to tasks (swap them such that in/out tokens + are inversed on each task) which will open the streams and use trio.StapledStream + to have a single bidirectional stream. + + Then take turns to send and receive messages. + + ''' + msg = generate_single_byte_msgs(100) + pair_0_msgs = [] + pair_1_msgs = [] + + pair_0_done = trio.Event() + pair_1_done = trio.Event() + + async def pair_0(token_in: RBToken, token_out: RBToken): + async with attach_to_ringbuf_pair( + token_in, + token_out, + cleanup_in=False, + cleanup_out=False + ) as stream: + # first turn to send + await stream.send_all(msg) + + # second turn to receive + while len(pair_0_msgs) != len(msg): + _msg = await stream.receive_some(max_bytes=1) + pair_0_msgs.append(_msg) + + pair_0_done.set() + await pair_1_done.wait() + + + async def pair_1(token_in: RBToken, token_out: RBToken): + async with attach_to_ringbuf_pair( + token_in, + token_out, + cleanup_in=False, + cleanup_out=False + ) as stream: + # first turn to receive + while len(pair_1_msgs) != len(msg): + _msg = await stream.receive_some(max_bytes=1) + pair_1_msgs.append(_msg) + + # second turn to send + await stream.send_all(msg) + + pair_1_done.set() + await pair_0_done.wait() + + + async def main(): + with tractor.ipc.open_ringbuf_pair( + 'test_stapled_ringbuf' + ) as (token_0, token_1): + async with trio.open_nursery() as n: + n.start_soon(pair_0, token_0, token_1) + n.start_soon(pair_1, token_1, token_0) + + + trio.run(main) + + assert msg == b''.join(pair_0_msgs) + assert msg == b''.join(pair_1_msgs) + + +@tractor.context +async def child_transport_sender( + ctx: tractor.Context, + msg_amount_min: int, + msg_amount_max: int, + token_in: RBToken, + token_out: RBToken +): + import random + msgs, _total_bytes = generate_sample_messages( + random.randint(msg_amount_min, msg_amount_max), + rand_min=256, + rand_max=1024, + ) + async with attach_to_ringbuf_stream( + token_in, + token_out + ) as transport: + await ctx.started(msgs) + + for msg in msgs: + await transport.send(msg) + + await transport.recv() + + +def test_ringbuf_transport(): + + msg_amount_min = 100 + msg_amount_max = 1000 + + async def main(): + with tractor.ipc.open_ringbuf_pair( + 'test_ringbuf_transport' + ) as (token_0, token_1): + async with ( + attach_to_ringbuf_stream(token_0, token_1) as transport, + tractor.open_nursery() as an + ): + recv_p = await an.start_actor( + 'test_ringbuf_transport_sender', + enable_modules=[__name__], + proc_kwargs={ + 'pass_fds': token_0.fds + token_1.fds + } + ) + async with ( + recv_p.open_context( + child_transport_sender, + msg_amount_min=msg_amount_min, + msg_amount_max=msg_amount_max, + token_in=token_1, + token_out=token_0 + ) as (ctx, msgs), + ): + recv_msgs = [] + while len(recv_msgs) < len(msgs): + recv_msgs.append(await transport.recv()) + + await transport.send(b'end') + await recv_p.cancel_actor() + assert recv_msgs == msgs + + trio.run(main) diff --git a/tractor/_testing/samples.py b/tractor/_testing/samples.py index a87a22c4..1454ee3d 100644 --- a/tractor/_testing/samples.py +++ b/tractor/_testing/samples.py @@ -2,6 +2,10 @@ import os import random +def generate_single_byte_msgs(amount: int) -> bytes: + return b''.join(str(i % 10).encode() for i in range(amount)) + + def generate_sample_messages( amount: int, rand_min: int = 0, diff --git a/tractor/ipc/__init__.py b/tractor/ipc/__init__.py index cd16a139..329dca1e 100644 --- a/tractor/ipc/__init__.py +++ b/tractor/ipc/__init__.py @@ -39,12 +39,19 @@ if platform.system() == 'Linux': write_eventfd as write_eventfd, read_eventfd as read_eventfd, close_eventfd as close_eventfd, + EFDReadCancelled as EFDReadCancelled, EventFD as EventFD, ) from ._ringbuf import ( RBToken as RBToken, + open_ringbuf as open_ringbuf, RingBuffSender as RingBuffSender, RingBuffReceiver as RingBuffReceiver, - open_ringbuf as open_ringbuf + open_ringbuf_pair as open_ringbuf_pair, + attach_to_ringbuf_receiver as attach_to_ringbuf_receiver, + attach_to_ringbuf_sender as attach_to_ringbuf_sender, + attach_to_ringbuf_pair as attach_to_ringbuf_pair, + attach_to_ringbuf_stream as attach_to_ringbuf_stream, + MsgpackRBStream as MsgpackRBStream ) diff --git a/tractor/ipc/_linux.py b/tractor/ipc/_linux.py index 88d80d1c..afce6bff 100644 --- a/tractor/ipc/_linux.py +++ b/tractor/ipc/_linux.py @@ -108,6 +108,10 @@ def close_eventfd(fd: int) -> int: raise OSError(errno.errorcode[ffi.errno], 'close failed') +class EFDReadCancelled(Exception): + ... + + class EventFD: ''' Use a previously opened eventfd(2), meant to be used in @@ -124,6 +128,7 @@ class EventFD: self._fd: int = fd self._omode: str = omode self._fobj = None + self._cscope: trio.CancelScope | None = None @property def fd(self) -> int | None: @@ -133,17 +138,38 @@ class EventFD: return write_eventfd(self._fd, value) async def read(self) -> int: - return await trio.to_thread.run_sync( - read_eventfd, self._fd, - abandon_on_cancel=True - ) + ''' + Async wrapper for `read_eventfd(self.fd)` + + `trio.to_thread.run_sync` is used, need to use a `trio.CancelScope` + in order to make it cancellable when `self.close()` is called. + + ''' + self._cscope = trio.CancelScope() + with self._cscope: + return await trio.to_thread.run_sync( + read_eventfd, self._fd, + abandon_on_cancel=True + ) + + if self._cscope.cancelled_caught: + raise EFDReadCancelled + + self._cscope = None def open(self): self._fobj = os.fdopen(self._fd, self._omode) def close(self): if self._fobj: - self._fobj.close() + try: + self._fobj.close() + + except OSError: + ... + + if self._cscope: + self._cscope.cancel() def __enter__(self): self.open() diff --git a/tractor/ipc/_ringbuf.py b/tractor/ipc/_ringbuf.py index 304454ed..7529c942 100644 --- a/tractor/ipc/_ringbuf.py +++ b/tractor/ipc/_ringbuf.py @@ -18,10 +18,22 @@ IPC Reliable RingBuffer implementation ''' from __future__ import annotations -from contextlib import contextmanager as cm +import struct +from collections.abc import ( + AsyncGenerator, + AsyncIterator +) +from contextlib import ( + contextmanager as cm, + asynccontextmanager as acm +) +from typing import ( + Any +) from multiprocessing.shared_memory import SharedMemory import trio +from tricycle import BufferedReceiveStream from msgspec import ( Struct, to_builtins @@ -30,10 +42,16 @@ from msgspec import ( from ._linux import ( open_eventfd, close_eventfd, + EFDReadCancelled, EventFD ) from ._mp_bs import disable_mantracker from tractor.log import get_logger +from tractor._exceptions import ( + TransportClosed, + InternalError +) +from tractor.ipc import MsgTransport log = get_logger(__name__) @@ -41,16 +59,21 @@ log = get_logger(__name__) disable_mantracker() +_DEFAULT_RB_SIZE = 10 * 1024 + class RBToken(Struct, frozen=True): ''' - RingBuffer token contains necesary info to open the two + RingBuffer token contains necesary info to open the three eventfds and the shared memory ''' shm_name: str - write_eventfd: int - wrap_eventfd: int + + 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 def as_msg(self): @@ -63,12 +86,29 @@ class RBToken(Struct, frozen=True): return RBToken(**msg) + @property + def fds(self) -> tuple[int, int, int]: + ''' + Useful for `pass_fds` params + + ''' + return ( + self.write_eventfd, + self.wrap_eventfd, + self.eof_eventfd + ) + @cm def open_ringbuf( shm_name: str, - buf_size: int = 10 * 1024, + buf_size: int = _DEFAULT_RB_SIZE, ) -> RBToken: + ''' + Handle resources for a ringbuf (shm, eventfd), yield `RBToken` to + be used with `attach_to_ringbuf_sender` and `attach_to_ringbuf_receiver` + + ''' shm = SharedMemory( name=shm_name, size=buf_size, @@ -79,11 +119,27 @@ def open_ringbuf( shm_name=shm_name, write_eventfd=open_eventfd(), wrap_eventfd=open_eventfd(), + eof_eventfd=open_eventfd(), buf_size=buf_size ) yield token - close_eventfd(token.write_eventfd) - close_eventfd(token.wrap_eventfd) + try: + close_eventfd(token.write_eventfd) + + except OSError: + ... + + try: + close_eventfd(token.wrap_eventfd) + + except OSError: + ... + + try: + close_eventfd(token.eof_eventfd) + + except OSError: + ... finally: shm.unlink() @@ -91,28 +147,36 @@ def open_ringbuf( 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. + +''' + class RingBuffSender(trio.abc.SendStream): ''' - IPC Reliable Ring Buffer sender side implementation + Ring Buffer sender side implementation - `eventfd(2)` is used for wrap around sync, and also to signal - writes to the reader. + Do not use directly! manage with `attach_to_ringbuf_sender` + after having opened a ringbuf context with `open_ringbuf`. ''' def __init__( self, token: RBToken, - start_ptr: int = 0, - is_ipc: bool = True + cleanup: bool = False ): self._token = RBToken.from_msg(token) self._shm: SharedMemory | None = None self._write_event = EventFD(self._token.write_eventfd, 'w') self._wrap_event = EventFD(self._token.wrap_eventfd, 'r') - self._ptr = start_ptr + self._eof_event = EventFD(self._token.eof_eventfd, 'w') + self._ptr = 0 - self._is_ipc = is_ipc + self._cleanup = cleanup self._send_lock = trio.StrictFIFOLock() @property @@ -170,13 +234,21 @@ class RingBuffSender(trio.abc.SendStream): ) self._write_event.open() self._wrap_event.open() + self._eof_event.open() - async def aclose(self): - if self._is_ipc: + 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() + async def aclose(self): + self.close() + async def __aenter__(self): self.open() return self @@ -184,25 +256,27 @@ class RingBuffSender(trio.abc.SendStream): class RingBuffReceiver(trio.abc.ReceiveStream): ''' - IPC Reliable Ring Buffer receiver side implementation + Ring Buffer receiver side implementation - `eventfd(2)` is used for wrap around sync, and also to signal - writes to the reader. + Do not use directly! manage with `attach_to_ringbuf_receiver` + after having opened a ringbuf context with `open_ringbuf`. ''' def __init__( self, token: RBToken, - start_ptr: int = 0, - is_ipc: bool = True + cleanup: bool = True, ): self._token = RBToken.from_msg(token) self._shm: SharedMemory | None = None self._write_event = EventFD(self._token.write_eventfd, 'w') self._wrap_event = EventFD(self._token.wrap_eventfd, 'r') - self._ptr = start_ptr - self._write_ptr = start_ptr - self._is_ipc = is_ipc + self._eof_event = EventFD(self._token.eof_eventfd, 'r') + self._ptr: int = 0 + self._write_ptr: int = 0 + self._end_ptr: int = -1 + + self._cleanup: bool = cleanup @property def name(self) -> str: @@ -226,21 +300,71 @@ class RingBuffReceiver(trio.abc.ReceiveStream): def wrap_fd(self) -> int: return self._wrap_event.fd - async def receive_some(self, max_bytes: int | None = None) -> memoryview: + 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() + self._write_event.close() + + except EFDReadCancelled: + ... + + async def receive_some(self, max_bytes: int | None = None) -> bytes: + ''' + Receive up to `max_bytes`, if no `max_bytes` is provided + a reasonable default is used. + + ''' + if max_bytes is None: + max_bytes: int = _DEFAULT_RB_SIZE + + if max_bytes < 1: + raise ValueError("max_bytes must be >= 1") + + # delta is remaining bytes we havent read delta = self._write_ptr - self._ptr if delta == 0: - delta = await self._write_event.read() - self._write_ptr += delta + # we have read all we can, see if new data is available + if self._end_ptr < 0: + # if we havent been signaled about EOF yet + try: + delta = await self._write_event.read() + self._write_ptr += delta - if isinstance(max_bytes, int): - if max_bytes == 0: - raise ValueError('if set, max_bytes must be > 0') - delta = min(delta, max_bytes) + except EFDReadCancelled: + # while waiting for new data `self._write_event` was closed + # this means writer signaled EOF + if self._end_ptr > 0: + # final self._write_ptr modification and recalculate delta + self._write_ptr = self._end_ptr + delta = self._end_ptr - self._ptr + + else: + # shouldnt happen cause self._eof_monitor_task always sets + # self._end_ptr before closing self._write_event + raise InternalError( + 'self._write_event.read cancelled but self._end_ptr is not set' + ) + + else: + # no more bytes to read and self._end_ptr set, EOF reached + return b'' + + # dont overflow caller + delta = min(delta, max_bytes) target_ptr = self._ptr + delta # fetch next segment and advance ptr - segment = self._shm.buf[self._ptr:target_ptr] + segment = bytes(self._shm.buf[self._ptr:target_ptr]) self._ptr = target_ptr if self._ptr == self.size: @@ -259,13 +383,284 @@ class RingBuffReceiver(trio.abc.ReceiveStream): ) self._write_event.open() self._wrap_event.open() + self._eof_event.open() - async def aclose(self): - if self._is_ipc: + def close(self): + if self._cleanup: self._write_event.close() self._wrap_event.close() + self._eof_event.close() self._shm.close() + async def aclose(self): + self.close() + async def __aenter__(self): self.open() return self + + +@acm +async def attach_to_ringbuf_receiver( + token: RBToken, + cleanup: bool = True +): + ''' + Instantiate a RingBuffReceiver from a previously opened + RBToken. + + Launches `receiver._eof_monitor_task` in a `trio.Nursery`. + ''' + async with ( + trio.open_nursery() as n, + RingBuffReceiver( + token, + cleanup=cleanup + ) as receiver + ): + n.start_soon(receiver._eof_monitor_task) + yield receiver + +@acm +async def attach_to_ringbuf_sender( + token: RBToken, + cleanup: bool = True +): + ''' + Instantiate a RingBuffSender from a previously opened + RBToken. + + ''' + async with RingBuffSender( + token, + cleanup=cleanup + ) as sender: + yield sender + + +@cm +def open_ringbuf_pair( + name: str, + buf_size: int = _DEFAULT_RB_SIZE +): + ''' + Handle resources for a ringbuf pair to be used for + bidirectional messaging. + + ''' + with ( + open_ringbuf( + name + '.pair0', + buf_size=buf_size + ) as token_0, + + open_ringbuf( + name + '.pair1', + buf_size=buf_size + ) as token_1 + ): + yield token_0, token_1 + + +@acm +async def attach_to_ringbuf_pair( + token_in: RBToken, + token_out: RBToken, + cleanup_in: bool = True, + cleanup_out: bool = True +): + ''' + Instantiate a trio.StapledStream from a previously opened + ringbuf pair. + + ''' + async with ( + attach_to_ringbuf_receiver( + token_in, + cleanup=cleanup_in + ) as receiver, + attach_to_ringbuf_sender( + token_out, + cleanup=cleanup_out + ) as sender, + ): + yield trio.StapledStream(sender, receiver) + + +class MsgpackRBStream(MsgTransport): + + def __init__( + self, + stream: trio.StapledStream + ): + self.stream = stream + + # create read loop intance + self._aiter_pkts = self._iter_packets() + self._send_lock = trio.StrictFIFOLock() + + self.drained: list[dict] = [] + + self.recv_stream = BufferedReceiveStream( + transport_stream=stream + ) + + async def _iter_packets(self) -> AsyncGenerator[dict, None]: + ''' + Yield `bytes`-blob decoded packets from the underlying TCP + stream using the current task's `MsgCodec`. + + This is a streaming routine implemented as an async generator + func (which was the original design, but could be changed?) + and is allocated by a `.__call__()` inside `.__init__()` where + it is assigned to the `._aiter_pkts` attr. + + ''' + + while True: + try: + header: bytes = await self.recv_stream.receive_exactly(4) + except ( + ValueError, + ConnectionResetError, + + # not sure entirely why we need this but without it we + # seem to be getting racy failures here on + # arbiter/registry name subs.. + trio.BrokenResourceError, + + ) as trans_err: + + loglevel = 'transport' + match trans_err: + # case ( + # ConnectionResetError() + # ): + # loglevel = 'transport' + + # peer actor (graceful??) TCP EOF but `tricycle` + # seems to raise a 0-bytes-read? + case ValueError() if ( + 'unclean EOF' in trans_err.args[0] + ): + pass + + # peer actor (task) prolly shutdown quickly due + # to cancellation + case trio.BrokenResourceError() if ( + 'Connection reset by peer' in trans_err.args[0] + ): + pass + + # unless the disconnect condition falls under "a + # normal operation breakage" we usualy console warn + # about it. + case _: + loglevel: str = 'warning' + + + raise TransportClosed( + message=( + f'IPC transport already closed by peer\n' + f'x)> {type(trans_err)}\n' + f' |_{self}\n' + ), + loglevel=loglevel, + ) from trans_err + + # XXX definitely can happen if transport is closed + # manually by another `trio.lowlevel.Task` in the + # same actor; we use this in some simulated fault + # testing for ex, but generally should never happen + # under normal operation! + # + # NOTE: as such we always re-raise this error from the + # RPC msg loop! + except trio.ClosedResourceError as closure_err: + raise TransportClosed( + message=( + f'IPC transport already manually closed locally?\n' + f'x)> {type(closure_err)} \n' + f' |_{self}\n' + ), + loglevel='error', + raise_on_report=( + closure_err.args[0] == 'another task closed this fd' + or + closure_err.args[0] in ['another task closed this fd'] + ), + ) from closure_err + + # graceful EOF disconnect + if header == b'': + raise TransportClosed( + message=( + f'IPC transport already gracefully closed\n' + f')>\n' + f'|_{self}\n' + ), + loglevel='transport', + # cause=??? # handy or no? + ) + + size: int + size, = struct.unpack(" None: + ''' + Send a msgpack encoded py-object-blob-as-msg. + + ''' + async with self._send_lock: + size: bytes = struct.pack(" Any: + return await self._aiter_pkts.asend(None) + + async def drain(self) -> AsyncIterator[dict]: + ''' + Drain the stream's remaining messages sent from + the far end until the connection is closed by + the peer. + + ''' + try: + async for msg in self._iter_packets(): + self.drained.append(msg) + except TransportClosed: + for msg in self.drained: + yield msg + + def __aiter__(self): + return self._aiter_pkts + + +@acm +async def attach_to_ringbuf_stream( + token_in: RBToken, + token_out: RBToken, + cleanup_in: bool = True, + cleanup_out: bool = True +): + ''' + Wrap a ringbuf trio.StapledStream in a MsgpackRBStream + + ''' + async with attach_to_ringbuf_pair( + token_in, + token_out, + cleanup_in=cleanup_in, + cleanup_out=cleanup_out + ) as stream: + yield MsgpackRBStream(stream) diff --git a/tractor/ipc/_tcp.py b/tractor/ipc/_tcp.py index 03185f82..3ce0b4ea 100644 --- a/tractor/ipc/_tcp.py +++ b/tractor/ipc/_tcp.py @@ -26,7 +26,6 @@ import struct from typing import ( Any, Callable, - Type, ) import msgspec diff --git a/tractor/ipc/_transport.py b/tractor/ipc/_transport.py index 24e03a90..64453c89 100644 --- a/tractor/ipc/_transport.py +++ b/tractor/ipc/_transport.py @@ -41,10 +41,10 @@ class MsgTransport(Protocol[MsgType]): # eventual msg definition/types? # - https://docs.python.org/3/library/typing.html#typing.Protocol - stream: trio.SocketStream + stream: trio.abc.Stream drained: list[MsgType] - def __init__(self, stream: trio.SocketStream) -> None: + def __init__(self, stream: trio.abc.Stream) -> None: ... # XXX: should this instead be called `.sendall()`? -- 2.34.1 From 2a9a78651b9da60fe52a1b072557d0839fffeaec Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Tue, 18 Mar 2025 13:19:40 -0300 Subject: [PATCH 15/19] Improve test_ringbuf test, drop MsgTransport ring buf impl for now in favour of a trio.abc.Channel[bytes] impl, add docstrings --- tests/test_ringbuf.py | 87 ++++++++---- tractor/_testing/samples.py | 46 +++++- tractor/ipc/__init__.py | 6 +- tractor/ipc/_ringbuf.py | 269 ++++++++++++++---------------------- 4 files changed, 210 insertions(+), 198 deletions(-) diff --git a/tests/test_ringbuf.py b/tests/test_ringbuf.py index 3aa32cf8..f987d4c8 100644 --- a/tests/test_ringbuf.py +++ b/tests/test_ringbuf.py @@ -1,4 +1,5 @@ import time +import hashlib import trio import pytest @@ -7,8 +8,8 @@ from tractor.ipc import ( open_ringbuf, attach_to_ringbuf_receiver, attach_to_ringbuf_sender, - attach_to_ringbuf_pair, attach_to_ringbuf_stream, + attach_to_ringbuf_channel, RBToken, ) from tractor._testing.samples import ( @@ -22,12 +23,26 @@ async def child_read_shm( ctx: tractor.Context, msg_amount: int, token: RBToken, -) -> None: - recvd_bytes = 0 +) -> str: + ''' + Sub-actor used in `test_ringbuf`. + + Attach to a ringbuf and receive all messages until end of stream. + Keep track of how many bytes received and also calculate + sha256 of the whole byte stream. + + Calculate and print performance stats, finally return calculated + hash. + + ''' await ctx.started() + print('reader started') + recvd_bytes = 0 + recvd_hash = hashlib.sha256() start_ts = time.time() async with attach_to_ringbuf_receiver(token) as receiver: async for msg in receiver: + recvd_hash.update(msg) recvd_bytes += len(msg) end_ts = time.time() @@ -37,7 +52,9 @@ async def child_read_shm( print(f'\n\telapsed ms: {elapsed_ms}') print(f'\tmsg/sec: {int(msg_amount / elapsed):,}') print(f'\tbytes/sec: {int(recvd_bytes / elapsed):,}') - print(f'\treceived bytes: {recvd_bytes}') + print(f'\treceived bytes: {recvd_bytes:,}') + + return recvd_hash.hexdigest() @tractor.context @@ -48,12 +65,26 @@ async def child_write_shm( rand_max: int, token: RBToken, ) -> None: - msgs, total_bytes = generate_sample_messages( + ''' + Sub-actor used in `test_ringbuf` + + Generate `msg_amount` payloads with + `random.randint(rand_min, rand_max)` random bytes at the end, + Calculate sha256 hash and send it to parent on `ctx.started`. + + Attach to ringbuf and send all generated messages. + + ''' + msgs, _total_bytes = generate_sample_messages( msg_amount, rand_min=rand_min, rand_max=rand_max, ) - await ctx.started(total_bytes) + print('writer hashing payload...') + sent_hash = hashlib.sha256(b''.join(msgs)).hexdigest() + print('writer done hashing.') + await ctx.started(sent_hash) + print('writer started') async with attach_to_ringbuf_sender(token, cleanup=False) as sender: for msg in msgs: await sender.send_all(msg) @@ -87,11 +118,12 @@ def test_ringbuf( ): ''' - Open a new ring buf on root actor - - Create a sender subactor and generate {msg_amount} messages - optionally with a random amount of bytes at the end of each, - return total_bytes on `ctx.started`, then send all messages - - Create a receiver subactor and receive until total_bytes are - read, print simple perf stats. + - Open `child_write_shm` ctx in sub-actor which will generate a + random payload and send its hash on `ctx.started`, finally sending + the payload through the stream. + - Open `child_read_shm` ctx in sub-actor which will receive the + payload, calculate perf stats and return the hash. + - Compare both hashes ''' async def main(): @@ -119,14 +151,16 @@ def test_ringbuf( msg_amount=msg_amount, rand_min=rand_min, rand_max=rand_max, - ) as (sctx, total_bytes), + ) as (_sctx, sent_hash), recv_p.open_context( child_read_shm, token=token, msg_amount=msg_amount - ) as (sctx, _sent), + ) as (rctx, _sent), ): - await recv_p.result() + recvd_hash = await rctx.result() + + assert sent_hash == recvd_hash await send_p.cancel_actor() await recv_p.cancel_actor() @@ -274,7 +308,7 @@ def test_stapled_ringbuf(): pair_1_done = trio.Event() async def pair_0(token_in: RBToken, token_out: RBToken): - async with attach_to_ringbuf_pair( + async with attach_to_ringbuf_stream( token_in, token_out, cleanup_in=False, @@ -293,7 +327,7 @@ def test_stapled_ringbuf(): async def pair_1(token_in: RBToken, token_out: RBToken): - async with attach_to_ringbuf_pair( + async with attach_to_ringbuf_stream( token_in, token_out, cleanup_in=False, @@ -327,7 +361,7 @@ def test_stapled_ringbuf(): @tractor.context -async def child_transport_sender( +async def child_channel_sender( ctx: tractor.Context, msg_amount_min: int, msg_amount_max: int, @@ -340,19 +374,17 @@ async def child_transport_sender( rand_min=256, rand_max=1024, ) - async with attach_to_ringbuf_stream( + async with attach_to_ringbuf_channel( token_in, token_out - ) as transport: + ) as chan: await ctx.started(msgs) for msg in msgs: - await transport.send(msg) - - await transport.recv() + await chan.send(msg) -def test_ringbuf_transport(): +def test_ringbuf_channel(): msg_amount_min = 100 msg_amount_max = 1000 @@ -362,7 +394,7 @@ def test_ringbuf_transport(): 'test_ringbuf_transport' ) as (token_0, token_1): async with ( - attach_to_ringbuf_stream(token_0, token_1) as transport, + attach_to_ringbuf_channel(token_0, token_1) as chan, tractor.open_nursery() as an ): recv_p = await an.start_actor( @@ -374,7 +406,7 @@ def test_ringbuf_transport(): ) async with ( recv_p.open_context( - child_transport_sender, + child_channel_sender, msg_amount_min=msg_amount_min, msg_amount_max=msg_amount_max, token_in=token_1, @@ -382,10 +414,9 @@ def test_ringbuf_transport(): ) as (ctx, msgs), ): recv_msgs = [] - while len(recv_msgs) < len(msgs): - recv_msgs.append(await transport.recv()) + async for msg in chan: + recv_msgs.append(msg) - await transport.send(b'end') await recv_p.cancel_actor() assert recv_msgs == msgs diff --git a/tractor/_testing/samples.py b/tractor/_testing/samples.py index 1454ee3d..4249bae9 100644 --- a/tractor/_testing/samples.py +++ b/tractor/_testing/samples.py @@ -3,6 +3,18 @@ import random def generate_single_byte_msgs(amount: int) -> bytes: + ''' + Generate a byte instance of len `amount` with: + + ``` + byte_at_index(i) = (i % 10).encode() + ``` + + this results in constantly repeating sequences of: + + b'0123456789' + + ''' return b''.join(str(i % 10).encode() for i in range(amount)) @@ -10,15 +22,39 @@ def generate_sample_messages( amount: int, rand_min: int = 0, rand_max: int = 0, - silent: bool = False + silent: bool = False, ) -> tuple[list[bytes], int]: + ''' + Generate bytes msgs for tests. + Messages will have the following format: + + ``` + b'[{i:08}]' + os.urandom(random.randint(rand_min, rand_max)) + ``` + + so for message index 25: + + b'[00000025]' + random_bytes + + ''' msgs = [] size = 0 + log_interval = None if not silent: print(f'\ngenerating {amount} messages...') + # calculate an apropiate log interval based on + # max message size + max_msg_size = 10 + rand_max + + if max_msg_size <= 32 * 1024: + log_interval = 10_000 + + else: + log_interval = 1000 + for i in range(amount): msg = f'[{i:08}]'.encode('utf-8') @@ -30,7 +66,13 @@ def generate_sample_messages( msgs.append(msg) - if not silent and i and i % 10_000 == 0: + if ( + not silent + and + i > 0 + and + i % log_interval == 0 + ): print(f'{i} generated') if not silent: diff --git a/tractor/ipc/__init__.py b/tractor/ipc/__init__.py index 329dca1e..689fc44b 100644 --- a/tractor/ipc/__init__.py +++ b/tractor/ipc/__init__.py @@ -51,7 +51,9 @@ if platform.system() == 'Linux': open_ringbuf_pair as open_ringbuf_pair, attach_to_ringbuf_receiver as attach_to_ringbuf_receiver, attach_to_ringbuf_sender as attach_to_ringbuf_sender, - attach_to_ringbuf_pair as attach_to_ringbuf_pair, attach_to_ringbuf_stream as attach_to_ringbuf_stream, - MsgpackRBStream as MsgpackRBStream + RingBuffBytesSender as RingBuffBytesSender, + RingBuffBytesReceiver as RingBuffBytesReceiver, + RingBuffChannel as RingBuffChannel, + attach_to_ringbuf_channel as attach_to_ringbuf_channel ) diff --git a/tractor/ipc/_ringbuf.py b/tractor/ipc/_ringbuf.py index 7529c942..038d9e73 100644 --- a/tractor/ipc/_ringbuf.py +++ b/tractor/ipc/_ringbuf.py @@ -19,17 +19,10 @@ IPC Reliable RingBuffer implementation ''' from __future__ import annotations import struct -from collections.abc import ( - AsyncGenerator, - AsyncIterator -) from contextlib import ( contextmanager as cm, asynccontextmanager as acm ) -from typing import ( - Any -) from multiprocessing.shared_memory import SharedMemory import trio @@ -48,10 +41,8 @@ from ._linux import ( from ._mp_bs import disable_mantracker from tractor.log import get_logger from tractor._exceptions import ( - TransportClosed, InternalError ) -from tractor.ipc import MsgTransport log = get_logger(__name__) @@ -147,6 +138,7 @@ def open_ringbuf( Buffer = bytes | bytearray | memoryview + ''' IPC Reliable Ring Buffer @@ -406,7 +398,7 @@ async def attach_to_ringbuf_receiver( cleanup: bool = True ): ''' - Instantiate a RingBuffReceiver from a previously opened + Attach a RingBuffReceiver from a previously opened RBToken. Launches `receiver._eof_monitor_task` in a `trio.Nursery`. @@ -421,13 +413,14 @@ async def attach_to_ringbuf_receiver( n.start_soon(receiver._eof_monitor_task) yield receiver + @acm async def attach_to_ringbuf_sender( token: RBToken, cleanup: bool = True ): ''' - Instantiate a RingBuffSender from a previously opened + Attach a RingBuffSender from a previously opened RBToken. ''' @@ -463,14 +456,14 @@ def open_ringbuf_pair( @acm -async def attach_to_ringbuf_pair( +async def attach_to_ringbuf_stream( token_in: RBToken, token_out: RBToken, cleanup_in: bool = True, cleanup_out: bool = True ): ''' - Instantiate a trio.StapledStream from a previously opened + Attach a trio.StapledStream from a previously opened ringbuf pair. ''' @@ -487,180 +480,124 @@ async def attach_to_ringbuf_pair( yield trio.StapledStream(sender, receiver) -class MsgpackRBStream(MsgTransport): +class RingBuffBytesSender(trio.abc.SendChannel[bytes]): + ''' + In order to guarantee full messages are received, all bytes + sent by `RingBuffBytesSender` are preceded with a 4 byte header + which decodes into a uint32 indicating the actual size of the + next payload. + + ''' def __init__( self, - stream: trio.StapledStream + sender: RingBuffSender ): - self.stream = stream - - # create read loop intance - self._aiter_pkts = self._iter_packets() + self._sender = sender self._send_lock = trio.StrictFIFOLock() - self.drained: list[dict] = [] - - self.recv_stream = BufferedReceiveStream( - transport_stream=stream - ) - - async def _iter_packets(self) -> AsyncGenerator[dict, None]: - ''' - Yield `bytes`-blob decoded packets from the underlying TCP - stream using the current task's `MsgCodec`. - - This is a streaming routine implemented as an async generator - func (which was the original design, but could be changed?) - and is allocated by a `.__call__()` inside `.__init__()` where - it is assigned to the `._aiter_pkts` attr. - - ''' - - while True: - try: - header: bytes = await self.recv_stream.receive_exactly(4) - except ( - ValueError, - ConnectionResetError, - - # not sure entirely why we need this but without it we - # seem to be getting racy failures here on - # arbiter/registry name subs.. - trio.BrokenResourceError, - - ) as trans_err: - - loglevel = 'transport' - match trans_err: - # case ( - # ConnectionResetError() - # ): - # loglevel = 'transport' - - # peer actor (graceful??) TCP EOF but `tricycle` - # seems to raise a 0-bytes-read? - case ValueError() if ( - 'unclean EOF' in trans_err.args[0] - ): - pass - - # peer actor (task) prolly shutdown quickly due - # to cancellation - case trio.BrokenResourceError() if ( - 'Connection reset by peer' in trans_err.args[0] - ): - pass - - # unless the disconnect condition falls under "a - # normal operation breakage" we usualy console warn - # about it. - case _: - loglevel: str = 'warning' - - - raise TransportClosed( - message=( - f'IPC transport already closed by peer\n' - f'x)> {type(trans_err)}\n' - f' |_{self}\n' - ), - loglevel=loglevel, - ) from trans_err - - # XXX definitely can happen if transport is closed - # manually by another `trio.lowlevel.Task` in the - # same actor; we use this in some simulated fault - # testing for ex, but generally should never happen - # under normal operation! - # - # NOTE: as such we always re-raise this error from the - # RPC msg loop! - except trio.ClosedResourceError as closure_err: - raise TransportClosed( - message=( - f'IPC transport already manually closed locally?\n' - f'x)> {type(closure_err)} \n' - f' |_{self}\n' - ), - loglevel='error', - raise_on_report=( - closure_err.args[0] == 'another task closed this fd' - or - closure_err.args[0] in ['another task closed this fd'] - ), - ) from closure_err - - # graceful EOF disconnect - if header == b'': - raise TransportClosed( - message=( - f'IPC transport already gracefully closed\n' - f')>\n' - f'|_{self}\n' - ), - loglevel='transport', - # cause=??? # handy or no? - ) - - size: int - size, = struct.unpack(" None: - ''' - Send a msgpack encoded py-object-blob-as-msg. - - ''' + async def send(self, value: bytes) -> None: async with self._send_lock: - size: bytes = struct.pack(" Any: - return await self._aiter_pkts.asend(None) + async def aclose(self) -> None: + async with self._send_lock: + await self._sender.aclose() - async def drain(self) -> AsyncIterator[dict]: + +class RingBuffBytesReceiver(trio.abc.ReceiveChannel[bytes]): + ''' + See `RingBuffBytesSender` docstring. + + A `tricycle.BufferedReceiveStream` is used for the + `receive_exactly` API. + ''' + def __init__( + self, + receiver: RingBuffReceiver + ): + self._receiver = receiver + + async def _receive_exactly(self, num_bytes: int) -> bytes: ''' - Drain the stream's remaining messages sent from - the far end until the connection is closed by - the peer. + Fetch bytes from receiver until we read exactly `num_bytes` + or end of stream is signaled. ''' - try: - async for msg in self._iter_packets(): - self.drained.append(msg) - except TransportClosed: - for msg in self.drained: - yield msg + payload = b'' + while len(payload) < num_bytes: + remaining = num_bytes - len(payload) - def __aiter__(self): - return self._aiter_pkts + new_bytes = await self._receiver.receive_some( + max_bytes=remaining + ) + + if new_bytes == b'': + raise trio.EndOfChannel + + payload += new_bytes + + return payload + + async def receive(self) -> bytes: + header: bytes = await self._receive_exactly(4) + size: int + size, = struct.unpack(" None: + await self._receiver.aclose() + + +class RingBuffChannel(trio.abc.Channel[bytes]): + ''' + Combine `RingBuffBytesSender` and `RingBuffBytesReceiver` + in order to expose the bidirectional `trio.abc.Channel` API. + + ''' + def __init__( + self, + sender: RingBuffBytesSender, + receiver: RingBuffBytesReceiver + ): + self._sender = sender + self._receiver = receiver + + async def send(self, value: bytes): + await self._sender.send(value) + + 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_stream( +async def attach_to_ringbuf_channel( token_in: RBToken, token_out: RBToken, cleanup_in: bool = True, cleanup_out: bool = True ): ''' - Wrap a ringbuf trio.StapledStream in a MsgpackRBStream + Attach to an already opened ringbuf pair and return + a `RingBuffChannel`. ''' - async with attach_to_ringbuf_pair( - token_in, - token_out, - cleanup_in=cleanup_in, - cleanup_out=cleanup_out - ) as stream: - yield MsgpackRBStream(stream) + async with ( + attach_to_ringbuf_receiver( + token_in, + cleanup=cleanup_in + ) as receiver, + attach_to_ringbuf_sender( + token_out, + cleanup=cleanup_out + ) as sender, + ): + yield RingBuffChannel( + RingBuffBytesSender(sender), + RingBuffBytesReceiver(receiver) + ) -- 2.34.1 From be7fc89ae9c6a15a86eafb64e3e871ccb5c1e022 Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Tue, 18 Mar 2025 13:47:41 -0300 Subject: [PATCH 16/19] Add direct ctx managers for RB channels --- tractor/ipc/__init__.py | 4 +++- tractor/ipc/_ringbuf.py | 39 +++++++++++++++++++++++++++++++++------ 2 files changed, 36 insertions(+), 7 deletions(-) diff --git a/tractor/ipc/__init__.py b/tractor/ipc/__init__.py index 689fc44b..4f0cd2b4 100644 --- a/tractor/ipc/__init__.py +++ b/tractor/ipc/__init__.py @@ -55,5 +55,7 @@ if platform.system() == 'Linux': RingBuffBytesSender as RingBuffBytesSender, RingBuffBytesReceiver as RingBuffBytesReceiver, RingBuffChannel as RingBuffChannel, - attach_to_ringbuf_channel as attach_to_ringbuf_channel + attach_to_ringbuf_schannel as attach_to_ringbuf_schannel, + attach_to_ringbuf_rchannel as attach_to_ringbuf_rchannel, + attach_to_ringbuf_channel as attach_to_ringbuf_channel, ) diff --git a/tractor/ipc/_ringbuf.py b/tractor/ipc/_ringbuf.py index 038d9e73..42403937 100644 --- a/tractor/ipc/_ringbuf.py +++ b/tractor/ipc/_ringbuf.py @@ -550,6 +550,36 @@ class RingBuffBytesReceiver(trio.abc.ReceiveChannel[bytes]): await self._receiver.aclose() +@acm +async def attach_to_ringbuf_rchannel( + token: RBToken, + cleanup: bool = True +): + ''' + Attach a RingBuffBytesReceiver from a previously opened + RBToken. + ''' + async with attach_to_ringbuf_receiver( + token, cleanup=cleanup + ) as receiver: + yield RingBuffBytesReceiver(receiver) + + +@acm +async def attach_to_ringbuf_schannel( + token: RBToken, + cleanup: bool = True +): + ''' + Attach a RingBuffBytesSender from a previously opened + RBToken. + ''' + async with attach_to_ringbuf_sender( + token, cleanup=cleanup + ) as sender: + yield RingBuffBytesSender(sender) + + class RingBuffChannel(trio.abc.Channel[bytes]): ''' Combine `RingBuffBytesSender` and `RingBuffBytesReceiver` @@ -588,16 +618,13 @@ async def attach_to_ringbuf_channel( ''' async with ( - attach_to_ringbuf_receiver( + attach_to_ringbuf_rchannel( token_in, cleanup=cleanup_in ) as receiver, - attach_to_ringbuf_sender( + attach_to_ringbuf_schannel( token_out, cleanup=cleanup_out ) as sender, ): - yield RingBuffChannel( - RingBuffBytesSender(sender), - RingBuffBytesReceiver(receiver) - ) + yield RingBuffChannel(sender, receiver) -- 2.34.1 From ea010ab46a920b16312f917758c8657b169c6155 Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Tue, 18 Mar 2025 22:48:12 -0300 Subject: [PATCH 17/19] Add direct read method on EventFD Type hint all ctx managers in _ringbuf.py Remove unnecesary send lock on ring chan sender Handle EOF on ring chan receiver Rename ringbuf tests to make it less redundant --- tests/test_ringbuf.py | 8 ++--- tractor/ipc/_linux.py | 8 +++++ tractor/ipc/_ringbuf.py | 75 ++++++++++++++++++----------------------- 3 files changed, 45 insertions(+), 46 deletions(-) diff --git a/tests/test_ringbuf.py b/tests/test_ringbuf.py index f987d4c8..8858215e 100644 --- a/tests/test_ringbuf.py +++ b/tests/test_ringbuf.py @@ -179,7 +179,7 @@ async def child_blocked_receiver( await receiver.receive_some() -def test_ring_reader_cancel(): +def test_reader_cancel(): ''' Test that a receiver blocked on eventfd(2) read responds to cancellation. @@ -222,7 +222,7 @@ async def child_blocked_sender( await sender.send_all(b'this will wrap') -def test_ring_sender_cancel(): +def test_sender_cancel(): ''' Test that a sender blocked on eventfd(2) read responds to cancellation. @@ -255,7 +255,7 @@ def test_ring_sender_cancel(): trio.run(main) -def test_ringbuf_max_bytes(): +def test_receiver_max_bytes(): ''' Test that RingBuffReceiver.receive_some's max_bytes optional argument works correctly, send a msg of size 100, then @@ -384,7 +384,7 @@ async def child_channel_sender( await chan.send(msg) -def test_ringbuf_channel(): +def test_channel(): msg_amount_min = 100 msg_amount_max = 1000 diff --git a/tractor/ipc/_linux.py b/tractor/ipc/_linux.py index afce6bff..0c05260e 100644 --- a/tractor/ipc/_linux.py +++ b/tractor/ipc/_linux.py @@ -157,6 +157,14 @@ class EventFD: self._cscope = None + def read_direct(self) -> int: + ''' + Direct call to `read_eventfd(self.fd)`, unless `eventfd` was + opened with `EFD_NONBLOCK` its gonna block the thread. + + ''' + return read_eventfd(self._fd) + def open(self): self._fobj = os.fdopen(self._fd, self._omode) diff --git a/tractor/ipc/_ringbuf.py b/tractor/ipc/_ringbuf.py index 42403937..09c955ac 100644 --- a/tractor/ipc/_ringbuf.py +++ b/tractor/ipc/_ringbuf.py @@ -19,6 +19,10 @@ IPC Reliable RingBuffer implementation ''' from __future__ import annotations import struct +from typing import ( + ContextManager, + AsyncContextManager +) from contextlib import ( contextmanager as cm, asynccontextmanager as acm @@ -26,7 +30,6 @@ from contextlib import ( from multiprocessing.shared_memory import SharedMemory import trio -from tricycle import BufferedReceiveStream from msgspec import ( Struct, to_builtins @@ -34,7 +37,6 @@ from msgspec import ( from ._linux import ( open_eventfd, - close_eventfd, EFDReadCancelled, EventFD ) @@ -94,7 +96,7 @@ class RBToken(Struct, frozen=True): def open_ringbuf( shm_name: str, buf_size: int = _DEFAULT_RB_SIZE, -) -> RBToken: +) -> ContextManager[RBToken]: ''' Handle resources for a ringbuf (shm, eventfd), yield `RBToken` to be used with `attach_to_ringbuf_sender` and `attach_to_ringbuf_receiver` @@ -106,31 +108,19 @@ def open_ringbuf( create=True ) try: - token = RBToken( - shm_name=shm_name, - write_eventfd=open_eventfd(), - wrap_eventfd=open_eventfd(), - eof_eventfd=open_eventfd(), - buf_size=buf_size - ) - yield token - try: - close_eventfd(token.write_eventfd) - - except OSError: - ... - - try: - close_eventfd(token.wrap_eventfd) - - except OSError: - ... - - try: - close_eventfd(token.eof_eventfd) - - except OSError: - ... + with ( + EventFD(open_eventfd(), 'r') as write_event, + EventFD(open_eventfd(), 'r') as wrap_event, + EventFD(open_eventfd(), 'r') as eof_event, + ): + token = RBToken( + shm_name=shm_name, + write_eventfd=write_event.fd, + wrap_eventfd=wrap_event.fd, + eof_eventfd=eof_event.fd, + buf_size=buf_size + ) + yield token finally: shm.unlink() @@ -232,6 +222,7 @@ class RingBuffSender(trio.abc.SendStream): self._eof_event.write( self._ptr if self._ptr > 0 else self.size ) + if self._cleanup: self._write_event.close() self._wrap_event.close() @@ -239,7 +230,8 @@ class RingBuffSender(trio.abc.SendStream): self._shm.close() async def aclose(self): - self.close() + async with self._send_lock: + self.close() async def __aenter__(self): self.open() @@ -396,7 +388,7 @@ class RingBuffReceiver(trio.abc.ReceiveStream): async def attach_to_ringbuf_receiver( token: RBToken, cleanup: bool = True -): +) -> AsyncContextManager[RingBuffReceiver]: ''' Attach a RingBuffReceiver from a previously opened RBToken. @@ -418,7 +410,7 @@ async def attach_to_ringbuf_receiver( async def attach_to_ringbuf_sender( token: RBToken, cleanup: bool = True -): +) -> AsyncContextManager[RingBuffSender]: ''' Attach a RingBuffSender from a previously opened RBToken. @@ -435,7 +427,7 @@ async def attach_to_ringbuf_sender( def open_ringbuf_pair( name: str, buf_size: int = _DEFAULT_RB_SIZE -): +) -> ContextManager[tuple(RBToken, RBToken)]: ''' Handle resources for a ringbuf pair to be used for bidirectional messaging. @@ -461,7 +453,7 @@ async def attach_to_ringbuf_stream( token_out: RBToken, cleanup_in: bool = True, cleanup_out: bool = True -): +) -> AsyncContextManager[trio.StapledStream]: ''' Attach a trio.StapledStream from a previously opened ringbuf pair. @@ -494,16 +486,13 @@ class RingBuffBytesSender(trio.abc.SendChannel[bytes]): sender: RingBuffSender ): self._sender = sender - self._send_lock = trio.StrictFIFOLock() async def send(self, value: bytes) -> None: - async with self._send_lock: - size: bytes = struct.pack(" None: - async with self._send_lock: - await self._sender.aclose() + await self._sender.aclose() class RingBuffBytesReceiver(trio.abc.ReceiveChannel[bytes]): @@ -544,6 +533,8 @@ class RingBuffBytesReceiver(trio.abc.ReceiveChannel[bytes]): header: bytes = await self._receive_exactly(4) size: int size, = struct.unpack(" None: @@ -554,7 +545,7 @@ class RingBuffBytesReceiver(trio.abc.ReceiveChannel[bytes]): async def attach_to_ringbuf_rchannel( token: RBToken, cleanup: bool = True -): +) -> AsyncContextManager[RingBuffBytesReceiver]: ''' Attach a RingBuffBytesReceiver from a previously opened RBToken. @@ -569,7 +560,7 @@ async def attach_to_ringbuf_rchannel( async def attach_to_ringbuf_schannel( token: RBToken, cleanup: bool = True -): +) -> AsyncContextManager[RingBuffBytesSender]: ''' Attach a RingBuffBytesSender from a previously opened RBToken. @@ -611,7 +602,7 @@ async def attach_to_ringbuf_channel( token_out: RBToken, cleanup_in: bool = True, cleanup_out: bool = True -): +) -> AsyncContextManager[RingBuffChannel]: ''' Attach to an already opened ringbuf pair and return a `RingBuffChannel`. -- 2.34.1 From 010874bed5f7d834281d1abd3e90aa39afbd6277 Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Thu, 20 Mar 2025 21:12:06 -0300 Subject: [PATCH 18/19] Catch trio cancellation on RingBuffReceiver bg eof listener task, add batched mode to RingBuffBytesSender --- tractor/ipc/_ringbuf.py | 39 ++++++++++++++++++++++++++++++++++----- 1 file changed, 34 insertions(+), 5 deletions(-) diff --git a/tractor/ipc/_ringbuf.py b/tractor/ipc/_ringbuf.py index 09c955ac..7d96eeda 100644 --- a/tractor/ipc/_ringbuf.py +++ b/tractor/ipc/_ringbuf.py @@ -301,6 +301,9 @@ class RingBuffReceiver(trio.abc.ReceiveStream): except EFDReadCancelled: ... + except trio.Cancelled: + ... + async def receive_some(self, max_bytes: int | None = None) -> bytes: ''' Receive up to `max_bytes`, if no `max_bytes` is provided @@ -480,16 +483,41 @@ class RingBuffBytesSender(trio.abc.SendChannel[bytes]): which decodes into a uint32 indicating the actual size of the next payload. + 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, - sender: RingBuffSender + sender: RingBuffSender, + batch_size: int = 1 ): self._sender = sender + self.batch_size = batch_size + self._batch_msg_len = 0 + self._batch: bytes = b'' + + async def flush(self) -> None: + await self._sender.send_all(self._batch) + self._batch = b'' + self._batch_msg_len = 0 async def send(self, value: bytes) -> None: - size: bytes = struct.pack(" None: await self._sender.aclose() @@ -559,7 +587,8 @@ async def attach_to_ringbuf_rchannel( @acm async def attach_to_ringbuf_schannel( token: RBToken, - cleanup: bool = True + cleanup: bool = True, + batch_size: int = 1, ) -> AsyncContextManager[RingBuffBytesSender]: ''' Attach a RingBuffBytesSender from a previously opened @@ -568,7 +597,7 @@ async def attach_to_ringbuf_schannel( async with attach_to_ringbuf_sender( token, cleanup=cleanup ) as sender: - yield RingBuffBytesSender(sender) + yield RingBuffBytesSender(sender, batch_size=batch_size) class RingBuffChannel(trio.abc.Channel[bytes]): -- 2.34.1 From bab265b2d894197137456e5dfa7e9f009797c12a Mon Sep 17 00:00:00 2001 From: Guillermo Rodriguez Date: Sat, 22 Mar 2025 16:54:00 -0300 Subject: [PATCH 19/19] Important RingBuffBytesSender fix on non batched mode! & downgrade nix-shell python to lowest supported --- default.nix | 2 +- tractor/ipc/_ringbuf.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/default.nix b/default.nix index 5a936971..31615def 100644 --- a/default.nix +++ b/default.nix @@ -13,6 +13,6 @@ pkgs.mkShell { shellHook = '' set -e - uv venv .venv --python=3.12 + uv venv .venv --python=3.11 ''; } diff --git a/tractor/ipc/_ringbuf.py b/tractor/ipc/_ringbuf.py index 7d96eeda..10975b7a 100644 --- a/tractor/ipc/_ringbuf.py +++ b/tractor/ipc/_ringbuf.py @@ -512,6 +512,7 @@ class RingBuffBytesSender(trio.abc.SendChannel[bytes]): msg: bytes = struct.pack("