Pipeline the websocket send path (#5051)

# Description of Changes

In keeping with the "pipeline everything" ethos, I've replaced `recv`
with `recv_many` on `ClientConnectionReceiver` so that a client
connection's websocket send path works on a batch of messages at a time,
if possible.

# API and ABI breaking changes

None

# Expected complexity level and risk

1.5

# Testing

Refactor

Author: joshua-spacetime

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

@@ -1064,13 +1064,13 @@ enum UnorderedWsMessage {
 /// Abstraction over [`ClientConnectionReceiver`], so tests can use a plain
 /// [`mpsc::Receiver`].
 trait Receiver<T> {
-    fn recv(&mut self) -> impl Future<Output = Option<T>> + Send;
+    fn recv_many(&mut self, buf: &mut Vec<T>, max: usize) -> impl Future<Output = usize> + Send;
     fn close(&mut self);
 }
 
 impl Receiver<OutboundMessage> for ClientConnectionReceiver {
-    async fn recv(&mut self) -> Option<OutboundMessage> {
-        ClientConnectionReceiver::recv(self).await
+    async fn recv_many(&mut self, buf: &mut Vec<OutboundMessage>, max: usize) -> usize {
+        ClientConnectionReceiver::recv_many(self, buf, max).await
     }
 
     fn close(&mut self) {
@@ -1079,8 +1079,8 @@ impl Receiver<OutboundMessage> for ClientConnectionReceiver {
 }
 
 impl<T: Send> Receiver<T> for mpsc::Receiver<T> {
-    async fn recv(&mut self) -> Option<T> {
-        mpsc::Receiver::recv(self).await
+    async fn recv_many(&mut self, buf: &mut Vec<T>, max: usize) -> usize {
+        mpsc::Receiver::recv_many(self, buf, max).await
     }
 
     fn close(&mut self) {
@@ -1148,6 +1148,8 @@ async fn ws_send_loop_inner<T, U, Encoder>(
     // The default frame size is 4KiB, hence we write in batches of 32KiB.
     const FRAME_BATCH_SIZE: usize = 8;
     let mut frames_batch = Vec::with_capacity(FRAME_BATCH_SIZE);
+    const MESSAGE_BATCH_SIZE: usize = ClientConnectionReceiver::DEFAULT_RECV_MANY_LIMIT;
+    let mut message_batch = Vec::new();
     let (frames_tx, mut frames_rx) = mpsc::unbounded_channel();
 
     let (encode_tx, encode_rx) = mpsc::unbounded_channel();
@@ -1262,12 +1264,15 @@ async fn ws_send_loop_inner<T, U, Encoder>(
             // Take on more work.
             //
             // Branch is disabled if we already sent a close frame.
-            Some(message) = messages.recv(), if !closed => {
-                encode_tx
-                    .send(message.into())
-                    // `ws_encode_task` shouldn't terminate until
-                    // `encode_tx` is dropped, except by panicking.
-                    .expect("encode task panicked");
+            n = messages.recv_many(&mut message_batch, MESSAGE_BATCH_SIZE), if !closed => {
+                log::trace!("encoding batch of {n} messages");
+                for message in message_batch.drain(..n) {
+                    encode_tx
+                        .send(message.into())
+                        // `ws_encode_task` shouldn't terminate until
+                        // `encode_tx` is dropped, except by panicking.
+                        .expect("encode task panicked");
+                }
             },
 
         }

crates/core/src/client/client_connection.rs

@@ -19,6 +19,7 @@ use bytes::Bytes;
 use bytestring::ByteString;
 use derive_more::From;
 use futures::prelude::*;
+use log::warn;
 use prometheus::{Histogram, IntCounter, IntGauge};
 use spacetimedb_auth::identity::{ConnectionAuthCtx, SpacetimeIdentityClaims};
 use spacetimedb_client_api_messages::websocket::{common as ws_common, v1 as ws_v1, v2 as ws_v2};
@@ -29,7 +30,7 @@ use spacetimedb_lib::Identity;
 use tokio::sync::mpsc::error::{SendError, TrySendError};
 use tokio::sync::{mpsc, oneshot, watch};
 use tokio::task::AbortHandle;
-use tracing::{trace, warn};
+use tracing::trace;
 
 #[derive(PartialEq, Eq, Clone, Copy, Hash, Debug)]
 pub enum Protocol {
@@ -154,11 +155,13 @@ impl DurableOffsetSupply for Arc<RelationalDB> {
 pub struct ClientConnectionReceiver {
     confirmed_reads: bool,
     channel: MeteredReceiver<ClientUpdate>,
-    current: Option<ClientUpdate>,
+    pending: Vec<ClientUpdate>,
     offset_supply: Box<dyn DurableOffsetSupply>,
 }
 
 impl ClientConnectionReceiver {
+    pub const DEFAULT_RECV_MANY_LIMIT: usize = 4096;
+
     fn new(
         confirmed_reads: bool,
         channel: MeteredReceiver<ClientUpdate>,
@@ -167,82 +170,115 @@ impl ClientConnectionReceiver {
         Self {
             confirmed_reads,
             channel,
-            current: None,
+            pending: Vec::new(),
             offset_supply: Box::new(offset_supply),
         }
     }
 
-    /// Receive the next message from this channel.
-    ///
-    /// If this method returns `None`, the channel is closed and no more messages
-    /// are in the internal buffers. No more messages can ever be received from
-    /// the channel.
+    #[cfg(test)]
+    pub(crate) async fn recv(&mut self) -> Option<OutboundMessage> {
+        let mut buf = Vec::with_capacity(1);
+        (self.recv_many(&mut buf, 1).await != 0).then(|| buf.remove(0))
+    }
+
+    /// Receive multiple messages from this channel.
     ///
     /// Messages are returned immediately if:
     ///
-    ///   - The (internal) [`ClientUpdate`] does not have a `tx_offset`
+    ///   - The [`ClientUpdate`] does not have a `tx_offset`
     ///     (such as for error messages).
-    ///   - The client hasn't requested confirmed reads (i.e.
-    ///     [`ClientConfig::confirmed_reads`] is `false`).
+    ///   - The client hasn't requested confirmed reads
+    ///     (i.e. [`ClientConfig::confirmed_reads`] is `false`).
     ///   - The database is configured to not persist transactions.
     ///
-    /// Otherwise, the update's `tx_offset` is compared against the module's
+    /// Otherwise, the last `tx_offset` in the batch is compared against the module's
     /// durable offset. If the durable offset is behind the `tx_offset`, the
     /// method waits until it catches up before returning the message.
     ///
     /// If the database is shut down while waiting for the durable offset,
-    /// `None` is returned. In this case, no more messages can ever be received
+    /// 0 is returned. In this case, no more messages can ever be received
     /// from the channel.
     ///
+    /// For non-zero values of `max`, this method will never return `0` unless the
+    /// input channel has been closed and there are no pending messages, or if the
+    /// database goes away. This indicates that no further values can ever be received
+    /// from this `Receiver`.
+    ///
     /// # Cancel safety
     ///
     /// This method is cancel safe, as long as `self` is not dropped.
     ///
-    /// If `recv` is used in a [`tokio::select!`] statement, it may get
+    /// If `recv_many` is used in a [`tokio::select!`] statement, it may get
     /// cancelled while waiting for the durable offset to catch up. At this
-    /// point, it has already received a value from the underlying channel.
-    /// This value is stored internally, so calling `recv` again will not lose
-    /// data.
-    //
-    // TODO: Can we make a cancel-safe `recv_many` with confirmed reads semantics?
-    pub async fn recv(&mut self) -> Option<OutboundMessage> {
-        let ClientUpdate { tx_offset, message } = match self.current.take() {
-            None => self.channel.recv().await?,
-            Some(update) => update,
-        };
+    /// point, it has already received values from the underlying channel.
+    /// These values are stored internally, so calling `recv_many` again will
+    /// not lose data.
+    pub async fn recv_many(&mut self, buf: &mut Vec<OutboundMessage>, max: usize) -> usize {
+        // If there are no pending updates and the input channel has been closed,
+        // no more messages can be received from this receiver.
+        if max == 0 || (self.pending.is_empty() && self.channel.recv_many(&mut self.pending, max).await == 0) {
+            return 0;
+        }
+
+        // If we don't have to wait for txns to be made durable,
+        // drain the pending updates.
         if !self.confirmed_reads {
-            return Some(message);
+            return self.drain_pending(buf, max);
+        }
+
+        // If we do have to wait for txns to be made durable,
+        // but the next client update doesn't have a tx offset,
+        // there's no reason to wait - just send it.
+        if !self.pending_update_has_offset() {
+            return self.drain_pending(buf, 1);
         }
 
-        if let Some(tx_offset) = tx_offset {
-            match self.offset_supply.durable_offset() {
-                Ok(Some(mut durable)) => {
-                    // Store the current update in case we get cancelled while
-                    // waiting for the durable offset.
-                    self.current = Some(ClientUpdate {
-                        tx_offset: Some(tx_offset),
-                        message,
-                    });
-                    trace!("waiting for offset {tx_offset} to become durable");
-                    durable
-                        .wait_for(tx_offset)
-                        .await
-                        .inspect_err(|_| {
-                            warn!("database went away while waiting for durable offset");
-                        })
-                        .ok()?;
-                    self.current.take().map(|update| update.message)
+        // Otherwise, grab the next offset that we should wait for.
+        let (n, wait_for_offset) = self.next_confirmed_reads_batch(max);
+
+        match self.offset_supply.durable_offset() {
+            Ok(Some(mut durable)) => {
+                trace!("waiting for offset {wait_for_offset} to become durable");
+                if durable.wait_for(wait_for_offset).await.is_err() {
+                    warn!("database went away while waiting for durable offset");
+                    return 0;
                 }
-                // Database shut down or crashed.
-                Err(NoSuchModule) => None,
-                // In-memory database.
-                Ok(None) => Some(message),
+                self.drain_pending(buf, n)
             }
-        } else {
-            Some(message)
+            // Database shut down or crashed.
+            Err(NoSuchModule) => 0,
+            // In-memory database.
+            Ok(None) => self.drain_pending(buf, max),
         }
     }
 
+    /// Compute the next batch of pending client updates that have a tx offset.
+    /// What is the size of the batch and what is the max offset?
+    fn next_confirmed_reads_batch(&self, max: usize) -> (usize, TxOffset) {
+        self.pending
+            .iter()
+            .take(max)
+            .map_while(|update| update.tx_offset)
+            .fold((0, 0), |(count, max_offset), tx_offset| {
+                (count + 1, max_offset.max(tx_offset))
+            })
+    }
+
+    /// Drain the pending [`ClientUpdate`]s, up to `max, into `buf`.
+    fn drain_pending(&mut self, buf: &mut Vec<OutboundMessage>, max: usize) -> usize {
+        let n = self.pending.len().min(max);
+        buf.reserve(n);
+        buf.extend(self.pending.drain(..n).map(|u| u.message));
+        n
+    }
+
+    /// Does the next pending update have a tx offset?
+    ///
+    /// Assumes that [`Self::pending`] is not empty.
+    fn pending_update_has_offset(&self) -> bool {
+        self.pending.first().is_some_and(|update| update.tx_offset.is_some())
+    }
+
     /// Close the receiver without dropping it.
     ///
     /// This is used to notify the [`ClientConnectionSender`] that the receiver

crates/testing/src/modules.rs

@@ -101,7 +101,8 @@ impl ModuleHandle {
     }
 
     pub async fn recv_message(&mut self) -> Option<OutboundMessage> {
-        self.receiver.recv().await
+        let mut buf = Vec::with_capacity(1);
+        (self.receiver.recv_many(&mut buf, 1).await != 0).then(|| buf.remove(0))
     }
 
     pub async fn recv_reducer_update(&mut self, request_id: RequestId) -> anyhow::Result<()> {