move building parts of websocket format to own module

Author: Centril

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

@@ -4,3 +4,4 @@ pub mod energy;
 pub mod http;
 pub mod name;
 pub mod websocket;
+mod websocket_building;

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

@@ -19,39 +19,29 @@ use bytes::Bytes;
 use bytestring::ByteString;
 use core::{
     fmt::Debug,
-    mem,
     ops::{Deref, Range},
 };
 use enum_as_inner::EnumAsInner;
 use smallvec::SmallVec;
 use spacetimedb_lib::{ConnectionId, Identity, TimeDuration, Timestamp};
 use spacetimedb_primitives::TableId;
 use spacetimedb_sats::{
-    bsatn::{self, ToBsatn},
+    bsatn,
     de::{Deserialize, Error},
     impl_deserialize, impl_serialize, impl_st,
-    ser::{serde::SerializeWrapper, Serialize},
+    ser::Serialize,
     AlgebraicType, SpacetimeType,
 };
 use std::{
-    io::{self, Read as _, Write as _},
+    io::{self, Read as _},
     sync::Arc,
 };
 
+pub use crate::websocket_building::*;
+
 pub const TEXT_PROTOCOL: &str = "v1.json.spacetimedb";
 pub const BIN_PROTOCOL: &str = "v1.bsatn.spacetimedb";
 
-/// A list of rows being built.
-pub trait RowListBuilder: Default {
-    type FinishedList;
-
-    /// Push a row to the list in a serialized format.
-    fn push(&mut self, row: impl ToBsatn + Serialize);
-
-    /// Finish the in flight list, throwing away the capability to mutate.
-    fn finish(self) -> Self::FinishedList;
-}
-
 pub trait RowListLen {
     /// Returns the length, in number of rows, not bytes, of the row list.
     fn len(&self) -> usize;
@@ -103,26 +93,6 @@ pub trait WebsocketFormat: Sized {
     type QueryUpdate: SpacetimeType + for<'de> Deserialize<'de> + Serialize + Debug + Clone + Send;
 }
 
-pub trait BuildableWebsocketFormat: WebsocketFormat {
-    /// The builder for [`Self::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) {
-        let mut num_rows = 0;
-        let mut list = Self::ListBuilder::default();
-        for elem in elems {
-            num_rows += 1;
-            list.push(elem);
-        }
-        (list.finish(), num_rows)
-    }
-
-    /// Convert a `QueryUpdate` into `Self::QueryUpdate`.
-    /// This allows some formats to e.g., compress the update.
-    fn into_query_update(qu: QueryUpdate<Self>, compression: Compression) -> Self::QueryUpdate;
-}
-
 /// Messages sent from the client to the server.
 ///
 /// Parametric over the reducer argument type to enable [`ClientMessage::map_args`].
@@ -785,25 +755,6 @@ impl WebsocketFormat for JsonFormat {
     type QueryUpdate = QueryUpdate<Self>;
 }
 
-impl BuildableWebsocketFormat for JsonFormat {
-    type ListBuilder = Self::List;
-
-    fn into_query_update(qu: QueryUpdate<Self>, _: Compression) -> Self::QueryUpdate {
-        qu
-    }
-}
-
-impl RowListBuilder for Vec<ByteString> {
-    type FinishedList = Self;
-    fn push(&mut self, row: impl ToBsatn + Serialize) {
-        let value = serde_json::to_string(&SerializeWrapper::new(row)).unwrap().into();
-        self.push(value);
-    }
-    fn finish(self) -> Self::FinishedList {
-        self
-    }
-}
-
 #[derive(Clone, Copy, Default, Debug, SpacetimeType)]
 #[sats(crate = spacetimedb_lib)]
 pub struct BsatnFormat;
@@ -814,30 +765,6 @@ impl WebsocketFormat for BsatnFormat {
     type QueryUpdate = CompressableQueryUpdate<Self>;
 }
 
-impl BuildableWebsocketFormat for BsatnFormat {
-    type ListBuilder = BsatnRowListBuilder;
-
-    fn into_query_update(qu: QueryUpdate<Self>, compression: Compression) -> Self::QueryUpdate {
-        let qu_len_would_have_been = bsatn::to_len(&qu).unwrap();
-
-        match decide_compression(qu_len_would_have_been, compression) {
-            Compression::None => CompressableQueryUpdate::Uncompressed(qu),
-            Compression::Brotli => {
-                let bytes = bsatn::to_vec(&qu).unwrap();
-                let mut out = Vec::new();
-                brotli_compress(&bytes, &mut out);
-                CompressableQueryUpdate::Brotli(out.into())
-            }
-            Compression::Gzip => {
-                let bytes = bsatn::to_vec(&qu).unwrap();
-                let mut out = Vec::new();
-                gzip_compress(&bytes, &mut out);
-                CompressableQueryUpdate::Gzip(out.into())
-            }
-        }
-    }
-}
-
 /// A specification of either a desired or decided compression algorithm.
 #[derive(serde::Deserialize, Default, PartialEq, Eq, Clone, Copy, Hash, Debug)]
 pub enum Compression {
@@ -850,54 +777,20 @@ pub enum Compression {
     Gzip,
 }
 
-pub fn decide_compression(len: usize, compression: Compression) -> Compression {
-    /// The threshold beyond which we start to compress messages.
-    /// 1KiB was chosen without measurement.
-    /// TODO(perf): measure!
-    const COMPRESS_THRESHOLD: usize = 1024;
-
-    if len > COMPRESS_THRESHOLD {
-        compression
-    } else {
-        Compression::None
-    }
-}
-
-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: i32 = 1;
-
-    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> {
     let mut decompressed = Vec::new();
     brotli::BrotliDecompress(&mut &bytes[..], &mut decompressed)?;
     Ok(decompressed)
 }
 
-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("should be able to gzip compress `bytes`");
-}
-
 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)
 }
 
