Add asynchronous handling
Signed-off-by: Nikolaos Karaolidis <nick@karaolidis.com>
This commit is contained in:
@@ -6,10 +6,7 @@ use crate::{
|
||||
data::authenticate_websocket,
|
||||
database,
|
||||
types::{
|
||||
alpaca::{
|
||||
api::{incoming, outgoing},
|
||||
websocket, Source,
|
||||
},
|
||||
alpaca::{api, websocket, Source},
|
||||
asset::{self, Asset},
|
||||
Bar, BarValidity, BroadcastMessage, Class,
|
||||
},
|
||||
@@ -63,16 +60,13 @@ pub async fn run(
|
||||
broadcast_sender.subscribe(),
|
||||
));
|
||||
|
||||
database::assets::select_where_class(&app_config.clickhouse_client, class)
|
||||
.await
|
||||
.into_iter()
|
||||
.for_each(|asset| {
|
||||
broadcast_sender
|
||||
.send(BroadcastMessage::Asset(asset::BroadcastMessage::Added(
|
||||
asset,
|
||||
)))
|
||||
.unwrap();
|
||||
});
|
||||
let assets = database::assets::select_where_class(&app_config.clickhouse_client, class).await;
|
||||
broadcast_sender
|
||||
.send(BroadcastMessage::Asset((
|
||||
asset::BroadcastMessage::Added,
|
||||
assets,
|
||||
)))
|
||||
.unwrap();
|
||||
|
||||
websocket_handler(app_config, class, stream, sink).await;
|
||||
|
||||
@@ -86,39 +80,41 @@ async fn broadcast_handler(
|
||||
) {
|
||||
loop {
|
||||
match broadcast_receiver.recv().await.unwrap() {
|
||||
BroadcastMessage::Asset(asset::BroadcastMessage::Added(asset))
|
||||
if asset.class == class =>
|
||||
{
|
||||
BroadcastMessage::Asset((action, assets)) => {
|
||||
let symbols = assets
|
||||
.into_iter()
|
||||
.filter(|asset| asset.class == class)
|
||||
.map(|asset| asset.symbol)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
if symbols.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
sink.write()
|
||||
.await
|
||||
.send(Message::Text(
|
||||
serde_json::to_string(&websocket::data::outgoing::Message::Subscribe(
|
||||
websocket::data::outgoing::subscribe::Message::new(
|
||||
asset.clone().symbol,
|
||||
),
|
||||
))
|
||||
serde_json::to_string(&match action {
|
||||
asset::BroadcastMessage::Added => {
|
||||
websocket::data::outgoing::Message::Subscribe(
|
||||
websocket::data::outgoing::subscribe::Message::new(
|
||||
symbols.clone(),
|
||||
),
|
||||
)
|
||||
}
|
||||
asset::BroadcastMessage::Deleted => {
|
||||
websocket::data::outgoing::Message::Unsubscribe(
|
||||
websocket::data::outgoing::subscribe::Message::new(
|
||||
symbols.clone(),
|
||||
),
|
||||
)
|
||||
}
|
||||
})
|
||||
.unwrap(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
BroadcastMessage::Asset(asset::BroadcastMessage::Deleted(asset))
|
||||
if asset.class == class =>
|
||||
{
|
||||
sink.write()
|
||||
.await
|
||||
.send(Message::Text(
|
||||
serde_json::to_string(&websocket::data::outgoing::Message::Unsubscribe(
|
||||
websocket::data::outgoing::subscribe::Message::new(
|
||||
asset.clone().symbol,
|
||||
),
|
||||
))
|
||||
.unwrap(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
BroadcastMessage::Asset(_) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -143,7 +139,12 @@ async fn websocket_handler(
|
||||
}
|
||||
|
||||
for message in parsed_data.unwrap_or_default() {
|
||||
websocket_handle_message(&app_config, class, &backfilled, message).await;
|
||||
spawn(websocket_handle_message(
|
||||
app_config.clone(),
|
||||
class,
|
||||
backfilled.clone(),
|
||||
message,
|
||||
));
|
||||
}
|
||||
}
|
||||
Some(Ok(Message::Ping(_))) => sink
|
||||
@@ -159,68 +160,75 @@ async fn websocket_handler(
|
||||
}
|
||||
|
||||
async fn websocket_handle_message(
|
||||
app_config: &Arc<Config>,
|
||||
app_config: Arc<Config>,
|
||||
class: Class,
|
||||
backfilled: &Arc<RwLock<HashMap<String, bool>>>,
|
||||
backfilled: Arc<RwLock<HashMap<String, bool>>>,
|
||||
message: websocket::data::incoming::Message,
|
||||
) {
|
||||
match message {
|
||||
websocket::data::incoming::Message::Subscription(subscription_message) => {
|
||||
let old_assets = backfilled
|
||||
.read()
|
||||
.await
|
||||
.keys()
|
||||
.cloned()
|
||||
.collect::<HashSet<_>>();
|
||||
let new_assets = subscription_message
|
||||
.bars
|
||||
.into_iter()
|
||||
.collect::<HashSet<_>>();
|
||||
websocket::data::incoming::Message::Subscription(message) => {
|
||||
let added_asset_symbols;
|
||||
let deleted_asset_symbols;
|
||||
|
||||
let added_assets = new_assets.difference(&old_assets).collect::<HashSet<_>>();
|
||||
let deleted_assets = old_assets.difference(&new_assets).collect::<HashSet<_>>();
|
||||
{
|
||||
let mut backfilled = backfilled.write().await;
|
||||
|
||||
for asset_symbol in &added_assets {
|
||||
let old_asset_sybols = backfilled.keys().cloned().collect::<HashSet<_>>();
|
||||
let new_asset_symbols = message.bars.into_iter().collect::<HashSet<_>>();
|
||||
|
||||
added_asset_symbols = new_asset_symbols
|
||||
.difference(&old_asset_sybols)
|
||||
.cloned()
|
||||
.collect::<HashSet<_>>();
|
||||
|
||||
for asset_symbol in &added_asset_symbols {
|
||||
backfilled.insert(asset_symbol.clone(), false);
|
||||
}
|
||||
|
||||
deleted_asset_symbols = old_asset_sybols
|
||||
.difference(&new_asset_symbols)
|
||||
.cloned()
|
||||
.collect::<HashSet<_>>();
|
||||
|
||||
for asset_symbol in &deleted_asset_symbols {
|
||||
backfilled.remove(asset_symbol);
|
||||
}
|
||||
|
||||
drop(backfilled);
|
||||
|
||||
info!(
|
||||
"Subscription update for {:?}: {:?} added, {:?} deleted.",
|
||||
class, added_asset_symbols, deleted_asset_symbols
|
||||
);
|
||||
}
|
||||
|
||||
for asset_symbol in added_asset_symbols {
|
||||
let asset = database::assets::select_where_symbol(
|
||||
&app_config.clickhouse_client,
|
||||
asset_symbol,
|
||||
&asset_symbol,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
backfilled.write().await.insert(asset.symbol.clone(), false);
|
||||
|
||||
let bar_validity = BarValidity::none(asset.symbol.clone());
|
||||
database::bars::insert_validity_if_not_exists(
|
||||
&app_config.clickhouse_client,
|
||||
&bar_validity,
|
||||
&BarValidity::none(asset.symbol.clone()),
|
||||
)
|
||||
.await;
|
||||
|
||||
spawn(backfill(
|
||||
app_config.clone(),
|
||||
backfilled.clone(),
|
||||
asset.clone(),
|
||||
));
|
||||
spawn(backfill(app_config.clone(), backfilled.clone(), asset));
|
||||
}
|
||||
|
||||
for asset_symbol in &deleted_assets {
|
||||
for asset_symbol in deleted_asset_symbols {
|
||||
database::bars::delete_validity_where_symbol(
|
||||
&app_config.clickhouse_client,
|
||||
asset_symbol,
|
||||
&asset_symbol,
|
||||
)
|
||||
.await;
|
||||
|
||||
database::bars::delete_where_symbol(&app_config.clickhouse_client, asset_symbol)
|
||||
database::bars::delete_where_symbol(&app_config.clickhouse_client, &asset_symbol)
|
||||
.await;
|
||||
|
||||
backfilled.write().await.remove(*asset_symbol);
|
||||
}
|
||||
|
||||
info!(
|
||||
"Subscription update for {:?}: {:?} added, {:?} deleted.",
|
||||
class, added_assets, deleted_assets
|
||||
);
|
||||
}
|
||||
websocket::data::incoming::Message::Bars(bar_message)
|
||||
| websocket::data::incoming::Message::UpdatedBars(bar_message) => {
|
||||
@@ -228,7 +236,7 @@ async fn websocket_handle_message(
|
||||
info!("websocket::Incoming bar for {}: {}", bar.symbol, bar.time);
|
||||
|
||||
database::bars::upsert(&app_config.clickhouse_client, &bar).await;
|
||||
if backfilled.read().await[&bar.symbol] {
|
||||
if *backfilled.read().await.get(&bar.symbol).unwrap() {
|
||||
database::bars::upsert_validity(&app_config.clickhouse_client, &bar.into()).await;
|
||||
}
|
||||
}
|
||||
@@ -255,7 +263,7 @@ pub async fn backfill(
|
||||
.send()
|
||||
.await
|
||||
.unwrap()
|
||||
.json::<incoming::clock::Clock>()
|
||||
.json::<api::incoming::clock::Clock>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -294,9 +302,9 @@ pub async fn backfill(
|
||||
Class::UsEquity => ALPACA_STOCK_DATA_URL,
|
||||
Class::Crypto => ALPACA_CRYPTO_DATA_URL,
|
||||
})
|
||||
.query(&outgoing::bar::Bar::new(
|
||||
.query(&api::outgoing::bar::Bar::new(
|
||||
vec![asset.symbol.clone()],
|
||||
String::from("1Min"),
|
||||
ONE_MINUTE,
|
||||
fetch_from,
|
||||
fetch_until,
|
||||
10000,
|
||||
@@ -305,7 +313,7 @@ pub async fn backfill(
|
||||
.send()
|
||||
.await
|
||||
.unwrap()
|
||||
.json::<incoming::bar::Message>()
|
||||
.json::<api::incoming::bar::Message>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
|
||||
Reference in New Issue
Block a user