Align Rust v3 websocket transport with TypeScript

Author: joshua-spacetime

crates/client-api-messages/DEVELOP.md

@@ -20,11 +20,7 @@ spacetime generate -p spacetimedb-cli --lang <SDK lang> \
   --module-def ws_schema_v2.json
 ```
 
-For the v3 WebSocket transport schema:
-
-```sh
-cargo run --example get_ws_schema_v3 > ws_schema_v3.json
-spacetime generate -p spacetimedb-cli --lang <SDK lang> \
-  --out-dir <sdk WebSocket schema bindings dir> \
-  --module-def ws_schema_v3.json
-```
+Note, the v3 WebSocket protocol does not have a separate schema.
+It reuses the v2 message schema and only changes the websocket binary framing.
+In v2 for example, a WebSocket frame contained a single BSATN-encoded v2 message,
+but in v3, a single WebSocket frame may contain a batch of one or more v2 messages.

crates/client-api-messages/examples/get_ws_schema_v3.rs

@@ -1,28 +1,13 @@
-use bytes::Bytes;
-pub use spacetimedb_sats::SpacetimeType;
+//! Binary framing for websocket protocol v3.
+//!
+//! Unlike v2, v3 does not define a new outer message schema.
+//! A single binary websocket payload contains one or more BSATN-encoded
+//! [`crate::websocket::v2::ClientMessage`] values from client to server,
+//! or one or more consecutive BSATN-encoded [`crate::websocket::v2::ServerMessage`]
+//! values from server to client.
+//!
+//! Client and server may coalesce multiple messages into one websocket payload,
+//! or send them separately, regardless of what the other one does,
+//! so long as logical order is preserved.
 
 pub const BIN_PROTOCOL: &str = "v3.bsatn.spacetimedb";
-
-/// Transport envelopes sent by the client over the v3 websocket protocol.
-///
-/// The inner bytes are BSATN-encoded v2 [`crate::websocket::v2::ClientMessage`] values.
-#[derive(SpacetimeType, Debug)]
-#[sats(crate = spacetimedb_lib)]
-pub enum ClientFrame {
-    /// A single logical client message.
-    Single(Bytes),
-    /// Multiple logical client messages that should be processed in-order.
-    Batch(Box<[Bytes]>),
-}
-
-/// Transport envelopes sent by the server over the v3 websocket protocol.
-///
-/// The inner bytes are BSATN-encoded v2 [`crate::websocket::v2::ServerMessage`] values.
-#[derive(SpacetimeType, Debug)]
-#[sats(crate = spacetimedb_lib)]
-pub enum ServerFrame {
-    /// A single logical server message.
-    Single(Bytes),
-    /// Multiple logical server messages that should be processed in-order.
-    Batch(Box<[Bytes]>),
-}

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

@@ -23,8 +23,8 @@ use prometheus::{Histogram, IntGauge};
 use scopeguard::{defer, ScopeGuard};
 use serde::Deserialize;
 use spacetimedb::client::messages::{
-    serialize, serialize_v2, serialize_v3, IdentityTokenMessage, InUseSerializeBuffer, SerializeBuffer,
-    SwitchedServerMessage, ToProtocol,
+    serialize, serialize_v2, IdentityTokenMessage, InUseSerializeBuffer, SerializeBuffer, SwitchedServerMessage,
+    ToProtocol,
 };
 use spacetimedb::client::{
     ClientActorId, ClientConfig, ClientConnection, ClientConnectionReceiver, DataMessage, MessageExecutionError,
@@ -1306,7 +1306,7 @@ async fn ws_encode_task(
     let binary_message_serializer = match config.version {
         WsVersion::V1 => None,
         WsVersion::V2 => Some(serialize_v2 as BinarySerializeFn),
-        WsVersion::V3 => Some(serialize_v3 as BinarySerializeFn),
+        WsVersion::V3 => Some(serialize_v2 as BinarySerializeFn),
     };
     let buf_pool = ArrayQueue::new(BUF_POOL_CAPACITY);
     let mut in_use_bufs: Vec<ScopeGuard<InUseSerializeBuffer, _>> = Vec::with_capacity(BUF_POOL_CAPACITY);

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

@@ -1,31 +1,33 @@
 use super::{ClientConnection, DataMessage, MessageHandleError};
 use serde::de::Error as _;
-use spacetimedb_client_api_messages::websocket::{v2 as ws_v2, v3 as ws_v3};
 use spacetimedb_lib::bsatn;
 use std::time::Instant;
 
+const EMPTY_V3_PAYLOAD_ERR: &str = "v3 websocket binary payload must contain at least one v2 client message";
+
 pub async fn handle(client: &ClientConnection, message: DataMessage, timer: Instant) -> Result<(), MessageHandleError> {
     client.observe_websocket_request_message(&message);
-    let frame = match message {
-        DataMessage::Binary(message_buf) => bsatn::from_slice::<ws_v3::ClientFrame>(&message_buf)?,
+    match message {
+        DataMessage::Binary(message_buf) => {
+            let mut remaining = &message_buf[..];
+
+            if remaining.is_empty() {
+                return Err(bsatn::DecodeError::Other(EMPTY_V3_PAYLOAD_ERR.into()).into());
+            }
+
+            loop {
+                let message = bsatn::from_reader(&mut remaining)?;
+                super::message_handlers_v2::handle_decoded_message(client, message, timer).await?;
+                if remaining.is_empty() {
+                    break;
+                }
+            }
+        }
         DataMessage::Text(_) => {
             return Err(MessageHandleError::TextDecode(serde_json::Error::custom(
                 "v3 websocket does not support text messages",
             )))
         }
-    };
-
-    match frame {
-        ws_v3::ClientFrame::Single(message) => {
-            let message = bsatn::from_slice::<ws_v2::ClientMessage>(&message)?;
-            super::message_handlers_v2::handle_decoded_message(client, message, timer).await?;
-        }
-        ws_v3::ClientFrame::Batch(messages) => {
-            for message in messages {
-                let message = bsatn::from_slice::<ws_v2::ClientMessage>(&message)?;
-                super::message_handlers_v2::handle_decoded_message(client, message, timer).await?;
-            }
-        }
     }
 
     Ok(())

crates/core/src/client/message_handlers_v3.rs

@@ -10,7 +10,6 @@ use derive_more::From;
 use spacetimedb_client_api_messages::websocket::common::{self as ws_common, RowListLen as _};
 use spacetimedb_client_api_messages::websocket::v1::{self as ws_v1};
 use spacetimedb_client_api_messages::websocket::v2 as ws_v2;
-use spacetimedb_client_api_messages::websocket::v3 as ws_v3;
 use spacetimedb_datastore::execution_context::WorkloadType;
 use spacetimedb_lib::identity::RequestId;
 use spacetimedb_lib::ser::serde::SerializeWrapper;
@@ -209,33 +208,6 @@ pub fn serialize_v2(
     finalize_binary_serialize_buffer(buffer, srv_msg_len, compression)
 }
 
-/// Serialize `msg` into a [`DataMessage`] containing a [`ws_v3::ServerFrame::Single`]
-/// whose payload is a BSATN-encoded [`ws_v2::ServerMessage`].
-///
-/// This mirrors the v2 framing by prepending the compression tag and applying
-/// conditional compression when configured.
-pub fn serialize_v3(
-    bsatn_rlb_pool: &BsatnRowListBuilderPool,
-    mut buffer: SerializeBuffer,
-    msg: ws_v2::ServerMessage,
-    compression: ws_v1::Compression,
-) -> (InUseSerializeBuffer, Bytes) {
-    let mut inner = BytesMut::with_capacity(SERIALIZE_BUFFER_INIT_CAP);
-    bsatn::to_writer((&mut inner).writer().into_inner(), &msg).expect("should be able to bsatn encode v2 message");
-
-    // 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));
-
-    let frame = ws_v3::ServerFrame::Single(inner.freeze());
-    let srv_msg = buffer.write_with_tag(ws_common::SERVER_MSG_COMPRESSION_TAG_NONE, |w| {
-        bsatn::to_writer(w.into_inner(), &frame).expect("should be able to bsatn encode v3 server frame");
-    });
-    let srv_msg_len = srv_msg.len();
-
-    finalize_binary_serialize_buffer(buffer, srv_msg_len, compression)
-}
-
 #[derive(Debug, From)]
 pub enum SerializableMessage {
     QueryBinary(OneOffQueryResponseMessage<ws_v1::BsatnFormat>),

crates/core/src/client/messages.rs

@@ -149,9 +149,7 @@ const V3_PREFERRED_PROTOCOL_HEADER: http::HeaderValue =
 #[cfg(not(feature = "browser"))]
 const MAX_V3_OUTBOUND_FRAME_BYTES: usize = 256 * 1024;
 #[cfg(not(feature = "browser"))]
-const BSATN_SUM_TAG_BYTES: usize = 1;
-#[cfg(not(feature = "browser"))]
-const BSATN_LENGTH_PREFIX_BYTES: usize = 4;
+const EMPTY_V3_SERVER_PAYLOAD_ERR: &str = "v3 websocket binary payload must contain at least one v2 server message";
 
 fn parse_scheme(scheme: Option<Scheme>) -> Result<Scheme, UriError> {
     Ok(match scheme {
@@ -305,48 +303,39 @@ fn decode_v2_server_message(bytes: &[u8]) -> Result<ws::v2::ServerMessage, WsErr
     bsatn::from_slice(bytes).map_err(|source| WsError::DeserializeMessage { source })
 }
 
-/// Expands a v3 server frame into the ordered sequence of encoded inner v2
-/// server messages it carries.
 #[cfg(not(feature = "browser"))]
-fn flatten_server_frame(frame: ws::v3::ServerFrame) -> Box<[Bytes]> {
-    match frame {
-        ws::v3::ServerFrame::Single(message) => Box::new([message]),
-        ws::v3::ServerFrame::Batch(messages) => messages,
+fn parse_v3_server_messages(bytes: &[u8]) -> Result<Vec<ws::v2::ServerMessage>, WsError> {
+    let mut remaining = bytes;
+    if remaining.is_empty() {
+        return Err(WsError::DeserializeMessage {
+            source: bsatn::DecodeError::Other(EMPTY_V3_SERVER_PAYLOAD_ERR.into()),
+        });
+    }
+
+    let mut messages = Vec::new();
+    while !remaining.is_empty() {
+        messages.push(bsatn::from_reader(&mut remaining).map_err(|source| WsError::DeserializeMessage { source })?);
     }
+    Ok(messages)
 }
 
 /// Encodes one logical v2 client message into raw BSATN bytes.
 fn encode_v2_client_message_bytes(msg: &ws::v2::ClientMessage) -> Bytes {
     Bytes::from(bsatn::to_vec(msg).expect("should be able to bsatn encode v2 client message"))
 }
 
-/// Wraps one or more encoded v2 client messages in a v3 transport frame.
-#[cfg(not(feature = "browser"))]
-fn encode_v3_client_frame(messages: Vec<Bytes>) -> Bytes {
-    let frame = if messages.len() == 1 {
-        ws::v3::ClientFrame::Single(messages.into_iter().next().unwrap())
-    } else {
-        ws::v3::ClientFrame::Batch(messages.into_boxed_slice())
-    };
-    Bytes::from(bsatn::to_vec(&frame).expect("should be able to bsatn encode v3 client frame"))
-}
-
-/// Returns the encoded size of a v3 `Single` frame carrying `message`.
-#[cfg(not(feature = "browser"))]
-fn encoded_v3_single_frame_size(message: &Bytes) -> usize {
-    BSATN_SUM_TAG_BYTES + BSATN_LENGTH_PREFIX_BYTES + message.len()
-}
-
-/// Returns the encoded size of a v3 `Batch` frame containing only its first logical message.
 #[cfg(not(feature = "browser"))]
-fn encoded_v3_batch_frame_size_for_first_message(message: &Bytes) -> usize {
-    BSATN_SUM_TAG_BYTES + BSATN_LENGTH_PREFIX_BYTES + BSATN_LENGTH_PREFIX_BYTES + message.len()
-}
+fn concatenate_v3_client_messages(messages: Vec<Bytes>) -> Bytes {
+    if messages.len() == 1 {
+        return messages.into_iter().next().unwrap();
+    }
 
-/// Returns the encoded contribution of one additional logical message inside a v3 `Batch` frame.
-#[cfg(not(feature = "browser"))]
-fn encoded_v3_batch_element_size(message: &Bytes) -> usize {
-    BSATN_LENGTH_PREFIX_BYTES + message.len()
+    let total_len = messages.iter().map(Bytes::len).sum();
+    let mut encoded = Vec::with_capacity(total_len);
+    for message in messages {
+        encoded.extend_from_slice(&message);
+    }
+    encoded.into()
 }
 
 /// Builds one bounded v3 transport frame from `first_message` and as many
@@ -363,25 +352,25 @@ where
     let first_message = encode_v2_client_message_bytes(&first_message);
     // Oversized logical messages are still sent alone so they cannot block the
     // queue forever behind the frame-size limit.
-    if encoded_v3_single_frame_size(&first_message) > MAX_V3_OUTBOUND_FRAME_BYTES {
-        if pending_outgoing.is_empty()
-            && let Some(next_message) = try_next_outgoing_now()
-        {
-            pending_outgoing.push_front(next_message);
+    if first_message.len() > MAX_V3_OUTBOUND_FRAME_BYTES {
+        if pending_outgoing.is_empty() {
+            if let Some(next_message) = try_next_outgoing_now() {
+                pending_outgoing.push_front(next_message);
+            }
         }
 
-        return encode_v3_client_frame(vec![first_message]);
+        return first_message;
     }
 
     let mut messages = vec![first_message];
-    let mut batch_size = encoded_v3_batch_frame_size_for_first_message(messages.first().unwrap());
+    let mut batch_size = messages.first().unwrap().len();
 
     loop {
         let Some(next_message) = pending_outgoing.pop_front().or_else(&mut try_next_outgoing_now) else {
             break;
         };
         let next_message_bytes = encode_v2_client_message_bytes(&next_message);
-        let next_batch_size = batch_size + encoded_v3_batch_element_size(&next_message_bytes);
+        let next_batch_size = batch_size + next_message_bytes.len();
         if next_batch_size > MAX_V3_OUTBOUND_FRAME_BYTES {
             pending_outgoing.push_front(next_message);
             break;
@@ -390,7 +379,7 @@ where
         messages.push(next_message_bytes);
     }
 
-    encode_v3_client_frame(messages)
+    concatenate_v3_client_messages(messages)
 }
 
 /// Encodes the next outbound logical message according to the negotiated
@@ -583,15 +572,7 @@ impl WsConnection {
         let bytes = &*decompress_server_message(bytes)?;
         match protocol {
             NegotiatedWsProtocol::V2 => Ok(vec![decode_v2_server_message(bytes)?]),
-            NegotiatedWsProtocol::V3 => {
-                let frame: ws::v3::ServerFrame =
-                    bsatn::from_slice(bytes).map_err(|source| WsError::DeserializeMessage { source })?;
-                flatten_server_frame(frame)
-                    .into_vec()
-                    .into_iter()
-                    .map(|message| decode_v2_server_message(&message))
-                    .collect()
-            }
+            NegotiatedWsProtocol::V3 => parse_v3_server_messages(bytes),
         }
     }
 
@@ -898,12 +879,22 @@ mod tests {
         encoded
     }
 
-    fn encode_server_frame(frame: &ws::v3::ServerFrame) -> Vec<u8> {
+    fn encode_server_messages(messages: &[ws::v2::ServerMessage]) -> Vec<u8> {
         let mut encoded = vec![ws::common::SERVER_MSG_COMPRESSION_TAG_NONE];
-        encoded.extend(bsatn::to_vec(frame).unwrap());
+        for message in messages {
+            encoded.extend(bsatn::to_vec(message).unwrap());
+        }
         encoded
     }
 
+    fn decode_client_messages(mut bytes: &[u8]) -> Vec<ws::v2::ClientMessage> {
+        let mut messages = Vec::new();
+        while !bytes.is_empty() {
+            messages.push(bsatn::from_reader(&mut bytes).unwrap());
+        }
+        messages
+    }
+
     #[test]
     fn negotiated_protocol_defaults_to_v2() {
         assert_eq!(
@@ -940,11 +931,14 @@ mod tests {
         assert!(!has_leftover_pending_outgoing);
         assert!(pending.is_empty());
 
-        let frame: ws::v3::ClientFrame = bsatn::from_slice(&raw).unwrap();
-        let ws::v3::ClientFrame::Batch(messages) = frame else {
-            panic!("expected batched v3 client frame");
-        };
+        let messages = decode_client_messages(&raw);
         assert_eq!(messages.len(), 2);
+        for (expected_request_id, message) in [1, 2].into_iter().zip(messages) {
+            match message {
+                ws::v2::ClientMessage::CallReducer(call) => assert_eq!(call.request_id, expected_request_id),
+                _ => panic!("expected CallReducer v3 client message"),
+            }
+        }
     }
 
     #[test]
@@ -960,12 +954,9 @@ mod tests {
         assert!(has_leftover_pending_outgoing);
         assert_eq!(pending.len(), 1);
 
-        let frame: ws::v3::ClientFrame = bsatn::from_slice(&raw).unwrap();
-        let ws::v3::ClientFrame::Single(message) = frame else {
-            panic!("expected single v3 client frame");
-        };
-        let inner: ws::v2::ClientMessage = bsatn::from_slice(&message).unwrap();
-        match inner {
+        let messages = decode_client_messages(&raw);
+        assert_eq!(messages.len(), 1);
+        match &messages[0] {
             ws::v2::ClientMessage::CallReducer(call) => assert_eq!(call.request_id, 1),
             _ => panic!("expected CallReducer inner message"),
         }
@@ -994,14 +985,7 @@ mod tests {
     fn parse_response_unwraps_v3_batches() {
         let first = procedure_result(1);
         let second = procedure_result(2);
-        let frame = ws::v3::ServerFrame::Batch(
-            vec![
-                Bytes::from(bsatn::to_vec(&first).unwrap()),
-                Bytes::from(bsatn::to_vec(&second).unwrap()),
-            ]
-            .into_boxed_slice(),
-        );
-        let encoded = encode_server_frame(&frame);
+        let encoded = encode_server_messages(&[first, second]);
 
         let messages = WsConnection::parse_responses(NegotiatedWsProtocol::V3, &encoded).unwrap();
         assert_eq!(messages.len(), 2);