move building and compression out of 'client-api-messages'

Centril · Centril · commit 6271604d8bf3 · 2025-07-15T19:02:28.000+02:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/client-api-messages/Cargo.toml b/crates/client-api-messages/Cargo.toml
@@ -12,10 +12,8 @@ spacetimedb-sats = { workspace = true, features = ["bytestring"] }
 
 bytes.workspace = true
 bytestring.workspace = true
-brotli.workspace = true
 chrono = { workspace = true, features = ["serde"] }
 enum-as-inner.workspace = true
-flate2.workspace = true
 serde = { workspace = true, features = ["derive"] }
 serde_json.workspace = true
 serde_with.workspace = true
diff --git a/crates/client-api-messages/src/lib.rs b/crates/client-api-messages/src/lib.rs
@@ -4,4 +4,3 @@ pub mod energy;
 pub mod http;
 pub mod name;
 pub mod websocket;
-mod websocket_building;
diff --git a/crates/client-api-messages/src/websocket.rs b/crates/client-api-messages/src/websocket.rs
@@ -26,18 +26,12 @@ use smallvec::SmallVec;
 use spacetimedb_lib::{ConnectionId, Identity, TimeDuration, Timestamp};
 use spacetimedb_primitives::TableId;
 use spacetimedb_sats::{
-    bsatn,
     de::{Deserialize, Error},
     impl_deserialize, impl_serialize, impl_st,
     ser::Serialize,
     AlgebraicType, SpacetimeType,
 };
-use std::{
-    io::{self, Read as _},
-    sync::Arc,
-};
-
-pub use crate::websocket_building::*;
+use std::sync::Arc;
 
 pub const TEXT_PROTOCOL: &str = "v1.json.spacetimedb";
 pub const BIN_PROTOCOL: &str = "v1.bsatn.spacetimedb";
@@ -661,22 +655,6 @@ pub enum CompressableQueryUpdate<F: WebsocketFormat> {
     Gzip(Bytes),
 }
 
-impl CompressableQueryUpdate<BsatnFormat> {
-    pub fn maybe_decompress(self) -> QueryUpdate<BsatnFormat> {
-        match self {
-            Self::Uncompressed(qu) => qu,
-            Self::Brotli(bytes) => {
-                let bytes = brotli_decompress(&bytes).unwrap();
-                bsatn::from_slice(&bytes).unwrap()
-            }
-            Self::Gzip(bytes) => {
-                let bytes = gzip_decompress(&bytes).unwrap();
-                bsatn::from_slice(&bytes).unwrap()
-            }
-        }
-    }
-}
-
 #[derive(SpacetimeType, Debug, Clone)]
 #[sats(crate = spacetimedb_lib)]
 pub struct QueryUpdate<F: WebsocketFormat> {
@@ -777,18 +755,6 @@ pub enum Compression {
     Gzip,
 }
 
-pub fn brotli_decompress(bytes: &[u8]) -> Result<Vec<u8>, io::Error> {
-    let mut decompressed = Vec::new();
-    brotli::BrotliDecompress(&mut &bytes[..], &mut decompressed)?;
-    Ok(decompressed)
-}
-
-pub fn gzip_decompress(bytes: &[u8]) -> Result<Vec<u8>, io::Error> {
-    let mut decompressed = Vec::new();
-    let _ = flate2::read::GzDecoder::new(bytes).read_to_end(&mut decompressed)?;
-    Ok(decompressed)
-}
-
 pub type RowSize = u16;
 pub type RowOffset = u64;
 
diff --git a/crates/core/src/client/messages.rs b/crates/core/src/client/messages.rs
@@ -3,6 +3,7 @@ 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::websocket_building::{brotli_compress, decide_compression, gzip_compress};
 use bytes::{BufMut, Bytes, BytesMut};
 use bytestring::ByteString;
 use derive_more::From;
@@ -148,10 +149,10 @@ pub fn serialize(
             });
 
             // Conditionally compress the message.
-            let (in_use, msg_bytes) = match ws::decide_compression(srv_msg.len(), config.compression) {
+            let (in_use, msg_bytes) = match decide_compression(srv_msg.len(), config.compression) {
                 Compression::None => buffer.uncompressed(),
-                Compression::Brotli => buffer.compress_with_tag(SERVER_MSG_COMPRESSION_TAG_BROTLI, ws::brotli_compress),
-                Compression::Gzip => buffer.compress_with_tag(SERVER_MSG_COMPRESSION_TAG_GZIP, ws::gzip_compress),
+                Compression::Brotli => buffer.compress_with_tag(SERVER_MSG_COMPRESSION_TAG_BROTLI, brotli_compress),
+                Compression::Gzip => buffer.compress_with_tag(SERVER_MSG_COMPRESSION_TAG_GZIP, gzip_compress),
             };
             (in_use, msg_bytes.into())
         }
