Rework `_FeedsBus` subscriptions mgmt using `set`

Allows using `set` ops for subscription management and guarantees no
duplicates per `brokerd` actor. New API is simpler for dynamic
pause/resume changes per `Feed`:
- `_FeedsBus.add_subs()`, `.get_subs()`, `.remove_subs()` all accept multi-sub
  `set` inputs.
- `Feed.pause()` / `.resume()` encapsulates management of *only* sending
  a msg on each unique underlying IPC msg stream.

Use new api in sampler task.
agg_feedz
Tyler Goodlet 2022-11-16 13:32:26 -05:00
parent 88870fdda7
commit 2a158aea2c
2 changed files with 73 additions and 60 deletions

View File

@ -24,7 +24,6 @@ from collections import Counter
import time import time
from typing import ( from typing import (
TYPE_CHECKING, TYPE_CHECKING,
Union,
) )
import tractor import tractor
@ -319,11 +318,10 @@ async def sample_and_broadcast(
sub_key: str = broker_symbol.lower() sub_key: str = broker_symbol.lower()
subs: list[ subs: list[
tuple[ tuple[
Union[tractor.MsgStream, trio.MemorySendChannel], tractor.MsgStream | trio.MemorySendChannel,
tractor.Context,
float | None, # tick throttle in Hz float | None, # tick throttle in Hz
] ]
] = bus._subscribers[sub_key] ] = bus.get_subs(sub_key)
# NOTE: by default the broker backend doesn't append # NOTE: by default the broker backend doesn't append
# it's own "name" into the fqsn schema (but maybe it # it's own "name" into the fqsn schema (but maybe it
@ -332,7 +330,7 @@ async def sample_and_broadcast(
fqsn = f'{broker_symbol}.{brokername}' fqsn = f'{broker_symbol}.{brokername}'
lags: int = 0 lags: int = 0
for (stream, ctx, tick_throttle) in subs: for (stream, tick_throttle) in subs.copy():
try: try:
with trio.move_on_after(0.2) as cs: with trio.move_on_after(0.2) as cs:
if tick_throttle: if tick_throttle:
@ -344,6 +342,7 @@ async def sample_and_broadcast(
) )
except trio.WouldBlock: except trio.WouldBlock:
overruns[sub_key] += 1 overruns[sub_key] += 1
ctx = stream._ctx
chan = ctx.chan chan = ctx.chan
log.warning( log.warning(
@ -399,9 +398,9 @@ async def sample_and_broadcast(
# so far seems like no since this should all # so far seems like no since this should all
# be single-threaded. Doing it anyway though # be single-threaded. Doing it anyway though
# since there seems to be some kinda race.. # since there seems to be some kinda race..
bus.remove_sub( bus.remove_subs(
sub_key, sub_key,
(stream, ctx, tick_throttle), {(stream, tick_throttle)},
) )

View File

@ -21,6 +21,7 @@ This module is enabled for ``brokerd`` daemons.
""" """
from __future__ import annotations from __future__ import annotations
from collections import defaultdict
from contextlib import asynccontextmanager as acm from contextlib import asynccontextmanager as acm
from datetime import datetime from datetime import datetime
from functools import partial from functools import partial
@ -111,16 +112,16 @@ class _FeedsBus(Struct):
task_lock: trio.StrictFIFOLock = trio.StrictFIFOLock() task_lock: trio.StrictFIFOLock = trio.StrictFIFOLock()
_subscribers: dict[ _subscribers: defaultdict[
str, str,
list[ set[
tuple[ tuple[
Union[tractor.MsgStream, trio.MemorySendChannel], tractor.MsgStream | trio.MemorySendChannel,
tractor.Context, # tractor.Context,
Optional[float], # tick throttle in Hz float | None, # tick throttle in Hz
] ]
] ]
] = {} ] = defaultdict(set)
async def start_task( async def start_task(
self, self,
@ -147,38 +148,53 @@ class _FeedsBus(Struct):
# task: trio.lowlevel.Task, # task: trio.lowlevel.Task,
# ) -> bool: # ) -> bool:
# ... # ...
def get_subs( def get_subs(
self, self,
key: str, key: str,
) -> list[ ) -> set[
tuple[ tuple[
Union[tractor.MsgStream, trio.MemorySendChannel], Union[tractor.MsgStream, trio.MemorySendChannel],
tractor.Context, # tractor.Context,
float | None, # tick throttle in Hz float | None, # tick throttle in Hz
] ]
]: ]:
'''
Get the ``set`` of consumer subscription entries for the given key.
'''
return self._subscribers[key] return self._subscribers[key]
def remove_sub( def add_subs(
self, self,
key: str, key: str,
sub: tuple, subs: set[tuple[
) -> bool: tractor.MsgStream | trio.MemorySendChannel,
# tractor.Context,
float | None, # tick throttle in Hz
]],
) -> set[tuple]:
''' '''
Remove a consumer's subscription entry for the given key. Add a ``set`` of consumer subscription entries for the given key.
''' '''
stream, ctx, tick_throttle = sub _subs = self._subscribers[key]
subs = self.get_subs(key) _subs.update(subs)
try: return _subs
subs.remove(sub)
except ValueError: def remove_subs(
chan = ctx.chan self,
log.error( key: str,
f'Stream was already removed from subs!?\n' subs: set[tuple],
f'{key}:'
f'{ctx.cid}@{chan.uid}' ) -> set[tuple]:
) '''
Remove a ``set`` of consumer subscription entries for key.
'''
_subs = self.get_subs(key)
_subs.difference_update(subs)
return _subs
_bus: _FeedsBus = None _bus: _FeedsBus = None
@ -969,12 +985,6 @@ class Flume(Struct):
else: else:
yield istream yield istream
async def pause(self) -> None:
await self.stream.send('pause')
async def resume(self) -> None:
await self.stream.send('resume')
def get_ds_info( def get_ds_info(
self, self,
) -> tuple[float, float, float]: ) -> tuple[float, float, float]:
@ -1308,7 +1318,7 @@ async def open_feed_bus(
# the sampler subscription since the backend isn't (yet) # the sampler subscription since the backend isn't (yet)
# expected to append it's own name to the fqsn, so we filter # expected to append it's own name to the fqsn, so we filter
# on keys which *do not* include that name (e.g .ib) . # on keys which *do not* include that name (e.g .ib) .
bus._subscribers.setdefault(bfqsn, []) bus._subscribers.setdefault(bfqsn, set())
# sync feed subscribers with flume handles # sync feed subscribers with flume handles
await ctx.started( await ctx.started(
@ -1324,7 +1334,7 @@ async def open_feed_bus(
ctx.open_stream() as stream, ctx.open_stream() as stream,
): ):
local_subs: list = [] local_subs: dict[str, set[tuple]] = {}
for fqsn, flume in flumes.items(): for fqsn, flume in flumes.items():
# re-send to trigger display loop cycle (necessary especially # re-send to trigger display loop cycle (necessary especially
# when the mkt is closed and no real-time messages are # when the mkt is closed and no real-time messages are
@ -1361,43 +1371,42 @@ async def open_feed_bus(
# stream it's the throttle task does the work of # stream it's the throttle task does the work of
# incrementally forwarding to the IPC stream at the throttle # incrementally forwarding to the IPC stream at the throttle
# rate. # rate.
sub = (send, ctx, tick_throttle) send._ctx = ctx # mock internal ``tractor.MsgStream`` ref
sub = (send, tick_throttle)
else: else:
sub = (stream, ctx, tick_throttle) sub = (stream, tick_throttle)
# TODO: add an api for this on the bus? # TODO: add an api for this on the bus?
# maybe use the current task-id to key the sub list that's # maybe use the current task-id to key the sub list that's
# added / removed? Or maybe we can add a general # added / removed? Or maybe we can add a general
# pause-resume by sub-key api? # pause-resume by sub-key api?
bfqsn = fqsn.removesuffix(f'.{brokername}') bfqsn = fqsn.removesuffix(f'.{brokername}')
bus_subs = bus._subscribers[bfqsn] local_subs.setdefault(bfqsn, set()).add(sub)
bus_subs.append(sub) bus.add_subs(bfqsn, {sub})
local_subs.append(sub)
# sync caller with all subs registered state
sub_registered.set() sub_registered.set()
uid = ctx.chan.uid
try: try:
uid = ctx.chan.uid
# ctrl protocol for start/stop of quote streams based on UI # ctrl protocol for start/stop of quote streams based on UI
# state (eg. don't need a stream when a symbol isn't being # state (eg. don't need a stream when a symbol isn't being
# displayed). # displayed).
async for msg in stream: async for msg in stream:
if msg == 'pause': if msg == 'pause':
for sub in local_subs: for bfqsn, subs in local_subs.items():
if sub in bus_subs: log.info(
log.info( f'Pausing {bfqsn} feed for {uid}')
f'Pausing {fqsn} feed for {uid}') bus.remove_subs(bfqsn, subs)
bus_subs.remove(sub)
elif msg == 'resume': elif msg == 'resume':
for sub in local_subs: for bfqsn, subs in local_subs.items():
if sub not in bus_subs: log.info(
log.info( f'Resuming {bfqsn} feed for {uid}')
f'Resuming {fqsn} feed for {uid}') bus.add_subs(bfqsn, subs)
bus_subs.append(sub)
else: else:
raise ValueError(msg) raise ValueError(msg)
finally: finally:
@ -1410,11 +1419,8 @@ async def open_feed_bus(
cs.cancel() cs.cancel()
# drop all subs for this task from the bus # drop all subs for this task from the bus
for sub in local_subs: for bfqsn, subs in local_subs.items():
try: bus.remove_subs(bfqsn, subs)
bus._subscribers[bfqsn].remove(sub)
except ValueError:
log.warning(f'{sub} for {symbol} was already removed?')
class Feed(Struct): class Feed(Struct):
@ -1492,6 +1498,14 @@ class Feed(Struct):
# def name(self) -> str: # def name(self) -> str:
# return self.mod.name # return self.mod.name
async def pause(self) -> None:
for stream in set(self.streams.values()):
await stream.send('pause')
async def resume(self) -> None:
for stream in set(self.streams.values()):
await stream.send('resume')
@acm @acm
async def install_brokerd_search( async def install_brokerd_search(