Refactor streaming logic to be less nested and readable
parent
8403d8a482
commit
9706803220
|
|
@ -32,6 +32,7 @@ import hmac
|
|||
import hashlib
|
||||
import wsproto
|
||||
from uuid import uuid4
|
||||
from functools import partial
|
||||
|
||||
import asks
|
||||
import tractor
|
||||
|
|
@ -443,6 +444,28 @@ async def open_symbol_search(
|
|||
await stream.send(await client.search_symbols(pattern))
|
||||
log.info('Kucoin symbol search opened')
|
||||
|
||||
@acm
|
||||
async def open_ping_task(ws: wsproto.WSConnection, ping_interval, connect_id):
|
||||
'''
|
||||
Spawn a non-blocking task that pings the ws
|
||||
server every ping_interval so Kucoin doesn't drop
|
||||
our connection
|
||||
|
||||
'''
|
||||
async with trio.open_nursery() as n:
|
||||
# TODO: cache this task so it's only called once
|
||||
async def ping_server():
|
||||
while True:
|
||||
await trio.sleep((ping_interval - 1000) / 1000)
|
||||
await ws.send_msg({'id': connect_id, 'type': 'ping'})
|
||||
|
||||
log.info(f'Starting ping task for kucoin ws connection')
|
||||
n.start_soon(ping_server)
|
||||
|
||||
yield ws
|
||||
|
||||
n.cancel_scope.cancel()
|
||||
|
||||
|
||||
async def stream_quotes(
|
||||
send_chan: trio.abc.SendChannel,
|
||||
|
|
@ -457,14 +480,21 @@ async def stream_quotes(
|
|||
Where the rubber hits the road baby
|
||||
|
||||
'''
|
||||
connect_id = str(uuid4())
|
||||
|
||||
async with open_cached_client('kucoin') as client:
|
||||
token, ping_interval = await client._get_ws_token()
|
||||
connect_id = str(uuid4())
|
||||
pairs = await client.cache_pairs()
|
||||
|
||||
# open ping task
|
||||
async with (
|
||||
open_autorecon_ws(
|
||||
f'wss://ws-api-spot.kucoin.com/?token={token}&[connectId={connect_id}]'
|
||||
) as ws,
|
||||
open_ping_task(ws, ping_interval, connect_id) as ws,
|
||||
):
|
||||
log.info('Starting up quote stream')
|
||||
# loop through symbols and sub to feedz
|
||||
for sym in symbols:
|
||||
token, ping_interval = await client._get_ws_token()
|
||||
pairs = await client.cache_pairs()
|
||||
pair: KucoinMktPair = pairs[sym]
|
||||
kucoin_sym = pair.symbol
|
||||
|
||||
|
|
@ -479,69 +509,15 @@ async def stream_quotes(
|
|||
},
|
||||
'shm_write_opts': {'sum_tick_vml': False},
|
||||
'fqsn': sym,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@acm
|
||||
async def subscribe(ws: wsproto.WSConnection):
|
||||
|
||||
@acm
|
||||
async def open_ping_task(ws: wsproto.WSConnection):
|
||||
'''
|
||||
Spawn a non-blocking task that pings the ws
|
||||
server every ping_interval so Kucoin doesn't drop
|
||||
our connection
|
||||
|
||||
'''
|
||||
async with trio.open_nursery() as n:
|
||||
# TODO: cache this task so it's only called once
|
||||
async def ping_server():
|
||||
while True:
|
||||
await trio.sleep((ping_interval - 1000) / 1000)
|
||||
await ws.send_msg({'id': connect_id, 'type': 'ping'})
|
||||
|
||||
log.info(f'Starting ping task for {sym}')
|
||||
n.start_soon(ping_server)
|
||||
|
||||
yield ws
|
||||
|
||||
n.cancel_scope.cancel()
|
||||
|
||||
# Spawn the ping task here
|
||||
async with open_ping_task(ws) as ws:
|
||||
tasks = []
|
||||
tasks.append(make_sub(kucoin_sym, connect_id, level='l3'))
|
||||
tasks.append(make_sub(kucoin_sym, connect_id, level='l1'))
|
||||
|
||||
for task in tasks:
|
||||
log.info(
|
||||
f'Subscribing to {task["topic"]} feed for {sym}')
|
||||
await ws.send_msg(task)
|
||||
|
||||
yield
|
||||
|
||||
# unsub
|
||||
if ws.connected():
|
||||
log.info(f'Unsubscribing to {kucoin_sym} feed')
|
||||
await ws.send_msg(
|
||||
{
|
||||
'id': connect_id,
|
||||
'type': 'unsubscribe',
|
||||
'topic': f'/market/ticker:{sym}',
|
||||
'privateChannel': False,
|
||||
'response': True,
|
||||
}
|
||||
)
|
||||
|
||||
async with (
|
||||
open_autorecon_ws(
|
||||
f'wss://ws-api-spot.kucoin.com/?token={token}&[connectId={connect_id}]',
|
||||
fixture=subscribe,
|
||||
) as ws,
|
||||
subscribe(ws, connect_id, kucoin_sym),
|
||||
stream_messages(ws, sym) as msg_gen,
|
||||
):
|
||||
typ, quote = await anext(msg_gen)
|
||||
|
||||
while typ != 'trade':
|
||||
# take care to not unblock here until we get a real trade quote
|
||||
typ, quote = await anext(msg_gen)
|
||||
|
|
@ -552,26 +528,41 @@ async def stream_quotes(
|
|||
async for typ, msg in msg_gen:
|
||||
await send_chan.send({sym: msg})
|
||||
|
||||
|
||||
def make_sub(sym, connect_id, level='l1') -> dict[str, str | bool] | None:
|
||||
match level:
|
||||
case 'l1':
|
||||
return {
|
||||
@acm
|
||||
async def subscribe(ws: wsproto.WSConnection, connect_id, sym):
|
||||
# breakpoint()
|
||||
# level 2 sub
|
||||
await ws.send_msg({
|
||||
'id': connect_id,
|
||||
'type': 'subscribe',
|
||||
'topic': f'/spotMarket/level2Depth5:{sym}',
|
||||
'privateChannel': False,
|
||||
'response': True,
|
||||
}
|
||||
})
|
||||
|
||||
case 'l3':
|
||||
return {
|
||||
# watch trades
|
||||
await ws.send_msg({
|
||||
'id': connect_id,
|
||||
'type': 'subscribe',
|
||||
'topic': f'/market/ticker:{sym}',
|
||||
'privateChannel': False,
|
||||
'response': True,
|
||||
})
|
||||
|
||||
yield
|
||||
|
||||
# unsub
|
||||
if ws.connected():
|
||||
log.info(f'Unsubscribing to {syn} feed')
|
||||
await ws.send_msg(
|
||||
{
|
||||
'id': connect_id,
|
||||
'type': 'unsubscribe',
|
||||
'topic': f'/market/ticker:{sym}',
|
||||
'privateChannel': False,
|
||||
'response': True,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@trio_async_generator
|
||||
|
|
|
|||
Loading…
Reference in New Issue