perf(1024): Remove serialization from tx execution thread (#1027)

Closes #1024.

Before this change,
we would serialize messages **before** inserting into the send queue.

Because we commit the tx only after inserting into the send queue,
this meant we were holding onto the database lock unnecessarily.

After this change,
we serialize messages **after** inserting into the send queue.
This means we serialize only after committing the tx.

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

@@ -9,9 +9,10 @@ use futures::{FutureExt, SinkExt, StreamExt};
 use http::{HeaderValue, StatusCode};
 use scopeguard::ScopeGuard;
 use serde::Deserialize;
-use spacetimedb::client::messages::{IdentityTokenMessage, ServerMessage};
+use spacetimedb::client::messages::{IdentityTokenMessage, SerializableMessage, ServerMessage};
 use spacetimedb::client::{ClientActorId, ClientConnection, DataMessage, MessageHandleError, Protocol};
 use spacetimedb::util::{also_poll, future_queue};
+use spacetimedb::worker_metrics::WORKER_METRICS;
 use spacetimedb_lib::address::AddressForUrl;
 use spacetimedb_lib::Address;
 use std::time::Instant;
@@ -168,7 +169,7 @@ where
 
 const LIVELINESS_TIMEOUT: Duration = Duration::from_secs(60);
 
-async fn ws_client_actor(client: ClientConnection, ws: WebSocketStream, sendrx: mpsc::Receiver<DataMessage>) {
+async fn ws_client_actor(client: ClientConnection, ws: WebSocketStream, sendrx: mpsc::Receiver<SerializableMessage>) {
     // ensure that even if this task gets cancelled, we always cleanup the connection
     let client = scopeguard::guard(client, |client| {
         tokio::spawn(client.disconnect());
@@ -182,7 +183,7 @@ async fn ws_client_actor(client: ClientConnection, ws: WebSocketStream, sendrx:
 async fn ws_client_actor_inner(
     client: &ClientConnection,
     mut ws: WebSocketStream,
-    mut sendrx: mpsc::Receiver<DataMessage>,
+    mut sendrx: mpsc::Receiver<SerializableMessage>,
 ) {
     let mut liveness_check_interval = tokio::time::interval(LIVELINESS_TIMEOUT);
     let mut got_pong = true;
@@ -240,7 +241,10 @@ async fn ws_client_actor_inner(
                     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).map(datamsg_to_wsmsg) {
+                        let id = client.id.identity;
+                        for msg in rx_buf.drain(..n).map(|msg| datamsg_to_wsmsg(msg.serialize(client.protocol))) {
+                            WORKER_METRICS.websocket_sent.with_label_values(&id).inc();
+                            WORKER_METRICS.websocket_sent_msg_size.with_label_values(&id).observe(msg.len() as f64);
                             // feed() buffers the message, but does not necessarily send it
                             ws.feed(msg).await?;
                         }

crates/core/src/client/client_connection.rs

@@ -3,7 +3,7 @@ use std::sync::atomic::{AtomicBool, Ordering::Relaxed};
 use std::sync::Arc;
 use std::time::Instant;
 
-use super::messages::{OneOffQueryResponseMessage, ServerMessage};
+use super::messages::{OneOffQueryResponseMessage, SerializableMessage};
 use super::{message_handlers, ClientActorId, MessageHandleError};
 use crate::error::DBError;
 use crate::host::{ModuleHost, ReducerArgs, ReducerCallError, ReducerCallResult};
@@ -26,7 +26,7 @@ pub enum Protocol {
 pub struct ClientConnectionSender {
     pub id: ClientActorId,
     pub protocol: Protocol,
-    sendtx: mpsc::Sender<DataMessage>,
+    sendtx: mpsc::Sender<SerializableMessage>,
     abort_handle: AbortHandle,
     cancelled: AtomicBool,
 }
@@ -56,13 +56,11 @@ impl ClientConnectionSender {
         }
     }
 
-    pub fn send_message(&self, message: impl ServerMessage) -> Result<(), ClientSendError> {
-        self.send(message.serialize(self.protocol))
+    pub fn send_message(&self, message: impl Into<SerializableMessage>) -> Result<(), ClientSendError> {
+        self.send(message.into())
     }
 
-    fn send(&self, message: DataMessage) -> Result<(), ClientSendError> {
-        let bytes_len = message.len();
-
+    fn send(&self, message: SerializableMessage) -> Result<(), ClientSendError> {
         if self.cancelled.load(Relaxed) {
             return Err(ClientSendError::Cancelled);
         }
@@ -77,13 +75,6 @@ impl ClientConnectionSender {
             mpsc::error::TrySendError::Closed(_) => ClientSendError::Disconnected,
         })?;
 
-        WORKER_METRICS.websocket_sent.with_label_values(&self.id.identity).inc();
-
-        WORKER_METRICS
-            .websocket_sent_msg_size
-            .with_label_values(&self.id.identity)
-            .observe(bytes_len as f64);
-
         Ok(())
     }
 }
@@ -137,7 +128,7 @@ impl ClientConnection {
         actor: F,
     ) -> Result<ClientConnection, ReducerCallError>
     where
-        F: FnOnce(ClientConnection, mpsc::Receiver<DataMessage>) -> Fut,
+        F: FnOnce(ClientConnection, mpsc::Receiver<SerializableMessage>) -> Fut,
         Fut: Future<Output = ()> + Send + 'static,
     {
         // Add this client as a subscriber
@@ -148,7 +139,7 @@ impl ClientConnection {
             .call_identity_connected_disconnected(id.identity, id.address, true)
             .await?;
 
-        let (sendtx, sendrx) = mpsc::channel::<DataMessage>(CLIENT_CHANNEL_CAPACITY);
+        let (sendtx, sendrx) = mpsc::channel::<SerializableMessage>(CLIENT_CHANNEL_CAPACITY);
 
         let db = module.info().address;
 

crates/core/src/client/message_handlers.rs

@@ -1,3 +1,4 @@
+use std::sync::Arc;
 use std::time::{Duration, Instant};
 
 use crate::energy::EnergyQuanta;
@@ -249,15 +250,15 @@ impl MessageExecutionError {
 impl ServerMessage for MessageExecutionError {
     fn serialize_text(self) -> crate::json::client_api::MessageJson {
         TransactionUpdateMessage::<DatabaseUpdate> {
-            event: &mut self.into_event(),
+            event: Arc::new(self.into_event()),
             database_update: Default::default(),
         }
         .serialize_text()
     }
 
     fn serialize_binary(self) -> Message {
         TransactionUpdateMessage::<DatabaseUpdate> {
-            event: &mut self.into_event(),
+            event: Arc::new(self.into_event()),
             database_update: Default::default(),
         }
         .serialize_binary()

crates/core/src/client/messages.rs

@@ -1,9 +1,11 @@
 use base64::Engine;
+use derive_more::From;
 use prost::Message as _;
 use spacetimedb_lib::identity::RequestId;
+use std::sync::Arc;
 use std::time::Instant;
 
-use crate::host::module_host::{EventStatus, ModuleEvent, ProtocolDatabaseUpdate};
+use crate::host::module_host::{DatabaseUpdate, EventStatus, ModuleEvent, ProtocolDatabaseUpdate};
 use crate::identity::Identity;
 use crate::json::client_api::{
     EventJson, FunctionCallJson, IdentityTokenJson, MessageJson, OneOffQueryResponseJson, OneOffTableJson,
@@ -14,6 +16,7 @@ use spacetimedb_client_api_messages::client_api::{OneOffQueryResponse, OneOffTab
 use spacetimedb_lib::Address;
 use spacetimedb_vm::relation::MemTable;
 
+use super::message_handlers::MessageExecutionError;
 use super::{DataMessage, Protocol};
 
 /// A message sent from the server to the client. Because clients can request either text or binary messages,
@@ -29,6 +32,41 @@ pub trait ServerMessage: Sized {
     fn serialize_binary(self) -> Message;
 }
 
+#[derive(Debug, From)]
+pub enum SerializableMessage {
+    Query(OneOffQueryResponseMessage),
+    Error(MessageExecutionError),
+    Identity(IdentityTokenMessage),
+    Subscribe(SubscriptionUpdateMessage),
+    DatabaseUpdate(TransactionUpdateMessage<DatabaseUpdate>),
+    ProtocolUpdate(TransactionUpdateMessage<ProtocolDatabaseUpdate>),
+}
+
+impl ServerMessage for SerializableMessage {
+    fn serialize_text(self) -> MessageJson {
+        match self {
+            SerializableMessage::Query(msg) => msg.serialize_text(),
+            SerializableMessage::Error(msg) => msg.serialize_text(),
+            SerializableMessage::Identity(msg) => msg.serialize_text(),
+            SerializableMessage::Subscribe(msg) => msg.serialize_text(),
+            SerializableMessage::DatabaseUpdate(msg) => msg.serialize_text(),
+            SerializableMessage::ProtocolUpdate(msg) => msg.serialize_text(),
+        }
+    }
+
+    fn serialize_binary(self) -> Message {
+        match self {
+            SerializableMessage::Query(msg) => msg.serialize_binary(),
+            SerializableMessage::Error(msg) => msg.serialize_binary(),
+            SerializableMessage::Identity(msg) => msg.serialize_binary(),
+            SerializableMessage::Subscribe(msg) => msg.serialize_binary(),
+            SerializableMessage::DatabaseUpdate(msg) => msg.serialize_binary(),
+            SerializableMessage::ProtocolUpdate(msg) => msg.serialize_binary(),
+        }
+    }
+}
+
+#[derive(Debug)]
 pub struct IdentityTokenMessage {
     pub identity: Identity,
     pub identity_token: String,
@@ -54,12 +92,13 @@ impl ServerMessage for IdentityTokenMessage {
     }
 }
 
-pub struct TransactionUpdateMessage<'a, U> {
-    pub event: &'a ModuleEvent,
+#[derive(Debug)]
+pub struct TransactionUpdateMessage<U> {
+    pub event: Arc<ModuleEvent>,
     pub database_update: SubscriptionUpdate<U>,
 }
 
-impl<U: Into<Vec<TableUpdate>> + Into<Vec<TableUpdateJson>>> ServerMessage for TransactionUpdateMessage<'_, U> {
+impl<U: Into<Vec<TableUpdate>> + Into<Vec<TableUpdateJson>>> ServerMessage for TransactionUpdateMessage<U> {
     fn serialize_text(self) -> MessageJson {
         let Self { event, database_update } = self;
         let (status_str, errmsg) = match &event.status {
@@ -123,25 +162,7 @@ impl<U: Into<Vec<TableUpdate>> + Into<Vec<TableUpdateJson>>> ServerMessage for T
     }
 }
 
-impl<U: Clone + Into<Vec<TableUpdate>> + Into<Vec<TableUpdateJson>>> ServerMessage
-    for &mut TransactionUpdateMessage<'_, U>
-{
-    fn serialize_text(self) -> MessageJson {
-        TransactionUpdateMessage {
-            event: self.event,
-            database_update: self.database_update.clone(),
-        }
-        .serialize_text()
-    }
-    fn serialize_binary(self) -> Message {
-        TransactionUpdateMessage {
-            event: self.event,
-            database_update: self.database_update.clone(),
-        }
-        .serialize_binary()
-    }
-}
-
+#[derive(Debug)]
 pub struct SubscriptionUpdateMessage {
     pub subscription_update: SubscriptionUpdate<ProtocolDatabaseUpdate>,
 }
@@ -185,42 +206,7 @@ impl<T: Into<Vec<TableUpdateJson>>> SubscriptionUpdate<T> {
     }
 }
 
-pub struct CachedMessage<M> {
-    msg: M,
-    text: Option<String>,
-    binary: Option<Vec<u8>>,
-}
-
-impl<M> CachedMessage<M> {
-    pub fn new(msg: M) -> Self {
-        Self {
-            msg,
-            text: None,
-            binary: None,
-        }
-    }
-}
-
-impl<M> CachedMessage<M>
-where
-    for<'b> &'b mut M: ServerMessage,
-{
-    pub fn serialize(&mut self, protocol: Protocol) -> DataMessage {
-        match protocol {
-            Protocol::Text => self
-                .text
-                .get_or_insert_with(|| self.msg.serialize_text().to_json())
-                .clone()
-                .into(),
-            Protocol::Binary => self
-                .binary
-                .get_or_insert_with(|| self.msg.serialize_binary().encode_to_vec())
-                .clone()
-                .into(),
-        }
-    }
-}
-
+#[derive(Debug)]
 pub struct OneOffQueryResponseMessage {
     pub message_id: Vec<u8>,
     pub error: Option<String>,

crates/core/src/subscription/module_subscription_actor.rs

@@ -140,7 +140,7 @@ impl ModuleSubscriptions {
             EventStatus::Failed(_) => {
                 if let Some(client) = client {
                     let message = TransactionUpdateMessage::<DatabaseUpdate> {
-                        event: &event,
+                        event,
                         database_update: <_>::default(),
                     };
                     let _ = client.send_message(message);
@@ -169,7 +169,7 @@ impl ModuleSubscriptions {
     /// it resolves, it's guaranteed that if you call `subscriber.send(x)` the client will receive
     /// x after they receive this subscription update).
     fn broadcast_commit_event(&self, subscriptions: &SubscriptionManager, event: Arc<ModuleEvent>) {
-        subscriptions.eval_updates(&self.relational_db, &event)
+        subscriptions.eval_updates(&self.relational_db, event)
     }
 }
 

crates/core/src/subscription/module_subscription_manager.rs

@@ -112,7 +112,7 @@ impl SubscriptionManager {
     /// evaluates only the necessary queries for those delta tables,
     /// and then sends the results to each client.
     #[tracing::instrument(skip_all)]
-    pub fn eval_updates(&self, db: &RelationalDB, event: &ModuleEvent) {
+    pub fn eval_updates(&self, db: &RelationalDB, event: Arc<ModuleEvent>) {
         let tables = &event.status.database_update().unwrap().tables;
 
         let tx = scopeguard::guard(db.begin_tx(), |tx| {
@@ -246,7 +246,10 @@ impl SubscriptionManager {
                     request_id: event.request_id,
                     timer: event.timer,
                 };
-                let message = TransactionUpdateMessage { event, database_update };
+                let message = TransactionUpdateMessage {
+                    event: event.clone(),
+                    database_update,
+                };
                 if let Err(e) = client.send_message(message) {
                     tracing::warn!(%client.id, "failed to send update message to client: {e}")
                 }