diff --git a/crates/core/src/host/module_host.rs b/crates/core/src/host/module_host.rs
@@ -18,6 +18,7 @@ use crate::sql::parser::RowLevelExpr;
 use crate::subscription::execute_plan;
 use crate::subscription::module_subscription_actor::ModuleSubscriptions;
 use crate::subscription::tx::DeltaTx;
+use crate::subscription::websocket_building::BuildableWebsocketFormat;
 use crate::util::asyncify;
 use crate::util::jobs::{JobCore, JobThread, JobThreadClosed, WeakJobThread};
 use crate::vm::check_row_limit;
@@ -28,9 +29,7 @@ use derive_more::From;
 use indexmap::IndexSet;
 use itertools::Itertools;
 use prometheus::{Histogram, IntGauge};
-use spacetimedb_client_api_messages::websocket::{
-    BuildableWebsocketFormat, ByteListLen, Compression, OneOffTable, QueryUpdate,
-};
+use spacetimedb_client_api_messages::websocket::{ByteListLen, Compression, OneOffTable, QueryUpdate};
 use spacetimedb_data_structures::error_stream::ErrorStream;
 use spacetimedb_data_structures::map::{HashCollectionExt as _, IntMap};
 use spacetimedb_execution::pipelined::PipelinedProject;
diff --git a/crates/core/src/subscription/execution_unit.rs b/crates/core/src/subscription/execution_unit.rs
@@ -6,11 +6,10 @@ use crate::error::DBError;
 use crate::estimation;
 use crate::host::module_host::{DatabaseTableUpdate, DatabaseTableUpdateRelValue, UpdatesRelValue};
 use crate::messages::websocket::TableUpdate;
+use crate::subscription::websocket_building::BuildableWebsocketFormat;
 use crate::util::slow::SlowQueryLogger;
 use crate::vm::{build_query, TxMode};
-use spacetimedb_client_api_messages::websocket::{
-    BuildableWebsocketFormat, Compression, QueryUpdate, RowListLen as _, SingleQueryUpdate,
-};
+use spacetimedb_client_api_messages::websocket::{Compression, QueryUpdate, RowListLen as _, SingleQueryUpdate};
 use spacetimedb_lib::Identity;
 use spacetimedb_primitives::TableId;
 use spacetimedb_sats::{u256, ProductValue};
diff --git a/crates/core/src/subscription/mod.rs b/crates/core/src/subscription/mod.rs
@@ -1,23 +1,21 @@
-use std::sync::Arc;
-
+use crate::subscription::websocket_building::{BuildableWebsocketFormat, RowListBuilder as _};
+use crate::{
+    db::{datastore::locking_tx_datastore::datastore::MetricsRecorder, db_metrics::DB_METRICS},
+    error::DBError,
+    execution_context::WorkloadType,
+    worker_metrics::WORKER_METRICS,
+};
 use anyhow::Result;
 use module_subscription_manager::Plan;
 use prometheus::IntCounter;
 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 use spacetimedb_client_api_messages::websocket::{
-    BuildableWebsocketFormat, ByteListLen, Compression, DatabaseUpdate, QueryUpdate, RowListBuilder as _,
-    SingleQueryUpdate, TableUpdate,
+    ByteListLen, Compression, DatabaseUpdate, QueryUpdate, SingleQueryUpdate, TableUpdate,
 };
 use spacetimedb_execution::{pipelined::PipelinedProject, Datastore, DeltaStore};
 use spacetimedb_lib::{metrics::ExecutionMetrics, Identity};
 use spacetimedb_primitives::TableId;
-
-use crate::{
-    db::{datastore::locking_tx_datastore::datastore::MetricsRecorder, db_metrics::DB_METRICS},
-    error::DBError,
-    execution_context::WorkloadType,
-    worker_metrics::WORKER_METRICS,
-};
+use std::sync::Arc;
 
 pub mod delta;
 pub mod execution_unit;
@@ -27,6 +25,7 @@ pub mod query;
 #[allow(clippy::module_inception)] // it's right this isn't ideal :/
 pub mod subscription;
 pub mod tx;
