diff --git a/piker/brokers/ib/feed.py b/piker/brokers/ib/feed.py index cfea2b5c..adeb6361 100644 --- a/piker/brokers/ib/feed.py +++ b/piker/brokers/ib/feed.py @@ -738,6 +738,10 @@ async def get_bars( ) +# per-actor cache of inter-eventloop-chans +_quote_streams: dict[str, trio.abc.ReceiveStream] = {} + + async def _setup_quote_stream( chan: tractor.to_asyncio.LinkedTaskChannel, symbol: str, @@ -766,6 +770,8 @@ async def _setup_quote_stream( and is thus run via `tractor.to_asyncio.open_channel_from()`. ''' + global _quote_streams + async with load_aio_clients( disconnect_on_exit=False, ) as accts2clients: @@ -828,6 +834,9 @@ async def _setup_quote_stream( client.ib.cancelMktData(contract) + # decouple broadcast mem chan + _quote_streams.pop(symbol, None) + def push( t: Ticker, tries_before_raise: int = 6, @@ -944,52 +953,54 @@ async def _setup_quote_stream( @acm -async def _open_aio_quote_channel( +async def open_aio_quote_stream( symbol: str, contract: Contract|None = None, -) -> tractor.to_asyncio.LinkedTaskChannel: +) -> ( + trio.abc.Channel| # iface + tractor.to_asyncio.LinkedTaskChannel # actually +): ''' - Open one raw inter-event-loop quote channel for `symbol`. + Open a real-time `Ticker` quote stream from an `asyncio.Task` + spawned via `tractor.to_asyncio.open_channel_from()`, deliver the + inter-event-loop channel to the `trio.Task` caller and cache it + globally for re-use. ''' - chan: tractor.to_asyncio.LinkedTaskChannel - started_contract: Contract + from tractor.trionics import broadcast_receiver + global _quote_streams + + from_aio = _quote_streams.get(symbol) + if from_aio: + + # if we already have a cached feed deliver a rx side clone + # to consumer + async with broadcast_receiver( + from_aio, + 2**6, + ) as from_aio: + yield from_aio + return + + from_aio: tractor.to_asyncio.LinkedTaskChannel async with tractor.to_asyncio.open_channel_from( _setup_quote_stream, symbol=symbol, contract=contract, - ) as (chan, started_contract): - assert started_contract - yield chan + ) as (from_aio, contract): + assert contract -@acm -async def open_aio_quote_stream( - symbol: str, - contract: Contract|None = None, + # TODO? de-reg on teardown of last consumer task? + # -> why aren't we using `.trionics.maybe_open_context()` + # here again?? (we are in `open_client_proxies()` tho?) + # + # cache feed for later consumers + _quote_streams[symbol] = from_aio -) -> tractor.trionics.BroadcastReceiver: - ''' - Open a caller-owned child of the cached quote channel for `symbol`. - - Keep the raw `LinkedTaskChannel` alive until its last cache user - exits, and retain exactly one broadcaster root on that channel. - - ''' - chan: tractor.to_asyncio.LinkedTaskChannel - async with tractor.trionics.maybe_open_context( - acm_func=_open_aio_quote_channel, - kwargs={ - 'symbol': symbol, - 'contract': contract, - }, - key=symbol, - ) as (_, chan): - bstream: tractor.trionics.BroadcastReceiver - async with chan.subscribe() as bstream: - yield bstream + yield from_aio # TODO: cython/mypyc/numba this! diff --git a/tests/test_ib_broadcast.py b/tests/test_ib_broadcast.py deleted file mode 100644 index 37250701..00000000 --- a/tests/test_ib_broadcast.py +++ /dev/null @@ -1,120 +0,0 @@ -''' -IB cached quote-channel ownership regressions. - -''' -from contextlib import asynccontextmanager as acm - -import tractor -import trio - -from piker.brokers.ib import feed - - -def test_quote_stream_caches_channel_and_owns_child( - monkeypatch, -) -> None: - ''' - IB quote callers must share one cached channel through child cursors. - - `open_aio_quote_stream()` previously stored the first caller's raw - `LinkedTaskChannel` in `_quote_streams`, yielded that root directly, - and wrapped the same source in a new broadcaster for each later - caller. The first caller therefore owned source lifetime while - competing roots split quote receives between concurrent users. - - Stub `maybe_open_context()` with one retained channel and enter a - miss owner plus concurrent hit. Prove both receive the same quote - through distinct children of one broadcaster, the hit closes - independently, and the owner keeps receiving until its own exit. - Also verify both calls request the same symbol cache key. - - ''' - class FakeLinkedTaskChannel: - ''' - Retain one real broadcaster and expose child subscriptions. - - ''' - def __init__(self) -> None: - self._tx: trio.MemorySendChannel - rx: trio.MemoryReceiveChannel - self._tx, rx = trio.open_memory_channel(8) - self._broadcaster: ( - tractor.trionics.BroadcastReceiver - ) = tractor.trionics.broadcast_receiver( - rx, - 8, - ) - - @acm - async def subscribe(self): - ''' - Yield one caller-owned child from the retained root. - - ''' - child: tractor.trionics.BroadcastReceiver - async with self._broadcaster.subscribe() as child: - yield child - - async def push(self, msg: dict) -> None: - ''' - Send one ticker through the retained source channel. - - ''' - await self._tx.send(msg) - - chan: FakeLinkedTaskChannel = FakeLinkedTaskChannel() - context_entries: int = 0 - keys: list[str] = [] - acm_funcs: list = [] - - @acm - async def maybe_open_context(**kwargs): - ''' - Return one retained channel as a miss followed by a hit. - - ''' - nonlocal context_entries - keys.append(kwargs['key']) - acm_funcs.append(kwargs['acm_func']) - cache_hit: bool = context_entries > 0 - context_entries += 1 - yield cache_hit, chan - - monkeypatch.setattr( - feed.tractor.trionics, - 'maybe_open_context', - maybe_open_context, - ) - - async def main() -> None: - owner: tractor.trionics.BroadcastReceiver - async with feed.open_aio_quote_stream( - 'NVDA', - ) as owner: - hit: tractor.trionics.BroadcastReceiver - async with feed.open_aio_quote_stream( - 'NVDA', - ) as hit: - assert hit is not owner - assert len(chan._broadcaster._state.subs) == 3 - - await chan.push({'value': 1}) - assert await owner.receive() == {'value': 1} - assert await hit.receive() == {'value': 1} - - assert hit._closed - assert not owner._closed - assert len(chan._broadcaster._state.subs) == 2 - - await chan.push({'value': 2}) - assert await owner.receive() == {'value': 2} - - assert owner._closed - assert len(chan._broadcaster._state.subs) == 1 - assert keys == ['NVDA', 'NVDA'] - assert acm_funcs == [ - feed._open_aio_quote_channel, - feed._open_aio_quote_channel, - ] - - trio.run(main) diff --git a/tests/test_ib_quote_restart.py b/tests/test_ib_quote_restart.py new file mode 100644 index 00000000..424a5d98 --- /dev/null +++ b/tests/test_ib_quote_restart.py @@ -0,0 +1,160 @@ +''' +IB quote restart regressions using the real Tractor asyncio bridge. + +''' +import asyncio +from contextlib import asynccontextmanager as acm +from types import SimpleNamespace + +from async_generator import aclosing +from ib_async import Contract, IB, RequestError, Ticker +import pytest +import tractor +import trio + +from piker.brokers.ib import feed +from piker.brokers.ib import _util +from piker.brokers.ib.api import Client, open_client_proxy + + +async def exercise_quote_restart() -> list[str]: + ''' + Cancel an active quote consumer, then reopen in the same actor. + + Only the external IB connection is replaced. The production quote + producer, eventkit callback, bridge, broadcaster and context cache + all run unchanged in an asyncio-infected Tractor child. The farm + OK event in stream_quotes() cancels its local scope after a reset; + reproduce that cancellation after an actual ticker delivery. + + A cached owner can turn this into a normal bridge exit, feeding + TrioTaskExited back into the consumer and killing brokerd. Two + generations must close their IB subscriptions and asyncio tasks. + Return callback failures separately so actor survival cannot be + mistaken for clean teardown. This proves local restart semantics, + not connectivity to a real Gateway. + + ''' + contract = Contract(symbol='QQQ') + tasks: list[asyncio.Task] = [] + tickers: list[Ticker] = [] + cancelled: list[Contract] = [] + callback_errors: list[dict] = [] + + class ReplayIB: + ''' + Deliver one real eventkit ticker per subscription. + + ''' + def reqMktData(self, con: Contract, opts: str) -> Ticker: + ''' + Emit once the production producer installs its callback. + + ''' + task = asyncio.current_task() + assert task is not None + tasks.append(task) + ticker = Ticker(contract=con) + tickers.append(ticker) + asyncio.get_running_loop().call_soon( + ticker.updateEvent.emit, + ticker, + ) + return ticker + + def cancelMktData(self, con: Contract) -> None: + ''' + Record the production producer's subscription teardown. + + ''' + cancelled.append(con) + + client = SimpleNamespace(ib=ReplayIB()) + + @acm + async def load_clients(**kwargs): + ''' + Supply an offline IB peer without opening broker connections. + + ''' + yield {'ib.replay': client} + + loop = asyncio.get_running_loop() + previous_handler = loop.get_exception_handler() + loop.set_exception_handler( + lambda loop, context: callback_errors.append(context) + ) + try: + with pytest.MonkeyPatch.context() as patch: + patch.setattr(feed, 'load_aio_clients', load_clients) + patch.setattr( + feed, + 'get_preferred_data_client', + lambda clients: ('ib.replay', client), + ) + for generation in range(2): + with trio.CancelScope() as scope: + async with feed.open_aio_quote_stream( + 'qqq.nasdaq', + contract=contract, + ) as quotes: + async with aclosing(quotes): + ticker = await quotes.receive() + assert ticker is tickers[generation] + scope.cancel() + await trio.lowlevel.checkpoint() + + assert scope.cancelled_caught + assert len(cancelled) == generation + 1 + assert tasks[generation].done() + assert len(tickers[generation].updateEvent) == 0 + + # Flush done callbacks before accepting a clean restart. + await tractor.to_asyncio.run_task(asyncio.sleep, delay=0) + return [ + repr(context['exception']) + for context in callback_errors + ] + finally: + loop.set_exception_handler(previous_handler) + + +@pytest.mark.parametrize( + 'check_callbacks', + [False, True], + ids=['actor-survival', 'callback-cleanliness'], +) +def test_quote_restart_preserves_actor( + reg_addr, + check_callbacks: bool, +) -> None: + ''' + A farm-style quote restart must leave its broker actor usable. + + Use a real TCP actor nursery and an asyncio-infected child. Run + two subscription generations, then invoke another RPC on the same + actor. Both cases reject a propagated bridge exception, leaked + quote task or lost actor. The strict case also rejects asyncio + callback failures; the survival case reports those independently. + The backout can restore restart survival while the strict case + still exposes Tractor's completion-callback defect. + + ''' + async def main() -> None: + with trio.fail_after(10): + async with ( + tractor.open_root_actor(registry_addrs=[reg_addr]), + tractor.open_nursery() as nursery, + ): + portal = await nursery.start_actor( + 'ib_quote_replay', + enable_modules=[__name__], + infect_asyncio=True, + ) + errors = await portal.run(exercise_quote_restart) + errors += await portal.run(exercise_quote_restart) + await portal.cancel_actor() + if check_callbacks: + assert not errors, errors + + trio.run(main)