subscriptions: reuse buffers in ServerMessage<BsatnFormat> via global pool

Author: Centril

crates/bench/benches/subscription.rs

@@ -1,11 +1,13 @@
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use spacetimedb::client::consume_each_list::ConsumeEachBuffer;
 use spacetimedb::error::DBError;
 use spacetimedb::execution_context::Workload;
 use spacetimedb::host::module_host::DatabaseTableUpdate;
 use spacetimedb::identity::AuthCtx;
 use spacetimedb::messages::websocket::BsatnFormat;
 use spacetimedb::sql::ast::SchemaViewer;
 use spacetimedb::subscription::query::compile_read_only_queryset;
+use spacetimedb::subscription::row_list_builder_pool::BsatnRowListBuilderPool;
 use spacetimedb::subscription::subscription::ExecutionSet;
 use spacetimedb::subscription::tx::DeltaTx;
 use spacetimedb::subscription::{collect_table_update, TableUpdateType};
@@ -119,6 +121,8 @@ fn eval(c: &mut Criterion) {
     let ins_rhs = insert_op(rhs, "location", new_rhs_row);
     let update = [&ins_lhs, &ins_rhs];
 
+    let bsatn_rlb_pool = BsatnRowListBuilderPool::new();
+
     // A benchmark runner for the new query engine
     let bench_query = |c: &mut Criterion, name, sql| {
         c.bench_function(name, |b| {
@@ -134,13 +138,17 @@ fn eval(c: &mut Criterion) {
             let tx = DeltaTx::from(&tx);
 
             b.iter(|| {
-                drop(black_box(collect_table_update::<_, BsatnFormat>(
+                let updates = black_box(collect_table_update::<_, BsatnFormat>(
                     &plans,
                     table_id,
                     table_name.clone(),
                     &tx,
                     TableUpdateType::Subscribe,
-                )))
+                    &bsatn_rlb_pool,
+                ));
+                if let Ok((updates, _)) = updates {
+                    updates.consume_each_list(&mut |buffer| black_box(bsatn_rlb_pool.try_put(buffer)));
+                }
             })
         });
     };
@@ -152,12 +160,9 @@ fn eval(c: &mut Criterion) {
             let query: ExecutionSet = query.into();
 
             b.iter(|| {
-                drop(black_box(query.eval::<BsatnFormat>(
-                    &raw.db,
-                    &tx,
-                    None,
-                    Compression::None,
-                )))
+                let updates =
+                    black_box(query.eval::<BsatnFormat>(&raw.db, &tx, &bsatn_rlb_pool, None, Compression::None));
+                updates.consume_each_list(&mut |buffer| black_box(bsatn_rlb_pool.try_put(buffer)));
             })
         });
     };

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

@@ -15,7 +15,7 @@
 //! rather than using an external mirror of this schema.
 
 use crate::energy::EnergyQuanta;
