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

Clean up error propagation

Browse Source 
Signed-off-by: Nikolaos Karaolidis <nick@karaolidis.com>
This commit is contained in:
Nick Karaolidis
2024-02-08 18:13:52 +00:00
parent 52e88f4bc9
commit 76bf2fddcb
24 changed files with 465 additions and 325 deletions
Download Patch File Download Diff File Expand all files Collapse all files

160
src/threads/data/backfill.rs
View File

@@ -9,7 +9,7 @@ use crate::{
Source,
},
news::Prediction,
Bar, Class, News,
Backfill, Bar, Class, News,
},
utils::{
duration_until, last_minute, remove_slash_from_pair, FIFTEEN_MINUTES, ONE_MINUTE,
@@ -18,14 +18,15 @@ use crate::{
};
use async_trait::async_trait;
use futures_util::future::join_all;
use log::{info, warn};
use log::{error, info, warn};
use std::{collections::HashMap, sync::Arc};
use time::OffsetDateTime;
use tokio::{
join, spawn,
spawn,
sync::{mpsc, oneshot, Mutex},
task::{block_in_place, JoinHandle},
time::sleep,
try_join,
};
pub enum Action {
@@ -64,9 +65,12 @@ impl Message {
#[async_trait]
pub trait Handler: Send + Sync {
async fn select_latest_backfill(&self, symbol: String) -> Option<crate::types::Backfill>;
async fn delete_backfills(&self, symbol: &[String]);
async fn delete_data(&self, symbol: &[String]);
async fn select_latest_backfill(
&self,
symbol: String,
) -> Result<Option<Backfill>, clickhouse::error::Error>;
async fn delete_backfills(&self, symbol: &[String]) -> Result<(), clickhouse::error::Error>;
async fn delete_data(&self, symbol: &[String]) -> Result<(), clickhouse::error::Error>;
async fn queue_backfill(&self, symbol: &str, fetch_to: OffsetDateTime);
async fn backfill(&self, symbol: String, fetch_from: OffsetDateTime, fetch_to: OffsetDateTime);
fn log_string(&self) -> &'static str;
@@ -111,13 +115,14 @@ async fn handle_backfill_message(
backfill_jobs.insert(
symbol.clone(),
spawn(async move {
let fetch_from = handler
let fetch_from = match handler
.select_latest_backfill(symbol.clone())
.await
.as_ref()
.map_or(OffsetDateTime::UNIX_EPOCH, |backfill| {
backfill.time + ONE_SECOND
});
.unwrap()
{
Some(latest_backfill) => latest_backfill.time + ONE_SECOND,
None => OffsetDateTime::UNIX_EPOCH,
};
let fetch_to = last_minute();
@@ -142,10 +147,11 @@ async fn handle_backfill_message(
}
}
join!(
try_join!(
handler.delete_backfills(&message.symbols),
handler.delete_data(&message.symbols)
);
)
.unwrap();
}
}
@@ -153,10 +159,10 @@ async fn handle_backfill_message(
}
struct BarHandler {
app_config: Arc<Config>,
config: Arc<Config>,
data_url: &'static str,
api_query_constructor: fn(
app_config: &Arc<Config>,
config: &Arc<Config>,
symbol: String,
fetch_from: OffsetDateTime,
fetch_to: OffsetDateTime,
@@ -165,7 +171,7 @@ struct BarHandler {
}
fn us_equity_query_constructor(
app_config: &Arc<Config>,
config: &Arc<Config>,
symbol: String,
fetch_from: OffsetDateTime,
fetch_to: OffsetDateTime,
@@ -179,7 +185,7 @@ fn us_equity_query_constructor(
limit: Some(10000),
adjustment: None,
asof: None,
feed: Some(app_config.alpaca_source),
feed: Some(config.alpaca_source),
currency: None,
page_token: next_page_token,
sort: Some(Sort::Asc),
@@ -206,30 +212,33 @@ fn crypto_query_constructor(
#[async_trait]
impl Handler for BarHandler {
async fn select_latest_backfill(&self, symbol: String) -> Option<crate::types::Backfill> {
async fn select_latest_backfill(
&self,
symbol: String,
) -> Result<Option<Backfill>, clickhouse::error::Error> {
database::backfills::select_latest_where_symbol(
&self.app_config.clickhouse_client,
&self.config.clickhouse_client,
&database::backfills::Table::Bars,
&symbol,
)
.await
}
async fn delete_backfills(&self, symbols: &[String]) {
async fn delete_backfills(&self, symbols: &[String]) -> Result<(), clickhouse::error::Error> {
database::backfills::delete_where_symbols(
&self.app_config.clickhouse_client,
&self.config.clickhouse_client,
&database::backfills::Table::Bars,
symbols,
)
.await;
.await
}
async fn delete_data(&self, symbols: &[String]) {
database::bars::delete_where_symbols(&self.app_config.clickhouse_client, symbols).await;
async fn delete_data(&self, symbols: &[String]) -> Result<(), clickhouse::error::Error> {
database::bars::delete_where_symbols(&self.config.clickhouse_client, symbols).await
}
async fn queue_backfill(&self, symbol: &str, fetch_to: OffsetDateTime) {
if self.app_config.alpaca_source == Source::Iex {
if self.config.alpaca_source == Source::Iex {
let run_delay = duration_until(fetch_to + FIFTEEN_MINUTES + ONE_MINUTE);
info!("Queing bar backfill for {} in {:?}.", symbol, run_delay);
sleep(run_delay).await;
@@ -243,18 +252,23 @@ impl Handler for BarHandler {
let mut next_page_token = None;
loop {
let message = alpaca::api::incoming::bar::get_historical(
&self.app_config,
let Ok(message) = alpaca::api::incoming::bar::get_historical(
&self.config,
self.data_url,
&(self.api_query_constructor)(
&self.app_config,
&self.config,
symbol.clone(),
fetch_from,
fetch_to,
next_page_token.clone(),
),
None,
)
.await;
.await
else {
error!("Failed to backfill bars for {}.", symbol);
return;
};
message.bars.into_iter().for_each(|(symbol, bar_vec)| {
for bar in bar_vec {
@@ -274,13 +288,17 @@ impl Handler for BarHandler {
}
let backfill = bars.last().unwrap().clone().into();
database::bars::upsert_batch(&self.app_config.clickhouse_client, bars).await;
database::bars::upsert_batch(&self.config.clickhouse_client, bars)
.await
.unwrap();
database::backfills::upsert(
&self.app_config.clickhouse_client,
&self.config.clickhouse_client,
&database::backfills::Table::Bars,
&backfill,
)
.await;
.await
.unwrap();
info!("Backfilled bars for {}.", symbol);
}
@@ -291,31 +309,34 @@ impl Handler for BarHandler {
}
struct NewsHandler {
app_config: Arc<Config>,
config: Arc<Config>,
}
#[async_trait]
impl Handler for NewsHandler {
async fn select_latest_backfill(&self, symbol: String) -> Option<crate::types::Backfill> {
async fn select_latest_backfill(
&self,
symbol: String,
) -> Result<Option<Backfill>, clickhouse::error::Error> {
database::backfills::select_latest_where_symbol(
&self.app_config.clickhouse_client,
&self.config.clickhouse_client,
&database::backfills::Table::News,
&symbol,
)
.await
}
async fn delete_backfills(&self, symbols: &[String]) {
async fn delete_backfills(&self, symbols: &[String]) -> Result<(), clickhouse::error::Error> {
database::backfills::delete_where_symbols(
&self.app_config.clickhouse_client,
&self.config.clickhouse_client,
&database::backfills::Table::News,
symbols,
)
.await;
.await
}
async fn delete_data(&self, symbols: &[String]) {
database::news::delete_where_symbols(&self.app_config.clickhouse_client, symbols).await;
async fn delete_data(&self, symbols: &[String]) -> Result<(), clickhouse::error::Error> {
database::news::delete_where_symbols(&self.config.clickhouse_client, symbols).await
}
async fn queue_backfill(&self, symbol: &str, fetch_to: OffsetDateTime) {
@@ -331,8 +352,8 @@ impl Handler for NewsHandler {
let mut next_page_token = None;
loop {
let message = alpaca::api::incoming::news::get_historical(
&self.app_config,
let Ok(message) = alpaca::api::incoming::news::get_historical(
&self.config,
&api::outgoing::news::News {
symbols: vec![remove_slash_from_pair(&symbol)],
start: Some(fetch_from),
@@ -343,8 +364,13 @@ impl Handler for NewsHandler {
page_token: next_page_token.clone(),
sort: Some(Sort::Asc),
},
None,
)
.await;
.await
else {
error!("Failed to backfill news for {}.", symbol);
return;
};
message.news.into_iter().for_each(|news_item| {
news.push(News::from(news_item));
@@ -366,23 +392,19 @@ impl Handler for NewsHandler {
.map(|news| format!("{}\n\n{}", news.headline, news.content))
.collect::<Vec<_>>();
let predictions = join_all(
inputs
.chunks(self.app_config.max_bert_inputs)
.map(|inputs| {
let sequence_classifier = self.app_config.sequence_classifier.clone();
async move {
let sequence_classifier = sequence_classifier.lock().await;
block_in_place(|| {
sequence_classifier
.predict(inputs.iter().map(String::as_str).collect::<Vec<_>>())
.into_iter()
.map(|label| Prediction::try_from(label).unwrap())
.collect::<Vec<_>>()
})
}
}),
)
let predictions = join_all(inputs.chunks(self.config.max_bert_inputs).map(|inputs| {
let sequence_classifier = self.config.sequence_classifier.clone();
async move {
let sequence_classifier = sequence_classifier.lock().await;
block_in_place(|| {
sequence_classifier
.predict(inputs.iter().map(String::as_str).collect::<Vec<_>>())
.into_iter()
.map(|label| Prediction::try_from(label).unwrap())
.collect::<Vec<_>>()
})
}
}))
.await
.into_iter()
.flatten();
@@ -398,13 +420,17 @@ impl Handler for NewsHandler {
.collect::<Vec<_>>();
let backfill = (news.last().unwrap().clone(), symbol.clone()).into();
database::news::upsert_batch(&self.app_config.clickhouse_client, news).await;
database::news::upsert_batch(&self.config.clickhouse_client, news)
.await
.unwrap();
database::backfills::upsert(
&self.app_config.clickhouse_client,
&self.config.clickhouse_client,
&database::backfills::Table::News,
&backfill,
)
.await;
.await
.unwrap();
info!("Backfilled news for {}.", symbol);
}
@@ -414,18 +440,18 @@ impl Handler for NewsHandler {
}
}
pub fn create_handler(thread_type: ThreadType, app_config: Arc<Config>) -> Box<dyn Handler> {
pub fn create_handler(thread_type: ThreadType, config: Arc<Config>) -> Box<dyn Handler> {
match thread_type {
ThreadType::Bars(Class::UsEquity) => Box::new(BarHandler {
app_config,
config,
data_url: ALPACA_STOCK_DATA_URL,
api_query_constructor: us_equity_query_constructor,
}),
ThreadType::Bars(Class::Crypto) => Box::new(BarHandler {
app_config,
config,
data_url: ALPACA_CRYPTO_DATA_URL,
api_query_constructor: crypto_query_constructor,
}),
ThreadType::News => Box::new(NewsHandler { app_config }),
ThreadType::News => Box::new(NewsHandler { config }),
}
}

83
src/threads/data/mod.rs
View File

@@ -6,9 +6,9 @@ use crate::{
config::{
Config, ALPACA_CRYPTO_WEBSOCKET_URL, ALPACA_NEWS_WEBSOCKET_URL, ALPACA_STOCK_WEBSOCKET_URL,
},
database,
create_send_await, database,
types::{alpaca, Asset, Class},
utils::{authenticate, cleanup},
utils::{authenticate, backoff, cleanup},
};
use futures_util::{future::join_all, StreamExt};
use itertools::{Either, Itertools};
@@ -52,22 +52,22 @@ pub enum ThreadType {
}
pub async fn run(
app_config: Arc<Config>,
config: Arc<Config>,
mut receiver: mpsc::Receiver<Message>,
mut clock_receiver: mpsc::Receiver<clock::Message>,
) {
let (bars_us_equity_websocket_sender, bars_us_equity_backfill_sender) =
init_thread(app_config.clone(), ThreadType::Bars(Class::UsEquity)).await;
init_thread(config.clone(), ThreadType::Bars(Class::UsEquity)).await;
let (bars_crypto_websocket_sender, bars_crypto_backfill_sender) =
init_thread(app_config.clone(), ThreadType::Bars(Class::Crypto)).await;
init_thread(config.clone(), ThreadType::Bars(Class::Crypto)).await;
let (news_websocket_sender, news_backfill_sender) =
init_thread(app_config.clone(), ThreadType::News).await;
init_thread(config.clone(), ThreadType::News).await;
loop {
select! {
Some(message) = receiver.recv() => {
spawn(handle_message(
app_config.clone(),
config.clone(),
bars_us_equity_websocket_sender.clone(),
bars_us_equity_backfill_sender.clone(),
bars_crypto_websocket_sender.clone(),
@@ -79,7 +79,7 @@ pub async fn run(
}
Some(_) = clock_receiver.recv() => {
spawn(handle_clock_message(
app_config.clone(),
config.clone(),
bars_us_equity_backfill_sender.clone(),
bars_crypto_backfill_sender.clone(),
news_backfill_sender.clone(),
@@ -91,34 +91,33 @@ pub async fn run(
}
async fn init_thread(
app_config: Arc<Config>,
config: Arc<Config>,
thread_type: ThreadType,
) -> (
mpsc::Sender<websocket::Message>,
mpsc::Sender<backfill::Message>,
) {
let websocket_url = match thread_type {
ThreadType::Bars(Class::UsEquity) => format!(
"{}/{}",
ALPACA_STOCK_WEBSOCKET_URL, &app_config.alpaca_source
),
ThreadType::Bars(Class::UsEquity) => {
format!("{}/{}", ALPACA_STOCK_WEBSOCKET_URL, &config.alpaca_source)
}
ThreadType::Bars(Class::Crypto) => ALPACA_CRYPTO_WEBSOCKET_URL.into(),
ThreadType::News => ALPACA_NEWS_WEBSOCKET_URL.into(),
};
let (websocket, _) = connect_async(websocket_url).await.unwrap();
let (mut websocket_sink, mut websocket_stream) = websocket.split();
authenticate(&app_config, &mut websocket_sink, &mut websocket_stream).await;
authenticate(&config, &mut websocket_sink, &mut websocket_stream).await;
let (backfill_sender, backfill_receiver) = mpsc::channel(100);
spawn(backfill::run(
Arc::new(backfill::create_handler(thread_type, app_config.clone())),
Arc::new(backfill::create_handler(thread_type, config.clone())),
backfill_receiver,
));
let (websocket_sender, websocket_receiver) = mpsc::channel(100);
spawn(websocket::run(
Arc::new(websocket::create_handler(thread_type, app_config.clone())),
Arc::new(websocket::create_handler(thread_type, config.clone())),
websocket_receiver,
websocket_stream,
websocket_sink,
@@ -127,17 +126,9 @@ async fn init_thread(
(websocket_sender, backfill_sender)
}
macro_rules! create_send_await {
($sender:expr, $action:expr, $($contents:expr),*) => {
let (message, receiver) = $action($($contents),*);
$sender.send(message).await.unwrap();
receiver.await.unwrap();
};
}
#[allow(clippy::too_many_arguments)]
async fn handle_message(
app_config: Arc<Config>,
config: Arc<Config>,
bars_us_equity_websocket_sender: mpsc::Sender<websocket::Message>,
bars_us_equity_backfill_sender: mpsc::Sender<backfill::Message>,
bars_crypto_websocket_sender: mpsc::Sender<websocket::Message>,
@@ -221,22 +212,30 @@ async fn handle_message(
match message.action {
Action::Add => {
let assets =
join_all(symbols.into_iter().map(|symbol| {
let app_config = app_config.clone();
async move {
alpaca::api::incoming::asset::get_by_symbol(&app_config, &symbol).await
}
}))
.await
.into_iter()
.map(|result| Asset::from(result.unwrap()))
.collect::<Vec<_>>();
let assets = join_all(symbols.into_iter().map(|symbol| {
let config = config.clone();
async move {
Asset::from(
alpaca::api::incoming::asset::get_by_symbol(
&config,
&symbol,
Some(backoff::infinite()),
)
.await
.unwrap(),
)
}
}))
.await;
database::assets::upsert_batch(&app_config.clickhouse_client, assets).await;
database::assets::upsert_batch(&config.clickhouse_client, assets)
.await
.unwrap();
}
Action::Remove => {
database::assets::delete_where_symbols(&app_config.clickhouse_client, &symbols).await;
database::assets::delete_where_symbols(&config.clickhouse_client, &symbols)
.await
.unwrap();
}
}
@@ -244,14 +243,16 @@ async fn handle_message(
}
async fn handle_clock_message(
app_config: Arc<Config>,
config: Arc<Config>,
bars_us_equity_backfill_sender: mpsc::Sender<backfill::Message>,
bars_crypto_backfill_sender: mpsc::Sender<backfill::Message>,
news_backfill_sender: mpsc::Sender<backfill::Message>,
) {
cleanup(&app_config.clickhouse_client).await;
cleanup(&config.clickhouse_client).await.unwrap();
let assets = database::assets::select(&app_config.clickhouse_client).await;
let assets = database::assets::select(&config.clickhouse_client)
.await
.unwrap();
let (us_equity_symbols, crypto_symbols): (Vec<_>, Vec<_>) = assets
.clone()

22
src/threads/data/websocket.rs
View File

@@ -221,7 +221,7 @@ async fn handle_websocket_message(
}
struct BarsHandler {
app_config: Arc<Config>,
config: Arc<Config>,
}
#[async_trait]
@@ -286,7 +286,10 @@ impl Handler for BarsHandler {
| websocket::incoming::Message::UpdatedBar(message) => {
let bar = Bar::from(message);
debug!("Received bar for {}: {}.", bar.symbol, bar.time);
database::bars::upsert(&self.app_config.clickhouse_client, &bar).await;
database::bars::upsert(&self.config.clickhouse_client, &bar)
.await
.unwrap();
}
websocket::incoming::Message::Success(_) => {}
websocket::incoming::Message::Error(message) => {
@@ -298,7 +301,7 @@ impl Handler for BarsHandler {
}
struct NewsHandler {
app_config: Arc<Config>,
config: Arc<Config>,
}
#[async_trait]
@@ -373,7 +376,7 @@ impl Handler for NewsHandler {
let input = format!("{}\n\n{}", news.headline, news.content);
let sequence_classifier = self.app_config.sequence_classifier.lock().await;
let sequence_classifier = self.config.sequence_classifier.lock().await;
let prediction = block_in_place(|| {
sequence_classifier
.predict(vec![input.as_str()])
@@ -388,7 +391,10 @@ impl Handler for NewsHandler {
confidence: prediction.confidence,
..news
};
database::news::upsert(&self.app_config.clickhouse_client, &news).await;
database::news::upsert(&self.config.clickhouse_client, &news)
.await
.unwrap();
}
websocket::incoming::Message::Success(_) => {}
websocket::incoming::Message::Error(message) => {
@@ -401,9 +407,9 @@ impl Handler for NewsHandler {
}
}
pub fn create_handler(thread_type: ThreadType, app_config: Arc<Config>) -> Box<dyn Handler> {
pub fn create_handler(thread_type: ThreadType, config: Arc<Config>) -> Box<dyn Handler> {
match thread_type {
ThreadType::Bars(_) => Box::new(BarsHandler { app_config }),
ThreadType::News => Box::new(NewsHandler { app_config }),
ThreadType::Bars(_) => Box::new(BarsHandler { config }),
ThreadType::News => Box::new(NewsHandler { config }),
}
}