Add debug gauge to web sockets (#28685)

This will help us triage which worker is getting stuck if it ever does.

GitOrigin-RevId: 4693b9c08ccf49dd591d807e2913042e22150645

Author: Preslav Le

crates/local_backend/src/subs/metrics.rs

@@ -119,3 +119,14 @@ register_convex_counter!(pub WEBSOCKET_CONNECTION_RESET_TOTAL, "Number of websoc
 pub fn log_websocket_connection_reset() {
     log_counter(&WEBSOCKET_CONNECTION_RESET_TOTAL, 1)
 }
+
+register_convex_gauge!(
+    DEBUG_SYNC_PROTOCOL_WEBSOCKETS_TOTAL,
+    "Number of WebSocket connected to a backend at a given state",
+    &["tag"]
+);
+pub fn log_debug_sync_protocol_websockets_total(tag: &'static str, delta: i8) {
+    DEBUG_SYNC_PROTOCOL_WEBSOCKETS_TOTAL
+        .with_label_values(&[tag])
+        .add(delta as f64)
+}

crates/local_backend/src/subs/mod.rs

@@ -60,6 +60,7 @@ use sync_types::IdentityVersion;
 mod metrics;
 
 use metrics::{
+    log_debug_sync_protocol_websockets_total,
     log_sync_protocol_websockets_total,
     log_websocket_client_timeout,
     log_websocket_closed,
@@ -96,6 +97,25 @@ impl Drop for SyncSocketDropToken {
     }
 }
 
+// TODO(presley): Remove. Used for debugging.
+struct DebugSyncSocketDropToken {
+    tag: &'static str,
+}
+
+/// Tracker that exists for the lifetime of a run_sync_socket.
+impl DebugSyncSocketDropToken {
+    fn new(tag: &'static str) -> Self {
+        log_debug_sync_protocol_websockets_total(tag, 1);
+        DebugSyncSocketDropToken { tag }
+    }
+}
+
+impl Drop for DebugSyncSocketDropToken {
+    fn drop(&mut self) {
+        log_debug_sync_protocol_websockets_total(self.tag, -1);
+    }
+}
+
 // The WebSocket layer for the sync protocol has three asynchronous processes:
 //
 // 1) A `receive_messages` loop that consumes messages from the WebSocket,
@@ -124,6 +144,7 @@ async fn run_sync_socket(
 
     let (client_tx, client_rx) = mpsc::unbounded();
     let receive_messages = async {
+        let _receive_message_drop_token = DebugSyncSocketDropToken::new("receive_message");
         while let Some(message_r) = rx.next().await {
             let message = match message_r {
                 Ok(message) => message,
@@ -175,6 +196,7 @@ async fn run_sync_socket(
 
     let (server_tx, mut server_rx) = measurable_unbounded_channel();
     let send_messages = async {
+        let _send_message_drop_token = DebugSyncSocketDropToken::new("send_message");
         let mut ping_ticker = tokio::time::interval(HEARTBEAT_INTERVAL);
         'top: loop {
             select_biased! {
@@ -209,6 +231,7 @@ async fn run_sync_socket(
     };
     let mut identity_version: Option<IdentityVersion> = None;
     let sync_worker_go = async {
+        let _sync_worker_drop_token = DebugSyncSocketDropToken::new("sync_worker");
         let mut sync_worker = SyncWorker::new(
             st.api.clone(),
             st.runtime.clone(),