-type RowSize = u16;
-type RowOffset = u64;
+pub type RowSize = u16;
+pub type RowOffset = u64;
 
 /// A packed list of BSATN-encoded rows.
 #[derive(SpacetimeType, Debug, Clone, Default)]
@@ -910,6 +803,17 @@ pub struct BsatnRowList {
     rows_data: Bytes,
 }
 
+impl BsatnRowList {
+    /// Returns a new row list where `rows_data` is the flattened byte array
+    /// containing the BSATN of each row, without any markers for where a row begins and end.
+    ///
+    /// The `size_hint` encodes the boundaries of each row in `rows_data`.
+    /// See [`RowSizeHint`] for more details on the encoding.
+    pub fn new(size_hint: RowSizeHint, rows_data: Bytes) -> Self {
+        Self { size_hint, rows_data }
+    }
+}
+
 /// NOTE(centril, 1.0): We might want to add a `None` variant to this
 /// where the client has to decode in a loop until `rows_data` has been exhausted.
 /// The use-case for this is clients who are bandwidth limited and where every byte counts.
@@ -1004,102 +908,3 @@ impl Iterator for BsatnRowListIter<'_> {
         self.list.get(index)
     }
 }
-
-/// A [`BsatnRowList`] that can be added to.
-#[derive(Default)]
-pub struct BsatnRowListBuilder {
-    /// A size hint about `rows_data`
-    /// 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>,
-}
-
-/// A [`RowSizeHint`] under construction.
-pub enum RowSizeHintBuilder {
-    /// We haven't seen any rows yet.
-    Empty,
-    /// Each row in `rows_data` is of the same fixed size as specified here
-    /// but we don't know whether the size fits in `RowSize`
-    /// and we don't know whether future rows will also have this size.
-    FixedSizeDyn(usize),
-    /// Each row in `rows_data` is of the same fixed size as specified here
-    /// and we know that this will be the case for future rows as well.
-    FixedSizeStatic(RowSize),
-    /// The offsets into `rows_data` defining the boundaries of each row.
-    /// Only stores the offset to the start of each row.
-    /// The ends of each row is inferred from the start of the next row, or `rows_data.len()`.
-    /// The behavior of this is identical to that of `PackedStr`.
-    RowOffsets(Vec<RowOffset>),
-}
-
-impl Default for RowSizeHintBuilder {
-    fn default() -> Self {
-        Self::Empty
-    }
-}
-
-impl RowListBuilder for BsatnRowListBuilder {
-    type FinishedList = BsatnRowList;
-
-    fn push(&mut self, row: impl ToBsatn + Serialize) {
-        use RowSizeHintBuilder::*;
-
-        // Record the length before. It will be the starting offset of `row`.
-        let len_before = self.rows_data.len();
-        // BSATN-encode the row directly to the buffer.
-        row.to_bsatn_extend(&mut self.rows_data).unwrap();
-
-        let encoded_len = || self.rows_data.len() - len_before;
-        let push_row_offset = |mut offsets: Vec<_>| {
-            offsets.push(len_before as u64);
-            RowOffsets(offsets)
-        };
-
-        let hint = mem::replace(&mut self.size_hint, Empty);
-        self.size_hint = match hint {
-            // Static size that is unchanging.
-            h @ FixedSizeStatic(_) => h,
-            // Dynamic size that is unchanging.
-            h @ FixedSizeDyn(size) if size == encoded_len() => h,
-            // Size mismatch for the dynamic fixed size.
-            // Now we must construct `RowOffsets` for all rows thus far.
-            // We know that `size != 0` here, as this was excluded when we had `Empty`.
-            FixedSizeDyn(size) => RowOffsets(collect_offsets_from_num_rows(1 + len_before / size, size)),
-            // Once there's a size for each row, we'll just add to it.
-            RowOffsets(offsets) => push_row_offset(offsets),
-            // First time a row is seen. Use `encoded_len()` as the hint.
-            // If we have a static layout, we'll always have a fixed size.
-            // Otherwise, let's start out with a potentially fixed size.
-            // In either case, if `encoded_len() == 0`, we have to store offsets,
-            // as we cannot recover the number of elements otherwise.
-            Empty => match row.static_bsatn_size() {
-                Some(0) => push_row_offset(Vec::new()),
-                Some(size) => FixedSizeStatic(size),
-                None => match encoded_len() {
-                    0 => push_row_offset(Vec::new()),
-                    size => FixedSizeDyn(size),
-                },
-            },
-        };
-    }
-
-    fn finish(self) -> Self::FinishedList {
-        let Self { size_hint, rows_data } = self;
-        let size_hint = match size_hint {
-            RowSizeHintBuilder::Empty => RowSizeHint::RowOffsets([].into()),
-            RowSizeHintBuilder::FixedSizeStatic(fs) => RowSizeHint::FixedSize(fs),
-            RowSizeHintBuilder::FixedSizeDyn(fs) => match u16::try_from(fs) {
-                Ok(fs) => RowSizeHint::FixedSize(fs),
-                Err(_) => RowSizeHint::RowOffsets(collect_offsets_from_num_rows(rows_data.len() / fs, fs).into()),
-            },
-            RowSizeHintBuilder::RowOffsets(ro) => RowSizeHint::RowOffsets(ro.into()),
-        };
-        let rows_data = rows_data.into();
-        BsatnRowList { 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()
-}