Make heartbeathandler shared struct

src/gateway/default/gateway.rs

@@ -7,7 +7,7 @@ use crate::types::{self, WebSocketEvent};
 #[derive(Debug)]
 pub struct DefaultGateway {
     events: Arc<Mutex<Events>>,
-    heartbeat_handler: DefaultHeartbeatHandler,
+    heartbeat_handler: HeartbeatHandler,
     websocket_send: Arc<
         Mutex<
             SplitSink<
@@ -28,10 +28,10 @@ impl
         WebSocketStream<MaybeTlsStream<TcpStream>>,
         WebSocketStream<MaybeTlsStream<TcpStream>>,
         DefaultGatewayHandle,
-        DefaultHeartbeatHandler,
+        HeartbeatHandler,
     > for DefaultGateway
 {
-    fn get_heartbeat_handler(&self) -> &DefaultHeartbeatHandler {
+    fn get_heartbeat_handler(&self) -> &HeartbeatHandler {
         &self.heartbeat_handler
     }
 
@@ -95,7 +95,7 @@ impl
 
         let mut gateway = DefaultGateway {
             events: shared_events.clone(),
-            heartbeat_handler: DefaultHeartbeatHandler::new(
+            heartbeat_handler: HeartbeatHandler::new(
                 Duration::from_millis(gateway_hello.heartbeat_interval),
                 shared_websocket_send.clone(),
                 kill_send.subscribe(),

src/gateway/default/heartbeat.rs

@@ -1,22 +1,7 @@
-use crate::types;
-
 use super::*;
 
-/// Handles sending heartbeats to the gateway in another thread
+#[async_trait]
-#[allow(dead_code)] // FIXME: Remove this, once HeartbeatHandler is used
+impl HeartbeatHandlerCapable<WebSocketStream<MaybeTlsStream<TcpStream>>> for HeartbeatHandler {
-#[derive(Debug)]
-pub struct DefaultHeartbeatHandler {
-    /// How ofter heartbeats need to be sent at a minimum
-    pub heartbeat_interval: Duration,
-    /// The send channel for the heartbeat thread
-    pub send: Sender<HeartbeatThreadCommunication>,
-    /// The handle of the thread
-    handle: JoinHandle<()>,
-}
-
-impl HeartbeatHandlerCapable<WebSocketStream<MaybeTlsStream<TcpStream>>>
-    for DefaultHeartbeatHandler
-{
     fn new(
         heartbeat_interval: Duration,
         websocket_tx: Arc<
@@ -28,12 +13,12 @@ impl HeartbeatHandlerCapable<WebSocketStream<MaybeTlsStream<TcpStream>>>
             >,
         >,
         kill_rc: tokio::sync::broadcast::Receiver<()>,
-    ) -> DefaultHeartbeatHandler {
+    ) -> HeartbeatHandler {
         let (send, receive) = tokio::sync::mpsc::channel(32);
         let kill_receive = kill_rc.resubscribe();
 
         let handle: JoinHandle<()> = task::spawn(async move {
-            DefaultHeartbeatHandler::heartbeat_task(
+            HeartbeatHandler::heartbeat_task(
                 websocket_tx,
                 heartbeat_interval,
                 receive,
@@ -57,92 +42,3 @@ impl HeartbeatHandlerCapable<WebSocketStream<MaybeTlsStream<TcpStream>>>
         self.heartbeat_interval
     }
 }
-
-impl DefaultHeartbeatHandler {
-    /// The main heartbeat task;
-    ///
-    /// Can be killed by the kill broadcast;
-    /// If the websocket is closed, will die out next time it tries to send a heartbeat;
-    pub async fn heartbeat_task(
-        websocket_tx: Arc<
-            Mutex<
-                SplitSink<
-                    WebSocketStream<MaybeTlsStream<TcpStream>>,
-                    tokio_tungstenite::tungstenite::Message,
-                >,
-            >,
-        >,
-        heartbeat_interval: Duration,
-        mut receive: tokio::sync::mpsc::Receiver<HeartbeatThreadCommunication>,
-        mut kill_receive: tokio::sync::broadcast::Receiver<()>,
-    ) {
-        let mut last_heartbeat_timestamp: Instant = time::Instant::now();
-        let mut last_heartbeat_acknowledged = true;
-        let mut last_seq_number: Option<u64> = None;
-
-        loop {
-            if kill_receive.try_recv().is_ok() {
-                trace!("GW: Closing heartbeat task");
-                break;
-            }
-
-            let timeout = if last_heartbeat_acknowledged {
-                heartbeat_interval
-            } else {
-                // If the server hasn't acknowledged our heartbeat we should resend it
-                Duration::from_millis(HEARTBEAT_ACK_TIMEOUT)
-            };
-
-            let mut should_send = false;
-
-            tokio::select! {
-                () = sleep_until(last_heartbeat_timestamp + timeout) => {
-                    should_send = true;
-                }
-                Some(communication) = receive.recv() => {
-                    // If we received a seq number update, use that as the last seq number
-                    if communication.sequence_number.is_some() {
-                        last_seq_number = communication.sequence_number;
-                    }
-
-                    if let Some(op_code) = communication.op_code {
-                        match op_code {
-                            GATEWAY_HEARTBEAT => {
-                                // As per the api docs, if the server sends us a Heartbeat, that means we need to respond with a heartbeat immediately
-                                should_send = true;
-                            }
-                            GATEWAY_HEARTBEAT_ACK => {
-                                // The server received our heartbeat
-                                last_heartbeat_acknowledged = true;
-                            }
-                            _ => {}
-                        }
-                    }
-                }
-            }
-
-            if should_send {
-                trace!("GW: Sending Heartbeat..");
-
-                let heartbeat = types::GatewayHeartbeat {
-                    op: GATEWAY_HEARTBEAT,
-                    d: last_seq_number,
-                };
-
-                let heartbeat_json = serde_json::to_string(&heartbeat).unwrap();
-
-                let msg = tokio_tungstenite::tungstenite::Message::text(heartbeat_json);
-
-                let send_result = websocket_tx.lock().await.send(msg).await;
-                if send_result.is_err() {
-                    // We couldn't send, the websocket is broken
-                    warn!("GW: Couldnt send heartbeat, websocket seems broken");
-                    break;
-                }
-
-                last_heartbeat_timestamp = time::Instant::now();
-                last_heartbeat_acknowledged = false;
-            }
-        }
-    }
-}

src/gateway/default/mod.rs

@@ -5,7 +5,6 @@ pub mod heartbeat;
 use super::*;
 pub use gateway::*;
 pub use handle::*;
-use heartbeat::*;
 use tokio_tungstenite::tungstenite::Message;
 
 use crate::errors::GatewayError;
@@ -15,18 +14,15 @@ use std::collections::HashMap;
 use std::fmt::Debug;
 use std::sync::{Arc, RwLock};
 use std::time::Duration;
-use tokio::time::sleep_until;
 
 use futures_util::stream::SplitSink;
 use futures_util::stream::SplitStream;
-use log::{info, trace, warn};
+use log::{info, warn};
 use tokio::net::TcpStream;
 use tokio::sync::mpsc::Sender;
 use tokio::sync::Mutex;
 use tokio::task;
 use tokio::task::JoinHandle;
-use tokio::time;
-use tokio::time::Instant;
 use tokio_tungstenite::MaybeTlsStream;
 use tokio_tungstenite::{connect_async_tls_with_config, Connector, WebSocketStream};
 

src/gateway/mod.rs

@@ -8,6 +8,8 @@ pub mod wasm;
 #[cfg(all(not(target_arch = "wasm32"), feature = "client"))]
 pub use default::*;
 pub use message::*;
+use safina_timer::sleep_until;
+use tokio::task::JoinHandle;
 #[cfg(all(target_arch = "wasm32", feature = "client"))]
 pub use wasm::*;
 
@@ -22,7 +24,7 @@ use crate::types::{
 use std::any::Any;
 use std::collections::HashMap;
 use std::sync::{Arc, RwLock};
-use std::time::Duration;
+use std::time::{self, Duration, Instant};
 
 use async_trait::async_trait;
 use futures_util::stream::SplitSink;
@@ -154,7 +156,7 @@ impl<T: WebSocketEvent> GatewayEvent<T> {
 pub trait GatewayCapable<R, S, G, H>
 where
     R: Stream,
-    S: Sink<Message>,
+    S: Sink<Message> + Send + 'static,
     G: GatewayHandleCapable<R, S>,
     H: HeartbeatHandlerCapable<S> + Send + Sync,
 {
@@ -441,6 +443,18 @@ where
     }
 }
 
+/// Handles sending heartbeats to the gateway in another thread
+#[allow(dead_code)] // FIXME: Remove this, once HeartbeatHandler is used
+#[derive(Debug)]
+pub struct HeartbeatHandler {
+    /// How ofter heartbeats need to be sent at a minimum
+    pub heartbeat_interval: Duration,
+    /// The send channel for the heartbeat thread
+    pub send: Sender<HeartbeatThreadCommunication>,
+    /// The handle of the thread
+    handle: JoinHandle<()>,
+}
+
 #[async_trait(?Send)]
 pub trait GatewayHandleCapable<R, S>
 where
@@ -541,7 +555,8 @@ where
     async fn close(&self);
 }
 
-pub trait HeartbeatHandlerCapable<S: Sink<Message>> {
+#[async_trait]
+pub trait HeartbeatHandlerCapable<S: Sink<Message> + Send + 'static> {
     fn new(
         heartbeat_interval: Duration,
         websocket_tx: Arc<Mutex<SplitSink<S, Message>>>,
@@ -550,4 +565,80 @@ pub trait HeartbeatHandlerCapable<S: Sink<Message>> {
 
     fn get_send(&self) -> &Sender<HeartbeatThreadCommunication>;
     fn get_heartbeat_interval(&self) -> Duration;
+    async fn heartbeat_task(
+        websocket_tx: Arc<Mutex<SplitSink<S, Message>>>,
+        heartbeat_interval: Duration,
+        mut receive: tokio::sync::mpsc::Receiver<HeartbeatThreadCommunication>,
+        mut kill_receive: tokio::sync::broadcast::Receiver<()>,
+    ) {
+        let mut last_heartbeat_timestamp: Instant = time::Instant::now();
+        let mut last_heartbeat_acknowledged = true;
+        let mut last_seq_number: Option<u64> = None;
+        safina_timer::start_timer_thread();
+
+        loop {
+            if kill_receive.try_recv().is_ok() {
+                trace!("GW: Closing heartbeat task");
+                break;
+            }
+
+            let timeout = if last_heartbeat_acknowledged {
+                heartbeat_interval
+            } else {
+                // If the server hasn't acknowledged our heartbeat we should resend it
+                Duration::from_millis(HEARTBEAT_ACK_TIMEOUT)
+            };
+
+            let mut should_send = false;
+
+            tokio::select! {
+                () = sleep_until(last_heartbeat_timestamp + timeout) => {
+                    should_send = true;
+                }
+                Some(communication) = receive.recv() => {
+                    // If we received a seq number update, use that as the last seq number
+                    if communication.sequence_number.is_some() {
+                        last_seq_number = communication.sequence_number;
+                    }
+
+                    if let Some(op_code) = communication.op_code {
+                        match op_code {
+                            GATEWAY_HEARTBEAT => {
+                                // As per the api docs, if the server sends us a Heartbeat, that means we need to respond with a heartbeat immediately
+                                should_send = true;
+                            }
+                            GATEWAY_HEARTBEAT_ACK => {
+                                // The server received our heartbeat
+                                last_heartbeat_acknowledged = true;
+                            }
+                            _ => {}
+                        }
+                    }
+                }
+            }
+
+            if should_send {
+                trace!("GW: Sending Heartbeat..");
+
+                let heartbeat = types::GatewayHeartbeat {
+                    op: GATEWAY_HEARTBEAT,
+                    d: last_seq_number,
+                };
+
+                let heartbeat_json = serde_json::to_string(&heartbeat).unwrap();
+
+                let msg = tokio_tungstenite::tungstenite::Message::text(heartbeat_json);
+
+                let send_result = websocket_tx.lock().await.send(msg).await;
+                if send_result.is_err() {
+                    // We couldn't send, the websocket is broken
+                    warn!("GW: Couldnt send heartbeat, websocket seems broken");
+                    break;
+                }
+
+                last_heartbeat_timestamp = time::Instant::now();
+                last_heartbeat_acknowledged = false;
+            }
+        }
+    }
 }