import time import hashlib import trio import pytest import tractor from tractor.ipc._ringbuf import ( open_ringbuf, open_ringbuf_pair, attach_to_ringbuf_receiver, attach_to_ringbuf_sender, attach_to_ringbuf_channel, RBToken, ) from tractor._testing.samples import ( generate_single_byte_msgs, RandomBytesGenerator ) # in case you don't want to melt your cores, uncomment dis! pytestmark = pytest.mark.skip @tractor.context async def child_read_shm( ctx: tractor.Context, token: RBToken, ) -> 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') msg_amount = 0 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: msg_amount += 1 recvd_hash.update(msg) recvd_bytes += len(msg) 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):,}') print(f'\treceived msgs: {msg_amount:,}') print(f'\treceived bytes: {recvd_bytes:,}') return recvd_hash.hexdigest() @tractor.context async def child_write_shm( ctx: tractor.Context, msg_amount: int, rand_min: int, rand_max: int, buf_size: int ) -> None: ''' 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. ''' rng = RandomBytesGenerator( msg_amount, rand_min=rand_min, rand_max=rand_max, ) async with ( open_ringbuf('test_ringbuf', buf_size=buf_size) as token, attach_to_ringbuf_sender(token) as sender ): await ctx.started(token) print('writer started') for msg in rng: await sender.send(msg) if rng.msgs_generated % rng.recommended_log_interval == 0: print(f'wrote {rng.msgs_generated} msgs') print('writer exit') return rng.hexdigest @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_ringbuf( msg_amount: int, rand_min: int, rand_max: int, buf_size: int ): ''' - Open a new ring buf on root actor - 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(): async with tractor.open_nursery() as an: send_p = await an.start_actor( 'ring_sender', enable_modules=[ __name__, 'tractor.linux._fdshare' ], ) recv_p = await an.start_actor( 'ring_receiver', enable_modules=[ __name__, 'tractor.linux._fdshare' ], ) async with ( send_p.open_context( child_write_shm, msg_amount=msg_amount, rand_min=rand_min, rand_max=rand_max, buf_size=buf_size ) as (sctx, token), recv_p.open_context( child_read_shm, token=token, ) as (rctx, _), ): sent_hash = await sctx.result() recvd_hash = await rctx.result() assert sent_hash == recvd_hash await an.cancel() trio.run(main) @tractor.context async def child_blocked_receiver(ctx: tractor.Context): async with ( open_ringbuf('test_ring_cancel_reader') as token, attach_to_ringbuf_receiver(token) as receiver ): await ctx.started(token) await receiver.receive_some() def test_reader_cancel(): ''' Test that a receiver blocked on eventfd(2) read responds to cancellation. ''' async def main(): async with tractor.open_nursery() as an: recv_p = await an.start_actor( 'ring_blocked_receiver', enable_modules=[ __name__, 'tractor.linux._fdshare' ], ) async with ( recv_p.open_context( child_blocked_receiver, ) as (sctx, token), attach_to_ringbuf_sender(token), ): await trio.sleep(.1) await an.cancel() with pytest.raises(tractor._exceptions.ContextCancelled): trio.run(main) @tractor.context async def child_blocked_sender(ctx: tractor.Context): async with ( open_ringbuf( 'test_ring_cancel_sender', buf_size=1 ) as token, attach_to_ringbuf_sender(token) as sender ): await ctx.started(token) await sender.send_all(b'this will wrap') def test_sender_cancel(): ''' Test that a sender blocked on eventfd(2) read responds to cancellation. ''' async def main(): async with tractor.open_nursery() as an: recv_p = await an.start_actor( 'ring_blocked_sender', enable_modules=[ __name__, 'tractor.linux._fdshare' ], ) async with ( recv_p.open_context( child_blocked_sender, ) as (sctx, token), attach_to_ringbuf_receiver(token) ): await trio.sleep(.1) await an.cancel() with pytest.raises(tractor._exceptions.ContextCancelled): trio.run(main) def test_receiver_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 = generate_single_byte_msgs(100) msgs = [] async def main(): async with ( tractor.open_nursery(), open_ringbuf( 'test_ringbuf_max_bytes', buf_size=10 ) as token, trio.open_nursery() as n, attach_to_ringbuf_sender(token, cleanup=False) as sender, attach_to_ringbuf_receiver(token, cleanup=False) as receiver ): 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) assert len(msg_part) == 1 msgs.append(msg_part) trio.run(main) assert msg == b''.join(msgs) @tractor.context async def child_channel_sender( ctx: tractor.Context, msg_amount_min: int, msg_amount_max: int, token_in: RBToken, token_out: RBToken ): import random rng = RandomBytesGenerator( random.randint(msg_amount_min, msg_amount_max), rand_min=256, rand_max=1024, ) async with attach_to_ringbuf_channel( token_in, token_out ) as chan: await ctx.started() for msg in rng: await chan.send(msg) await chan.send(b'bye') await chan.receive() return rng.hexdigest def test_channel(): msg_amount_min = 100 msg_amount_max = 1000 mods = [ __name__, 'tractor.linux._fdshare' ] async def main(): async with ( tractor.open_nursery(enable_modules=mods) as an, open_ringbuf_pair( 'test_ringbuf_transport' ) as (send_token, recv_token), attach_to_ringbuf_channel(send_token, recv_token) as chan, ): sender = await an.start_actor( 'test_ringbuf_transport_sender', enable_modules=mods, ) async with ( sender.open_context( child_channel_sender, msg_amount_min=msg_amount_min, msg_amount_max=msg_amount_max, token_in=recv_token, token_out=send_token ) as (ctx, _), ): recvd_hash = hashlib.sha256() async for msg in chan: if msg == b'bye': await chan.send(b'bye') break recvd_hash.update(msg) sent_hash = await ctx.result() assert recvd_hash.hexdigest() == sent_hash await an.cancel() trio.run(main)