fix: feed rpc respones and notifications instead of send and flush

Author: ali-behjati

config/config.toml

@@ -6,6 +6,18 @@
 # connection is not exposed for unauthorized access.
 listen_address = "127.0.0.1:8910"
 
+# Size of the buffer of each Server's channel on which `notify_price` events are
+# received from the Price state.
+# notify_price_tx_buffer = 10000
+
+# Size of the buffer of each Server's channel on which `notify_price_sched` events are
+# received from the Price state.
+# notify_price_sched_tx_buffer = 10000
+
+# Flush interval for the notifications. This is the maximum time the server
+# will wait before flushing the notifications to the client.
+# notify_flush_interval_duration = "100ms"
+
 # Configuration for the primary network this agent will publish data to. In most cases this should be a Pythnet endpoint.
 [primary_network]
 ### Required fields ###

src/agent/pyth/rpc.rs

@@ -48,6 +48,7 @@ use {
         fmt::Debug,
         net::SocketAddr,
         sync::Arc,
+        time::Duration,
     },
     tokio::sync::mpsc,
     tracing::instrument,
@@ -115,6 +116,7 @@ async fn handle_connection<S>(
     state: Arc<S>,
     notify_price_tx_buffer: usize,
     notify_price_sched_tx_buffer: usize,
+    notify_flush_interval_duration: Duration,
 ) where
     S: state::Prices,
     S: Send,
@@ -127,6 +129,8 @@ async fn handle_connection<S>(
     let (mut notify_price_sched_tx, mut notify_price_sched_rx) =
         mpsc::channel(notify_price_sched_tx_buffer);
 
+    let mut notify_flush_interval = tokio::time::interval(notify_flush_interval_duration);
+
     loop {
         if let Err(err) = handle_next(
             &*state,
@@ -136,6 +140,7 @@ async fn handle_connection<S>(
             &mut notify_price_rx,
             &mut notify_price_sched_tx,
             &mut notify_price_sched_rx,
+            &mut notify_flush_interval,
         )
         .await
         {
@@ -159,6 +164,7 @@ async fn handle_next<S>(
     notify_price_rx: &mut mpsc::Receiver<NotifyPrice>,
     notify_price_sched_tx: &mut mpsc::Sender<NotifyPriceSched>,
     notify_price_sched_rx: &mut mpsc::Receiver<NotifyPriceSched>,
+    notify_flush_interval: &mut tokio::time::Interval,
 ) -> Result<()>
 where
     S: state::Prices,
@@ -183,13 +189,16 @@ where
             }
         }
         Some(notify_price) = notify_price_rx.recv() => {
-            send_notification(ws_tx, Method::NotifyPrice, Some(notify_price))
+            feed_notification(ws_tx, Method::NotifyPrice, Some(notify_price))
                 .await
         }
         Some(notify_price_sched) = notify_price_sched_rx.recv() => {
-            send_notification(ws_tx, Method::NotifyPriceSched, Some(notify_price_sched))
+            feed_notification(ws_tx, Method::NotifyPriceSched, Some(notify_price_sched))
                 .await
         }
+        _ = notify_flush_interval.tick() => {
+            flush(ws_tx).await
+        }
     }
 }
 
@@ -357,18 +366,18 @@ async fn send_error(
     send_text(ws_tx, &response.to_string()).await
 }
 
-async fn send_notification<T>(
+async fn feed_notification<T>(
     ws_tx: &mut SplitSink<WebSocket, Message>,
     method: Method,
     params: Option<T>,
 ) -> Result<()>
 where
     T: Sized + Serialize + DeserializeOwned,
 {
-    send_request(ws_tx, IdReq::Notification, method, params).await
+    feed_request(ws_tx, IdReq::Notification, method, params).await
 }
 
-async fn send_request<I, T>(
+async fn feed_request<I, T>(
     ws_tx: &mut SplitSink<WebSocket, Message>,
     id: I,
     method: Method,
@@ -379,7 +388,14 @@ where
     T: Sized + Serialize + DeserializeOwned,
 {
     let request = Request::with_params(id, method, params);
-    send_text(ws_tx, &request.to_string()).await
+    feed_text(ws_tx, &request.to_string()).await
+}
+
+async fn feed_text(ws_tx: &mut SplitSink<WebSocket, Message>, msg: &str) -> Result<()> {
+    ws_tx
+        .feed(Message::text(msg.to_string()))
+        .await
+        .map_err(|e| e.into())
 }
 
 async fn send_text(ws_tx: &mut SplitSink<WebSocket, Message>, msg: &str) -> Result<()> {
@@ -389,25 +405,33 @@ async fn send_text(ws_tx: &mut SplitSink<WebSocket, Message>, msg: &str) -> Resu
         .map_err(|e| e.into())
 }
 
+async fn flush(ws_tx: &mut SplitSink<WebSocket, Message>) -> Result<()> {
+    ws_tx.flush().await.map_err(|e| e.into())
+}
+
 #[derive(Clone, Debug, Serialize, Deserialize)]
 #[serde(default)]
 pub struct Config {
     /// The address which the websocket API server will listen on.
-    pub listen_address:               String,
+    pub listen_address:                 String,
     /// Size of the buffer of each Server's channel on which `notify_price` events are
     /// received from the Price state.
-    pub notify_price_tx_buffer:       usize,
+    pub notify_price_tx_buffer:         usize,
     /// Size of the buffer of each Server's channel on which `notify_price_sched` events are
     /// received from the Price state.
-    pub notify_price_sched_tx_buffer: usize,
+    pub notify_price_sched_tx_buffer:   usize,
+    /// Flush interval duration for the notifications.
+    #[serde(with = "humantime_serde")]
+    pub notify_flush_interval_duration: Duration,
 }
 
 impl Default for Config {
     fn default() -> Self {
         Self {
-            listen_address:               "127.0.0.1:8910".to_string(),
-            notify_price_tx_buffer:       10000,
-            notify_price_sched_tx_buffer: 10000,
+            listen_address:                 "127.0.0.1:8910".to_string(),
+            notify_price_tx_buffer:         10000,
+            notify_price_sched_tx_buffer:   10000,
+            notify_flush_interval_duration: Duration::from_millis(50),
         }
     }
 }
@@ -448,6 +472,7 @@ where
                         state,
                         config.notify_price_tx_buffer,
                         config.notify_price_sched_tx_buffer,
+                        config.notify_flush_interval_duration,
                     )
                     .await
                 })