Generalize the publisher/fan-out system
Start working toward a more general (on-demand) pub-sub system which can be brought into ``tractor``. Right now this just means making the code in the `fan_out_to_ctxs()` less specific but, eventually I think this function should be coupled with a decorator and shipped as a standard "message pattern". Additionally, - try out making `BrokerFeed` a `@dataclass` - strip out all the `trio.Event` / uneeded nursery / extra task crap from `start_quote_stream()`kivy_mainline_and_py3.8
Tyler Goodlet
parent
c94ce47aa6
commit
22670afe58
|
|
@ -3,12 +3,13 @@ Live data feed machinery
|
|||
"""
|
||||
import time
|
||||
from functools import partial
|
||||
from dataclasses import dataclass, field
|
||||
from itertools import cycle
|
||||
import socket
|
||||
import json
|
||||
from types import ModuleType
|
||||
import typing
|
||||
from typing import Coroutine, Callable, Dict, List, Any
|
||||
from typing import Coroutine, Callable, Dict, List, Any, Tuple
|
||||
import contextlib
|
||||
from operator import itemgetter
|
||||
|
||||
|
|
@ -44,8 +45,9 @@ async def wait_for_network(net_func: Callable, sleep: int = 1) -> dict:
|
|||
|
||||
|
||||
async def stream_quotes(
|
||||
get_topics: typing.Callable,
|
||||
get_quotes: Coroutine,
|
||||
brokermod: ModuleType,
|
||||
request_quotes: Coroutine,
|
||||
rate: int = 5, # delay between quote requests
|
||||
diff_cached: bool = True, # only deliver "new" quotes to the queue
|
||||
) -> None:
|
||||
|
|
@ -58,8 +60,14 @@ async def stream_quotes(
|
|||
sleeptime = round(1. / rate, 3)
|
||||
_cache = {} # ticker to quote caching
|
||||
|
||||
while True: # use an event here to trigger exit?
|
||||
async def request_quotes():
|
||||
"""Get quotes for current symbol subscription set.
|
||||
"""
|
||||
symbols = get_topics()
|
||||
# subscription can be changed at any time
|
||||
return await get_quotes(symbols) if symbols else ()
|
||||
|
||||
while True: # use an event here to trigger exit?
|
||||
prequote_start = time.time()
|
||||
|
||||
with trio.move_on_after(3) as cancel_scope:
|
||||
|
|
@ -71,14 +79,13 @@ async def stream_quotes(
|
|||
if cancelled:
|
||||
log.warn("Quote query timed out after 3 seconds, retrying...")
|
||||
# handle network outages by idling until response is received
|
||||
# quotes = await wait_for_network(partial(get_quotes, tickers))
|
||||
quotes = await wait_for_network(request_quotes)
|
||||
|
||||
new_quotes = []
|
||||
new_quotes = {}
|
||||
if diff_cached:
|
||||
# If cache is enabled then only deliver "new" changes.
|
||||
# Useful for polling setups but obviously should be
|
||||
# disabled if you're rx-ing event data.
|
||||
# disabled if you're rx-ing per-tick data.
|
||||
for quote in quotes:
|
||||
symbol = quote['symbol']
|
||||
last = _cache.setdefault(symbol, {})
|
||||
|
|
@ -87,10 +94,11 @@ async def stream_quotes(
|
|||
log.info(
|
||||
f"New quote {quote['symbol']}:\n{new}")
|
||||
_cache[symbol] = quote
|
||||
new_quotes.append(quote)
|
||||
new_quotes[symbol] = quote
|
||||
else:
|
||||
new_quotes = quotes
|
||||
log.info(f"Delivering quotes:\n{quotes}")
|
||||
for quote in quotes:
|
||||
newquotes[quote['symbol']] = quote
|
||||
|
||||
yield new_quotes
|
||||
|
||||
|
|
@ -112,7 +120,8 @@ async def stream_quotes(
|
|||
# TODO: at this point probably just just make this a class and
|
||||
# a lot of these functions should be methods. It will definitely
|
||||
# make stateful UI apps easier to implement
|
||||
class BrokerFeed(typing.NamedTuple):
|
||||
@dataclass
|
||||
class BrokerFeed:
|
||||
"""A per broker "client feed" container.
|
||||
|
||||
A structure to keep track of components used by
|
||||
|
|
@ -124,20 +133,26 @@ class BrokerFeed(typing.NamedTuple):
|
|||
mod: ModuleType
|
||||
client: object
|
||||
exit_stack: contextlib.AsyncExitStack
|
||||
quoter_keys: List[str] = ['stock', 'option']
|
||||
tasks: Dict[str, trio.Event] = dict.fromkeys(
|
||||
quoter_keys, False)
|
||||
quoters: Dict[str, typing.Coroutine] = {}
|
||||
subscriptions: Dict[str, Dict[str, set]] = {'option': {}, 'stock': {}}
|
||||
quoter_keys: Tuple[str] = ('stock', 'option')
|
||||
locks: Dict[str, trio.StrictFIFOLock] = field(
|
||||
default_factory=lambda:
|
||||
{'stock': trio.StrictFIFOLock(), 'option': trio.StrictFIFOLock()}
|
||||
)
|
||||
quoters: Dict[str, typing.Coroutine] = field(default_factory=dict)
|
||||
subscriptions: Dict[str, Dict[str, set]] = field(
|
||||
default_factory=partial(dict, **{'option': {}, 'stock': {}})
|
||||
)
|
||||
|
||||
|
||||
async def fan_out_to_chans(
|
||||
async def fan_out_to_ctxs(
|
||||
pub_gen: typing.AsyncGenerator,
|
||||
feed: BrokerFeed,
|
||||
get_quotes: Coroutine,
|
||||
symbols2chans: Dict[str, tractor.Channel],
|
||||
topics2ctxs: Dict[str, tractor.Context],
|
||||
topic_key: str = 'key',
|
||||
packet_key: str = 'symbol',
|
||||
rate: int = 5, # delay between quote requests
|
||||
diff_cached: bool = True, # only deliver "new" quotes to the queue
|
||||
cid: str = None,
|
||||
) -> None:
|
||||
"""Request and fan out quotes to each subscribed actor channel.
|
||||
"""
|
||||
|
|
@ -146,45 +161,44 @@ async def fan_out_to_chans(
|
|||
rate = broker_limit
|
||||
log.warn(f"Limiting {feed.mod.__name__} query rate to {rate}/sec")
|
||||
|
||||
async def request():
|
||||
"""Get quotes for current symbol subscription set.
|
||||
"""
|
||||
symbols = list(symbols2chans.keys())
|
||||
# subscription can be changed at any time
|
||||
return await get_quotes(symbols) if symbols else ()
|
||||
def get_topics():
|
||||
return tuple(topics2ctxs.keys())
|
||||
|
||||
async for quotes in stream_quotes(
|
||||
feed.mod, request, rate,
|
||||
async for published in pub_gen(
|
||||
get_topics,
|
||||
get_quotes,
|
||||
feed.mod,
|
||||
rate,
|
||||
diff_cached=diff_cached,
|
||||
):
|
||||
chan_payloads = {}
|
||||
for quote in quotes:
|
||||
packet = {quote['symbol']: quote}
|
||||
for chan, cid in symbols2chans.get(quote['key'], set()):
|
||||
chan_payloads.setdefault(
|
||||
(chan, cid),
|
||||
{'yield': {}, 'cid': cid}
|
||||
)['yield'].update(packet)
|
||||
ctx_payloads = {}
|
||||
for packet_key, data in published.items():
|
||||
# grab each suscription topic using provided key for lookup
|
||||
topic = data[topic_key]
|
||||
# build a new dict packet for passing to multiple underlying channels
|
||||
packet = {packet_key: data}
|
||||
for ctx in topics2ctxs.get(topic, set()):
|
||||
ctx_payloads.setdefault(ctx, {}).update(packet),
|
||||
|
||||
# deliver to each subscriber (fan out)
|
||||
if chan_payloads:
|
||||
for (chan, cid), payload in chan_payloads.items():
|
||||
if ctx_payloads:
|
||||
for ctx, payload in ctx_payloads.items():
|
||||
try:
|
||||
await chan.send(payload)
|
||||
await ctx.send_yield(payload)
|
||||
except (
|
||||
# That's right, anything you can think of...
|
||||
trio.ClosedStreamError, ConnectionResetError,
|
||||
ConnectionRefusedError,
|
||||
):
|
||||
log.warn(f"{chan} went down?")
|
||||
for chanset in symbols2chans.values():
|
||||
chanset.discard((chan, cid))
|
||||
log.warn(f"{ctx.chan} went down?")
|
||||
for ctx_set in topics2ctxs.values():
|
||||
ctx_set.discard(ctx)
|
||||
|
||||
if not any(symbols2chans.values()):
|
||||
if not any(topics2ctxs.values()):
|
||||
log.warn(f"No subs left for broker {feed.mod.name}, exiting task")
|
||||
break
|
||||
|
||||
log.info(f"Terminating stream quoter task for {feed.mod.name}")
|
||||
log.info(f"Terminating stream quoter task for {pub_gen.__name__}")
|
||||
|
||||
|
||||
async def symbol_data(broker: str, tickers: List[str]):
|
||||
|
|
@ -241,12 +255,12 @@ async def smoke_quote(get_quotes, tickers, broker):
|
|||
###########################################
|
||||
|
||||
|
||||
def modify_quote_stream(broker, feed_type, symbols, chan, cid):
|
||||
def modify_quote_stream(broker, feed_type, symbols, ctx):
|
||||
"""Absolute symbol subscription list for each quote stream.
|
||||
|
||||
Effectively a symbol subscription api.
|
||||
"""
|
||||
log.info(f"{chan} changed symbol subscription to {symbols}")
|
||||
log.info(f"{ctx.chan} changed symbol subscription to {symbols}")
|
||||
ss = tractor.current_actor().statespace
|
||||
feed = ss['feeds'].get(broker)
|
||||
if feed is None:
|
||||
|
|
@ -254,26 +268,23 @@ def modify_quote_stream(broker, feed_type, symbols, chan, cid):
|
|||
"`get_cached_feed()` must be called before modifying its stream"
|
||||
)
|
||||
|
||||
symbols2chans = feed.subscriptions[feed_type]
|
||||
symbols2ctxs = feed.subscriptions[feed_type]
|
||||
# update map from each symbol to requesting client's chan
|
||||
for ticker in symbols:
|
||||
symbols2chans.setdefault(ticker, set()).add((chan, cid))
|
||||
symbols2ctxs.setdefault(ticker, set()).add(ctx)
|
||||
|
||||
# remove any existing symbol subscriptions if symbol is not
|
||||
# found in ``symbols``
|
||||
# TODO: this can likely be factored out into the pub-sub api
|
||||
for ticker in filter(
|
||||
lambda ticker: ticker not in symbols, symbols2chans.copy()
|
||||
lambda ticker: ticker not in symbols, symbols2ctxs.copy()
|
||||
):
|
||||
chanset = symbols2chans.get(ticker)
|
||||
# XXX: cid will be different on unsub call
|
||||
for item in chanset.copy():
|
||||
if (chan, cid) == item:
|
||||
chanset.discard(item)
|
||||
ctx_set = symbols2ctxs.get(ticker)
|
||||
ctx_set.discard(ctx)
|
||||
|
||||
if not chanset:
|
||||
if not ctx_set:
|
||||
# pop empty sets which will trigger bg quoter task termination
|
||||
symbols2chans.pop(ticker)
|
||||
symbols2ctxs.pop(ticker)
|
||||
|
||||
|
||||
async def get_cached_feed(
|
||||
|
|
@ -310,8 +321,7 @@ async def start_quote_stream(
|
|||
symbols: List[Any],
|
||||
feed_type: str = 'stock',
|
||||
diff_cached: bool = True,
|
||||
chan: tractor.Channel = None,
|
||||
cid: str = None,
|
||||
ctx: tractor.Context = None,
|
||||
rate: int = 3,
|
||||
) -> None:
|
||||
"""Handle per-broker quote stream subscriptions using a "lazy" pub-sub
|
||||
|
|
@ -321,16 +331,17 @@ async def start_quote_stream(
|
|||
Since most brokers seems to support batch quote requests we
|
||||
limit to one task per process for now.
|
||||
"""
|
||||
actor = tractor.current_actor()
|
||||
# set log level after fork
|
||||
get_console_log(actor.loglevel)
|
||||
# XXX: why do we need this again?
|
||||
get_console_log(tractor.current_actor().loglevel)
|
||||
|
||||
# pull global vars from local actor
|
||||
symbols = list(symbols)
|
||||
log.info(
|
||||
f"{chan.uid} subscribed to {broker} for symbols {symbols}")
|
||||
f"{ctx.chan.uid} subscribed to {broker} for symbols {symbols}")
|
||||
# another actor task may have already created it
|
||||
feed = await get_cached_feed(broker)
|
||||
symbols2chans = feed.subscriptions[feed_type]
|
||||
symbols2ctxs = feed.subscriptions[feed_type]
|
||||
task_is_dead = None
|
||||
|
||||
if feed_type == 'stock':
|
||||
get_quotes = feed.quoters.setdefault(
|
||||
|
|
@ -341,7 +352,7 @@ async def start_quote_stream(
|
|||
# out bad symbols for now)
|
||||
payload = await smoke_quote(get_quotes, symbols, broker)
|
||||
# push initial smoke quote response for client initialization
|
||||
await chan.send({'yield': payload, 'cid': cid})
|
||||
await ctx.send_yield(payload)
|
||||
elif feed_type == 'option':
|
||||
# FIXME: yeah we need maybe a more general way to specify
|
||||
# the arg signature for the option feed beasides a symbol
|
||||
|
|
@ -355,77 +366,51 @@ async def start_quote_stream(
|
|||
for quote in await get_quotes(symbols)
|
||||
}
|
||||
# push initial smoke quote response for client initialization
|
||||
await chan.send({'yield': payload, 'cid': cid})
|
||||
await ctx.send_yield(payload)
|
||||
try:
|
||||
# update map from each symbol to requesting client's chan
|
||||
modify_quote_stream(broker, feed_type, symbols, chan, cid)
|
||||
modify_quote_stream(broker, feed_type, symbols, ctx)
|
||||
|
||||
# event indicating that task was started and then killed
|
||||
task_is_dead = feed.tasks.get(feed_type)
|
||||
if task_is_dead is False:
|
||||
task_is_dead = trio.Event()
|
||||
task_is_dead.set()
|
||||
feed.tasks[feed_type] = task_is_dead
|
||||
# prevents more then one broker feed task from spawning
|
||||
lock = feed.locks.get(feed_type)
|
||||
|
||||
if not task_is_dead.is_set():
|
||||
# block and let existing feed task deliver
|
||||
# stream data until it is cancelled in which case
|
||||
# we'll take over and spawn it again
|
||||
await task_is_dead.wait()
|
||||
# client channel was likely disconnected
|
||||
# but we still want to keep the broker task
|
||||
# alive if there are other consumers (including
|
||||
# ourselves)
|
||||
if any(symbols2chans.values()):
|
||||
log.warn(
|
||||
f"Data feed task for {feed.mod.name} was cancelled but"
|
||||
f" there are still active clients, respawning")
|
||||
|
||||
async with lock:
|
||||
# no data feeder task yet; so start one
|
||||
respawn = True
|
||||
while respawn:
|
||||
respawn = False
|
||||
log.info(f"Spawning data feed task for {feed.mod.name}")
|
||||
try:
|
||||
async with trio.open_nursery() as nursery:
|
||||
nursery.start_soon(
|
||||
partial(
|
||||
fan_out_to_chans, feed, get_quotes,
|
||||
symbols2chans,
|
||||
# unblocks when no more symbols subscriptions exist and the
|
||||
# quote streamer task terminates
|
||||
await fan_out_to_ctxs(
|
||||
stream_quotes,
|
||||
feed,
|
||||
get_quotes,
|
||||
symbols2ctxs,
|
||||
diff_cached=diff_cached,
|
||||
cid=cid,
|
||||
rate=rate,
|
||||
)
|
||||
)
|
||||
# it's alive!
|
||||
task_is_dead.clear()
|
||||
|
||||
except trio.BrokenResourceError:
|
||||
log.exception("Respawning failed data feed task")
|
||||
respawn = True
|
||||
|
||||
# unblocks when no more symbols subscriptions exist and the
|
||||
# quote streamer task terminates (usually because another call
|
||||
# was made to `modify_quoter` to unsubscribe from streaming
|
||||
# symbols)
|
||||
finally:
|
||||
log.info(f"Terminated {feed_type} quoter task for {feed.mod.name}")
|
||||
task_is_dead.set()
|
||||
|
||||
# if we're cancelled externally unsubscribe our quote feed
|
||||
modify_quote_stream(broker, feed_type, [], chan, cid)
|
||||
modify_quote_stream(broker, feed_type, [], ctx)
|
||||
|
||||
# if there are truly no more subscriptions with this broker
|
||||
# drop from broker subs dict
|
||||
if not any(symbols2chans.values()):
|
||||
log.info(f"No more subscriptions for {broker}")
|
||||
# broker2symbolsubs.pop(broker, None)
|
||||
if not any(symbols2ctxs.values()):
|
||||
log.info(f"No more subscriptions for broker {broker}")
|
||||
|
||||
# destroy the API client
|
||||
await feed.exit_stack.aclose()
|
||||
|
||||
|
||||
class DataFeed(object):
|
||||
class DataFeed:
|
||||
"""Data feed client for streaming symbol data from a (remote)
|
||||
``brokerd`` data daemon.
|
||||
"""
|
||||
|
|
@ -472,7 +457,7 @@ class DataFeed(object):
|
|||
filename=test
|
||||
)
|
||||
else:
|
||||
log.info(f"Starting new stream for {self._symbols}")
|
||||
log.info(f"Starting new stream for {symbols}")
|
||||
# start live streaming from broker daemon
|
||||
quote_gen = await self.portal.run(
|
||||
"piker.brokers.data",
|
||||
|
|
|
|||
Loading…
Reference in New Issue