fix(hermes): disconnect slow WS/SSE consumers to prevent OOM

apps/hermes/server/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name        = "hermes"
-version     = "0.10.8"
+version     = "0.11.0"
 description = "Hermes is an agent that provides Verified Prices from the Pythnet Pyth Oracle."
 edition     = "2021"
 

apps/hermes/server/src/api.rs

@@ -20,10 +20,17 @@ pub mod token;
 pub mod types;
 mod ws;
 
+#[derive(Clone, Debug)]
+pub struct StreamingConfig {
+    pub disconnect_slow_consumers: bool,
+    pub ws_max_write_buffer_bytes: usize,
+}
+
 pub struct ApiState<S> {
     pub state: Arc<S>,
     pub ws: Arc<ws::WsState>,
     pub metrics: Arc<metrics_middleware::ApiMetrics>,
+    pub streaming: StreamingConfig,
 }
 
 /// Manually implement `Clone` as the derive macro will try and slap `Clone` on
@@ -34,12 +41,18 @@ impl<S> Clone for ApiState<S> {
             state: self.state.clone(),
             ws: self.ws.clone(),
             metrics: self.metrics.clone(),
+            streaming: self.streaming.clone(),
         }
     }
 }
 
 impl<S> ApiState<S> {
-    pub fn new(state: Arc<S>, ws_whitelist: Vec<IpNet>, requester_ip_header_name: String) -> Self
+    pub fn new(
+        state: Arc<S>,
+        ws_whitelist: Vec<IpNet>,
+        requester_ip_header_name: String,
+        streaming: StreamingConfig,
+    ) -> Self
     where
         S: Metrics,
         S: Send + Sync + 'static,
@@ -52,6 +65,7 @@ impl<S> ApiState<S> {
                 state.clone(),
             )),
             state,
+            streaming,
         }
     }
 }
@@ -71,6 +85,10 @@ where
             state,
             opts.rpc.ws_whitelist,
             opts.rpc.requester_ip_header_name,
+            StreamingConfig {
+                disconnect_slow_consumers: opts.rpc.disconnect_slow_consumers,
+                ws_max_write_buffer_bytes: opts.rpc.ws_max_write_buffer_bytes,
+            },
         )
     };
 

apps/hermes/server/src/api/metrics_middleware.rs

@@ -9,16 +9,24 @@ use {
     },
     prometheus_client::{
         encoding::EncodeLabelSet,
-        metrics::{counter::Counter, family::Family, histogram::Histogram},
+        metrics::{counter::Counter, family::Family, gauge::Gauge, histogram::Histogram},
     },
     std::sync::Arc,
     tokio::time::Instant,
 };
 
+#[derive(Clone, Debug, PartialEq, Eq, Hash, EncodeLabelSet)]
+pub struct StreamProtocolLabel {
+    pub protocol: String,
+}
+
 pub struct ApiMetrics {
     pub requests: Family<Labels, Counter>,
     pub latencies: Family<LatencyLabels, Histogram>,
     pub sse_broadcast_latency: Histogram,
+    pub stream_active_connections: Family<StreamProtocolLabel, Gauge>,
+    pub stream_slow_consumer_disconnects: Family<StreamProtocolLabel, Counter>,
+    pub sse_broadcast_lagged: Counter,
 }
 
 impl ApiMetrics {
@@ -43,12 +51,18 @@ impl ApiMetrics {
                 ]
                 .into_iter(),
             ),
+            stream_active_connections: Family::default(),
+            stream_slow_consumer_disconnects: Family::default(),
+            sse_broadcast_lagged: Counter::default(),
         };
 
         {
             let requests = new.requests.clone();
             let latencies = new.latencies.clone();
             let sse_broadcast_latency = new.sse_broadcast_latency.clone();
+            let stream_active_connections = new.stream_active_connections.clone();
+            let stream_slow_consumer_disconnects = new.stream_slow_consumer_disconnects.clone();
+            let sse_broadcast_lagged = new.sse_broadcast_lagged.clone();
 
             tokio::spawn(async move {
                 Metrics::register(
@@ -76,6 +90,36 @@ impl ApiMetrics {
                     ),
                 )
                 .await;
+
+                Metrics::register(
+                    &*state,
+                    (
+                        "stream_active_connections",
+                        "Number of active WebSocket and SSE streaming connections",
+                        stream_active_connections,
+                    ),
+                )
+                .await;
+
+                Metrics::register(
+                    &*state,
+                    (
+                        "stream_slow_consumer_disconnects_total",
+                        "Total streaming connections closed due to slow consumption",
+                        stream_slow_consumer_disconnects,
+                    ),
+                )
+                .await;
+
+                Metrics::register(
+                    &*state,
+                    (
+                        "sse_broadcast_lagged_total",
+                        "Total SSE broadcast lag events (client missed updates)",
+                        sse_broadcast_lagged,
+                    ),
+                )
+                .await;
             });
         }
 
@@ -141,3 +185,9 @@ pub async fn track_metrics<B, S>(
 
     response
 }
+
+pub fn stream_protocol_label(protocol: &str) -> StreamProtocolLabel {
+    StreamProtocolLabel {
+        protocol: protocol.to_string(),
+    }
+}

apps/hermes/server/src/api/rest.rs

@@ -127,6 +127,7 @@ where
 mod tests {
     use {
         super::*,
+        crate::api::StreamingConfig,
         crate::state::{
             aggregate::{
                 AggregationEvent, PriceFeedsWithUpdateData, PublisherStakeCapsWithUpdateData,
@@ -215,7 +216,15 @@ mod tests {
         available_ids.insert(id2);
 
         let mock_state = MockAggregates { available_ids };
-        let api_state = ApiState::new(Arc::new(mock_state), vec![], String::new());
+        let api_state = ApiState::new(
+            Arc::new(mock_state),
+            vec![],
+            String::new(),
+            StreamingConfig {
+                disconnect_slow_consumers: true,
+                ws_max_write_buffer_bytes: 2 * 1024 * 1024,
+            },
+        );
 
         let input_ids = vec![id1, id2];
         let result = validate_price_ids(&api_state, &input_ids, false).await;
@@ -234,7 +243,15 @@ mod tests {
         available_ids.insert(id2);
 
         let mock_state = MockAggregates { available_ids };
-        let api_state = ApiState::new(Arc::new(mock_state), vec![], String::new());
+        let api_state = ApiState::new(
+            Arc::new(mock_state),
+            vec![],
+            String::new(),
+            StreamingConfig {
+                disconnect_slow_consumers: true,
+                ws_max_write_buffer_bytes: 2 * 1024 * 1024,
+            },
+        );
 
         let input_ids = vec![id1, id2, id3];
         let result = validate_price_ids(&api_state, &input_ids, true).await;
@@ -253,7 +270,15 @@ mod tests {
         available_ids.insert(id2);
 
         let mock_state = MockAggregates { available_ids };
-        let api_state = ApiState::new(Arc::new(mock_state), vec![], String::new());
+        let api_state = ApiState::new(
+            Arc::new(mock_state),
+            vec![],
+            String::new(),
+            StreamingConfig {
+                disconnect_slow_consumers: true,
+                ws_max_write_buffer_bytes: 2 * 1024 * 1024,
+            },
+        );
 
         let input_ids = vec![id1, id2, id3];
         let result = validate_price_ids(&api_state, &input_ids, false).await;