Organize websocket codebase

Signed-off-by: Nikolaos Karaolidis <nick@karaolidis.com>
This commit is contained in:
2024-01-14 10:29:32 +00:00
parent 687fbb909f
commit 63a9ca950f
18 changed files with 138 additions and 128 deletions

View File

@@ -3,7 +3,7 @@ RUN apk add --no-cache pkgconf musl-dev openssl-dev
WORKDIR /usr/src/qrust
ENV SQLX_OFFLINE true
ENV SQLX_OFFLINE=true
RUN mkdir src && echo "fn main() {}" > src/main.rs
COPY Cargo.toml .sqlx ./

View File

@@ -3,6 +3,7 @@ use crate::{
Config, ALPACA_CRYPTO_DATA_URL, ALPACA_CRYPTO_WEBSOCKET_URL, ALPACA_STOCK_DATA_URL,
ALPACA_STOCK_WEBSOCKET_URL, ALPACA_TIMESTAMP_FORMAT,
},
data::authenticate_websocket,
database,
time::{duration_until, last_minute, next_30s, next_minute, ONE_MINUTE, THIRTY_SECONDS},
types::{
@@ -19,7 +20,7 @@ use futures_util::{
use http::StatusCode;
use indexmap::IndexMap;
use log::{error, info, warn};
use serde_json::{from_str, to_string};
use serde_json::from_str;
use std::{
collections::{HashMap, HashSet},
sync::Arc,
@@ -83,50 +84,6 @@ pub async fn run(
unreachable!()
}
async fn authenticate_websocket(
app_config: &Arc<Config>,
stream: &mut SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
sink: &mut SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>,
) {
match stream.next().await {
Some(Ok(Message::Text(data)))
if from_str::<Vec<websocket::incoming::Message>>(&data)
.unwrap()
.get(0)
== Some(&websocket::incoming::Message::Success(
websocket::incoming::success::Message {
msg: websocket::incoming::success::MessageType::Connected,
},
)) => {}
_ => panic!(),
}
sink.send(Message::Text(
to_string(&websocket::outgoing::Message::Auth(
websocket::outgoing::auth::Message::new(
app_config.alpaca_api_key.clone(),
app_config.alpaca_api_secret.clone(),
),
))
.unwrap(),
))
.await
.unwrap();
match stream.next().await {
Some(Ok(Message::Text(data)))
if from_str::<Vec<websocket::incoming::Message>>(&data)
.unwrap()
.get(0)
== Some(&websocket::incoming::Message::Success(
websocket::incoming::success::Message {
msg: websocket::incoming::success::MessageType::Authenticated,
},
)) => {}
_ => panic!(),
};
}
async fn websocket_broadcast_handler(
class: Class,
sink: Arc<RwLock<SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>>>,
@@ -140,8 +97,10 @@ async fn websocket_broadcast_handler(
sink.write()
.await
.send(Message::Text(
serde_json::to_string(&websocket::outgoing::Message::Subscribe(
websocket::outgoing::subscribe::Message::new(asset.clone().symbol),
serde_json::to_string(&websocket::data::outgoing::Message::Subscribe(
websocket::data::outgoing::subscribe::Message::new(
asset.clone().symbol,
),
))
.unwrap(),
))
@@ -154,8 +113,10 @@ async fn websocket_broadcast_handler(
sink.write()
.await
.send(Message::Text(
serde_json::to_string(&websocket::outgoing::Message::Unsubscribe(
websocket::outgoing::subscribe::Message::new(asset.clone().symbol),
serde_json::to_string(&websocket::data::outgoing::Message::Unsubscribe(
websocket::data::outgoing::subscribe::Message::new(
asset.clone().symbol,
),
))
.unwrap(),
))
@@ -177,9 +138,12 @@ async fn websocket_message_handler(
loop {
match stream.next().await {
Some(Ok(Message::Text(data))) => {
let parsed_data = from_str::<Vec<websocket::incoming::Message>>(&data);
let parsed_data = from_str::<Vec<websocket::data::incoming::Message>>(&data);
if let Err(e) = &parsed_data {
warn!("Unparsed websocket::incoming message: {:?}: {}", data, e);
warn!(
"Unparsed websocket::data::incoming message: {:?}: {}",
data, e
);
}
for message in parsed_data.unwrap_or_default() {
@@ -192,7 +156,7 @@ async fn websocket_message_handler(
.send(Message::Pong(vec![]))
.await
.unwrap(),
Some(unknown) => error!("Unknown websocket::incoming message: {:?}", unknown),
Some(unknown) => error!("Unknown websocket::data::incoming message: {:?}", unknown),
None => panic!(),
}
}
@@ -201,11 +165,11 @@ async fn websocket_message_handler(
async fn websocket_handle_text_message(
app_config: &Arc<Config>,
class: Class,
message: websocket::incoming::Message,
message: websocket::data::incoming::Message,
backfilled: &Arc<RwLock<HashMap<String, bool>>>,
) {
match message {
websocket::incoming::Message::Subscription(subscription_message) => {
websocket::data::incoming::Message::Subscription(subscription_message) => {
let old_assets = backfilled
.read()
.await
@@ -243,7 +207,7 @@ async fn websocket_handle_text_message(
class, added_assets, deleted_assets
);
}
websocket::incoming::Message::Bars(bar_message) => {
websocket::data::incoming::Message::Bars(bar_message) => {
let bar = Bar::from(bar_message);
info!(
"websocket::Incoming bar for {}: {}",
@@ -256,7 +220,7 @@ async fn websocket_handle_text_message(
)
.await;
}
websocket::incoming::Message::UpdatedBars(bar_message) => {
websocket::data::incoming::Message::UpdatedBars(bar_message) => {
let bar = Bar::from(bar_message);
info!(
"websocket::Incoming bar for {}: {}",
@@ -276,7 +240,7 @@ async fn websocket_handle_text_message(
}
transaction.commit().await.unwrap();
}
websocket::incoming::Message::Success(_) => {}
_ => {}
}
}

View File

@@ -1 +1,56 @@
pub mod market;
use crate::{config::Config, types::websocket};
use core::panic;
use futures_util::{
stream::{SplitSink, SplitStream},
SinkExt, StreamExt,
};
use serde_json::{from_str, to_string};
use std::sync::Arc;
use tokio::net::TcpStream;
use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream, WebSocketStream};
async fn authenticate_websocket(
app_config: &Arc<Config>,
stream: &mut SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
sink: &mut SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>,
) {
match stream.next().await {
Some(Ok(Message::Text(data)))
if from_str::<Vec<websocket::data::incoming::Message>>(&data)
.unwrap()
.get(0)
== Some(&websocket::data::incoming::Message::Success(
websocket::data::incoming::success::Message {
msg: websocket::data::incoming::success::MessageType::Connected,
},
)) => {}
_ => panic!(),
}
sink.send(Message::Text(
to_string(&websocket::data::outgoing::Message::Auth(
websocket::data::outgoing::auth::Message::new(
app_config.alpaca_api_key.clone(),
app_config.alpaca_api_secret.clone(),
),
))
.unwrap(),
))
.await
.unwrap();
match stream.next().await {
Some(Ok(Message::Text(data)))
if from_str::<Vec<websocket::data::incoming::Message>>(&data)
.unwrap()
.get(0)
== Some(&websocket::data::incoming::Message::Success(
websocket::data::incoming::success::Message {
msg: websocket::data::incoming::success::MessageType::Authenticated,
},
)) => {}
_ => panic!(),
};
}

View File

@@ -1,8 +1,56 @@
use super::{api::incoming, class::Class, exchange::Exchange};
use serde::Serialize;
use sqlx::FromRow;
use super::api;
use serde::{Deserialize, Serialize};
use sqlx::{FromRow, Type};
use time::OffsetDateTime;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize, Type)]
pub enum Class {
#[sqlx(rename = "us_equity")]
#[serde(rename = "us_equity")]
UsEquity,
#[sqlx(rename = "crypto")]
#[serde(rename = "crypto")]
Crypto,
}
#[derive(Clone, Copy, Debug, Serialize, Deserialize, Type)]
pub enum Exchange {
#[sqlx(rename = "AMEX")]
#[serde(rename = "AMEX")]
Amex,
#[sqlx(rename = "ARCA")]
#[serde(rename = "ARCA")]
Arca,
#[sqlx(rename = "BATS")]
#[serde(rename = "BATS")]
Bats,
#[sqlx(rename = "NYSE")]
#[serde(rename = "NYSE")]
Nyse,
#[sqlx(rename = "NASDAQ")]
#[serde(rename = "NASDAQ")]
Nasdaq,
#[sqlx(rename = "NYSEARCA")]
#[serde(rename = "NYSEARCA")]
Nysearca,
#[sqlx(rename = "OTC")]
#[serde(rename = "OTC")]
Otc,
#[sqlx(rename = "CRYPTO")]
#[serde(rename = "CRYPTO")]
Crypto,
}
#[derive(PartialEq, Eq, Deserialize, Type)]
pub enum Status {
#[sqlx(rename = "active")]
#[serde(rename = "active")]
Active,
#[sqlx(rename = "inactive")]
#[serde(rename = "inactive")]
Inactive,
}
#[derive(Clone, Debug, FromRow, Serialize)]
pub struct Asset {
pub symbol: String,
@@ -14,8 +62,10 @@ pub struct Asset {
pub timestamp_last: OffsetDateTime,
}
impl From<(incoming::Asset, bool, OffsetDateTime)> for Asset {
fn from((asset, trading, timestamp_first): (incoming::Asset, bool, OffsetDateTime)) -> Self {
impl From<(api::incoming::Asset, bool, OffsetDateTime)> for Asset {
fn from(
(asset, trading, timestamp_first): (api::incoming::Asset, bool, OffsetDateTime),
) -> Self {
Self {
symbol: asset.symbol,
class: asset.class,

View File

@@ -39,8 +39,8 @@ impl Bar {
}
}
impl From<websocket::incoming::bar::Message> for Bar {
fn from(bar_message: websocket::incoming::bar::Message) -> Self {
impl From<websocket::data::incoming::bar::Message> for Bar {
fn from(bar_message: websocket::data::incoming::bar::Message) -> Self {
Self {
timestamp: bar_message.timestamp,
asset_symbol: bar_message.symbol,

View File

@@ -1,12 +0,0 @@
use serde::{Deserialize, Serialize};
use sqlx::Type;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize, Type)]
pub enum Class {
#[sqlx(rename = "us_equity")]
#[serde(rename = "us_equity")]
UsEquity,
#[sqlx(rename = "crypto")]
#[serde(rename = "crypto")]
Crypto,
}

View File

@@ -1,30 +0,0 @@
use serde::{Deserialize, Serialize};
use sqlx::Type;
#[derive(Clone, Copy, Debug, Serialize, Deserialize, Type)]
pub enum Exchange {
#[sqlx(rename = "AMEX")]
#[serde(rename = "AMEX")]
Amex,
#[sqlx(rename = "ARCA")]
#[serde(rename = "ARCA")]
Arca,
#[sqlx(rename = "BATS")]
#[serde(rename = "BATS")]
Bats,
#[sqlx(rename = "NYSE")]
#[serde(rename = "NYSE")]
Nyse,
#[sqlx(rename = "NASDAQ")]
#[serde(rename = "NASDAQ")]
Nasdaq,
#[sqlx(rename = "NYSEARCA")]
#[serde(rename = "NYSEARCA")]
Nysearca,
#[sqlx(rename = "OTC")]
#[serde(rename = "OTC")]
Otc,
#[sqlx(rename = "CRYPTO")]
#[serde(rename = "CRYPTO")]
Crypto,
}

View File

@@ -1,18 +1,12 @@
pub mod api;
pub mod asset;
pub mod bar;
pub mod class;
pub mod exchange;
pub mod source;
pub mod status;
pub mod websocket;
pub use asset::Asset;
pub use asset::{Asset, Class, Exchange, Status};
pub use bar::Bar;
pub use class::Class;
pub use exchange::Exchange;
pub use source::Source;
pub use status::Status;
#[derive(Clone, Debug)]
pub enum BroadcastMessage {

View File

@@ -1,12 +0,0 @@
use serde::Deserialize;
use sqlx::Type;
#[derive(PartialEq, Eq, Deserialize, Type)]
pub enum Status {
#[sqlx(rename = "active")]
#[serde(rename = "active")]
Active,
#[sqlx(rename = "inactive")]
#[serde(rename = "inactive")]
Inactive,
}

View File

@@ -0,0 +1,2 @@
pub mod incoming;
pub mod outgoing;

View File

@@ -1,2 +1 @@
pub mod incoming;
pub mod outgoing;
pub mod data;