Refactor threads to use trait implementations

Signed-off-by: Nikolaos Karaolidis <nick@karaolidis.com>
This commit is contained in:
2024-02-05 13:47:43 +00:00
parent a796feb299
commit 85eef2bf0b
9 changed files with 524 additions and 439 deletions

View File

@@ -1,4 +1,4 @@
use super::{backfill, Guard, ThreadType};
use super::{backfill, Guard};
use crate::{
config::Config,
database,
@@ -23,22 +23,18 @@ use tokio_tungstenite::{tungstenite, MaybeTlsStream, WebSocketStream};
pub async fn run(
app_config: Arc<Config>,
thread_type: ThreadType,
guard: Arc<RwLock<Guard>>,
websocket_sender: Arc<
Mutex<SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, tungstenite::Message>>,
>,
mut websocket_receiver: SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
sender: Arc<Mutex<SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, tungstenite::Message>>>,
mut receiver: SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
backfill_sender: mpsc::Sender<backfill::Message>,
) {
loop {
let message = websocket_receiver.next().await.unwrap().unwrap();
let message = receiver.next().await.unwrap().unwrap();
spawn(handle_websocket_message(
app_config.clone(),
thread_type,
guard.clone(),
websocket_sender.clone(),
sender.clone(),
backfill_sender.clone(),
message,
));
@@ -47,11 +43,8 @@ pub async fn run(
async fn handle_websocket_message(
app_config: Arc<Config>,
thread_type: ThreadType,
guard: Arc<RwLock<Guard>>,
websocket_sender: Arc<
Mutex<SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, tungstenite::Message>>,
>,
sender: Arc<Mutex<SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, tungstenite::Message>>>,
backfill_sender: mpsc::Sender<backfill::Message>,
message: tungstenite::Message,
) {
@@ -63,31 +56,24 @@ async fn handle_websocket_message(
for message in message {
spawn(handle_parsed_websocket_message(
app_config.clone(),
thread_type,
guard.clone(),
backfill_sender.clone(),
message,
));
}
} else {
error!(
"{:?} - Failed to deserialize websocket message: {:?}",
thread_type, message
);
error!("Failed to deserialize websocket message: {:?}", message);
}
}
tungstenite::Message::Ping(_) => {
websocket_sender
sender
.lock()
.await
.send(tungstenite::Message::Pong(vec![]))
.await
.unwrap();
}
_ => error!(
"{:?} - Unexpected websocket message: {:?}",
thread_type, message
),
_ => error!("Unexpected websocket message: {:?}", message),
}
}
@@ -95,19 +81,20 @@ async fn handle_websocket_message(
#[allow(clippy::too_many_lines)]
async fn handle_parsed_websocket_message(
app_config: Arc<Config>,
thread_type: ThreadType,
guard: Arc<RwLock<Guard>>,
backfill_sender: mpsc::Sender<backfill::Message>,
message: websocket::incoming::Message,
) {
match message {
websocket::incoming::Message::Subscription(message) => {
let symbols = match message {
websocket::incoming::subscription::Message::Market { bars, .. } => bars,
websocket::incoming::subscription::Message::News { news } => news
.into_iter()
.map(|symbol| add_slash_to_pair(&symbol))
.collect(),
let (symbols, log_string) = match message {
websocket::incoming::subscription::Message::Market { bars, .. } => (bars, "bars"),
websocket::incoming::subscription::Message::News { news } => (
news.into_iter()
.map(|symbol| add_slash_to_pair(&symbol))
.collect(),
"news",
),
};
let mut guard = guard.write().await;
@@ -127,8 +114,8 @@ async fn handle_parsed_websocket_message(
let newly_subscribed_future = async {
if !newly_subscribed.is_empty() {
info!(
"{:?} - Subscribed to {:?}.",
thread_type,
"Subscribed to {} for {:?}.",
log_string,
newly_subscribed
.iter()
.map(|asset| asset.symbol.clone())
@@ -148,8 +135,8 @@ async fn handle_parsed_websocket_message(
let newly_unsubscribed_future = async {
if !newly_unsubscribed.is_empty() {
info!(
"{:?} - Unsubscribed from {:?}.",
thread_type,
"Unsubscribed from {} for {:?}.",
log_string,
newly_unsubscribed
.iter()
.map(|asset| asset.symbol.clone())
@@ -175,16 +162,13 @@ async fn handle_parsed_websocket_message(
let guard = guard.read().await;
if !guard.assets.contains_right(&bar.symbol) {
warn!(
"{:?} - Race condition: received bar for unsubscribed symbol: {:?}.",
thread_type, bar.symbol
"Race condition: received bar for unsubscribed symbol: {:?}.",
bar.symbol
);
return;
}
debug!(
"{:?} - Received bar for {}: {}.",
thread_type, bar.symbol, bar.time
);
debug!("Received bar for {}: {}.", bar.symbol, bar.time);
database::bars::upsert(&app_config.clickhouse_client, &bar).await;
}
websocket::incoming::Message::News(message) => {
@@ -197,15 +181,15 @@ async fn handle_parsed_websocket_message(
.any(|symbol| guard.assets.contains_right(symbol))
{
warn!(
"{:?} - Race condition: received news for unsubscribed symbols: {:?}.",
thread_type, news.symbols
"Race condition: received news for unsubscribed symbols: {:?}.",
news.symbols
);
return;
}
debug!(
"{:?} - Received news for {:?}: {}.",
thread_type, news.symbols, news.time_created
"Received news for {:?}: {}.",
news.symbols, news.time_created
);
let input = format!("{}\n\n{}", news.headline, news.content);
@@ -229,10 +213,7 @@ async fn handle_parsed_websocket_message(
}
websocket::incoming::Message::Success(_) => {}
websocket::incoming::Message::Error(message) => {
error!(
"{:?} - Received error message: {}.",
thread_type, message.message
);
error!("Received error message: {}.", message.message);
}
}
}