Add news data support

- Refactor everything in the process, oops

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

src/threads/data/mod.rs

@@ -0,0 +1,233 @@
+pub mod asset_status;
+pub mod backfill;
+pub mod websocket;
+
+use super::clock;
+use crate::{
+    config::{
+        Config, ALPACA_CRYPTO_WEBSOCKET_URL, ALPACA_NEWS_WEBSOCKET_URL, ALPACA_STOCK_WEBSOCKET_URL,
+    },
+    types::{Asset, Class, Subset},
+    utils::authenticate,
+};
+use futures_util::StreamExt;
+use std::{
+    collections::{HashMap, HashSet},
+    sync::Arc,
+};
+use tokio::{
+    join, select, spawn,
+    sync::{mpsc, Mutex, RwLock},
+};
+use tokio_tungstenite::connect_async;
+
+pub struct Guard {
+    pub symbols: HashSet<String>,
+    pub pending_subscriptions: HashMap<String, Asset>,
+    pub pending_unsubscriptions: HashMap<String, Asset>,
+}
+
+impl Guard {
+    pub fn new() -> Self {
+        Self {
+            symbols: HashSet::new(),
+            pending_subscriptions: HashMap::new(),
+            pending_unsubscriptions: HashMap::new(),
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug)]
+pub enum ThreadType {
+    Bars(Class),
+    News,
+}
+
+pub async fn run(
+    app_config: Arc<Config>,
+    mut asset_receiver: mpsc::Receiver<asset_status::Message>,
+    mut clock_receiver: mpsc::Receiver<clock::Message>,
+) {
+    let (bars_us_equity_asset_status_sender, bars_us_equity_backfill_sender) =
+        init_thread(app_config.clone(), ThreadType::Bars(Class::UsEquity)).await;
+    let (bars_crypto_asset_status_sender, bars_crypto_backfill_sender) =
+        init_thread(app_config.clone(), ThreadType::Bars(Class::Crypto)).await;
+    let (news_asset_status_sender, news_backfill_sender) =
+        init_thread(app_config.clone(), ThreadType::News).await;
+
+    loop {
+        select! {
+            Some(asset_message) = asset_receiver.recv() => {
+                let bars_us_equity_asset_status_sender = bars_us_equity_asset_status_sender.clone();
+                let bars_crypto_asset_status_sender = bars_crypto_asset_status_sender.clone();
+                let news_asset_status_sender = news_asset_status_sender.clone();
+
+                spawn(handle_asset_message(
+                    bars_us_equity_asset_status_sender,
+                    bars_crypto_asset_status_sender,
+                    news_asset_status_sender,
+                    asset_message,
+                ));
+            }
+            Some(_) = clock_receiver.recv() => {
+                let bars_us_equity_backfill_sender =  bars_us_equity_backfill_sender.clone();
+                let bars_crypto_backfill_sender =  bars_crypto_backfill_sender.clone();
+                let news_backfill_sender =  news_backfill_sender.clone();
+
+                spawn(handle_clock_message(
+                    bars_us_equity_backfill_sender,
+                    bars_crypto_backfill_sender,
+                    news_backfill_sender,
+                ));
+            }
+            else => {
+                panic!("Communication channel unexpectedly closed.")
+            }
+        }
+    }
+}
+
+async fn init_thread(
+    app_config: Arc<Config>,
+    thread_type: ThreadType,
+) -> (
+    mpsc::Sender<asset_status::Message>,
+    mpsc::Sender<backfill::Message>,
+) {
+    let guard = Arc::new(RwLock::new(Guard::new()));
+
+    let websocket_url = match thread_type {
+        ThreadType::Bars(Class::UsEquity) => format!(
+            "{}/{}",
+            ALPACA_STOCK_WEBSOCKET_URL, &app_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_sender, mut websocket_receiver) = websocket.split();
+    authenticate(&app_config, &mut websocket_sender, &mut websocket_receiver).await;
+    let websocket_sender = Arc::new(Mutex::new(websocket_sender));
+
+    let (asset_status_sender, asset_status_receiver) = mpsc::channel(100);
+    spawn(asset_status::run(
+        app_config.clone(),
+        thread_type,
+        guard.clone(),
+        asset_status_receiver,
+        websocket_sender.clone(),
+    ));
+
+    let (backfill_sender, backfill_receiver) = mpsc::channel(100);
+    spawn(backfill::run(
+        app_config.clone(),
+        thread_type,
+        guard.clone(),
+        backfill_receiver,
+    ));
+
+    spawn(websocket::run(
+        app_config.clone(),
+        thread_type,
+        guard.clone(),
+        websocket_sender,
+        websocket_receiver,
+        backfill_sender.clone(),
+    ));
+
+    (asset_status_sender, backfill_sender)
+}
+
+async fn handle_asset_message(
+    bars_us_equity_asset_status_sender: mpsc::Sender<asset_status::Message>,
+    bars_crypto_asset_status_sender: mpsc::Sender<asset_status::Message>,
+    news_asset_status_sender: mpsc::Sender<asset_status::Message>,
+    asset_status_message: asset_status::Message,
+) {
+    let (us_equity_assets, crypto_assets): (Vec<_>, Vec<_>) = asset_status_message
+        .assets
+        .clone()
+        .into_iter()
+        .partition(|asset| asset.class == Class::UsEquity);
+
+    let bars_us_equity_future = async {
+        if !us_equity_assets.is_empty() {
+            let (bars_us_equity_asset_status_message, bars_us_equity_asset_status_receiver) =
+                asset_status::Message::new(asset_status_message.action.clone(), us_equity_assets);
+            bars_us_equity_asset_status_sender
+                .send(bars_us_equity_asset_status_message)
+                .await
+                .unwrap();
+            bars_us_equity_asset_status_receiver.await.unwrap();
+        }
+    };
+
+    let bars_crypto_future = async {
+        if !crypto_assets.is_empty() {
+            let (crypto_asset_status_message, crypto_asset_status_receiver) =
+                asset_status::Message::new(asset_status_message.action.clone(), crypto_assets);
+            bars_crypto_asset_status_sender
+                .send(crypto_asset_status_message)
+                .await
+                .unwrap();
+            crypto_asset_status_receiver.await.unwrap();
+        }
+    };
+
+    let news_future = async {
+        if !asset_status_message.assets.is_empty() {
+            let (news_asset_status_message, news_asset_status_receiver) =
+                asset_status::Message::new(
+                    asset_status_message.action.clone(),
+                    asset_status_message.assets,
+                );
+            news_asset_status_sender
+                .send(news_asset_status_message)
+                .await
+                .unwrap();
+            news_asset_status_receiver.await.unwrap();
+        }
+    };
+
+    join!(bars_us_equity_future, bars_crypto_future, news_future);
+    asset_status_message.response.send(()).unwrap();
+}
+
+async fn handle_clock_message(
+    bars_us_equity_backfill_sender: mpsc::Sender<backfill::Message>,
+    bars_crypto_backfill_sender: mpsc::Sender<backfill::Message>,
+    news_backfill_sender: mpsc::Sender<backfill::Message>,
+) {
+    let bars_us_equity_future = async {
+        let (bars_us_equity_backfill_message, bars_us_equity_backfill_receiver) =
+            backfill::Message::new(backfill::Action::Backfill, Subset::All);
+        bars_us_equity_backfill_sender
+            .send(bars_us_equity_backfill_message)
+            .await
+            .unwrap();
+        bars_us_equity_backfill_receiver.await.unwrap();
+    };
+
+    let bars_crypto_future = async {
+        let (bars_crypto_backfill_message, bars_crypto_backfill_receiver) =
+            backfill::Message::new(backfill::Action::Backfill, Subset::All);
+        bars_crypto_backfill_sender
+            .send(bars_crypto_backfill_message)
+            .await
+            .unwrap();
+        bars_crypto_backfill_receiver.await.unwrap();
+    };
+
+    let news_future = async {
+        let (news_backfill_message, news_backfill_receiver) =
+            backfill::Message::new(backfill::Action::Backfill, Subset::All);
+        news_backfill_sender
+            .send(news_backfill_message)
+            .await
+            .unwrap();
+        news_backfill_receiver.await.unwrap();
+    };
+
+    join!(bars_us_equity_future, bars_crypto_future, news_future);
+}