feat: small ws improvements (#59)

* feat: doc example and timestamp parsing

* feat: cors/logging layer + ping msg

* Update ohlc.rs

* feat: msg handling

* fix: clippy

* fix: taplo

Author: EvolveArt

Cargo.lock

@@ -30,6 +30,7 @@ strum = { version = "0.25.0", features = ["derive"] }
 thiserror = "1.0.49"
 time = "0.3.29"
 tokio = { version = "1.0", features = ["sync", "macros", "rt-multi-thread"] }
+tower-http = { version = "0.4.0", features = ["fs", "trace", "cors"] }
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 url = "2.5.0"

pragma-node/Cargo.toml

@@ -34,7 +34,7 @@ pub async fn get_entry(
     // Construct pair id
     let pair_id = currency_pair_to_pair_id(&pair.0, &pair.1);
 
-    let now = chrono::Utc::now().timestamp() as u64;
+    let now = chrono::Utc::now().timestamp();
 
     let timestamp = if let Some(timestamp) = params.timestamp {
         timestamp

pragma-node/src/handlers/entries/get_entry.rs

@@ -31,7 +31,7 @@ pub async fn get_ohlc(
     // Construct pair id
     let pair_id = currency_pair_to_pair_id(&pair.0, &pair.1);
 
-    let now = chrono::Utc::now().timestamp() as u64;
+    let now = chrono::Utc::now().timestamp();
 
     let timestamp = if let Some(timestamp) = params.timestamp {
         timestamp

pragma-node/src/handlers/entries/get_ohlc.rs

@@ -73,12 +73,26 @@ async fn handle_channel(mut socket: WebSocket, state: AppState) {
     let mut ohlc_to_compute = 10;
     let mut ohlc_data: Vec<OHLCEntry> = Vec::new();
 
+    //send a ping (unsupported by some browsers) just to kick things off and get a response
+    if socket.send(Message::Ping(vec![])).await.is_ok() {
+        tracing::info!("Pinged ...");
+    } else {
+        tracing::info!("Could not send ping !");
+        // no Error here since the only thing we can do is to close the connection.
+        // If we can not send messages, there is no way to salvage the statemachine anyway.
+        return;
+    }
+
     loop {
         tokio::select! {
-            Some(msg) = socket.recv() => {
-                if let Ok(Message::Text(text)) = msg {
-                    handle_message_received(&mut socket, &state, &mut subscribed_pair, &mut network, &mut interval, text).await;
-                }
+            Some(maybe_msg) = socket.recv() => {
+                // TODO: remove once we have proper top-level error handling
+                let msg = if let Ok(msg) = maybe_msg {
+                    msg
+                } else {
+                    break;
+                };
+                handle_message_received(&mut socket, &state, &mut subscribed_pair, &mut network, &mut interval, msg).await;
             },
             _ = update_interval.tick() => {
                 match send_ohlc_data(&mut socket, &state, &subscribed_pair, &mut ohlc_data, network, interval, ohlc_to_compute).await {
@@ -103,9 +117,27 @@ async fn handle_message_received(
     subscribed_pair: &mut Option<String>,
     network: &mut Network,
     interval: &mut Interval,
-    message: String,
+    message: Message,
 ) {
-    if let Ok(subscription_msg) = serde_json::from_str::<SubscriptionRequest>(&message) {
+    let maybe_client_message = match message {
+        Message::Close(_) => {
+            // TODO: Send the close message to gracefully shut down the connection
+            // Otherwise the client might get an abnormal Websocket closure
+            // error.
+            return;
+        }
+        Message::Text(text) => serde_json::from_str::<SubscriptionRequest>(&text),
+        Message::Binary(data) => serde_json::from_slice::<SubscriptionRequest>(&data),
+        Message::Ping(_) => {
+            // Axum will send Pong automatically
+            return;
+        }
+        Message::Pong(_) => {
+            return;
+        }
+    };
+
+    if let Ok(subscription_msg) = maybe_client_message {
         match subscription_msg.msg_type {
             SubscriptionType::Subscribe => {
                 let pair_exists = is_onchain_existing_pair(
@@ -196,5 +228,7 @@ async fn send_ohlc_data(
 /// (Does not close the connection)
 async fn send_error_message(socket: &mut WebSocket, error: &str) {
     let error_msg = json!({ "error": error }).to_string();
-    socket.send(Message::Text(error_msg)).await.unwrap();
+    if socket.send(Message::Text(error_msg)).await.is_err() {
+        tracing::error!("Client already disconnected. Could not send error message.");
+    }
 }

pragma-node/src/handlers/entries/get_onchain/ohlc.rs

@@ -10,7 +10,10 @@ pub use get_ohlc::get_ohlc;
 pub use get_onchain::get_onchain;
 pub use get_volatility::get_volatility;
 
-use crate::infra::repositories::entry_repository::OHLCEntry;
+use crate::{
+    infra::repositories::entry_repository::OHLCEntry,
+    utils::{doc_examples, UnixTimestamp},
+};
 
 pub mod create_entry;
 pub mod get_entry;
@@ -134,7 +137,11 @@ pub struct GetOnchainCheckpointsResponse(pub Vec<Checkpoint>);
 
 #[derive(Debug, Deserialize, IntoParams, ToSchema)]
 pub struct GetEntryParams {
-    pub timestamp: Option<u64>,
+    /// The unix timestamp in seconds. This endpoint will return the first update whose
+    /// timestamp is <= the provided value.
+    #[param(value_type = i64)]
+    #[param(example = doc_examples::timestamp_example)]
+    pub timestamp: Option<UnixTimestamp>,
     pub interval: Option<Interval>,
     pub routing: Option<bool>,
     pub aggregation: Option<AggregationMode>,
@@ -143,7 +150,7 @@ pub struct GetEntryParams {
 impl Default for GetEntryParams {
     fn default() -> Self {
         Self {
-            timestamp: Some(chrono::Utc::now().timestamp() as u64),
+            timestamp: Some(chrono::Utc::now().timestamp()),
             interval: Some(Interval::default()),
             routing: Some(false),
             aggregation: Some(AggregationMode::default()),

pragma-node/src/handlers/entries/mod.rs

@@ -87,7 +87,7 @@ pub async fn routing(
     pool: &deadpool_diesel::postgres::Pool,
     pair_id: String,
     interval: Interval,
-    timestamp: u64,
+    timestamp: i64,
     is_routing: bool,
     agg_mode: AggregationMode,
 ) -> Result<(MedianEntry, u32), InfraError> {
@@ -164,7 +164,7 @@ async fn find_alternative_pair_price(
     base: &str,
     quote: &str,
     interval: Interval,
-    timestamp: u64,
+    timestamp: i64,
     agg_mode: AggregationMode,
 ) -> Result<(MedianEntry, u32), InfraError> {
     let conn = pool.get().await.map_err(adapt_infra_error)?;
@@ -213,7 +213,7 @@ async fn get_price_decimals(
     pool: &deadpool_diesel::postgres::Pool,
     pair_id: String,
     interval: Interval,
-    timestamp: u64,
+    timestamp: i64,
     agg_mode: AggregationMode,
 ) -> Result<(MedianEntry, u32), InfraError> {
     let entry = match agg_mode {
@@ -251,7 +251,7 @@ pub async fn get_twap_price(
     pool: &deadpool_diesel::postgres::Pool,
     pair_id: String,
     interval: Interval,
-    time: u64,
+    time: i64,
 ) -> Result<MedianEntry, InfraError> {
     let conn = pool.get().await.map_err(adapt_infra_error)?;
 
@@ -330,7 +330,7 @@ pub async fn get_twap_price(
         }
     };
 
-    let date_time = DateTime::from_timestamp(time as i64, 0).ok_or(InfraError::InvalidTimeStamp)?;
+    let date_time = DateTime::from_timestamp(time, 0).ok_or(InfraError::InvalidTimeStamp)?;
 
     let raw_entry = conn
         .interact(move |conn| {
@@ -357,7 +357,7 @@ pub async fn get_median_price(
     pool: &deadpool_diesel::postgres::Pool,
     pair_id: String,
     interval: Interval,
-    time: u64,
+    time: i64,
 ) -> Result<MedianEntry, InfraError> {
     let conn = pool.get().await.map_err(adapt_infra_error)?;
 
@@ -436,7 +436,7 @@ pub async fn get_median_price(
         }
     };
 
-    let date_time = DateTime::from_timestamp(time as i64, 0).ok_or(InfraError::InvalidTimeStamp)?;
+    let date_time = DateTime::from_timestamp(time, 0).ok_or(InfraError::InvalidTimeStamp)?;
 
     let raw_entry = conn
         .interact(move |conn| {
@@ -595,7 +595,7 @@ pub async fn get_ohlc(
     pool: &deadpool_diesel::postgres::Pool,
     pair_id: String,
     interval: Interval,
-    time: u64,
+    time: i64,
 ) -> Result<Vec<OHLCEntry>, InfraError> {
     let conn = pool.get().await.map_err(adapt_infra_error)?;
 
@@ -682,7 +682,7 @@ pub async fn get_ohlc(
         }
     };
 
-    let date_time = DateTime::from_timestamp(time as i64, 0).ok_or(InfraError::InvalidTimeStamp)?;
+    let date_time = DateTime::from_timestamp(time, 0).ok_or(InfraError::InvalidTimeStamp)?;
 
     let raw_entries = conn
         .interact(move |conn| {

pragma-node/src/infra/repositories/entry_repository.rs

@@ -1,6 +1,8 @@
 use deadpool_diesel::postgres::Pool;
 use pragma_entities::connection::{ENV_POSTGRES_DATABASE_URL, ENV_TS_DATABASE_URL};
 use std::net::SocketAddr;
+use tower_http::cors::CorsLayer;
+use tower_http::trace::{DefaultMakeSpan, TraceLayer};
 use utoipa::openapi::security::{ApiKey, ApiKeyValue, SecurityScheme};
 use utoipa::Modify;
 use utoipa::OpenApi;
@@ -110,7 +112,15 @@ async fn main() {
         postgres_pool,
     };
 
-    let app = app_router::<ApiDoc>(state.clone()).with_state(state);
+    let app = app_router::<ApiDoc>(state.clone())
+        .with_state(state)
+        // Logging so we can see whats going on
+        .layer(
+            TraceLayer::new_for_http()
+                .make_span_with(DefaultMakeSpan::default().include_headers(true)),
+        )
+        // Permissive CORS layer to allow all origins
+        .layer(CorsLayer::permissive());
 
     let host = config.server_host();
     let port = config.server_port();

pragma-node/src/main.rs

@@ -0,0 +1,7 @@
+use crate::utils::types::UnixTimestamp;
+
+/// Example value for a unix timestamp
+pub fn timestamp_example() -> UnixTimestamp {
+    const STATIC_UNIX_TIMESTAMP: UnixTimestamp = 1717632000; // Thursday, 6 June 2024 00:00:00 GMT
+    STATIC_UNIX_TIMESTAMP
+}

pragma-node/src/utils/doc_examples.rs

@@ -2,7 +2,10 @@ pub use conversion::{convert_via_quote, format_bigdecimal_price, normalize_to_de
 pub use custom_extractors::json_extractor::JsonExtractor;
 pub use custom_extractors::path_extractor::PathExtractor;
 pub use signing::typed_data::TypedData;
+pub use types::UnixTimestamp;
 
 mod conversion;
 mod custom_extractors;
+pub mod doc_examples;
 mod signing;
+mod types;

pragma-node/src/utils/mod.rs

@@ -0,0 +1,4 @@
+/// The number of seconds since the Unix epoch (00:00:00 UTC on 1 Jan 1970). The timestamp is
+/// always positive, but represented as a signed integer because that's the standard on Unix
+/// systems and allows easy subtraction to compute durations.
+pub type UnixTimestamp = i64;