fix incoming message queue length metric (#3172)

# Description of Changes

<!-- Please describe your change, mention any related tickets, and so on
here. -->

We stopped incrementing the incoming queue length metric. This patch
increments it again and adds a regression test.

# API and ABI breaking changes

<!-- If this is an API or ABI breaking change, please apply the
corresponding GitHub label. -->

None

# Expected complexity level and risk

<!--
How complicated do you think these changes are? Grade on a scale from 1
to 5,
where 1 is a trivial change, and 5 is a deep-reaching and complex
change.

This complexity rating applies not only to the complexity apparent in
the diff,
but also to its interactions with existing and future code.

If you answered more than a 2, explain what is complex about the PR,
and what other components it interacts with in potentially concerning
ways. -->

1

# Testing

<!-- Describe any testing you've done, and any testing you'd like your
reviewers to do,
so that you're confident that all the changes work as expected! -->

- [x] Regression test

Author: joshua-spacetime

crates/client-api/src/routes/subscribe.rs

@@ -26,7 +26,7 @@ use spacetimedb::client::messages::{
 };
 use spacetimedb::client::{
     ClientActorId, ClientConfig, ClientConnection, DataMessage, MessageExecutionError, MessageHandleError,
-    MeteredReceiver, Protocol,
+    MeteredReceiver, MeteredSender, Protocol,
 };
 use spacetimedb::host::module_host::ClientConnectedError;
 use spacetimedb::host::NoSuchModule;
@@ -676,7 +676,10 @@ async fn ws_recv_task<MessageHandler>(
 ) where
     MessageHandler: Future<Output = Result<(), MessageHandleError>>,
 {
-    let recv_queue = ws_recv_queue(state.clone(), unordered_tx.clone(), ws);
+    let recv_queue_gauge = WORKER_METRICS
+        .total_incoming_queue_length
+        .with_label_values(&state.database);
+    let recv_queue = ws_recv_queue(state.clone(), unordered_tx.clone(), recv_queue_gauge, ws);
     let recv_loop = pin!(ws_recv_loop(state.clone(), idle_tx, recv_queue));
     let recv_handler = ws_client_message_handler(state.clone(), client_closed_metric, recv_loop);
     pin_mut!(recv_handler);
@@ -816,6 +819,7 @@ fn ws_recv_loop(
 fn ws_recv_queue(
     state: Arc<ActorState>,
     unordered_tx: mpsc::UnboundedSender<UnorderedWsMessage>,
+    recv_queue_gauge: IntGauge,
     mut ws: impl Stream<Item = Result<WsMessage, WsError>> + Unpin + Send + 'static,
 ) -> impl Stream<Item = Result<WsMessage, WsError>> {
     const CLOSE: UnorderedWsMessage = UnorderedWsMessage::Close(CloseFrame {
@@ -829,14 +833,10 @@ fn ws_recv_queue(
     let max_incoming_queue_length = state.config.incoming_queue_length.get();
 
     let (tx, rx) = mpsc::channel(max_incoming_queue_length);
-    let rx = MeteredReceiverStream {
-        inner: MeteredReceiver::with_gauge(
-            rx,
-            WORKER_METRICS
-                .total_incoming_queue_length
-                .with_label_values(&state.database),
-        ),
-    };
+
+    let mut tx = MeteredSender::with_gauge(tx, recv_queue_gauge.clone());
+    let rx = MeteredReceiver::with_gauge(rx, recv_queue_gauge);
+    let rx = MeteredReceiverStream { inner: rx };
 
     tokio::spawn(async move {
         while let Some(item) = ws.next().await {
@@ -1692,8 +1692,14 @@ mod tests {
         let (unordered_tx, mut unordered_rx) = mpsc::unbounded_channel();
         let input = stream::iter((0..20).map(|i| Ok(WsMessage::text(format!("message {i}")))));
 
-        let received = ws_recv_queue(state, unordered_tx, input).collect::<Vec<_>>().await;
+        let metric = IntGauge::new("bleep", "unhelpful").unwrap();
+        let received = ws_recv_queue(state, unordered_tx, metric.clone(), input)
+            .collect::<Vec<_>>()
+            .await;
+
         assert_matches!(unordered_rx.recv().await, Some(UnorderedWsMessage::Close(_)));
+        // Queue length metric should be zero
+        assert_eq!(metric.get(), 0);
         // Should have received all of the input.
         assert_eq!(received.len(), 20);
     }
@@ -1708,10 +1714,14 @@ mod tests {
         let (unordered_tx, _) = mpsc::unbounded_channel();
         let input = stream::iter((0..20).map(|i| Ok(WsMessage::text(format!("message {i}")))));
 
-        let received = ws_recv_queue(state.clone(), unordered_tx, input)
+        let metric = IntGauge::new("bleep", "unhelpful").unwrap();
+        let received = ws_recv_queue(state.clone(), unordered_tx, metric.clone(), input)
             .collect::<Vec<_>>()
             .await;
+
         assert!(state.closed());
+        // Queue length metric should be zero
+        assert_eq!(metric.get(), 0);
         // Should have received up to capacity.
         assert_eq!(received.len(), 10);
     }

crates/core/src/client.rs

@@ -8,7 +8,7 @@ pub mod messages;
 
 pub use client_connection::{
     ClientConfig, ClientConnection, ClientConnectionSender, ClientSendError, DataMessage, MeteredDeque,
-    MeteredReceiver, Protocol,
+    MeteredReceiver, MeteredSender, Protocol,
 };
 pub use client_connection_index::ClientActorIndex;
 pub use message_handlers::{MessageExecutionError, MessageHandleError};

crates/core/src/client/client_connection.rs

@@ -28,6 +28,7 @@ use spacetimedb_client_api_messages::websocket::{
 use spacetimedb_lib::identity::RequestId;
 use spacetimedb_lib::metrics::ExecutionMetrics;
 use spacetimedb_lib::Identity;
+use tokio::sync::mpsc::error::{SendError, TrySendError};
 use tokio::sync::{mpsc, oneshot, watch};
 use tokio::task::AbortHandle;
 
@@ -386,6 +387,41 @@ impl<T> Drop for MeteredReceiver<T> {
     }
 }
 
+/// Wraps the transmitting end of a channel with a gauge for tracking the size of the channel.
+pub struct MeteredSender<T> {
+    inner: mpsc::Sender<T>,
+    gauge: Option<IntGauge>,
+}
+
+impl<T> MeteredSender<T> {
+    pub fn new(inner: mpsc::Sender<T>) -> Self {
+        Self { inner, gauge: None }
+    }
+
+    pub fn with_gauge(inner: mpsc::Sender<T>, gauge: IntGauge) -> Self {
+        Self {
+            inner,
+            gauge: Some(gauge),
+        }
+    }
+
+    pub async fn send(&mut self, value: T) -> Result<(), SendError<T>> {
+        self.inner.send(value).await?;
+        if let Some(gauge) = &self.gauge {
+            gauge.inc();
+        }
+        Ok(())
+    }
+
+    pub fn try_send(&mut self, value: T) -> Result<(), TrySendError<T>> {
+        self.inner.try_send(value)?;
+        if let Some(gauge) = &self.gauge {
+            gauge.inc();
+        }
+        Ok(())
+    }
+}
+
 // if a client racks up this many messages in the queue without ACK'ing
 // anything, we boot 'em.
 const CLIENT_CHANNEL_CAPACITY: usize = 16 * KB;