Merge pull request #15 from drift-labs/master

add pub/sub heartbeat channel

Author: moosecat2

src/swift_server.rs

@@ -231,6 +231,75 @@ pub async fn process_order(
     }
 }
 
+pub async fn send_heartbeat(server_params: &mut ServerParams) -> () {
+    let hearbeat_time = unix_now_ms();
+    let log_prefix = format!("[hearbeat: {hearbeat_time}]");
+
+    if let Some(kafka_producer) = &server_params.kafka_producer {
+        match kafka_producer
+            .send(
+                FutureRecord::<String, String>::to(&"hearbeat").payload(&"love you".to_string()),
+                Timeout::After(Duration::ZERO),
+            )
+            .await
+        {
+            Ok(_) => {
+                log::trace!(target: "kafka", "{log_prefix}: Sent heartbeat");
+                server_params
+                    .metrics
+                    .order_type_counter
+                    .with_label_values(&["_", "heartbeat"])
+                    .inc();
+
+                server_params
+                    .metrics
+                    .response_time_histogram
+                    .observe((unix_now_ms() - hearbeat_time) as f64);
+            }
+            Err((e, _)) => {
+                log::error!(
+                    target: "kafka",
+                    "{log_prefix}: Failed to deliver heartbeat, error: {e:?}"
+                );
+                server_params.metrics.kafka_forward_fail_counter.inc();
+            }
+        }
+    } else {
+        let mut conn = match server_params.redis_pool.as_ref().unwrap().get().await {
+            Ok(conn) => conn,
+            Err(_) => {
+                log::error!(target: "redis", "{log_prefix}: Obtaining redis connection failed");
+                return;
+            }
+        };
+
+        match conn
+            .publish::<String, String, i64>("heartbeat".to_string(), "love you".to_string())
+            .await
+        {
+            Ok(_) => {
+                log::trace!(target: "redis", "{log_prefix}: Sent redis heartbeat");
+                server_params
+                    .metrics
+                    .order_type_counter
+                    .with_label_values(&["_", "heartbeat"])
+                    .inc();
+
+                server_params
+                    .metrics
+                    .response_time_histogram
+                    .observe((unix_now_ms() - hearbeat_time) as f64);
+            }
+            Err(e) => {
+                log::error!(
+                    target: "redis",
+                    "{log_prefix}: Failed to deliver heartbeat, error: {e:?}"
+                );
+            }
+        }
+    }
+}
+
 pub async fn health_check() -> impl axum::response::IntoResponse {
     axum::http::StatusCode::OK
 }
@@ -408,10 +477,25 @@ pub async fn start_server() {
         }
     });
 
+    let mut state_clone = state.clone();
+    let send_heartbeat_loop = tokio::spawn(async move {
+        let mut interval = tokio::time::interval(Duration::from_secs(2));
+
+        loop {
+            interval.tick().await;
+            send_heartbeat(&mut state_clone).await;
+        }
+    });
+
     let axum_server = tokio::spawn(async { axum::serve(listener, app).await });
     let metrics_server = tokio::spawn(async { axum::serve(listener_metrics, metrics_app).await });
 
-    let _ = tokio::try_join!(rpc_sim_loop, axum_server, metrics_server);
+    let _ = tokio::try_join!(
+        rpc_sim_loop,
+        axum_server,
+        metrics_server,
+        send_heartbeat_loop
+    );
 }
 
 /// Simple validation from program's `handle_signed_order_ix`

src/types/messages.rs

@@ -10,7 +10,7 @@ use base64::Engine;
 use drift_rs::types::{MarketType, SignedMsgOrderParamsMessage};
 use ed25519_dalek::{PublicKey, Signature, Verifier};
 use serde_json::json;
