1393 lines
43 KiB
Python
1393 lines
43 KiB
Python
"""
|
|
Broadcast channels for fan-out to local tasks.
|
|
|
|
"""
|
|
from contextlib import (
|
|
asynccontextmanager as acm,
|
|
)
|
|
from functools import partial
|
|
from itertools import cycle
|
|
import time
|
|
from types import SimpleNamespace
|
|
from typing import Optional
|
|
import warnings
|
|
|
|
import pytest
|
|
import trio
|
|
from trio.lowlevel import current_task
|
|
import tractor
|
|
from tractor.to_asyncio import LinkedTaskChannel
|
|
from tractor.trionics import (
|
|
broadcast_receiver,
|
|
BroadcastReceiveError,
|
|
Lagged,
|
|
collapse_eg,
|
|
)
|
|
|
|
|
|
@tractor.context
|
|
async def echo_sequences(
|
|
|
|
ctx: tractor.Context,
|
|
|
|
) -> None:
|
|
'''Bidir streaming endpoint which will stream
|
|
back any sequence it is sent item-wise.
|
|
|
|
'''
|
|
await ctx.started()
|
|
|
|
async with ctx.open_stream() as stream:
|
|
async for sequence in stream:
|
|
seq = list(sequence)
|
|
for value in seq:
|
|
await stream.send(value)
|
|
print(f'producer sent {value}')
|
|
|
|
|
|
async def ensure_sequence(
|
|
|
|
stream: tractor.MsgStream,
|
|
sequence: list,
|
|
delay: Optional[float] = None,
|
|
|
|
) -> None:
|
|
|
|
name = current_task().name
|
|
async with stream.subscribe() as bcaster:
|
|
assert not isinstance(bcaster, type(stream))
|
|
async for value in bcaster:
|
|
print(f'{name} rx: {value}')
|
|
assert value == sequence[0]
|
|
sequence.remove(value)
|
|
|
|
if delay:
|
|
await trio.sleep(delay)
|
|
|
|
if not sequence:
|
|
# fully consumed
|
|
break
|
|
|
|
|
|
@acm
|
|
async def open_sequence_streamer(
|
|
|
|
sequence: list[int],
|
|
reg_addr: tuple[str, int],
|
|
start_method: str,
|
|
|
|
) -> tractor.MsgStream:
|
|
|
|
async with tractor.open_nursery(
|
|
registry_addrs=[reg_addr],
|
|
start_method=start_method,
|
|
) as an:
|
|
|
|
portal = await an.start_actor(
|
|
'sequence_echoer',
|
|
enable_modules=[__name__],
|
|
)
|
|
|
|
async with portal.open_context(
|
|
echo_sequences,
|
|
) as (ctx, first):
|
|
|
|
assert first is None
|
|
async with ctx.open_stream(allow_overruns=True) as stream:
|
|
yield stream
|
|
|
|
await portal.cancel_actor()
|
|
|
|
|
|
def test_stream_fan_out_to_local_subscriptions(
|
|
reg_addr,
|
|
start_method,
|
|
):
|
|
|
|
sequence = list(range(1000))
|
|
|
|
async def main():
|
|
|
|
async with open_sequence_streamer(
|
|
sequence,
|
|
reg_addr,
|
|
start_method,
|
|
) as stream:
|
|
|
|
async with trio.open_nursery() as n:
|
|
for i in range(10):
|
|
n.start_soon(
|
|
ensure_sequence,
|
|
stream,
|
|
sequence.copy(),
|
|
name=f'consumer_{i}',
|
|
)
|
|
|
|
await stream.send(tuple(sequence))
|
|
|
|
async for value in stream:
|
|
print(f'source stream rx: {value}')
|
|
assert value == sequence[0]
|
|
sequence.remove(value)
|
|
|
|
if not sequence:
|
|
# fully consumed
|
|
break
|
|
|
|
trio.run(main)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
'task_delays',
|
|
[
|
|
(0.01, 0.001),
|
|
(0.001, 0.01),
|
|
]
|
|
)
|
|
def test_consumer_and_parent_maybe_lag(
|
|
reg_addr,
|
|
start_method,
|
|
task_delays,
|
|
):
|
|
|
|
async def main():
|
|
|
|
sequence = list(range(300))
|
|
parent_delay, sub_delay = task_delays
|
|
|
|
async with open_sequence_streamer(
|
|
sequence,
|
|
reg_addr,
|
|
start_method,
|
|
) as stream:
|
|
|
|
try:
|
|
async with (
|
|
collapse_eg(),
|
|
trio.open_nursery() as tn,
|
|
):
|
|
|
|
tn.start_soon(
|
|
ensure_sequence,
|
|
stream,
|
|
sequence.copy(),
|
|
sub_delay,
|
|
name='consumer_task',
|
|
)
|
|
|
|
await stream.send(tuple(sequence))
|
|
|
|
# async for value in stream:
|
|
lagged = False
|
|
lag_count = 0
|
|
|
|
while True:
|
|
try:
|
|
value = await stream.receive()
|
|
print(f'source stream rx: {value}')
|
|
|
|
if lagged:
|
|
# re set the sequence starting at our last
|
|
# value
|
|
sequence = sequence[sequence.index(value) + 1:]
|
|
else:
|
|
assert value == sequence[0]
|
|
sequence.remove(value)
|
|
|
|
lagged = False
|
|
|
|
except Lagged:
|
|
lagged = True
|
|
print(f'source stream lagged after {value}')
|
|
lag_count += 1
|
|
continue
|
|
|
|
# lag the parent
|
|
await trio.sleep(parent_delay)
|
|
|
|
if not sequence:
|
|
# fully consumed
|
|
break
|
|
print(f'parent + source stream lagged: {lag_count}')
|
|
|
|
if parent_delay > sub_delay:
|
|
assert lag_count > 0
|
|
|
|
except Lagged:
|
|
# child was lagged
|
|
assert parent_delay < sub_delay
|
|
|
|
trio.run(main)
|
|
|
|
|
|
def test_faster_task_to_recv_is_cancelled_by_slower(
|
|
reg_addr,
|
|
start_method,
|
|
):
|
|
'''
|
|
Ensure that if a faster task consuming from a stream is cancelled
|
|
the slower task can continue to receive all expected values.
|
|
|
|
'''
|
|
async def main():
|
|
|
|
sequence = list(range(1000))
|
|
|
|
async with open_sequence_streamer(
|
|
sequence,
|
|
reg_addr,
|
|
start_method,
|
|
|
|
) as stream:
|
|
|
|
async with trio.open_nursery() as tn:
|
|
tn.start_soon(
|
|
ensure_sequence,
|
|
stream,
|
|
sequence.copy(),
|
|
0,
|
|
name='consumer_task',
|
|
)
|
|
|
|
await stream.send(tuple(sequence))
|
|
|
|
# pull 3 values, cancel the subtask, then
|
|
# expect to be able to pull all values still
|
|
for i in range(20):
|
|
try:
|
|
value = await stream.receive()
|
|
print(f'source stream rx: {value}')
|
|
await trio.sleep(0.01)
|
|
except Lagged:
|
|
print(f'parent overrun after {value}')
|
|
continue
|
|
|
|
print('cancelling faster subtask')
|
|
tn.cancel_scope.cancel()
|
|
|
|
try:
|
|
value = await stream.receive()
|
|
print(f'source stream after cancel: {value}')
|
|
except Lagged:
|
|
print(f'parent overrun after {value}')
|
|
|
|
# expect to see all remaining values
|
|
with trio.fail_after(0.5):
|
|
async for value in stream:
|
|
assert stream._broadcaster._state.recv_ready is None
|
|
print(f'source stream rx: {value}')
|
|
if value == 999:
|
|
# fully consumed and we missed no values once
|
|
# the faster subtask was cancelled
|
|
break
|
|
|
|
# await tractor.pause()
|
|
# await stream.receive()
|
|
print(f'final value: {value}')
|
|
|
|
trio.run(main)
|
|
|
|
|
|
def test_subscribe_errors_after_close():
|
|
|
|
async def main():
|
|
|
|
size = 1
|
|
tx, rx = trio.open_memory_channel(size)
|
|
async with broadcast_receiver(rx, size) as brx:
|
|
pass
|
|
|
|
try:
|
|
# open and close
|
|
async with brx.subscribe():
|
|
pass
|
|
|
|
except trio.ClosedResourceError:
|
|
assert brx.key not in brx._state.subs
|
|
|
|
else:
|
|
assert 0
|
|
|
|
trio.run(main)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
('size', 'sent', 'dropped'),
|
|
[
|
|
(1, 2, 1),
|
|
(3, 5, 2),
|
|
],
|
|
)
|
|
def test_lagged_reports_exact_drop_count(
|
|
size: int,
|
|
sent: int,
|
|
dropped: int,
|
|
) -> None:
|
|
'''
|
|
`Lagged` must report every value outside the retained window.
|
|
|
|
`BroadcastReceiver.receive_nowait()` previously subtracted the
|
|
queue length from an already-invalid deque index without counting
|
|
that first displaced value. A one-slot queue therefore claimed it
|
|
dropped zero values after two sends. Keep one root receiver idle
|
|
while a child subscriber drains every produced value, then prove
|
|
the lag error reports the exact overrun and positions the root at
|
|
the oldest value still retained by `BroadcastState.queue`.
|
|
|
|
'''
|
|
async def main() -> None:
|
|
tx, rx = trio.open_memory_channel(size)
|
|
brx = broadcast_receiver(rx, size)
|
|
|
|
async with brx.subscribe() as fast:
|
|
for value in range(sent):
|
|
await tx.send(value)
|
|
assert await fast.receive() == value
|
|
|
|
match = rf'dropped `{dropped}` values'
|
|
with pytest.raises(Lagged, match=match):
|
|
await brx.receive()
|
|
|
|
assert await brx.receive() == sent - size
|
|
|
|
trio.run(main)
|
|
|
|
|
|
def test_broadcast_statistics_report_queued_counts() -> None:
|
|
'''
|
|
`BroadcastState.statistics()` must report counts, not indexes.
|
|
|
|
Each `BroadcastState.subs` value is the deque index of a
|
|
receiver's next unread value, with `-1` meaning caught up. The
|
|
statistics API returned these indexes directly, so one queued
|
|
value appeared as zero and every positive count was one short.
|
|
Keep one root receiver idle while a child synchronously receives
|
|
four produced values. Prove the root count advances through one
|
|
and three retained values, then remains clamped to the three-slot
|
|
retention window after lagging.
|
|
|
|
Finally install an actual unwaited `trio.Event` in
|
|
`BroadcastState.recv_ready` while treating deprecations as errors.
|
|
This proves statistics checks `None` explicitly instead of using
|
|
deprecated `trio.Event` truthiness.
|
|
|
|
'''
|
|
async def main() -> None:
|
|
tx, rx = trio.open_memory_channel(3)
|
|
brx = broadcast_receiver(rx, 3)
|
|
|
|
async with brx.subscribe() as child:
|
|
state = brx._state
|
|
assert state.statistics()['queued_len_by_task'] == {
|
|
brx.key: 0,
|
|
child.key: 0,
|
|
}
|
|
|
|
await tx.send(0)
|
|
assert await child.receive() == 0
|
|
assert state.statistics()['queued_len_by_task'] == {
|
|
brx.key: 1,
|
|
child.key: 0,
|
|
}
|
|
|
|
for value in range(1, 4):
|
|
await tx.send(value)
|
|
assert await child.receive() == value
|
|
|
|
state.recv_ready = (child.key, trio.Event())
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter('error', DeprecationWarning)
|
|
stats = state.statistics()
|
|
|
|
assert stats['queued_len_by_task'] == {
|
|
brx.key: 3,
|
|
child.key: 0,
|
|
}
|
|
assert stats['tasks_waiting'] == 0
|
|
state.recv_ready = None
|
|
|
|
trio.run(main)
|
|
|
|
|
|
def test_cancelled_reader_diagnostics_are_transient() -> None:
|
|
'''
|
|
Cancelled-reader diagnostics must not retain stale `Task`s.
|
|
|
|
`BroadcastState.cancelled` previously accumulated every source
|
|
owner cancelled during `BroadcastReceiver.receive()`. Even after
|
|
that receiver successfully read again or its subscription closed,
|
|
`BroadcastState.statistics()` retained the old `Task`, reporting
|
|
stale state and keeping the completed task alive.
|
|
|
|
Cancel one child's source read under a receiver-local scope and
|
|
verify its task is reported. Reuse that same receiver for one
|
|
successful read to prove progress clears the entry. Cancel it once
|
|
more, then leave the subscription and prove close also removes the
|
|
diagnostic while the root receiver remains registered.
|
|
|
|
'''
|
|
async def main() -> None:
|
|
tx, rx = trio.open_memory_channel(1)
|
|
brx = broadcast_receiver(rx, 1)
|
|
cancel_scope = trio.CancelScope()
|
|
child_key: int
|
|
child_task = None
|
|
|
|
async with brx.subscribe() as child:
|
|
child_key = child.key
|
|
|
|
async def cancel_source_read() -> None:
|
|
nonlocal child_task
|
|
child_task = current_task()
|
|
with cancel_scope:
|
|
await child.receive()
|
|
assert cancel_scope.cancelled_caught
|
|
|
|
async with trio.open_nursery() as nursery:
|
|
nursery.start_soon(cancel_source_read)
|
|
while brx._state.recv_ready is None:
|
|
await trio.lowlevel.checkpoint()
|
|
cancel_scope.cancel()
|
|
|
|
stats = brx._state.statistics()
|
|
assert child_task is not None
|
|
assert stats['tasks_cancelled'] == {
|
|
child_key: child_task,
|
|
}
|
|
|
|
await tx.send(1)
|
|
assert await child.receive() == 1
|
|
assert not brx._state.cancelled
|
|
|
|
cancel_scope = trio.CancelScope()
|
|
async with trio.open_nursery() as nursery:
|
|
nursery.start_soon(cancel_source_read)
|
|
while brx._state.recv_ready is None:
|
|
await trio.lowlevel.checkpoint()
|
|
cancel_scope.cancel()
|
|
|
|
assert child_key in brx._state.cancelled
|
|
|
|
assert child_key not in brx._state.cancelled
|
|
assert brx.key in brx._state.subs
|
|
|
|
trio.run(main)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
'terminal_exc',
|
|
[
|
|
trio.EndOfChannel(),
|
|
RuntimeError('terminal source failure'),
|
|
],
|
|
ids=['end-of-channel', 'receive-error'],
|
|
)
|
|
def test_terminal_broadcast_clears_cancelled_tasks(
|
|
terminal_exc: Exception,
|
|
) -> None:
|
|
'''
|
|
Terminal broadcast state must release every cancelled `Task`.
|
|
|
|
A receiver which owned and cancelled a source read can leave its
|
|
task in `BroadcastState.cancelled`. If another receiver later gets
|
|
EOC or a terminal source failure, no subscriber can make source
|
|
progress to clear that stale diagnostic. Clearing only the terminal
|
|
owner's key therefore retained the first receiver's completed task.
|
|
|
|
Cancel a child during the first controlled source read, then let
|
|
the root own a second read which raises EOC or `RuntimeError`.
|
|
Prove each terminal path clears the other receiver's diagnostic
|
|
before propagating its exact source outcome.
|
|
|
|
'''
|
|
class TerminalReceiver:
|
|
'''
|
|
Block one cancellable read, then raise a terminal outcome.
|
|
|
|
'''
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
self.first_started = trio.Event()
|
|
|
|
async def receive(self) -> None:
|
|
'''
|
|
Drive cancellation followed by terminal source state.
|
|
|
|
'''
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
self.first_started.set()
|
|
await trio.sleep_forever()
|
|
|
|
raise terminal_exc
|
|
|
|
async def main() -> None:
|
|
source = TerminalReceiver()
|
|
brx = broadcast_receiver(source, 1)
|
|
cancel_scope = trio.CancelScope()
|
|
|
|
async with brx.subscribe() as child:
|
|
async def cancel_child_read() -> None:
|
|
with cancel_scope:
|
|
await child.receive()
|
|
assert cancel_scope.cancelled_caught
|
|
|
|
async with trio.open_nursery() as nursery:
|
|
nursery.start_soon(cancel_child_read)
|
|
await source.first_started.wait()
|
|
cancel_scope.cancel()
|
|
|
|
assert child.key in brx._state.cancelled
|
|
with pytest.raises(type(terminal_exc)) as exc_info:
|
|
await brx.receive()
|
|
assert exc_info.value is terminal_exc
|
|
assert not brx._state.cancelled
|
|
|
|
trio.run(main)
|
|
|
|
|
|
def test_end_of_channel_is_terminal_for_waiting_peer() -> None:
|
|
'''
|
|
EOC must not let an awakened peer re-enter the closed source.
|
|
|
|
`BroadcastState.eoc` was set when one source owner received EOC,
|
|
but neither receive path consulted it. A peer waiting behind that
|
|
owner therefore woke, saw no queued value, and started a second
|
|
source read. Cancellation at that checkpoint could repopulate
|
|
`BroadcastState.cancelled` after the broadcast became terminal.
|
|
|
|
Block one child in the sole source read while the root waits on its
|
|
event, then release EOC. Both receivers must terminate from that
|
|
one source call, and the root's later receive must replay EOC
|
|
immediately without retaining cancellation diagnostics.
|
|
|
|
'''
|
|
class EOCReceiver:
|
|
'''
|
|
Publish one controlled EOC and reject any second source read.
|
|
|
|
'''
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
self.started = trio.Event()
|
|
self.release = trio.Event()
|
|
|
|
async def receive(self) -> None:
|
|
'''
|
|
Block the only valid source read until EOC release.
|
|
|
|
'''
|
|
self.calls += 1
|
|
assert self.calls == 1
|
|
self.started.set()
|
|
await self.release.wait()
|
|
raise trio.EndOfChannel
|
|
|
|
async def main() -> None:
|
|
source = EOCReceiver()
|
|
brx = broadcast_receiver(source, 1)
|
|
outcomes: list[str] = []
|
|
|
|
async with brx.subscribe() as child:
|
|
async def receive_eoc(
|
|
receiver,
|
|
name: str,
|
|
) -> None:
|
|
with pytest.raises(trio.EndOfChannel):
|
|
await receiver.receive()
|
|
outcomes.append(name)
|
|
|
|
async with trio.open_nursery() as nursery:
|
|
nursery.start_soon(receive_eoc, child, 'child')
|
|
await source.started.wait()
|
|
nursery.start_soon(receive_eoc, brx, 'root')
|
|
|
|
_, event = brx._state.recv_ready
|
|
while not event.statistics().tasks_waiting:
|
|
await trio.lowlevel.checkpoint()
|
|
source.release.set()
|
|
|
|
with pytest.raises(trio.EndOfChannel):
|
|
await brx.receive()
|
|
|
|
assert sorted(outcomes) == ['child', 'root']
|
|
assert source.calls == 1
|
|
assert not brx._state.cancelled
|
|
|
|
trio.run(main)
|
|
|
|
|
|
def test_msgstream_eoc_close_preserves_aclose_override() -> None:
|
|
'''
|
|
Internal EOC cleanup must preserve the public `aclose()` contract.
|
|
|
|
Passing a new private keyword from `MsgStream.receive()` to
|
|
`self.aclose()` broke subclasses whose compatible override kept
|
|
the original zero-argument signature. Use a minimal subclass which
|
|
records virtual dispatch and delegates to the base implementation.
|
|
Drive graceful EOC through the real root broadcaster and prove the
|
|
override runs without closing that active root re-entrantly.
|
|
|
|
'''
|
|
class Stream(tractor.MsgStream):
|
|
'''
|
|
Record public close dispatch with the established signature.
|
|
|
|
'''
|
|
close_calls = 0
|
|
|
|
async def aclose(self):
|
|
'''
|
|
Delegate closure without accepting private arguments.
|
|
|
|
'''
|
|
self.close_calls += 1
|
|
return await super().aclose()
|
|
|
|
class PldRx:
|
|
'''
|
|
Delegate source receive and terminate the close drain.
|
|
|
|
'''
|
|
def __init__(self, rx) -> None:
|
|
self._rx = rx
|
|
|
|
async def recv_pld(self, **kwargs):
|
|
'''
|
|
Receive directly from the test source channel.
|
|
|
|
'''
|
|
return await self._rx.receive()
|
|
|
|
def recv_msg_nowait(self, **kwargs):
|
|
'''
|
|
Report EOC to finish `MsgStream.aclose()` draining.
|
|
|
|
'''
|
|
raise trio.EndOfChannel
|
|
|
|
async def main() -> None:
|
|
tx, rx = trio.open_memory_channel(1)
|
|
ctx = SimpleNamespace(
|
|
cid='test-context',
|
|
_pld_rx=PldRx(rx),
|
|
send_stop=lambda: trio.lowlevel.checkpoint(),
|
|
side='caller',
|
|
peer_side='callee',
|
|
maybe_raise=lambda **kwargs: None,
|
|
)
|
|
stream = Stream(ctx, rx)
|
|
|
|
async with stream.subscribe():
|
|
await tx.aclose()
|
|
with pytest.raises(trio.EndOfChannel):
|
|
await stream.receive()
|
|
assert stream.close_calls == 1
|
|
assert not stream._broadcaster._closed
|
|
|
|
trio.run(main)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
'close_wrapper',
|
|
[
|
|
tractor.MsgStream.aclose,
|
|
LinkedTaskChannel.aclose,
|
|
],
|
|
ids=['msg-stream', 'linked-task-channel'],
|
|
)
|
|
def test_wrapper_close_clears_root_cancelled_task(
|
|
close_wrapper,
|
|
) -> None:
|
|
'''
|
|
Public stream close must release root cancellation diagnostics.
|
|
|
|
Root broadcasters allocated by `MsgStream.subscribe()` and
|
|
`LinkedTaskChannel.subscribe()` are private implementation state.
|
|
If their source receive was cancelled, callers had no public way
|
|
to close the root, so wrapper teardown retained the completed
|
|
`Task` in `BroadcastState.cancelled` indefinitely.
|
|
|
|
Cancel a root source read, attach that broadcaster to a minimal
|
|
public wrapper, and close it through each real `aclose()` method.
|
|
The root receiver and its task diagnostic must both be removed;
|
|
for `MsgStream`, pre-close the source to cover its idempotent early
|
|
return path.
|
|
|
|
'''
|
|
async def main() -> None:
|
|
_, rx = trio.open_memory_channel(1)
|
|
brx = broadcast_receiver(rx, 1)
|
|
cancel_scope = trio.CancelScope()
|
|
|
|
async def cancel_source_read() -> None:
|
|
with cancel_scope:
|
|
await brx.receive()
|
|
assert cancel_scope.cancelled_caught
|
|
|
|
async with trio.open_nursery() as nursery:
|
|
nursery.start_soon(cancel_source_read)
|
|
while brx._state.recv_ready is None:
|
|
await trio.lowlevel.checkpoint()
|
|
cancel_scope.cancel()
|
|
|
|
assert brx.key in brx._state.cancelled
|
|
|
|
if close_wrapper is tractor.MsgStream.aclose:
|
|
ctx = SimpleNamespace(cid='test-context')
|
|
wrapper = tractor.MsgStream(ctx, rx)
|
|
wrapper._broadcaster = brx
|
|
await rx.aclose()
|
|
else:
|
|
wrapper = SimpleNamespace(
|
|
_broadcaster=brx,
|
|
_from_aio=rx,
|
|
)
|
|
|
|
await close_wrapper(wrapper)
|
|
assert brx.key not in brx._state.subs
|
|
assert brx.key not in brx._state.cancelled
|
|
|
|
trio.run(main)
|
|
|
|
|
|
def test_broadcast_rejects_zero_buffer_size() -> None:
|
|
'''
|
|
A broadcaster must retain at least one value for peer fan-out.
|
|
|
|
`collections.deque(maxlen=0)` silently discards every appended
|
|
value, so `broadcast_receiver(..., 0)` allowed the source owner to
|
|
receive while peer cursors advanced into an always-empty queue.
|
|
Their lag recovery then reset to index `-1` and recursively retried
|
|
without any retained value to consume.
|
|
|
|
Construct a rendezvous memory channel and prove broadcaster setup
|
|
rejects its zero capacity synchronously with a clear public error,
|
|
before any receiver is registered or source receive can begin.
|
|
|
|
'''
|
|
_, rx = trio.open_memory_channel(0)
|
|
with pytest.raises(
|
|
ValueError,
|
|
match='`max_buffer_size` must be greater than zero',
|
|
):
|
|
broadcast_receiver(rx, 0)
|
|
|
|
|
|
def test_underlying_receive_failure_wakes_all_subscribers() -> None:
|
|
'''
|
|
A shared receive failure must terminate every broadcast receiver.
|
|
|
|
Previously, only `EndOfChannel` and receiver cancellation woke
|
|
peer tasks waiting on `BroadcastState.recv_ready`. If the shared
|
|
underlying receiver raised another error, its owner propagated
|
|
the failure and cleared the event while every peer remained
|
|
blocked forever.
|
|
|
|
Script one successful receive followed by a controlled
|
|
`RuntimeError`. Let a fast child own both underlying receives
|
|
while the root first drains its retained value and then waits on
|
|
the child's second receive. Release the failure only after both
|
|
tasks have reached those positions. Both exact errors prove the
|
|
peer was awakened without losing buffered data. A later
|
|
subscriber proves the terminal failure remains published for new
|
|
receivers instead of retrying the failed underlying channel.
|
|
|
|
'''
|
|
class FailingReceiver:
|
|
'''
|
|
Return one value, then fail after deterministic release.
|
|
|
|
'''
|
|
def __init__(self) -> None:
|
|
self.calls: int = 0
|
|
self.failure_started = trio.Event()
|
|
self.release_failure = trio.Event()
|
|
|
|
async def receive(self) -> int:
|
|
'''
|
|
Drive the scripted success-then-failure sequence.
|
|
|
|
'''
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
return 1
|
|
|
|
self.failure_started.set()
|
|
await self.release_failure.wait()
|
|
raise RuntimeError('underlying receive failed')
|
|
|
|
async def main() -> None:
|
|
source = FailingReceiver()
|
|
brx = broadcast_receiver(source, 3)
|
|
child_error: list[RuntimeError] = []
|
|
root_error: list[BroadcastReceiveError] = []
|
|
late_error: list[BroadcastReceiveError] = []
|
|
root_drained = trio.Event()
|
|
|
|
async def receive_child() -> None:
|
|
async with brx.subscribe() as child:
|
|
assert await child.receive() == 1
|
|
try:
|
|
await child.receive()
|
|
except RuntimeError as exc:
|
|
child_error.append(exc)
|
|
|
|
async def receive_root() -> None:
|
|
assert await brx.receive() == 1
|
|
root_drained.set()
|
|
try:
|
|
await brx.receive()
|
|
except BroadcastReceiveError as exc:
|
|
root_error.append(exc)
|
|
|
|
with trio.fail_after(1):
|
|
async with trio.open_nursery() as nursery:
|
|
nursery.start_soon(receive_child)
|
|
await source.failure_started.wait()
|
|
|
|
nursery.start_soon(receive_root)
|
|
await root_drained.wait()
|
|
|
|
source.release_failure.set()
|
|
|
|
assert source.calls == 2
|
|
assert [str(exc) for exc in child_error] == [
|
|
'underlying receive failed',
|
|
]
|
|
assert [str(exc) for exc in root_error] == [
|
|
'Shared broadcast receiver failed',
|
|
]
|
|
assert child_error[0] is not root_error[0]
|
|
assert root_error[0].__cause__ is child_error[0]
|
|
|
|
async with brx.subscribe() as late:
|
|
with pytest.raises(
|
|
BroadcastReceiveError,
|
|
match='Shared broadcast receiver failed',
|
|
) as exc_info:
|
|
await late.receive()
|
|
late_error.append(exc_info.value)
|
|
assert late_error[0] is not child_error[0]
|
|
assert late_error[0] is not root_error[0]
|
|
assert late_error[0].__cause__ is child_error[0]
|
|
assert source.calls == 2
|
|
|
|
trio.run(main)
|
|
|
|
|
|
def test_control_flow_exit_wakes_broadcast_peer() -> None:
|
|
'''
|
|
Non-terminal control flow must wake peers without being retained.
|
|
|
|
Process-control and cancellation-like `BaseException` values
|
|
should remain local to the task which receives them, but the old
|
|
owner still has to wake subscribers blocked on its shared event.
|
|
Make one child own a controlled `BaseException` receive while the
|
|
root waits behind it. After release, prove the child gets that
|
|
exact exit and the root takes ownership of the next underlying
|
|
receive instead of hanging or replaying the control-flow event.
|
|
|
|
'''
|
|
class ReceiveExit(BaseException):
|
|
'''
|
|
Model a non-terminal process-control receive exit.
|
|
|
|
'''
|
|
|
|
class ControlFlowReceiver:
|
|
'''
|
|
Raise one controlled exit, then return a value.
|
|
|
|
'''
|
|
def __init__(self) -> None:
|
|
self.calls: int = 0
|
|
self.exit_started = trio.Event()
|
|
self.release_exit = trio.Event()
|
|
|
|
async def receive(self) -> int:
|
|
'''
|
|
Drive the scripted control-flow-then-value sequence.
|
|
|
|
'''
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
self.exit_started.set()
|
|
await self.release_exit.wait()
|
|
raise ReceiveExit
|
|
|
|
return 2
|
|
|
|
async def main() -> None:
|
|
source = ControlFlowReceiver()
|
|
brx = broadcast_receiver(source, 3)
|
|
child_exit: list[ReceiveExit] = []
|
|
root_value: list[int] = []
|
|
|
|
async def receive_child() -> None:
|
|
async with brx.subscribe() as child:
|
|
try:
|
|
await child.receive()
|
|
except ReceiveExit as exc:
|
|
child_exit.append(exc)
|
|
|
|
async def receive_root() -> None:
|
|
root_value.append(await brx.receive())
|
|
|
|
with trio.fail_after(1):
|
|
async with trio.open_nursery() as nursery:
|
|
nursery.start_soon(receive_child)
|
|
await source.exit_started.wait()
|
|
nursery.start_soon(receive_root)
|
|
|
|
while True:
|
|
_, event = brx._state.recv_ready
|
|
if event.statistics().tasks_waiting:
|
|
break
|
|
await trio.lowlevel.checkpoint()
|
|
|
|
source.release_exit.set()
|
|
|
|
assert len(child_exit) == 1
|
|
assert root_value == [2]
|
|
assert source.calls == 2
|
|
assert brx._state.receive_exc is None
|
|
|
|
trio.run(main)
|
|
|
|
|
|
def test_closing_non_owner_preserves_source_wait() -> None:
|
|
'''
|
|
Closing one subscriber must not wake another receiver's peers.
|
|
|
|
`BroadcastReceiver.aclose()` previously set the one shared
|
|
`BroadcastState.recv_ready` event even when a different receiver
|
|
owned the source read. Waiting peers then repeatedly awaited an
|
|
already-set event until the source produced another value,
|
|
creating a runnable hot loop on idle streams.
|
|
|
|
Block one child in the source receive, then place both the root
|
|
and a closing child behind its event. Close only that waiting
|
|
child and prove it gets `ClosedResourceError` without setting the
|
|
shared event. Both remaining receivers must still get the same
|
|
value after the source is released.
|
|
|
|
'''
|
|
async def main() -> None:
|
|
tx, rx = trio.open_memory_channel(1)
|
|
brx = broadcast_receiver(rx, 3)
|
|
owner_value: list[int] = []
|
|
root_value: list[int] = []
|
|
closing_closed = trio.Event()
|
|
|
|
async with (
|
|
brx.subscribe() as owner,
|
|
brx.subscribe() as closing,
|
|
):
|
|
async def receive_owner() -> None:
|
|
owner_value.append(await owner.receive())
|
|
|
|
async def receive_root() -> None:
|
|
root_value.append(await brx.receive())
|
|
|
|
async def receive_closing() -> None:
|
|
with pytest.raises(trio.ClosedResourceError):
|
|
await closing.receive()
|
|
closing_closed.set()
|
|
|
|
with trio.fail_after(1):
|
|
async with trio.open_nursery() as nursery:
|
|
nursery.start_soon(receive_owner)
|
|
while brx._state.recv_ready is None:
|
|
await trio.lowlevel.checkpoint()
|
|
|
|
nursery.start_soon(receive_root)
|
|
nursery.start_soon(receive_closing)
|
|
_, event = brx._state.recv_ready
|
|
while event.statistics().tasks_waiting < 2:
|
|
await trio.lowlevel.checkpoint()
|
|
|
|
await closing.aclose()
|
|
await closing_closed.wait()
|
|
assert not event.is_set()
|
|
await tx.send(1)
|
|
|
|
assert owner_value == [1]
|
|
assert root_value == [1]
|
|
|
|
trio.run(main)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
'first_outcome',
|
|
[
|
|
1,
|
|
RuntimeError('discarded source error'),
|
|
trio.EndOfChannel(),
|
|
],
|
|
)
|
|
def test_closing_source_owner_hands_read_to_peer(
|
|
first_outcome: int|Exception,
|
|
) -> None:
|
|
'''
|
|
Closing the source-read owner must transfer ownership to a peer.
|
|
|
|
Merely suppressing the old shared-event wake would leave peers
|
|
blocked behind an externally closed receiver that still owned an
|
|
idle source read. Script a first receive which blocks until its
|
|
private scope is cancelled and a second which returns immediately.
|
|
Close that owner only after the root is waiting behind it. Cover
|
|
a shielded value, ordinary error and EOC from the cancelled source
|
|
read. The owner must always get `ClosedResourceError`, while the
|
|
awakened root takes the second source read without publishing the
|
|
discarded source outcome.
|
|
|
|
'''
|
|
class HandoffReceiver:
|
|
'''
|
|
Block the first source read and satisfy the second.
|
|
|
|
'''
|
|
def __init__(self) -> None:
|
|
self.calls: int = 0
|
|
self.first_started = trio.Event()
|
|
self.release_first = trio.Event()
|
|
|
|
async def receive(self) -> int:
|
|
'''
|
|
Drive one cancelled read followed by one value.
|
|
|
|
'''
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
self.first_started.set()
|
|
with trio.CancelScope(shield=True):
|
|
await self.release_first.wait()
|
|
if isinstance(first_outcome, BaseException):
|
|
raise first_outcome
|
|
return first_outcome
|
|
|
|
return 2
|
|
|
|
async def main() -> None:
|
|
source = HandoffReceiver()
|
|
brx = broadcast_receiver(source, 3)
|
|
owner_closed = trio.Event()
|
|
root_value: list[int] = []
|
|
|
|
async with brx.subscribe() as owner:
|
|
async def receive_owner() -> None:
|
|
with pytest.raises(trio.ClosedResourceError):
|
|
await owner.receive()
|
|
owner_closed.set()
|
|
|
|
async def receive_root() -> None:
|
|
root_value.append(await brx.receive())
|
|
|
|
with trio.fail_after(1):
|
|
async with trio.open_nursery() as nursery:
|
|
nursery.start_soon(receive_owner)
|
|
await source.first_started.wait()
|
|
nursery.start_soon(receive_root)
|
|
|
|
_, event = brx._state.recv_ready
|
|
while not event.statistics().tasks_waiting:
|
|
await trio.lowlevel.checkpoint()
|
|
|
|
await owner.aclose()
|
|
source.release_first.set()
|
|
await owner_closed.wait()
|
|
|
|
assert source.calls == 2
|
|
assert root_value == [2]
|
|
assert brx._state.receive_exc is None
|
|
assert not brx._state.eoc
|
|
|
|
trio.run(main)
|
|
|
|
|
|
def test_ensure_slow_consumers_lag_out(
|
|
reg_addr,
|
|
start_method,
|
|
):
|
|
'''This is a pure local task test; no tractor
|
|
machinery is really required.
|
|
|
|
'''
|
|
async def main():
|
|
|
|
# make sure it all works within the runtime
|
|
async with tractor.open_root_actor():
|
|
|
|
num_laggers = 4
|
|
laggers: dict[str, int] = {}
|
|
retries = 3
|
|
size = 100
|
|
tx, rx = trio.open_memory_channel(size)
|
|
brx = broadcast_receiver(rx, size)
|
|
|
|
async def sub_and_print(
|
|
delay: float,
|
|
) -> None:
|
|
|
|
task = current_task()
|
|
start = time.time()
|
|
|
|
async with brx.subscribe() as lbrx:
|
|
while True:
|
|
print(f'{task.name}: starting consume loop')
|
|
try:
|
|
async for value in lbrx:
|
|
print(f'{task.name}: {value}')
|
|
await trio.sleep(delay)
|
|
|
|
if task.name == 'sub_1':
|
|
# trigger checkpoint to clean out other subs
|
|
await trio.sleep(0.01)
|
|
|
|
# the non-lagger got
|
|
# a ``trio.EndOfChannel``
|
|
# because the ``tx`` below was closed
|
|
assert len(lbrx._state.subs) == 1
|
|
|
|
await lbrx.aclose()
|
|
|
|
assert len(lbrx._state.subs) == 0
|
|
|
|
except trio.ClosedResourceError:
|
|
# only the fast sub will try to re-enter
|
|
# iteration on the now closed bcaster
|
|
assert task.name == 'sub_1'
|
|
return
|
|
|
|
except Lagged:
|
|
lag_time = time.time() - start
|
|
lags = laggers[task.name]
|
|
print(
|
|
f'restarting slow task {task.name} '
|
|
f'that bailed out on {lags}:{value} '
|
|
f'after {lag_time:.3f}')
|
|
if lags <= retries:
|
|
laggers[task.name] += 1
|
|
continue
|
|
else:
|
|
print(
|
|
f'{task.name} was too slow and terminated '
|
|
f'on {lags}:{value}')
|
|
return
|
|
|
|
async with trio.open_nursery() as tn:
|
|
|
|
for i in range(1, num_laggers):
|
|
|
|
task_name = f'sub_{i}'
|
|
laggers[task_name] = 0
|
|
tn.start_soon(
|
|
partial(
|
|
sub_and_print,
|
|
delay=i*0.001,
|
|
),
|
|
name=task_name,
|
|
)
|
|
|
|
# allow subs to sched
|
|
await trio.sleep(0.1)
|
|
|
|
async with tx:
|
|
for i in cycle(range(size)):
|
|
await tx.send(i)
|
|
if len(brx._state.subs) == 2:
|
|
# only one, the non lagger, sub is left
|
|
break
|
|
|
|
# the non-lagger
|
|
assert laggers.pop('sub_1') == 0
|
|
|
|
for n, v in laggers.items():
|
|
assert v == 4
|
|
|
|
assert tx._closed
|
|
assert not tx._state.open_send_channels
|
|
|
|
# check that "first" bcaster that we created
|
|
# above, never was iterated and is thus overrun
|
|
try:
|
|
await brx.receive()
|
|
except Lagged:
|
|
# expect tokio style index truncation
|
|
seq = brx._state.subs[brx.key]
|
|
assert seq == len(brx._state.queue) - 1
|
|
|
|
# all no_overruns entries in the underlying
|
|
# channel should have been copied into the bcaster
|
|
# queue trailing-window
|
|
async for i in rx:
|
|
print(f'bped: {i}')
|
|
assert i in brx._state.queue
|
|
|
|
# should be noop
|
|
await brx.aclose()
|
|
|
|
trio.run(main)
|
|
|
|
|
|
def test_first_recver_is_cancelled():
|
|
|
|
async def main():
|
|
|
|
# make sure it all works within the runtime
|
|
async with tractor.open_root_actor():
|
|
|
|
tx, rx = trio.open_memory_channel(1)
|
|
brx = broadcast_receiver(rx, 1)
|
|
cs = trio.CancelScope()
|
|
|
|
async def sub_and_recv():
|
|
with cs:
|
|
async with brx.subscribe() as bc:
|
|
async for value in bc:
|
|
print(value)
|
|
assert cs.cancelled_caught
|
|
|
|
async def cancel_and_send():
|
|
await trio.sleep(0.2)
|
|
cs.cancel()
|
|
await tx.send(1)
|
|
|
|
async with trio.open_nursery() as n:
|
|
|
|
n.start_soon(sub_and_recv)
|
|
await trio.sleep(0.1)
|
|
assert brx._state.recv_ready
|
|
|
|
n.start_soon(cancel_and_send)
|
|
|
|
# ensure that we don't hang because no-task is now
|
|
# waiting on the underlying receive..
|
|
with trio.fail_after(0.5):
|
|
value = await brx.receive()
|
|
print(f'parent: {value}')
|
|
assert value == 1
|
|
|
|
trio.run(main)
|
|
|
|
|
|
def test_no_raise_on_lag():
|
|
'''
|
|
Run a simple 2-task broadcast where one task is slow but configured
|
|
so that it does not raise `Lagged` on overruns using
|
|
`raise_on_lasg=False` and verify that the task does not raise.
|
|
|
|
'''
|
|
size = 100
|
|
tx, rx = trio.open_memory_channel(size)
|
|
brx = broadcast_receiver(rx, size)
|
|
|
|
async def slow():
|
|
async with brx.subscribe(
|
|
raise_on_lag=False,
|
|
) as br:
|
|
async for msg in br:
|
|
print(f'slow task got: {msg}')
|
|
await trio.sleep(0.1)
|
|
|
|
async def fast():
|
|
async with brx.subscribe() as br:
|
|
async for msg in br:
|
|
print(f'fast task got: {msg}')
|
|
|
|
async def main():
|
|
async with (
|
|
tractor.open_root_actor(
|
|
# NOTE: so we see the warning msg emitted by the bcaster
|
|
# internals when the no raise flag is set.
|
|
loglevel='warning',
|
|
),
|
|
collapse_eg(),
|
|
trio.open_nursery() as n,
|
|
):
|
|
n.start_soon(slow)
|
|
n.start_soon(fast)
|
|
|
|
for i in range(1000):
|
|
await tx.send(i)
|
|
|
|
# simulate user nailing ctl-c after realizing
|
|
# there's a lag in the slow task.
|
|
await trio.sleep(1)
|
|
raise KeyboardInterrupt
|
|
|
|
with pytest.raises(KeyboardInterrupt):
|
|
trio.run(main)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
('subscribe', 'chan_attr'),
|
|
[
|
|
(tractor.MsgStream.subscribe, '_rx_chan'),
|
|
(LinkedTaskChannel.subscribe, '_from_aio'),
|
|
],
|
|
ids=['msg-stream', 'linked-task-channel'],
|
|
)
|
|
def test_stream_subscribe_forwards_lag_policy(
|
|
subscribe,
|
|
chan_attr: str,
|
|
) -> None:
|
|
'''
|
|
Stream wrappers must expose per-subscriber lag policy.
|
|
|
|
`MsgStream.subscribe()` and `LinkedTaskChannel.subscribe()`
|
|
previously omitted `BroadcastReceiver.raise_on_lag`, forcing
|
|
downstream users to mutate a private receiver attribute. Invoke
|
|
each public wrapper against a minimal receive-compatible handle.
|
|
Prove the first non-raising subscription configures both the
|
|
irreversible root broadcaster and its child, while a later strict
|
|
child selects its own policy without changing that root.
|
|
|
|
'''
|
|
class StreamHandle:
|
|
'''
|
|
Provide the wrapper fields needed for local fan-out.
|
|
|
|
'''
|
|
def __init__(self) -> None:
|
|
self._broadcaster = None
|
|
setattr(
|
|
self,
|
|
chan_attr,
|
|
SimpleNamespace(
|
|
_state=SimpleNamespace(max_buffer_size=1),
|
|
),
|
|
)
|
|
|
|
async def receive(self):
|
|
'''
|
|
Block if a regression unexpectedly enters source receive.
|
|
|
|
'''
|
|
await trio.sleep_forever()
|
|
|
|
async def send(self, value) -> None:
|
|
'''
|
|
Satisfy `MsgStream` duplex-handle patching.
|
|
|
|
'''
|
|
|
|
async def main() -> None:
|
|
stream = StreamHandle()
|
|
async with subscribe(
|
|
stream,
|
|
raise_on_lag=False,
|
|
) as first:
|
|
assert not stream._broadcaster._raise_on_lag
|
|
assert not first._raise_on_lag
|
|
|
|
async with subscribe(
|
|
stream,
|
|
raise_on_lag=True,
|
|
) as second:
|
|
assert not stream._broadcaster._raise_on_lag
|
|
assert second._raise_on_lag
|
|
|
|
trio.run(main)
|