117 lines
3.6 KiB
Rust
117 lines
3.6 KiB
Rust
use super::State;
|
|
use crate::config::{Config, CLICKHOUSE_BATCH_NEWS_SIZE};
|
|
use async_trait::async_trait;
|
|
use clickhouse::inserter::Inserter;
|
|
use log::{debug, error, info};
|
|
use nonempty::NonEmpty;
|
|
use qrust::{
|
|
types::{alpaca::websocket, News},
|
|
utils::ONE_SECOND,
|
|
};
|
|
use std::{collections::HashMap, sync::Arc};
|
|
use tokio::sync::{Mutex, RwLock};
|
|
|
|
pub struct Handler {
|
|
pub config: Arc<Config>,
|
|
pub inserter: Arc<Mutex<Inserter<News>>>,
|
|
}
|
|
|
|
#[async_trait]
|
|
impl super::Handler for Handler {
|
|
fn create_subscription_message(
|
|
&self,
|
|
symbols: NonEmpty<String>,
|
|
) -> websocket::data::outgoing::subscribe::Message {
|
|
websocket::data::outgoing::subscribe::Message::new_news(symbols)
|
|
}
|
|
|
|
async fn handle_websocket_message(
|
|
&self,
|
|
state: Arc<RwLock<State>>,
|
|
message: websocket::data::incoming::Message,
|
|
) {
|
|
match message {
|
|
websocket::data::incoming::Message::Subscription(message) => {
|
|
let websocket::data::incoming::subscription::Message::News { news: symbols } =
|
|
message
|
|
else {
|
|
unreachable!()
|
|
};
|
|
|
|
let mut state = state.write().await;
|
|
|
|
let newly_subscribed = state
|
|
.pending_subscriptions
|
|
.extract_if(|symbol, _| symbols.contains(symbol))
|
|
.collect::<HashMap<_, _>>();
|
|
|
|
let newly_unsubscribed = state
|
|
.pending_unsubscriptions
|
|
.extract_if(|symbol, _| !symbols.contains(symbol))
|
|
.collect::<HashMap<_, _>>();
|
|
|
|
state
|
|
.active_subscriptions
|
|
.extend(newly_subscribed.keys().cloned());
|
|
|
|
drop(state);
|
|
|
|
if !newly_subscribed.is_empty() {
|
|
info!(
|
|
"Subscribed to news for {:?}.",
|
|
newly_subscribed.keys().collect::<Vec<_>>()
|
|
);
|
|
|
|
for sender in newly_subscribed.into_values() {
|
|
sender.send(()).unwrap();
|
|
}
|
|
}
|
|
|
|
if !newly_unsubscribed.is_empty() {
|
|
info!(
|
|
"Unsubscribed from news for {:?}.",
|
|
newly_unsubscribed.keys().collect::<Vec<_>>()
|
|
);
|
|
|
|
for sender in newly_unsubscribed.into_values() {
|
|
sender.send(()).unwrap();
|
|
}
|
|
}
|
|
}
|
|
websocket::data::incoming::Message::News(message) => {
|
|
let news = News::from(message);
|
|
debug!(
|
|
"Received news for {:?}: {}.",
|
|
news.symbols, news.time_created
|
|
);
|
|
self.inserter.lock().await.write(&news).await.unwrap();
|
|
}
|
|
websocket::data::incoming::Message::Error(message) => {
|
|
error!("Received error message: {}.", message.message);
|
|
}
|
|
_ => unreachable!(),
|
|
}
|
|
}
|
|
|
|
fn log_string(&self) -> &'static str {
|
|
"news"
|
|
}
|
|
|
|
async fn run_inserter(&self) {
|
|
super::run_inserter(self.inserter.clone()).await;
|
|
}
|
|
}
|
|
|
|
pub fn create_handler(config: Arc<Config>) -> Box<dyn super::Handler> {
|
|
let inserter = Arc::new(Mutex::new(
|
|
config
|
|
.clickhouse_client
|
|
.inserter("news")
|
|
.unwrap()
|
|
.with_period(Some(ONE_SECOND))
|
|
.with_max_entries((*CLICKHOUSE_BATCH_NEWS_SIZE).try_into().unwrap()),
|
|
));
|
|
|
|
Box::new(Handler { config, inserter })
|
|
}
|