Use same approach for client incoming message queue

Author: joshua-spacetime

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

@@ -1,4 +1,3 @@
-use std::collections::VecDeque;
 use std::mem;
 use std::pin::{pin, Pin};
 use std::time::Duration;
@@ -16,7 +15,8 @@ use scopeguard::ScopeGuard;
 use serde::Deserialize;
 use spacetimedb::client::messages::{serialize, IdentityTokenMessage, SerializableMessage, SerializeBuffer};
 use spacetimedb::client::{
-    ClientActorId, ClientConfig, ClientConnection, DataMessage, MessageHandleError, MeteredReceiver, Protocol,
+    ClientActorId, ClientConfig, ClientConnection, DataMessage, MessageHandleError, MeteredDeque, MeteredReceiver,
+    Protocol,
 };
 use spacetimedb::execution_context::WorkloadType;
 use spacetimedb::host::module_host::ClientConnectedError;
@@ -210,6 +210,8 @@ async fn ws_client_actor_inner(
     let mut liveness_check_interval = tokio::time::interval(LIVELINESS_TIMEOUT);
     let mut got_pong = true;
 
+    let addr = client.module.info().database_identity;
+
     // Build a queue of incoming messages to handle, to be processed one at a time,
     // in the order they're received.
     //
@@ -223,29 +225,14 @@ async fn ws_client_actor_inner(
     //       `select!` for examples of how to do this.
     //
     // TODO: do we want this to have a fixed capacity? or should it be unbounded
-    let mut message_queue = VecDeque::<(DataMessage, Instant)>::new();
+    let mut message_queue = MeteredDeque::<(DataMessage, Instant)>::new(
+        WORKER_METRICS.total_incoming_queue_length.with_label_values(&addr),
+    );
     let mut current_message = pin!(MaybeDone::Gone);
 
     let mut closed = false;
     let mut rx_buf = Vec::new();
 
-    let addr = client.module.info().database_identity;
-
-    // Grab handles on the total incoming and outgoing queue length metrics,
-    // which we'll increment and decrement as we push into and pull out of those queues.
-    // Note that `total_outgoing_queue_length` is incremented separately,
-    // by `ClientConnectionSender::send` in core/src/client/client_connection.rs;
-    // we're only responsible for decrementing that one.
-    // Also note that much care must be taken to clean up these metrics when the connection closes!
-    // Any path which exits this function must decrement each of these metrics
-    // by the number of messages still waiting in this client's queue,
-    // or else they will grow without bound as clients disconnect, and be useless.
-    let incoming_queue_length_metric = WORKER_METRICS.total_incoming_queue_length.with_label_values(&addr);
-
-    let clean_up_metrics = |message_queue: &VecDeque<(DataMessage, Instant)>| {
-        incoming_queue_length_metric.sub(message_queue.len() as _);
-    };
-
     let mut msg_buffer = SerializeBuffer::new(client.config);
     loop {
         rx_buf.clear();
@@ -255,7 +242,6 @@ async fn ws_client_actor_inner(
         }
         if let MaybeDone::Gone = *current_message {
             if let Some((message, timer)) = message_queue.pop_front() {
-                incoming_queue_length_metric.dec();
                 let client = client.clone();
                 let fut = async move { client.handle_message(message, timer).await };
                 current_message.set(MaybeDone::Future(fut));
@@ -284,7 +270,6 @@ async fn ws_client_actor_inner(
                 }
                 // the client sent us a close frame
                 None => {
-                    clean_up_metrics(&message_queue);
                     break
                 },
             },
@@ -376,7 +361,6 @@ async fn ws_client_actor_inner(
                 } else {
                     // the client never responded to our ping; drop them without trying to send them a Close
                     log::warn!("client {} timed out", client.id);
-                    clean_up_metrics(&message_queue);
                     break;
                 }
             }
@@ -391,7 +375,6 @@ async fn ws_client_actor_inner(
         match message {
             Item::Message(ClientMessage::Message(message)) => {
                 let timer = Instant::now();
-                incoming_queue_length_metric.inc();
                 message_queue.push_back((message, timer))
             }
             Item::HandleResult(res) => {

crates/core/src/client.rs

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

crates/core/src/client/client_connection.rs

@@ -1,3 +1,4 @@
+use std::collections::VecDeque;
 use std::ops::Deref;
 use std::sync::atomic::{AtomicBool, Ordering::Relaxed};
 use std::sync::Arc;
@@ -258,6 +259,59 @@ impl DataMessage {
     }
 }
 
+/// Wraps a [VecDeque] with a gauge for tracking its size.
+/// We subtract its size from the gauge on drop to avoid leaking the metric.
+pub struct MeteredDeque<T> {
+    inner: VecDeque<T>,
+    gauge: IntGauge,
+}
+
+impl<T> MeteredDeque<T> {
+    pub fn new(gauge: IntGauge) -> Self {
+        Self {
+            inner: VecDeque::new(),
+            gauge,
+        }
+    }
+
+    pub fn pop_front(&mut self) -> Option<T> {
+        self.inner.pop_front().inspect(|_| {
+            self.gauge.dec();
+        })
+    }
+
+    pub fn pop_back(&mut self) -> Option<T> {
+        self.inner.pop_back().inspect(|_| {
+            self.gauge.dec();
+        })
+    }
+
+    pub fn push_front(&mut self, value: T) {
+        self.gauge.inc();
+        self.inner.push_front(value);
+    }
+
+    pub fn push_back(&mut self, value: T) {
+        self.gauge.inc();
+        self.inner.push_back(value);
+    }
+
+    pub fn len(&self) -> usize {
+        self.inner.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.inner.is_empty()
+    }
+}
+
+impl<T> Drop for MeteredDeque<T> {
+    fn drop(&mut self) {
+        // Record the number of elements still in the deque on drop
+        self.gauge.sub(self.inner.len() as _);
+    }
+}
+
 /// Wraps the receiving end of a channel with a gauge for tracking the size of the channel.
 /// We subtract the size of the channel from the gauge on drop to avoid leaking the metric.
 pub struct MeteredReceiver<T> {