karaolidis/qrust
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Separate clock handler

Browse Source 
Signed-off-by: Nikolaos Karaolidis <nick@karaolidis.com>
This commit is contained in:
Nick Karaolidis
2024-01-20 18:22:36 +00:00
parent 2d14fe35c8
commit 178a062c25
15 changed files with 262 additions and 221 deletions
Download Patch File Download Diff File Expand all files Collapse all files

352
src/data/market.rs
View File

@@ -1,13 +1,10 @@
use crate::{
config::{
Config, ALPACA_CLOCK_API_URL, ALPACA_CRYPTO_WEBSOCKET_URL, ALPACA_STOCK_WEBSOCKET_URL,
},
config::{Config, ALPACA_CRYPTO_WEBSOCKET_URL, ALPACA_STOCK_WEBSOCKET_URL},
data::authenticate_websocket,
database,
state::{self, BroadcastMessage},
types::{
alpaca::{api, websocket, Source},
Asset, Backfill, Bar, Class,
state, Asset, Backfill, Bar, BroadcastMessage, Class,
},
utils::{duration_until, last_minute, FIFTEEN_MINUTES, ONE_MINUTE},
};
@@ -26,10 +23,7 @@ use time::OffsetDateTime;
use tokio::{
net::TcpStream,
spawn,
sync::{
broadcast::{Receiver, Sender},
Mutex, RwLock,
},
sync::{broadcast::Sender, Mutex, RwLock},
task::JoinHandle,
time::sleep,
};
@@ -47,7 +41,7 @@ pub async fn run(
class: Class,
broadcast_bus_sender: Sender<BroadcastMessage>,
) {
info!("Running live data threads for {:?}.", class);
info!("Running live threads for {:?}.", class);
let websocket_url = match class {
Class::UsEquity => format!(
@@ -73,7 +67,7 @@ pub async fn run(
app_config.clone(),
class,
sink.clone(),
broadcast_bus_sender.subscribe(),
broadcast_bus_sender.clone(),
guard.clone(),
));
@@ -85,12 +79,6 @@ pub async fn run(
guard.clone(),
));
spawn(clock_handler(
app_config.clone(),
class,
broadcast_bus_sender.clone(),
));
let assets = database::assets::select_where_class(&app_config.clickhouse_client, &class).await;
broadcast_bus_sender
.send(BroadcastMessage::Asset((
@@ -100,184 +88,159 @@ pub async fn run(
.unwrap();
}
#[allow(clippy::too_many_lines)]
#[allow(clippy::significant_drop_tightening)]
pub async fn broadcast_bus_handler(
app_config: Arc<Config>,
class: Class,
sink: Arc<Mutex<SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>>>,
mut broadcast_bus_receiver: Receiver<BroadcastMessage>,
broadcast_bus_sender: Sender<BroadcastMessage>,
guard: Arc<RwLock<Guard>>,
) {
let mut broadcast_bus_receiver = broadcast_bus_sender.subscribe();
loop {
match broadcast_bus_receiver.recv().await.unwrap() {
BroadcastMessage::Asset((action, mut assets)) => {
assets.retain(|asset| asset.class == class);
let app_config = app_config.clone();
let sink = sink.clone();
let broadcast_bus_sender = broadcast_bus_sender.clone();
let guard = guard.clone();
let message = broadcast_bus_receiver.recv().await.unwrap();
if assets.is_empty() {
continue;
}
let symbols = assets
.iter()
.map(|asset| asset.symbol.clone())
.collect::<Vec<_>>();
match action {
state::asset::BroadcastMessage::Add => {
database::assets::upsert_batch(&app_config.clickhouse_client, &assets)
.await;
info!("Added {:?}.", symbols);
let mut guard = guard.write().await;
guard.pending_subscriptions.extend(
assets
.into_iter()
.map(|asset| (asset.symbol.clone(), asset)),
);
guard.symbols.extend(symbols.clone());
sink.lock()
.await
.send(Message::Text(
to_string(&websocket::data::outgoing::Message::Subscribe(
websocket::data::outgoing::subscribe::Message::new(symbols),
))
.unwrap(),
))
.await
.unwrap();
}
state::asset::BroadcastMessage::Delete => {
database::assets::delete_where_symbols(
&app_config.clickhouse_client,
&symbols,
)
.await;
info!("Deleted {:?}.", symbols);
let mut guard = guard.write().await;
guard.pending_unsubscriptions.extend(
assets
.into_iter()
.map(|asset| (asset.symbol.clone(), asset)),
);
guard.symbols.retain(|symbol| !symbols.contains(symbol));
sink.lock()
.await
.send(Message::Text(
to_string(&websocket::data::outgoing::Message::Unsubscribe(
websocket::data::outgoing::subscribe::Message::new(symbols),
))
.unwrap(),
))
.await
.unwrap();
}
state::asset::BroadcastMessage::Backfill => {
let guard_clone = guard.clone();
let mut guard = guard.write().await;
info!("Creating backfill jobs for {:?}.", symbols);
for asset in assets {
if let Some(backfill_job) = guard.backfill_jobs.remove(&asset.symbol) {
backfill_job.abort();
backfill_job.await.unwrap_err();
}
guard.backfill_jobs.insert(asset.symbol.clone(), {
let guard = guard_clone.clone();
let app_config = app_config.clone();
spawn(async move {
backfill(app_config, class, asset.clone()).await;
let mut guard = guard.write().await;
guard.backfill_jobs.remove(&asset.symbol);
})
});
}
}
state::asset::BroadcastMessage::Purge => {
let mut guard = guard.write().await;
info!("Purging {:?}.", symbols);
for asset in assets {
if let Some(backfill_job) = guard.backfill_jobs.remove(&asset.symbol) {
backfill_job.abort();
backfill_job.await.unwrap_err();
}
}
database::backfills::delete_where_symbols(
&app_config.clickhouse_client,
&symbols,
)
.await;
database::bars::delete_where_symbols(
&app_config.clickhouse_client,
&symbols,
)
.await;
}
}
}
}
spawn(broadcast_bus_handle_message(
app_config,
class,
sink,
broadcast_bus_sender,
guard,
message,
));
}
}
pub async fn clock_handler(
#[allow(clippy::significant_drop_tightening)]
#[allow(clippy::too_many_lines)]
async fn broadcast_bus_handle_message(
app_config: Arc<Config>,
class: Class,
sink: Arc<Mutex<SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>>>,
broadcast_bus_sender: Sender<BroadcastMessage>,
guard: Arc<RwLock<Guard>>,
message: BroadcastMessage,
) {
loop {
let clock = retry(ExponentialBackoff::default(), || async {
app_config.alpaca_rate_limit.until_ready().await;
app_config
.alpaca_client
.get(ALPACA_CLOCK_API_URL)
.send()
.await?
.error_for_status()?
.json::<api::incoming::clock::Clock>()
.await
.map_err(backoff::Error::Permanent)
})
.await
.unwrap();
let sleep_until = duration_until(if clock.is_open {
if class == Class::UsEquity {
info!("Market is open, will close at {}.", clock.next_close);
match message {
BroadcastMessage::Asset((action, mut assets)) => {
assets.retain(|asset| asset.class == class);
if assets.is_empty() {
return;
}
clock.next_close
} else {
if class == Class::UsEquity {
info!("Market is closed, will reopen at {}.", clock.next_open);
let assets = assets
.into_iter()
.map(|asset| (asset.symbol.clone(), asset))
.collect::<HashMap<_, _>>();
let symbols = assets.keys().cloned().collect::<Vec<_>>();
match action {
state::asset::BroadcastMessage::Add => {
database::assets::upsert_batch(
&app_config.clickhouse_client,
assets.clone().into_values(),
)
.await;
let mut guard = guard.write().await;
guard.symbols.extend(symbols.clone());
guard.pending_subscriptions.extend(assets);
info!("Added {:?}.", symbols);
sink.lock()
.await
.send(Message::Text(
to_string(&websocket::data::outgoing::Message::Subscribe(
websocket::data::outgoing::subscribe::Message::new(symbols),
))
.unwrap(),
))
.await
.unwrap();
}
state::asset::BroadcastMessage::Delete => {
database::assets::delete_where_symbols(&app_config.clickhouse_client, &symbols)
.await;
let mut guard = guard.write().await;
guard.symbols.retain(|symbol| !assets.contains_key(symbol));
guard.pending_unsubscriptions.extend(assets);
info!("Deleted {:?}.", symbols);
sink.lock()
.await
.send(Message::Text(
to_string(&websocket::data::outgoing::Message::Unsubscribe(
websocket::data::outgoing::subscribe::Message::new(symbols),
))
.unwrap(),
))
.await
.unwrap();
}
state::asset::BroadcastMessage::Backfill => {
let guard_clone = guard.clone();
let mut guard = guard.write().await;
info!("Creating backfill jobs for {:?}.", symbols);
for (symbol, asset) in assets {
if let Some(backfill_job) = guard.backfill_jobs.remove(&symbol) {
backfill_job.abort();
backfill_job.await.unwrap_err();
}
guard.backfill_jobs.insert(symbol.clone(), {
let guard = guard_clone.clone();
let app_config = app_config.clone();
spawn(async move {
backfill(app_config, class, asset.clone()).await;
let mut guard = guard.write().await;
guard.backfill_jobs.remove(&symbol);
})
});
}
}
state::asset::BroadcastMessage::Purge => {
let mut guard = guard.write().await;
info!("Purging {:?}.", symbols);
for (symbol, _) in assets {
if let Some(backfill_job) = guard.backfill_jobs.remove(&symbol) {
backfill_job.abort();
backfill_job.await.unwrap_err();
}
}
database::backfills::delete_where_symbols(
&app_config.clickhouse_client,
&symbols,
)
.await;
database::bars::delete_where_symbols(&app_config.clickhouse_client, &symbols)
.await;
}
}
clock.next_open
});
sleep(sleep_until).await;
let assets = database::assets::select(&app_config.clickhouse_client).await;
broadcast_bus_sender
.send(BroadcastMessage::Asset((
state::asset::BroadcastMessage::Backfill,
assets,
)))
.unwrap();
}
BroadcastMessage::Clock(_) => {
broadcast_bus_sender
.send(BroadcastMessage::Asset((
state::asset::BroadcastMessage::Backfill,
database::assets::select(&app_config.clickhouse_client).await,
)))
.unwrap();
}
}
}
@@ -343,22 +306,18 @@ async fn websocket_handle_message(
let newly_subscribed_assets = guard
.pending_subscriptions
.extract_if(|symbol, _| symbols.contains(symbol))
.map(|(_, asset)| asset)
.collect::<Vec<_>>();
.collect::<HashMap<_, _>>();
if !newly_subscribed_assets.is_empty() {
info!(
"Subscribed to {:?}.",
newly_subscribed_assets
.iter()
.map(|asset| asset.symbol.clone())
.collect::<Vec<_>>()
newly_subscribed_assets.keys().collect::<Vec<_>>()
);
broadcast_bus_sender
.send(BroadcastMessage::Asset((
state::asset::BroadcastMessage::Backfill,
newly_subscribed_assets,
newly_subscribed_assets.into_values().collect::<Vec<_>>(),
)))
.unwrap();
}
@@ -366,22 +325,18 @@ async fn websocket_handle_message(
let newly_unsubscribed_assets = guard
.pending_unsubscriptions
.extract_if(|symbol, _| !symbols.contains(symbol))
.map(|(_, asset)| asset)
.collect::<Vec<_>>();
.collect::<HashMap<_, _>>();
if !newly_unsubscribed_assets.is_empty() {
info!(
"Unsubscribed from {:?}.",
newly_unsubscribed_assets
.iter()
.map(|asset| asset.symbol.clone())
.collect::<Vec<_>>()
newly_unsubscribed_assets.keys().collect::<Vec<_>>()
);
broadcast_bus_sender
.send(BroadcastMessage::Asset((
state::asset::BroadcastMessage::Purge,
newly_unsubscribed_assets,
newly_unsubscribed_assets.into_values().collect::<Vec<_>>(),
)))
.unwrap();
}
@@ -429,13 +384,13 @@ pub async fn backfill(app_config: Arc<Config>, class: Class, asset: Asset) {
if app_config.alpaca_source == Source::Iex {
let task_run_delay = duration_until(fetch_until + FIFTEEN_MINUTES + ONE_MINUTE);
info!(
"Queing historical data backfill for {} in {:?}.",
"Queing backfill for {} in {:?}.",
asset.symbol, task_run_delay
);
sleep(task_run_delay).await;
}
info!("Running historical data backfill for {}.", asset.symbol);
info!("Running backfill for {}.", asset.symbol);
let mut bars = Vec::new();
let mut next_page_token = None;
@@ -466,10 +421,7 @@ pub async fn backfill(app_config: Arc<Config>, class: Class, asset: Asset) {
let message = match message {
Ok(message) => message,
Err(e) => {
error!(
"Failed to backfill historical data for {}: {}.",
asset.symbol, e
);
error!("Failed to backfill data for {}: {}.", asset.symbol, e);
return;
}
};
@@ -486,12 +438,12 @@ pub async fn backfill(app_config: Arc<Config>, class: Class, asset: Asset) {
next_page_token = message.next_page_token;
}
database::bars::upsert_batch(&app_config.clickhouse_client, &bars).await;
database::bars::upsert_batch(&app_config.clickhouse_client, bars).await;
database::backfills::upsert(
&app_config.clickhouse_client,
&Backfill::new(asset.symbol.clone(), fetch_until),
)
.await;
info!("Backfilled historical data for {}.", asset.symbol);
info!("Backfilled data for {}.", asset.symbol);
}