+pub mod websocket_building;
 
 #[derive(Debug)]
 pub struct ExecutionCounters {
diff --git a/crates/core/src/subscription/module_subscription_manager.rs b/crates/core/src/subscription/module_subscription_manager.rs
@@ -10,15 +10,15 @@ use crate::error::DBError;
 use crate::host::module_host::{DatabaseTableUpdate, ModuleEvent, UpdatesRelValue};
 use crate::messages::websocket::{self as ws, TableUpdate};
 use crate::subscription::delta::eval_delta;
+use crate::subscription::websocket_building::BuildableWebsocketFormat;
 use crate::worker_metrics::WORKER_METRICS;
 use core::mem;
 use hashbrown::hash_map::OccupiedError;
 use hashbrown::{HashMap, HashSet};
 use parking_lot::RwLock;
 use prometheus::IntGauge;
 use spacetimedb_client_api_messages::websocket::{
-    BsatnFormat, BuildableWebsocketFormat, CompressableQueryUpdate, FormatSwitch, JsonFormat, QueryId, QueryUpdate,
-    SingleQueryUpdate,
+    BsatnFormat, CompressableQueryUpdate, FormatSwitch, JsonFormat, QueryId, QueryUpdate, SingleQueryUpdate,
 };
 use spacetimedb_data_structures::map::{Entry, IntMap};
 use spacetimedb_lib::metrics::ExecutionMetrics;
diff --git a/crates/core/src/subscription/query.rs b/crates/core/src/subscription/query.rs
@@ -160,7 +160,7 @@ mod tests {
     use crate::vm::tests::create_table_with_rows;
     use crate::vm::DbProgram;
     use itertools::Itertools;
-    use spacetimedb_client_api_messages::websocket::{BsatnFormat, Compression};
+    use spacetimedb_client_api_messages::websocket::{BsatnFormat, CompressableQueryUpdate, Compression};
     use spacetimedb_lib::bsatn;
     use spacetimedb_lib::db::auth::{StAccess, StTableType};
     use spacetimedb_lib::error::ResultTest;
@@ -353,7 +353,7 @@ mod tests {
         total_tables: usize,
         rows: &[ProductValue],
     ) -> ResultTest<()> {
-        let result = s.eval::<BsatnFormat>(db, tx, None, Compression::Brotli).tables;
+        let result = s.eval::<BsatnFormat>(db, tx, None, Compression::None).tables;
         assert_eq!(
             result.len(),
             total_tables,
@@ -363,7 +363,10 @@ mod tests {
         let result = result
             .into_iter()
             .flat_map(|x| x.updates)
-            .map(|x| x.maybe_decompress())
+            .map(|x| match x {
+                CompressableQueryUpdate::Uncompressed(x) => x,
+                _ => unreachable!(),
+            })
             .flat_map(|x| {
                 (&x.deletes)
                     .into_iter()
diff --git a/crates/core/src/subscription/subscription.rs b/crates/core/src/subscription/subscription.rs
@@ -29,10 +29,11 @@ use crate::error::{DBError, SubscriptionError};
 use crate::host::module_host::{DatabaseTableUpdate, DatabaseUpdateRelValue, UpdatesRelValue};
 use crate::messages::websocket as ws;
 use crate::sql::ast::SchemaViewer;
+use crate::subscription::websocket_building::BuildableWebsocketFormat;
 use crate::vm::{build_query, TxMode};
 use anyhow::Context;
 use itertools::Either;
-use spacetimedb_client_api_messages::websocket::{BuildableWebsocketFormat, Compression};
+use spacetimedb_client_api_messages::websocket::Compression;
 use spacetimedb_data_structures::map::HashSet;
 use spacetimedb_lib::db::auth::{StAccess, StTableType};
 use spacetimedb_lib::identity::AuthCtx;
diff --git a/crates/core/src/subscription/websocket_building.rs b/crates/core/src/subscription/websocket_building.rs
@@ -1,9 +1,9 @@
-use crate::websocket::{
+use bytestring::ByteString;
+use core::mem;
+use spacetimedb_client_api_messages::websocket::{
     BsatnFormat, BsatnRowList, CompressableQueryUpdate, Compression, JsonFormat, QueryUpdate, RowOffset, RowSize,
     RowSizeHint, WebsocketFormat,
 };
-use bytestring::ByteString;
-use core::mem;
 use spacetimedb_sats::bsatn::{self, ToBsatn};
 use spacetimedb_sats::ser::serde::SerializeWrapper;
 use spacetimedb_sats::Serialize;
diff --git a/crates/sdk/Cargo.toml b/crates/sdk/Cargo.toml
@@ -19,6 +19,7 @@ anymap.workspace = true
 base64.workspace = true
 brotli.workspace = true
 bytes.workspace = true
+flate2.workspace = true
 futures.workspace = true
 futures-channel.workspace = true
 home.workspace = true
diff --git a/crates/sdk/src/compression.rs b/crates/sdk/src/compression.rs
@@ -0,0 +1,57 @@
+use crate::websocket::WsError;
+use spacetimedb_client_api_messages::websocket::{
+    BsatnFormat, CompressableQueryUpdate, QueryUpdate, SERVER_MSG_COMPRESSION_TAG_BROTLI,
+    SERVER_MSG_COMPRESSION_TAG_GZIP, SERVER_MSG_COMPRESSION_TAG_NONE,
+};
+use spacetimedb_sats::bsatn;
+use std::borrow::Cow;
+use std::io::{self, Read as _};
+use std::sync::Arc;
+
+fn brotli_decompress(bytes: &[u8]) -> Result<Vec<u8>, io::Error> {
+    let mut decompressed = Vec::new();
+    brotli::BrotliDecompress(&mut &bytes[..], &mut decompressed)?;
+    Ok(decompressed)
+}
+
+fn gzip_decompress(bytes: &[u8]) -> Result<Vec<u8>, io::Error> {
+    let mut decompressed = Vec::new();
+    let _ = flate2::read::GzDecoder::new(bytes).read_to_end(&mut decompressed)?;
+    Ok(decompressed)
+}
+
+pub(crate) fn maybe_decompress_cqu(cqu: CompressableQueryUpdate<BsatnFormat>) -> QueryUpdate<BsatnFormat> {
+    match cqu {
+        CompressableQueryUpdate::Uncompressed(qu) => qu,
+        CompressableQueryUpdate::Brotli(bytes) => {
+            let bytes = brotli_decompress(&bytes).unwrap();
+            bsatn::from_slice(&bytes).unwrap()
+        }
+        CompressableQueryUpdate::Gzip(bytes) => {
+            let bytes = gzip_decompress(&bytes).unwrap();
+            bsatn::from_slice(&bytes).unwrap()
+        }
+    }
+}
+
+/// Decompresses a `ServerMessage` encoded in BSATN into the raw BSATN
+/// for further deserialization.
+pub(crate) fn decompress_server_message(raw: &[u8]) -> Result<Cow<'_, [u8]>, WsError> {
+    let err_decompress = |scheme| {
+        move |source| WsError::Decompress {
+            scheme,
+            source: Arc::new(source),
+        }
+    };
+    match raw {
+        [] => Err(WsError::EmptyMessage),
+        [SERVER_MSG_COMPRESSION_TAG_NONE, bytes @ ..] => Ok(Cow::Borrowed(bytes)),
+        [SERVER_MSG_COMPRESSION_TAG_BROTLI, bytes @ ..] => brotli_decompress(bytes)
+            .map(Cow::Owned)
+            .map_err(err_decompress("brotli")),
+        [SERVER_MSG_COMPRESSION_TAG_GZIP, bytes @ ..] => {
+            gzip_decompress(bytes).map(Cow::Owned).map_err(err_decompress("gzip"))
+        }
+        [c, ..] => Err(WsError::UnknownCompressionScheme { scheme: *c }),
+    }
+}
diff --git a/crates/sdk/src/lib.rs b/crates/sdk/src/lib.rs
@@ -12,6 +12,7 @@
 
 mod callbacks;
 mod client_cache;
+mod compression;
 mod db_connection;
 mod metrics;
 mod spacetime_module;
diff --git a/crates/sdk/src/spacetime_module.rs b/crates/sdk/src/spacetime_module.rs
@@ -9,6 +9,7 @@ use crate::{
     subscription::{OnEndedCallback, SubscriptionHandleImpl},
     Event, ReducerEvent,
     __codegen::InternalError,
+    compression::maybe_decompress_cqu,
 };
 use bytes::Bytes;
 use spacetimedb_client_api_messages::websocket::{self as ws, RowListLen as _};
@@ -228,7 +229,7 @@ impl<Row: DeserializeOwned + Debug> TableUpdate<Row> {
         let mut inserts = Vec::new();
         let mut deletes = Vec::new();
         for update in raw_updates.updates {
-            let update = update.maybe_decompress();
+            let update = maybe_decompress_cqu(update);
             Self::parse_from_row_list(&mut deletes, &update.deletes)?;
             Self::parse_from_row_list(&mut inserts, &update.inserts)?;
         }
diff --git a/crates/sdk/src/websocket.rs b/crates/sdk/src/websocket.rs
