sdk: Basic connection metrics (#2299)

crates/sdk/Cargo.toml

@@ -10,8 +10,9 @@ description = "A Rust SDK for clients to interface with SpacetimeDB"
 [dependencies]
 spacetimedb-data-structures.workspace = true
 spacetimedb-sats.workspace = true
-spacetimedb-lib = { workspace = true, features = ["serde"]}
+spacetimedb-lib = { workspace = true, features = ["serde", "metrics_impls"]}
 spacetimedb-client-api-messages.workspace = true
+spacetimedb-metrics.workspace = true
 
 thiserror.workspace = true
 anymap.workspace = true
@@ -23,6 +24,8 @@ futures-channel.workspace = true
 home.workspace = true
 http.workspace = true
 log.workspace = true
+once_cell.workspace = true
+prometheus.workspace = true
 rand.workspace = true
 tokio.workspace = true
 tokio-tungstenite.workspace = true

crates/sdk/src/lib.rs

@@ -13,6 +13,7 @@
 mod callbacks;
 mod client_cache;
 mod db_connection;
+mod metrics;
 mod spacetime_module;
 mod subscription;
 mod websocket;
@@ -67,5 +68,6 @@ pub mod unstable {
     //!
     //! These may change incompatibly without a major version bump.
     pub use crate::db_connection::set_connection_id;
+    pub use crate::metrics::{ClientMetrics, CLIENT_METRICS};
     pub use spacetimedb_client_api_messages::websocket::CallReducerFlags;
 }

crates/sdk/src/metrics.rs

@@ -0,0 +1,20 @@
+use once_cell::sync::Lazy;
+use prometheus::{HistogramVec, IntCounterVec};
+use spacetimedb_lib::ConnectionId;
+use spacetimedb_metrics::metrics_group;
+
+metrics_group!(
+    pub struct ClientMetrics {
+        #[name = spacetime_client_received_total]
+        #[help = "The cumulative number of received websocket messages"]
+        #[labels(db: Box<str>, connection_id: ConnectionId)]
+        pub websocket_received: IntCounterVec,
+
+        #[name = spacetime_client_received_msg_size]
+        #[help = "The size of received websocket messages"]
+        #[labels(db: Box<str>, connection_id: ConnectionId)]
+        pub websocket_received_msg_size: HistogramVec,
+    }
+);
+
+pub static CLIENT_METRICS: Lazy<ClientMetrics> = Lazy::new(ClientMetrics::new);

crates/sdk/src/websocket.rs

@@ -24,6 +24,8 @@ use tokio_tungstenite::{
     MaybeTlsStream, WebSocketStream,
 };
 
+use crate::metrics::CLIENT_METRICS;
+
 #[derive(Error, Debug, Clone)]
 pub enum UriError {
     #[error("Unknown URI scheme {scheme}, expected http, https, ws or wss")]
@@ -80,6 +82,8 @@ pub enum WsError {
 }
 
 pub(crate) struct WsConnection {
+    db_name: Box<str>,
+    connection_id: ConnectionId,
     sock: WebSocketStream<MaybeTlsStream<TcpStream>>,
 }
 
@@ -223,7 +227,11 @@ impl WsConnection {
             uri,
             source: Arc::new(source),
         })?;
-        Ok(WsConnection { sock })
+        Ok(WsConnection {
+            db_name: db_name.into(),
+            connection_id,
+            sock,
+        })
     }
 
     pub(crate) fn parse_response(bytes: &[u8]) -> Result<ServerMessage<BsatnFormat>, WsError> {
@@ -268,6 +276,17 @@ impl WsConnection {
         incoming_messages: mpsc::UnboundedSender<ServerMessage<BsatnFormat>>,
         outgoing_messages: mpsc::UnboundedReceiver<ClientMessage<Bytes>>,
     ) {
+        let record_metrics = |msg_size: usize| {
+            CLIENT_METRICS
+                .websocket_received
+                .with_label_values(&self.db_name, &self.connection_id)
+                .inc();
+            CLIENT_METRICS
+                .websocket_received_msg_size
+                .with_label_values(&self.db_name, &self.connection_id)
+                .observe(msg_size as f64);
+        };
+
         let mut outgoing_messages = Some(outgoing_messages);
         loop {
             tokio::select! {
@@ -280,6 +299,7 @@ impl WsConnection {
                     ),
 
                     Ok(Some(WebSocketMessage::Binary(bytes))) => {
+                        record_metrics(bytes.len());
                         match Self::parse_response(&bytes) {
                             Err(e) => Self::maybe_log_error::<(), _>(
                                 "Error decoding WebSocketMessage::Binary payload",
@@ -292,9 +312,14 @@ impl WsConnection {
                         }
                     }
 
-                    Ok(Some(WebSocketMessage::Ping(_))) => {}
+                    Ok(Some(WebSocketMessage::Ping(payload))) => {
+                        record_metrics(payload.len());
+                    }
 
-                    Ok(Some(other)) => log::warn!("Unexpected WebSocket message {:?}", other),
+                    Ok(Some(other)) => {
+                        log::warn!("Unexpected WebSocket message {:?}", other);
+                        record_metrics(other.len());
+                    },
                 },
 
                 // this is stupid. we want to handle the channel close *once*, and then disable this branch