-use solana_sdk::{hash::Hash, pubkey::Pubkey};
+use solana_sdk::pubkey::Pubkey;
 
 #[derive(AnchorDeserialize, AnchorSerialize, Debug, Clone, InitSpace, Copy, Default, PartialEq)]
 pub struct SignedOrderParamsMessageWithPrefix {
@@ -241,10 +241,6 @@ pub struct BorshBuf<const N: usize> {
 }
 
 impl<const N: usize> BorshBuf<N> {
-    /// Get the SHA256 digest of the buffer
-    pub fn sha256_digest(&self) -> Hash {
-        solana_sdk::hash::hash(self.data())
-    }
     /// Deserialize the buffer as `T`
     pub fn deserialize<T: anchor_lang::AnchorDeserialize>(&self) -> Result<T, std::io::Error> {
         T::deserialize(&mut self.data())
@@ -290,7 +286,9 @@ where
 }
 
 /// Deserialize base64 str as fixed size byte array with Borsh helpers
-pub fn base64_to_borsh_buf<'de, D, const N: usize>(deserializer: D) -> Result<BorshBuf<N>, D::Error>
+pub fn _base64_to_borsh_buf<'de, D, const N: usize>(
+    deserializer: D,
+) -> Result<BorshBuf<N>, D::Error>
 where
     D: serde::Deserializer<'de>,
 {

src/ws_server.rs

@@ -599,7 +599,7 @@ async fn subscribe_kafka_consumer(
     topics_prefix: &str,
 ) {
     kafka_consumer
-        .subscribe(&[topics_prefix])
+        .subscribe(&[topics_prefix, &"heartbeat"])
         .context("Failed to subscribe to topics")
         .expect("subscribed topic prefix");
 
@@ -626,9 +626,11 @@ async fn subscribe_kafka_consumer(
                     .expect("topic channel exists");
 
                 if (tx.receiver_count() as u32) < 1 {
-                    // TODO: consumer got the message but it won't be forwarded since no one's listening??
-                    // do we want to have something where it replays recent message on connect?
-                    log::error!(target: "kafka", "No receiver found for topic: {topic}, order message lost!");
+                    if topic != "heartbeat" {
+                        log::warn!(target: "kafka", "No receiver found for topic: {topic}, order message lost!");
+                    } else {
+                        log::debug!(target: "kafka", "Received heartbeat message");
+                    }
                     continue;
                 }
                 let payload: &[u8] = message.payload().context("Failed to get payload").unwrap();
@@ -728,9 +730,11 @@ async fn subscribe_redis_pubsub(
             .expect("topic channel exists");
 
         if (tx.receiver_count() as u32) < 1 {
-            // TODO: consumer got the message but it won't be forwarded since no one's listening??
-            // do we want to have something where it replays recent message on connect?
-            log::error!(target: "ws", "No receiver found for topic: {topic}, order message lost!");
+            if topic != "heartbeat" {
+                log::warn!(target: "ws", "No receiver found for topic: {topic}, order message lost!");
+            } else {
+                log::debug!(target: "ws", "Received heartbeat message");
+            }
             continue;
         }
         let payload: &[u8] = message.get_payload_bytes();
@@ -793,8 +797,12 @@ pub async fn start_server() {
 
     // Set up the server with the server params
     let subscriptions = DashMap::new();
-    let mut topic_names: Vec<String> = vec![];
+    let mut topic_names: Vec<String> = vec!["heartbeat".to_string()];
     for market in &perp_market_accounts {
+        if market.symbol().contains("BET") {
+            log::info!("Skipping BET market");
+            continue;
+        }
         let topic = format!(
             "swift_orders_{}_{}",
             market.market_type(),
@@ -1299,7 +1307,7 @@ mod test {
         let mut ws_conn = WsConnection::new(Pubkey::new_unique());
         ws_conn.authenticated = true;
 
-        ws_conn
+        let _ = ws_conn
             .spawn_handler(&mut client_rx, &mut client_tx, Box::leak(Box::default()))
             .await;
 
@@ -1330,7 +1338,7 @@ mod test {
         let mut client_rx = vec![];
         let mut client_tx = stream::iter([Ok(Message::Close(None))]);
         let ws_conn = WsConnection::new(Pubkey::new_unique());
-        ws_conn
+        let _ = ws_conn
             .spawn_handler(&mut client_rx, &mut client_tx, Box::leak(Box::default()))
             .await;
     }
@@ -1344,7 +1352,7 @@ mod test {
         let _ = ws_conn.send_message(WsMessage::heartbeat());
         let _ = ws_conn.send_message(WsMessage::heartbeat());
         let _ = ws_conn.send_message(WsMessage::heartbeat());
-        ws_conn
+        let _ = ws_conn
             .spawn_handler(&mut client_rx, &mut client_tx, Box::leak(Box::default()))
             .await;
         assert!(client_rx.len() >= 3);