Merge branch 'master' into kim/commitlog/index-log-noise

Author: kim

crates/client-api-messages/src/websocket.rs

@@ -622,13 +622,19 @@ pub struct TableUpdate<F: WebsocketFormat> {
     pub updates: SmallVec<[F::QueryUpdate; 1]>,
 }
 
+/// Computed update for a single query, annotated with the number of matching rows.
+pub struct SingleQueryUpdate<F: WebsocketFormat> {
+    pub update: F::QueryUpdate,
+    pub num_rows: u64,
+}
+
 impl<F: WebsocketFormat> TableUpdate<F> {
-    pub fn new(table_id: TableId, table_name: Box<str>, (update, num_rows): (F::QueryUpdate, u64)) -> Self {
+    pub fn new(table_id: TableId, table_name: Box<str>, update: SingleQueryUpdate<F>) -> Self {
         Self {
             table_id,
             table_name,
-            num_rows,
-            updates: [update].into(),
+            num_rows: update.num_rows,
+            updates: [update.update].into(),
         }
     }
 
@@ -641,9 +647,9 @@ impl<F: WebsocketFormat> TableUpdate<F> {
         }
     }
 
-    pub fn push(&mut self, (update, num_rows): (F::QueryUpdate, u64)) {
-        self.updates.push(update);
-        self.num_rows += num_rows;
+    pub fn push(&mut self, update: SingleQueryUpdate<F>) {
+        self.updates.push(update.update);
+        self.num_rows += update.num_rows;
     }
 
     pub fn num_rows(&self) -> usize {

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

@@ -228,6 +228,24 @@ async fn ws_client_actor_inner(
 
     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 outgoing_queue_length_metric = WORKER_METRICS.total_outgoing_queue_length.with_label_values(&addr);
+
+    let clean_up_metrics = |message_queue: &VecDeque<(DataMessage, Instant)>,
+                            sendrx: &mpsc::Receiver<SerializableMessage>| {
+        incoming_queue_length_metric.sub(message_queue.len() as _);
+        outgoing_queue_length_metric.sub(sendrx.len() as _);
+    };
+
     loop {
         rx_buf.clear();
         enum Item {
@@ -236,6 +254,7 @@ 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));
@@ -263,43 +282,47 @@ async fn ws_client_actor_inner(
                     continue;
                 }
                 // the client sent us a close frame
-                None => break,
+                None => {
+                    clean_up_metrics(&message_queue, &sendrx);
+                    break
+                },
             },
 
             // If we have an outgoing message to send, send it off.
             // No incoming `message` to handle, so `continue`.
             Some(n) = sendrx.recv_many(&mut rx_buf, 32).map(|n| (n != 0).then_some(n)) => {
+                outgoing_queue_length_metric.sub(n as _);
                 if closed {
                     // TODO: this isn't great. when we receive a close request from the peer,
                     //       tungstenite doesn't let us send any new messages on the socket,
                     //       even though the websocket RFC allows it. should we fork tungstenite?
                     log::info!("dropping messages due to ws already being closed: {:?}", &rx_buf[..n]);
                 } else {
-                    let send_all = async {
-                        for msg in rx_buf.drain(..n) {
-                            let workload = msg.workload();
-                            let num_rows = msg.num_rows();
-
-                            let msg = datamsg_to_wsmsg(serialize(msg, client.config));
-
-                            // These metrics should be updated together,
-                            // or not at all.
-                            if let (Some(workload), Some(num_rows)) = (workload, num_rows) {
-                                WORKER_METRICS
-                                    .websocket_sent_num_rows
-                                    .with_label_values(&addr, &workload)
-                                    .observe(num_rows as f64);
-                                WORKER_METRICS
-                                    .websocket_sent_msg_size
-                                    .with_label_values(&addr, &workload)
-                                    .observe(msg.len() as f64);
-                            }
-                            // feed() buffers the message, but does not necessarily send it
-                            ws.feed(msg).await?;
+                let send_all = async {
+                    for msg in rx_buf.drain(..n) {
+                        let workload = msg.workload();
+                        let num_rows = msg.num_rows();
+
+                        let msg = datamsg_to_wsmsg(serialize(msg, client.config));
+
+                        // These metrics should be updated together,
+                        // or not at all.
+                        if let (Some(workload), Some(num_rows)) = (workload, num_rows) {
+                            WORKER_METRICS
+                                .websocket_sent_num_rows
+                                .with_label_values(&addr, &workload)
+                                .observe(num_rows as f64);
+                            WORKER_METRICS
+                                .websocket_sent_msg_size
+                                .with_label_values(&addr, &workload)
+                                .observe(msg.len() as f64);
                         }
-                        // now we flush all the messages to the socket
-                        ws.flush().await
-                    };
+                        // feed() buffers the message, but does not necessarily send it
+                        ws.feed(msg).await?;
+                    }
+                    // now we flush all the messages to the socket
+                     ws.flush().await
+                 };
                     // Flush the websocket while continuing to poll the `handle_queue`,
                     // to avoid deadlocks or delays due to enqueued futures holding resources.
                     let send_all = also_poll(send_all, make_progress(&mut current_message));
@@ -348,6 +371,7 @@ 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, &sendrx);
                     break;
                 }
             }
@@ -362,6 +386,7 @@ 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/commitlog/src/index/indexfile.rs

@@ -5,7 +5,7 @@ use std::{
     mem,
 };
 
-use log::debug;
+use log::{debug, trace};
 use memmap2::MmapMut;
 use spacetimedb_paths::server::OffsetIndexFile;
 
@@ -190,19 +190,32 @@ impl<Key: Into<u64> + From<u64>> IndexFileMut<Key> {
         self.inner.flush_async()
     }
 
-    /// Truncates the index file starting from the entry with a key greater than or equal to the given key.
+    /// Truncates the index file starting from the entry with a key greater than
+    /// or equal to the given key.
+    ///
+    /// If successful, `key` will no longer be in the index.
     pub(crate) fn truncate(&mut self, key: Key) -> Result<(), IndexError> {
         let key = key.into();
-        let (found_key, index) = self.find_index(Key::from(key))?;
+        let (found_key, index) = self
+            .find_index(Key::from(key))
+            .map(|(found, index)| (found.into(), index))?;
 
-        // If returned key is smalled than asked key, truncate from next entry
-        self.num_entries = if found_key.into() == key {
+        // If returned key is smaller than asked key, truncate from next entry
+        self.num_entries = if found_key == key {
             index as usize
         } else {
             index as usize + 1
         };
 
         let start = self.num_entries * ENTRY_SIZE;
+        trace!(
+            "truncate key={} found={} index={} num-entries={} start={}",
+            key,
+            found_key,
+            index,
+            self.num_entries,
+            start
+        );
 
         if start < self.inner.len() {
             self.inner[start..].fill(0);

crates/commitlog/src/segment.rs

@@ -697,6 +697,7 @@ mod tests {
     use super::*;
     use crate::{payload::ArrayDecoder, repo, Options};
     use itertools::Itertools;
+    use pretty_assertions::assert_matches;
     use proptest::prelude::*;
     use spacetimedb_paths::server::CommitLogDir;
     use tempfile::tempdir;
@@ -926,16 +927,28 @@ mod tests {
 
         // Truncating to any offset in the written range or larger
         // retains that offset - 1, or the max offset written.
-        let truncate_to: TxOffset = rand::random_range(1..=32);
-        let retained_key = truncate_to.saturating_sub(1).min(10);
-        let retained_val = retained_key * 128;
-        let retained = (retained_key, retained_val);
-
-        writer.ftruncate(truncate_to, rand::random()).unwrap();
-        assert_eq!(writer.head.key_lookup(truncate_to).unwrap(), retained);
-        // Make sure this also holds after reopen.
-        drop(writer);
-        let index = TxOffsetIndex::open_index_file(&index_path).unwrap();
-        assert_eq!(index.key_lookup(truncate_to).unwrap(), retained);
+        for truncate_to in (2..=10u64).rev() {
+            let retained_key = truncate_to.saturating_sub(1).min(10);
+            let retained_val = retained_key * 128;
+            let retained = (retained_key, retained_val);
+
+            writer.ftruncate(truncate_to, rand::random()).unwrap();
+            assert_matches!(
+                writer.head.key_lookup(truncate_to),
+                Ok(x) if x == retained,
+                "truncate to {truncate_to} should retain {retained:?}"
+            );
+            // Make sure this also holds after reopen.
+            let index = TxOffsetIndex::open_index_file(&index_path).unwrap();
+            assert_matches!(
+                index.key_lookup(truncate_to),
+                Ok(x) if x == retained,
+                "truncate to {truncate_to} should retain {retained:?} after reopen"
+            );
+        }
+
+        // Truncating to 1 leaves no entries in the index
+        writer.ftruncate(1, rand::random()).unwrap();
+        assert_matches!(writer.head.key_lookup(1), Err(IndexError::KeyNotFound));
     }
 }

crates/core/src/client/client_connection.rs

@@ -16,7 +16,7 @@ use bytes::Bytes;
 use bytestring::ByteString;
 use derive_more::From;
 use futures::prelude::*;
-use prometheus::{Histogram, IntCounter};
+use prometheus::{Histogram, IntCounter, IntGauge};
 use spacetimedb_client_api_messages::websocket::{
     BsatnFormat, CallReducerFlags, Compression, FormatSwitch, JsonFormat, SubscribeMulti, SubscribeSingle, Unsubscribe,
     UnsubscribeMulti, WebsocketFormat,
@@ -91,6 +91,15 @@ pub struct ClientConnectionSender {
 pub struct ClientConnectionMetrics {
     pub websocket_request_msg_size: Histogram,
     pub websocket_requests: IntCounter,
+
+    /// The `total_outgoing_queue_length` metric labeled with this database's `Identity`,
+    /// which we'll increment whenever sending a message.
+    ///
+    /// This metric will be decremented, and cleaned up,
+    /// by `ws_client_actor_inner` in client-api/src/routes/subscribe.rs.
+    /// Care must be taken not to increment it after the client has disconnected
+    /// and performed its clean-up.
+    pub sendtx_queue_size: IntGauge,
 }
 
 impl ClientConnectionMetrics {
@@ -102,10 +111,14 @@ impl ClientConnectionMetrics {
         let websocket_requests = WORKER_METRICS
             .websocket_requests
             .with_label_values(&database_identity, message_kind);
+        let sendtx_queue_size = WORKER_METRICS
+            .total_outgoing_queue_length
+            .with_label_values(&database_identity);
 
         Self {
             websocket_request_msg_size,
             websocket_requests,
+            sendtx_queue_size,
         }
     }
 }
@@ -126,6 +139,7 @@ impl ClientConnectionSender {
             Ok(h) => h.spawn(async {}).abort_handle(),
             Err(_) => tokio::runtime::Runtime::new().unwrap().spawn(async {}).abort_handle(),
         };
+
         let cancelled = AtomicBool::new(false);
         let sender = Self {
             id,
@@ -150,17 +164,27 @@ impl ClientConnectionSender {
         if self.cancelled.load(Relaxed) {
             return Err(ClientSendError::Cancelled);
         }
-        self.sendtx.try_send(message).map_err(|e| match e {
-            mpsc::error::TrySendError::Full(_) => {
+
+        match self.sendtx.try_send(message) {
+            Err(mpsc::error::TrySendError::Full(_)) => {
                 // we've hit CLIENT_CHANNEL_CAPACITY messages backed up in
                 // the channel, so forcibly kick the client
                 tracing::warn!(identity = %self.id.identity, connection_id = %self.id.connection_id, "client channel capacity exceeded");
                 self.abort_handle.abort();
                 self.cancelled.store(true, Relaxed);
-                ClientSendError::Cancelled
+                return Err(ClientSendError::Cancelled);
             }
-            mpsc::error::TrySendError::Closed(_) => ClientSendError::Disconnected,
-        })?;
+            Err(mpsc::error::TrySendError::Closed(_)) => return Err(ClientSendError::Disconnected),
+            Ok(()) => {
+                // If we successfully pushed a message into the queue, increment the queue size metric.
+                // Don't do this before pushing because, if the client has disconnected,
+                // it will already have performed its clean-up,
+                // and so would never perform the corresponding `dec` to this `inc`.
+                if let Some(metrics) = &self.metrics {
+                    metrics.sendtx_queue_size.inc();
+                }
+            }
+        }
 
         Ok(())
     }