-use bytes::Bytes;
+use bytes::{Bytes, BytesMut};
 use bytestring::ByteString;
 use core::{
     fmt::Debug,
@@ -41,8 +41,15 @@ use std::{
 pub const TEXT_PROTOCOL: &str = "v1.json.spacetimedb";
 pub const BIN_PROTOCOL: &str = "v1.bsatn.spacetimedb";
 
+/// A source of row list builders for a given [`WebsocketFormat`].
+pub trait RowListBuilderSource<F: WebsocketFormat> {
+    /// Returns a row list builder from the source `self`.
+    fn take_row_list_builder(&self) -> F::ListBuilder;
+}
+
 /// A list of rows being built.
-pub trait RowListBuilder: Default {
+pub trait RowListBuilder {
+    /// The type of a finished list returned by [`RowListBuilder::finish`].
     type FinishedList;
 
     /// Push a row to the list in a serialized format.
@@ -98,13 +105,17 @@ pub trait WebsocketFormat: Sized {
         + Clone
         + Default;
 
-    /// The builder for [`Self::List`].
+    /// The builder for [`WebsocketFormat::List`].
     type ListBuilder: RowListBuilder<FinishedList = Self::List>;
 
     /// Encodes the `elems` to a list in the format and also returns the length of the list.
-    fn encode_list<R: ToBsatn + Serialize>(elems: impl Iterator<Item = R>) -> (Self::List, u64) {
+    ///
+    /// Needs to be provided with an empty [`WebsocketFormat::ListBuilder`].
+    fn encode_list<R: ToBsatn + Serialize>(
+        mut list: Self::ListBuilder,
+        elems: impl Iterator<Item = R>,
+    ) -> (Self::List, u64) {
         let mut num_rows = 0;
-        let mut list = Self::ListBuilder::default();
         for elem in elems {
             num_rows += 1;
             list.push(elem);
@@ -116,7 +127,7 @@ pub trait WebsocketFormat: Sized {
     /// This type exists so that some formats, e.g., BSATN, can compress an update.
     type QueryUpdate: SpacetimeType + for<'de> Deserialize<'de> + Serialize + Debug + Clone + Send;
 
-    /// Convert a `QueryUpdate` into `Self::QueryUpdate`.
+    /// Convert a `QueryUpdate` into [`WebsocketFormat::QueryUpdate`].
     /// This allows some formats to e.g., compress the update.
     fn into_query_update(qu: QueryUpdate<Self>, compression: Compression) -> Self::QueryUpdate;
 }
@@ -976,6 +987,11 @@ impl BsatnRowList {
         let data_range = self.size_hint.index_to_range(index, data_end)?;
         Some(self.rows_data.slice(data_range))
     }
+
+    /// Consumes the list and returns the parts.
+    pub fn into_inner(self) -> (RowSizeHint, Bytes) {
+        (self.size_hint, self.rows_data)
+    }
 }
 
 /// An iterator over all the elements in a [`BsatnRowList`].
@@ -1008,7 +1024,7 @@ pub struct BsatnRowListBuilder {
     /// intended to facilitate parallel decode purposes on large initial updates.
     size_hint: RowSizeHintBuilder,
     /// The flattened byte array for a list of rows.
-    rows_data: Vec<u8>,
+    rows_data: BytesMut,
 }
 
 /// A [`RowSizeHint`] under construction.
@@ -1096,6 +1112,14 @@ impl RowListBuilder for BsatnRowListBuilder {
     }
 }
 
+impl BsatnRowListBuilder {
+    /// Returns a new builder using an empty [`BytesMut`] for the `rows_data` buffer.
+    pub fn new_from_bytes(rows_data: BytesMut) -> Self {
+        let size_hint = <_>::default();
+        Self { size_hint, rows_data }
+    }
+}
+
 fn collect_offsets_from_num_rows(num_rows: usize, size: usize) -> Vec<u64> {
     (0..num_rows).map(|i| i * size).map(|o| o as u64).collect()
 }

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

@@ -234,6 +234,8 @@ async fn ws_client_actor_inner(
     let mut closed = false;
     let mut rx_buf = Vec::new();
 
+    let bsatn_rlb_pool = &client.module.subscriptions().bsatn_rlb_pool.clone();
+
     let mut msg_buffer = SerializeBuffer::new(client.config);
     loop {
         rx_buf.clear();
@@ -292,7 +294,7 @@ async fn ws_client_actor_inner(
 
                             // Serialize the message, report metrics,
                             // and keep a handle to the buffer.
-                            let (msg_alloc, msg_data) = serialize(msg_buffer, msg, client.config);
+                            let (msg_alloc, msg_data) = serialize(bsatn_rlb_pool, msg_buffer, msg, client.config);
                             report_ws_sent_metrics(&addr, workload, num_rows, &msg_data);
 
                             // Buffer the message without necessarily sending it.
@@ -440,7 +442,7 @@ async fn ws_client_actor_inner(
                     if let MessageHandleError::Execution(err) = e {
                         log::error!("reducer execution error: {err:#}");
                         // Serialize the message and keep a handle to the buffer.
-                        let (msg_alloc, msg_data) = serialize(msg_buffer, err, client.config);
+                        let (msg_alloc, msg_data) = serialize(bsatn_rlb_pool, msg_buffer, err, client.config);
 
                         let send = async { ws.send(datamsg_to_wsmsg(msg_data)).await };
                         let send = tokio::time::timeout(SEND_TIMEOUT, send);

crates/core/src/client.rs

@@ -3,6 +3,7 @@ use std::fmt;
 
 mod client_connection;
 mod client_connection_index;
+pub mod consume_each_list;
 mod message_handlers;
 pub mod messages;
 

crates/core/src/client/client_connection.rs

@@ -11,6 +11,7 @@ use crate::error::DBError;
 use crate::host::module_host::ClientConnectedError;
 use crate::host::{ModuleHost, NoSuchModule, ReducerArgs, ReducerCallError, ReducerCallResult};
 use crate::messages::websocket::Subscribe;
+use crate::subscription::row_list_builder_pool::JsonRowListBuilderFakePool;
 use crate::util::asyncify;
 use crate::util::prometheus_handle::IntGaugeExt;
 use crate::worker_metrics::WORKER_METRICS;
@@ -586,6 +587,7 @@ impl ClientConnection {
                 self.sender.clone(),
                 message_id.to_owned(),
                 timer,
+                JsonRowListBuilderFakePool,
                 |msg: OneOffQueryResponseMessage<JsonFormat>| msg.into(),
             )
             .await
@@ -597,13 +599,15 @@ impl ClientConnection {
         message_id: &[u8],
         timer: Instant,
     ) -> Result<(), anyhow::Error> {
+        let bsatn_rlb_pool = self.module.replica_ctx().subscriptions.bsatn_rlb_pool.clone();
         self.module
             .one_off_query::<BsatnFormat>(
                 self.id.identity,
                 query.to_owned(),
                 self.sender.clone(),
                 message_id.to_owned(),
                 timer,
+                bsatn_rlb_pool,
                 |msg: OneOffQueryResponseMessage<BsatnFormat>| msg.into(),
             )
             .await

crates/core/src/client/consume_each_list.rs

@@ -0,0 +1,80 @@
+use bytes::Bytes;
+use spacetimedb_client_api_messages::websocket::{
+    BsatnFormat, BsatnRowList, CompressableQueryUpdate, DatabaseUpdate, OneOffQueryResponse, QueryUpdate,
+    ServerMessage, TableUpdate, UpdateStatus,
+};
+
+/// Moves each buffer in `self` into a closure.
+pub trait ConsumeEachBuffer {
+    /// Consumes `self`, moving each `Bytes` buffer in `self` into the closure `each`.
+    fn consume_each_list(self, each: &mut impl FnMut(Bytes));
+}
+
+impl ConsumeEachBuffer for ServerMessage<BsatnFormat> {
+    fn consume_each_list(self, each: &mut impl FnMut(Bytes)) {
+        use ServerMessage::*;
+        match self {
+            InitialSubscription(x) => x.database_update.consume_each_list(each),
+            TransactionUpdate(x) => x.status.consume_each_list(each),
+            TransactionUpdateLight(x) => x.update.consume_each_list(each),
+            IdentityToken(_) | SubscriptionError(_) => {}
+            OneOffQueryResponse(x) => x.consume_each_list(each),
+            SubscribeApplied(x) => x.rows.table_rows.consume_each_list(each),
+            UnsubscribeApplied(x) => x.rows.table_rows.consume_each_list(each),
+            SubscribeMultiApplied(x) => x.update.consume_each_list(each),
+            UnsubscribeMultiApplied(x) => x.update.consume_each_list(each),
+        }
+    }
+}
+
+impl ConsumeEachBuffer for OneOffQueryResponse<BsatnFormat> {
+    fn consume_each_list(self, each: &mut impl FnMut(Bytes)) {
+        Vec::from(self.tables)
+            .into_iter()
+            .for_each(|x| x.rows.consume_each_list(each));
+    }
+}
+
+impl ConsumeEachBuffer for UpdateStatus<BsatnFormat> {
+    fn consume_each_list(self, each: &mut impl FnMut(Bytes)) {
+        match self {
+            Self::Committed(x) => x.consume_each_list(each),
+            Self::Failed(_) | UpdateStatus::OutOfEnergy => {}
+        }
+    }
+}
+
+impl ConsumeEachBuffer for DatabaseUpdate<BsatnFormat> {
+    fn consume_each_list(self, each: &mut impl FnMut(Bytes)) {
+        self.tables.into_iter().for_each(|x| x.consume_each_list(each));
+    }
+}
+
+impl ConsumeEachBuffer for TableUpdate<BsatnFormat> {
+    fn consume_each_list(self, each: &mut impl FnMut(Bytes)) {
+        self.updates.into_iter().for_each(|x| x.consume_each_list(each));
+    }
+}
+
+impl ConsumeEachBuffer for CompressableQueryUpdate<BsatnFormat> {
+    fn consume_each_list(self, each: &mut impl FnMut(Bytes)) {
+        match self {
+            Self::Uncompressed(x) => x.consume_each_list(each),
+            Self::Brotli(bytes) | Self::Gzip(bytes) => each(bytes),
+        }
+    }
+}
+
+impl ConsumeEachBuffer for QueryUpdate<BsatnFormat> {
+    fn consume_each_list(self, each: &mut impl FnMut(Bytes)) {
+        self.deletes.consume_each_list(each);
+        self.inserts.consume_each_list(each);
+    }
+}
+
+impl ConsumeEachBuffer for BsatnRowList {
+    fn consume_each_list(self, each: &mut impl FnMut(Bytes)) {
+        let (_, buffer) = self.into_inner();
+        each(buffer);
+    }
+}

crates/core/src/client/messages.rs

@@ -1,8 +1,10 @@
 use super::{ClientConfig, DataMessage, Protocol};
+use crate::client::consume_each_list::ConsumeEachBuffer;
 use crate::execution_context::WorkloadType;
 use crate::host::module_host::{EventStatus, ModuleEvent};
 use crate::host::ArgsTuple;
 use crate::messages::websocket as ws;
+use crate::subscription::row_list_builder_pool::BsatnRowListBuilderPool;
 use bytes::{BufMut, Bytes, BytesMut};
 use bytestring::ByteString;
 use derive_more::From;
@@ -125,6 +127,7 @@ impl InUseSerializeBuffer {
 /// If `protocol` is [`Protocol::Binary`],
 /// the message will be conditionally compressed by this method according to `compression`.
 pub fn serialize(
+    bsatn_rlb_pool: &BsatnRowListBuilderPool,
     mut buffer: SerializeBuffer,
     msg: impl ToProtocol<Encoded = SwitchedServerMessage>,
     config: ClientConfig,
@@ -147,6 +150,10 @@ pub fn serialize(
                 bsatn::to_writer(w.into_inner(), &msg).unwrap()
             });
 
+            // At this point, we no longer have a use for `msg`,
+            // so try to reclaim its buffers.
+            msg.consume_each_list(&mut |buffer| bsatn_rlb_pool.try_put(buffer));
+
             // Conditionally compress the message.
             let (in_use, msg_bytes) = match ws::decide_compression(srv_msg.len(), config.compression) {
                 Compression::None => buffer.uncompressed(),