Add news data support

- Refactor everything in the process, oops

Signed-off-by: Nikolaos Karaolidis <nick@karaolidis.com>

src/threads/data/websocket.rs

@@ -0,0 +1,217 @@
+use super::{backfill, Guard, ThreadType};
+use crate::{
+    config::Config,
+    database,
+    types::{alpaca::websocket, Bar, News, Subset},
+};
+use futures_util::{
+    stream::{SplitSink, SplitStream},
+    SinkExt, StreamExt,
+};
+use log::{error, info, warn};
+use serde_json::from_str;
+use std::{
+    collections::{HashMap, HashSet},
+    sync::Arc,
+};
+use tokio::{
+    join,
+    net::TcpStream,
+    spawn,
+    sync::{mpsc, Mutex, RwLock},
+};
+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>>>,
+    backfill_sender: mpsc::Sender<backfill::Message>,
+) {
+    loop {
+        let app_config = app_config.clone();
+        let guard = guard.clone();
+        let websocket_sender = websocket_sender.clone();
+        let backfill_sender = backfill_sender.clone();
+
+        let message = websocket_receiver.next().await.unwrap().unwrap();
+
+        spawn(handle_websocket_message(
+            app_config,
+            thread_type,
+            guard,
+            websocket_sender,
+            backfill_sender,
+            message,
+        ));
+    }
+}
+
+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>>,
+    >,
+    backfill_sender: mpsc::Sender<backfill::Message>,
+    message: tungstenite::Message,
+) {
+    match message {
+        tungstenite::Message::Text(message) => {
+            let message = from_str::<Vec<websocket::incoming::Message>>(&message);
+
+            if let Ok(message) = message {
+                for message in message {
+                    let app_config = app_config.clone();
+                    let guard = guard.clone();
+                    let backfill_sender = backfill_sender.clone();
+
+                    spawn(handle_parsed_websocket_message(
+                        app_config,
+                        thread_type,
+                        guard,
+                        backfill_sender,
+                        message,
+                    ));
+                }
+            } else {
+                error!(
+                    "{:?} - Failed to deserialize websocket message: {:?}",
+                    thread_type, message
+                );
+            }
+        }
+        tungstenite::Message::Ping(_) => {
+            websocket_sender
+                .lock()
+                .await
+                .send(tungstenite::Message::Pong(vec![]))
+                .await
+                .unwrap();
+        }
+        _ => error!(
+            "{:?} - Unexpected websocket message: {:?}",
+            thread_type, message
+        ),
+    }
+}
+
+#[allow(clippy::significant_drop_tightening)]
+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(message) => message.bars,
+                websocket::incoming::subscription::Message::News(message) => message.news,
+            };
+
+            let mut guard = guard.write().await;
+
+            let newly_subscribed = guard
+                .pending_subscriptions
+                .extract_if(|symbol, _| symbols.contains(symbol))
+                .collect::<HashMap<_, _>>();
+
+            let newly_unsubscribed = guard
+                .pending_unsubscriptions
+                .extract_if(|symbol, _| !symbols.contains(symbol))
+                .collect::<HashMap<_, _>>();
+
+            drop(guard);
+
+            let newly_subscribed_future = async {
+                if !newly_subscribed.is_empty() {
+                    info!(
+                        "{:?} - Subscribed to {:?}.",
+                        thread_type,
+                        newly_subscribed.keys().collect::<Vec<_>>()
+                    );
+
+                    let (backfill_message, backfill_receiver) = backfill::Message::new(
+                        backfill::Action::Backfill,
+                        Subset::Some(newly_subscribed.into_values().collect::<Vec<_>>()),
+                    );
+
+                    backfill_sender.send(backfill_message).await.unwrap();
+                    backfill_receiver.await.unwrap();
+                }
+            };
+
+            let newly_unsubscribed_future = async {
+                if !newly_unsubscribed.is_empty() {
+                    info!(
+                        "{:?} - Unsubscribed from {:?}.",
+                        thread_type,
+                        newly_unsubscribed.keys().collect::<Vec<_>>()
+                    );
+
+                    let (purge_message, purge_receiver) = backfill::Message::new(
+                        backfill::Action::Purge,
+                        Subset::Some(newly_unsubscribed.into_values().collect::<Vec<_>>()),
+                    );
+
+                    backfill_sender.send(purge_message).await.unwrap();
+                    purge_receiver.await.unwrap();
+                }
+            };
+
+            join!(newly_subscribed_future, newly_unsubscribed_future);
+        }
+        websocket::incoming::Message::Bar(message)
+        | websocket::incoming::Message::UpdatedBar(message) => {
+            let bar = Bar::from(message);
+
+            let guard = guard.read().await;
+            if guard.symbols.get(&bar.symbol).is_none() {
+                warn!(
+                    "{:?} - Race condition: received bar for unsubscribed symbol: {:?}.",
+                    thread_type, bar.symbol
+                );
+                return;
+            }
+
+            info!(
+                "{:?} - Received bar for {}: {}.",
+                thread_type, bar.symbol, bar.time
+            );
+            database::bars::upsert(&app_config.clickhouse_client, &bar).await;
+        }
+        websocket::incoming::Message::News(message) => {
+            let news = News::from(message);
+            let symbols = news.symbols.clone().into_iter().collect::<HashSet<_>>();
+
+            let guard = guard.read().await;
+            if !guard.symbols.iter().any(|symbol| symbols.contains(symbol)) {
+                warn!(
+                    "{:?} - Race condition: received news for unsubscribed symbols: {:?}.",
+                    thread_type, news.symbols
+                );
+                return;
+            }
+
+            info!(
+                "{:?} - Received news for {:?}: {}.",
+                thread_type, news.symbols, news.time_created
+            );
+            database::news::upsert(&app_config.clickhouse_client, &news).await;
+        }
+        websocket::incoming::Message::Success(_) => {}
+        websocket::incoming::Message::Error(message) => {
+            error!(
+                "{:?} - Received error message: {}.",
+                thread_type, message.message
+            );
+        }
+    }
+}