messages::serialize: take/put buffers from/into a SerializeBufferPool

Author: Centril

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

@@ -859,25 +859,19 @@ pub fn decide_compression(len: usize, compression: Compression) -> Compression {
     }
 }
 
-pub fn brotli_compress(bytes: &[u8], out: &mut Vec<u8>) {
-    let reader = &mut &bytes[..];
-
-    // The default Brotli buffer size.
-    const BUFFER_SIZE: usize = 4096;
+pub fn brotli_compress(bytes: &[u8], out: &mut impl io::Write) {
     // We are optimizing for compression speed,
     // so we choose the lowest (fastest) level of compression.
     // Experiments on internal workloads have shown compression ratios between 7:1 and 10:1
     // for large `SubscriptionUpdate` messages at this level.
-    const COMPRESSION_LEVEL: u32 = 1;
-    // The default value for an internal compression parameter.
-    // See `BrotliEncoderParams` for more details.
-    const LG_WIN: u32 = 22;
+    const COMPRESSION_LEVEL: i32 = 1;
 
-    let mut encoder = brotli::CompressorReader::new(reader, BUFFER_SIZE, COMPRESSION_LEVEL, LG_WIN);
-
-    encoder
-        .read_to_end(out)
-        .expect("Failed to Brotli compress `SubscriptionUpdateMessage`");
+    let params = brotli::enc::BrotliEncoderParams {
+        quality: COMPRESSION_LEVEL,
+        ..<_>::default()
+    };
+    let reader = &mut &bytes[..];
+    brotli::BrotliCompress(reader, out, &params).expect("should be able to BrotliCompress");
 }
 
 pub fn brotli_decompress(bytes: &[u8]) -> Result<Vec<u8>, io::Error> {
@@ -886,10 +880,10 @@ pub fn brotli_decompress(bytes: &[u8]) -> Result<Vec<u8>, io::Error> {
     Ok(decompressed)
 }
 
-pub fn gzip_compress(bytes: &[u8], out: &mut Vec<u8>) {
+pub fn gzip_compress(bytes: &[u8], out: &mut impl io::Write) {
     let mut encoder = flate2::write::GzEncoder::new(out, flate2::Compression::fast());
     encoder.write_all(bytes).unwrap();
-    encoder.finish().expect("Failed to gzip compress `bytes`");
+    encoder.finish().expect("should be able to gzip compress `bytes`");
 }
 
 pub fn gzip_decompress(bytes: &[u8]) -> Result<Vec<u8>, io::Error> {

crates/client-api/src/lib.rs

@@ -5,6 +5,7 @@ use async_trait::async_trait;
 use axum::response::ErrorResponse;
 use http::StatusCode;
 
+use spacetimedb::client::messages::SerializeBufferPool;
 use spacetimedb::client::ClientActorIndex;
 use spacetimedb::energy::{EnergyBalance, EnergyQuanta};
 use spacetimedb::host::{HostController, ModuleHost, NoSuchModule, UpdateDatabaseResult};
@@ -39,6 +40,8 @@ pub trait NodeDelegate: Send + Sync {
     /// The [`Host`] is spawned implicitly if not already running.
     async fn leader(&self, database_id: u64) -> anyhow::Result<Option<Host>>;
     fn module_logs_dir(&self, replica_id: u64) -> ModuleLogsDir;
+
+    fn websocket_send_serialize_buffer_pool(&self) -> &Arc<SerializeBufferPool>;
 }
 
 /// Client view of a running module.
@@ -371,6 +374,10 @@ impl<T: NodeDelegate + ?Sized> NodeDelegate for Arc<T> {
     fn module_logs_dir(&self, replica_id: u64) -> ModuleLogsDir {
         (**self).module_logs_dir(replica_id)
     }
+
+    fn websocket_send_serialize_buffer_pool(&self) -> &Arc<SerializeBufferPool> {
+        (**self).websocket_send_serialize_buffer_pool()
+    }
 }
 
 pub fn log_and_500(e: impl std::fmt::Display) -> ErrorResponse {

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

@@ -1,6 +1,7 @@
 use std::collections::VecDeque;
 use std::mem;
 use std::pin::{pin, Pin};
+use std::sync::Arc;
 use std::time::Duration;
 
 use axum::extract::{Path, Query, State};
@@ -14,12 +15,14 @@ use futures::{Future, FutureExt, SinkExt, StreamExt};
 use http::{HeaderValue, StatusCode};
 use scopeguard::ScopeGuard;
 use serde::Deserialize;
-use spacetimedb::client::messages::{serialize, IdentityTokenMessage, SerializableMessage};
+use spacetimedb::client::messages::{serialize, IdentityTokenMessage, SerializableMessage, SerializeBufferPool};
 use spacetimedb::client::{ClientActorId, ClientConfig, ClientConnection, DataMessage, MessageHandleError, Protocol};
+use spacetimedb::execution_context::WorkloadType;
 use spacetimedb::host::module_host::ClientConnectedError;
 use spacetimedb::host::NoSuchModule;
 use spacetimedb::util::also_poll;
 use spacetimedb::worker_metrics::WORKER_METRICS;
+use spacetimedb::Identity;
 use spacetimedb_client_api_messages::websocket::{self as ws_api, Compression};
 use spacetimedb_lib::connection_id::{ConnectionId, ConnectionIdForUrl};
 use std::time::Instant;
@@ -125,6 +128,8 @@ where
         name: ctx.client_actor_index().next_client_name(),
     };
 
+    let serialize_buffer_pool = ctx.websocket_send_serialize_buffer_pool().clone();
+
     let ws_config = WebSocketConfig::default()
         .max_message_size(Some(0x2000000))
         .max_frame_size(None)
@@ -146,7 +151,7 @@ where
             None => log::debug!("New client connected from unknown ip"),
         }
 
-        let actor = |client, sendrx| ws_client_actor(client, ws, sendrx);
+        let actor = |client, sendrx| ws_client_actor(client, ws, sendrx, serialize_buffer_pool);
         let client = match ClientConnection::spawn(client_id, client_config, leader.replica_id, module_rx, actor).await
         {
             Ok(s) => s,
@@ -180,13 +185,18 @@ where
 
 const LIVELINESS_TIMEOUT: Duration = Duration::from_secs(60);
 
-async fn ws_client_actor(client: ClientConnection, ws: WebSocketStream, sendrx: mpsc::Receiver<SerializableMessage>) {
+async fn ws_client_actor(
+    client: ClientConnection,
+    ws: WebSocketStream,
+    sendrx: mpsc::Receiver<SerializableMessage>,
+    serialize_buffer_pool: Arc<SerializeBufferPool>,
+) {
     // ensure that even if this task gets cancelled, we always cleanup the connection
     let mut client = scopeguard::guard(client, |client| {
         tokio::spawn(client.disconnect());
     });
 
-    ws_client_actor_inner(&mut client, ws, sendrx).await;
+    ws_client_actor_inner(&mut client, ws, sendrx, &serialize_buffer_pool).await;
 
     ScopeGuard::into_inner(client).disconnect().await;
 }
@@ -203,6 +213,7 @@ async fn ws_client_actor_inner(
     client: &mut ClientConnection,
     mut ws: WebSocketStream,
     mut sendrx: mpsc::Receiver<SerializableMessage>,
+    serialize_buffer_pool: &SerializeBufferPool,
 ) {
     let mut liveness_check_interval = tokio::time::interval(LIVELINESS_TIMEOUT);
     let mut got_pong = true;
@@ -298,31 +309,31 @@ async fn ws_client_actor_inner(
                     //       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);
+                    let send_all = async {
+                        for msg in rx_buf.drain(..n) {
+                            let workload = msg.workload();
+                            let num_rows = msg.num_rows();
+
+                            // Serialize the message, report metrics,
+                            // and keep a handle to the buffer.
+                            let msg_data = serialize(serialize_buffer_pool, msg, client.config);
+                            report_ws_sent_metrics(&addr, workload, num_rows, &msg_data);
+                            let msg_alloc = msg_data.allocation();
+
+                            // Buffer the message without necessarily sending it.
+                            ws.feed(datamsg_to_wsmsg(msg_data)).await?;
+
+                            // At this point,
+                            // the underlying allocation of `msg_data` should have a single referent
+                            // and this should be `msg_alloc`.
+                            // We can put this back into our pool.
+                            let msg_alloc = msg_alloc.try_into_mut()
+                                .expect("should have a unique referent to `msg_alloc`");
+                            serialize_buffer_pool.put(msg_alloc);
                         }
-                        // 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
-                 };
+                        // 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));
@@ -393,10 +404,24 @@ async fn ws_client_actor_inner(
                 if let Err(e) = res {
                     if let MessageHandleError::Execution(err) = e {
                         log::error!("{err:#}");
-                        let msg = serialize(err, client.config);
-                        if let Err(error) = ws.send(datamsg_to_wsmsg(msg)).await {
+                        // Serialize the message and keep a handle to the buffer.
+                        let msg_data = serialize(serialize_buffer_pool, err, client.config);
+                        let msg_alloc = msg_data.allocation();
+
+                        // Buffer the message without necessarily sending it.
+                        if let Err(error) = ws.send(datamsg_to_wsmsg(msg_data)).await {
                             log::warn!("Websocket send error: {error}")
                         }
+
+                        // At this point,
+                        // the underlying allocation of `msg_data` should have a single referent
+                        // and this should be `msg_alloc`.
+                        // We can put this back into our pool.
+                        let msg_alloc = msg_alloc
+                            .try_into_mut()
+                            .expect("should have a unique referent to `msg_alloc`");
+                        serialize_buffer_pool.put(msg_alloc);
+
                         continue;
                     }
                     log::debug!("Client caused error on text message: {}", e);
@@ -460,6 +485,27 @@ impl ClientMessage {
     }
 }
 
+/// Report metrics on sent rows and message sizes to a websocket client.
+fn report_ws_sent_metrics(
+    addr: &Identity,
+    workload: Option<WorkloadType>,
+    num_rows: Option<usize>,
+    msg_ws: &DataMessage,
+) {
+    // 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_ws.len() as f64);
+    }
+}
+
 fn datamsg_to_wsmsg(msg: DataMessage) -> WsMessage {
     match msg {
         DataMessage::Text(text) => WsMessage::Text(bytestring_to_utf8bytes(text)),

crates/core/src/client/client_connection.rs

@@ -234,17 +234,27 @@ impl From<Vec<u8>> for DataMessage {
 }
 
 impl DataMessage {
+    /// Returns the number of bytes this message consists of.
     pub fn len(&self) -> usize {
         match self {
-            DataMessage::Text(s) => s.len(),
-            DataMessage::Binary(b) => b.len(),
+            Self::Text(s) => s.len(),
+            Self::Binary(b) => b.len(),
         }
     }
 
+    /// Is the message empty?
     #[must_use]
     pub fn is_empty(&self) -> bool {
         self.len() == 0
     }
+
+    /// Returns a handle to the underlying allocation of the message without consuming it.
+    pub fn allocation(&self) -> Bytes {
+        match self {
+            DataMessage::Text(alloc) => alloc.as_bytes().clone(),
+            DataMessage::Binary(alloc) => alloc.clone(),
+        }
+    }
 }
 
 // if a client racks up this many messages in the queue without ACK'ing