Add client and server ts information to websocket messages to surface websocket latency issues (#40827)

This PR adds the `clientTs` to the initial websocket connect message, which the sync worker uses to calculate the `client_clock_skew`, which is the difference between the client clock and server clock, plus whatever latency between client sending and server receiving the message.

We send the latest possible `serverTs` before the server `Transition` message gets sent down to the client. On the client, we can then understand how long it took for the `Transition` message to download by comparing the current client time with the `serverTs` and subtracting the `clientClockSkew`.

GitOrigin-RevId: 9b8d2693a40196a78ee08aa3e2f49c695d0172d0

Author: emmaling27

crates/convex/src/base_client/mod.rs

@@ -277,6 +277,8 @@ impl RemoteQuerySet {
             start_version,
             end_version,
             modifications,
+            client_clock_skew: _,
+            server_ts: _,
         } = transition
         else {
             panic!("not transition");

crates/convex/src/client/mod.rs

@@ -565,6 +565,8 @@ pub mod tests {
                         log_lines: LogLinesMessage(vec![]),
                     })
                     .collect(),
+                client_clock_skew: None,
+                server_ts: None,
             },
             end_version,
         )
@@ -719,6 +721,7 @@ pub mod tests {
                     connection_count: 0,
                     last_close_reason: "InitialConnect".to_string(),
                     max_observed_timestamp: None,
+                    client_ts: None,
                 },
                 ClientMessage::ModifyQuerySet {
                     base_version: 0,
@@ -852,6 +855,7 @@ pub mod tests {
                     connection_count: 0,
                     last_close_reason: "InitialConnect".to_string(),
                     max_observed_timestamp: None,
+                    client_ts: None,
                 },
                 ClientMessage::ModifyQuerySet {
                     base_version: 0,
@@ -912,6 +916,7 @@ pub mod tests {
                     connection_count: 0,
                     last_close_reason: "InitialConnect".to_string(),
                     max_observed_timestamp: None,
+                    client_ts: None,
                 },
                 ClientMessage::ModifyQuerySet {
                     base_version: 0,

crates/convex/src/sync/testing.rs

@@ -83,6 +83,7 @@ impl SyncProtocol for TestProtocolManager {
                 connection_count,
                 last_close_reason: "InitialConnect".to_string(),
                 max_observed_timestamp: None,
+                client_ts: None,
             })
             .await?;
 

crates/convex/src/sync/web_socket_manager.rs

@@ -341,6 +341,7 @@ impl WebSocketInternal {
             connection_count,
             last_close_reason,
             max_observed_timestamp,
+            client_ts: Some(0),
         };
         let msg = Message::Text(
             serde_json::Value::try_from(message)

crates/convex/sync_types/src/json.rs

@@ -173,6 +173,10 @@ enum ClientMessageJson {
         #[serde(default)]
         #[serde(skip_serializing_if = "Option::is_none")]
         max_observed_timestamp: Option<String>,
+
+        #[serde(default)]
+        #[serde(skip_serializing_if = "Option::is_none")]
+        client_ts: Option<i64>,
     },
     #[serde(rename_all = "camelCase")]
     ModifyQuerySet {
@@ -219,11 +223,13 @@ impl TryFrom<ClientMessage> for JsonValue {
                 connection_count,
                 last_close_reason,
                 max_observed_timestamp,
+                client_ts,
             } => ClientMessageJson::Connect {
                 session_id: format!("{}", session_id.as_hyphenated()),
                 connection_count,
                 last_close_reason: Some(last_close_reason),
                 max_observed_timestamp: max_observed_timestamp.map(|ts| u64_to_string(ts.into())),
+                client_ts: client_ts.map(|ts| ts as i64),
             },
             ClientMessage::ModifyQuerySet {
                 base_version,
@@ -303,6 +309,7 @@ impl TryFrom<JsonValue> for ClientMessage {
                 connection_count,
                 last_close_reason,
                 max_observed_timestamp,
+                client_ts,
             } => ClientMessage::Connect {
                 session_id: session_id.parse()?,
                 connection_count,
@@ -312,6 +319,7 @@ impl TryFrom<JsonValue> for ClientMessage {
                     .transpose()?
                     .map(Timestamp::try_from)
                     .transpose()?,
+                client_ts: client_ts.map(|ts| ts as u64),
             },
             ClientMessageJson::ModifyQuerySet {
                 base_version,
@@ -523,11 +531,15 @@ impl<V: Into<JsonValue>> From<ServerMessage<V>> for JsonValue {
                 start_version,
                 end_version,
                 modifications,
+                client_clock_skew,
+                server_ts,
             } => json!({
                 "type": "Transition",
                 "startVersion": JsonValue::from(start_version),
                 "endVersion": JsonValue::from(end_version),
                 "modifications": modifications.into_iter().map(JsonValue::from).collect::<Vec<JsonValue>>(),
+                "clientClockSkew": JsonValue::from(client_clock_skew),
+                "serverTs": JsonValue::from(server_ts),
             }),
             ServerMessage::MutationResponse {
                 request_id,
@@ -634,6 +646,8 @@ impl<V: TryFrom<JsonValue, Error = anyhow::Error>> TryFrom<JsonValue> for Server
                 start_version: JsonValue,
                 end_version: JsonValue,
                 modifications: Vec<JsonValue>,
+                client_clock_skew: Option<i64>,
+                server_ts: Option<i64>,
             },
             #[serde(rename_all = "camelCase")]
             MutationResponse {
@@ -671,13 +685,17 @@ impl<V: TryFrom<JsonValue, Error = anyhow::Error>> TryFrom<JsonValue> for Server
                 start_version,
                 end_version,
                 modifications,
+                client_clock_skew,
+                server_ts,
             } => ServerMessage::Transition {
                 start_version: start_version.try_into()?,
                 end_version: end_version.try_into()?,
                 modifications: modifications
                     .into_iter()
                     .map(|sm: JsonValue| sm.try_into())
                     .collect::<anyhow::Result<Vec<StateModification<V>>>>()?,
+                client_clock_skew,
+                server_ts: server_ts.map(Timestamp::try_from).transpose()?,
             },
             ServerMessageJson::MutationResponse {
                 request_id,
@@ -1054,6 +1072,8 @@ mod tests {
                 journal: None,
                 error_data: Some(TestValue(JsonValue::Null)),
             }],
+            client_clock_skew: Some(1),
+            server_ts: Some(Timestamp::must(1)),
         });
     }
 }

crates/convex/sync_types/src/types.rs

@@ -108,6 +108,7 @@ pub enum ClientMessage {
         connection_count: u32,
         last_close_reason: String,
         max_observed_timestamp: Option<Timestamp>,
+        client_ts: Option<u64>,
     },
     ModifyQuerySet {
         base_version: QuerySetVersion,
@@ -308,6 +309,12 @@ pub enum ServerMessage<V: 'static> {
             proptest(strategy = "prop::collection::vec(any::<StateModification<V>>(), 0..8)")
         )]
         modifications: Vec<StateModification<V>>,
+        /// The difference between the timestamp in `ClientMessage::Connect` and
+        /// the timestamp when the client message was received by the server.
+        client_clock_skew: Option<i64>,
+        /// The timestamp right before this message was sent back to the
+        /// client.
+        server_ts: Option<Timestamp>,
     },
     MutationResponse {
         request_id: SessionRequestSeqNumber,
@@ -333,6 +340,17 @@ pub enum ServerMessage<V: 'static> {
     Ping,
 }
 
+impl<V: 'static> ServerMessage<V> {
+    pub fn inject_server_ts(&mut self, ts: Timestamp) {
+        match self {
+            Self::Transition {
+                ref mut server_ts, ..
+            } => *server_ts = Some(ts),
+            _ => {},
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 #[cfg_attr(any(test, feature = "testing"), derive(proptest_derive::Arbitrary))]
 pub enum ErrorPayload<V: 'static> {

crates/local_backend/src/subs/mod.rs

@@ -214,12 +214,13 @@ async fn run_sync_socket(
                     }
                 },
                 maybe_message = server_rx.next().fuse() => {
-                    let (message, send_time) = match maybe_message {
+                    let (mut message, send_time) = match maybe_message {
                         Some(m) => m,
                         None => break 'top,
                     };
                     let delay = st.runtime.monotonic_now() - send_time;
                     log_websocket_message_out(&message, delay);
+                    message.inject_server_ts(st.runtime.generate_timestamp()?);
                     let serialized = serde_json::to_string(&JsonValue::from(message))?;
                     if tx.send(Message::Text(serialized.into())).await.is_err() {
                         break 'top;

crates/sync/src/tests.rs

@@ -174,6 +174,7 @@ impl<RT: Runtime> SyncTest<RT> {
                 connection_count: 0,
                 last_close_reason: "InitialConnect".to_string(),
                 max_observed_timestamp,
+                client_ts: None,
             },
             self.rt.monotonic_now(),
         ))?;
@@ -303,6 +304,7 @@ async fn test_basic_account(rt: TestRuntime) -> anyhow::Result<()> {
         start_version: start,
         end_version: end,
         modifications,
+        ..
     } = sync_worker.receive().await?);
     assert_eq!(start.query_set, 0);
     assert_eq!(end.query_set, 1);
@@ -328,6 +330,7 @@ async fn test_basic_account(rt: TestRuntime) -> anyhow::Result<()> {
         start_version: start,
         end_version: end,
         modifications,
+        ..
     } = sync_worker.receive().await?);
     assert_eq!(start.query_set, 1);
     assert_eq!(end.query_set, 1);
@@ -355,7 +358,8 @@ async fn test_basic_account(rt: TestRuntime) -> anyhow::Result<()> {
     must_let!(let ServerMessage::Transition {
         start_version: start,
         end_version: end,
-        modifications
+        modifications,
+        ..
     } = sync_worker.receive().await?);
     assert_eq!(start.query_set, 1);
     assert_eq!(end.query_set, 2);
@@ -380,6 +384,7 @@ async fn test_basic_account(rt: TestRuntime) -> anyhow::Result<()> {
         start_version: start,
         end_version: end,
         modifications,
+        ..
     } = sync_worker.receive().await?);
     assert_eq!(start.query_set, 2);
     assert_eq!(end.query_set, 2);
@@ -406,6 +411,7 @@ async fn test_basic_account(rt: TestRuntime) -> anyhow::Result<()> {
         start_version: start,
         end_version: end,
         modifications,
+        ..
     } = sync_worker.receive().await?);
     assert_eq!(start.query_set, 2);
     assert_eq!(end.query_set, 3);
@@ -450,6 +456,8 @@ async fn test_remove_in_progress_query(rt: TestRuntime) -> anyhow::Result<()> {
         start_version: start,
         end_version: end,
         modifications,
+        client_clock_skew: _,
+        server_ts: _,
     } = sync_worker.receive().await?);
     assert_eq!(start.query_set, 0);
     assert_eq!(end.query_set, 1);

crates/sync/src/worker.rs

@@ -247,6 +247,10 @@ pub struct SyncWorker<RT: Runtime> {
 
     on_connect: Option<(StatusTimer, Box<dyn FnOnce(SessionId) + Send>)>,
     partition_id: u64,
+
+    /// The difference between the client's clock and the server's clock, in
+    /// milliseconds. Includes latency between the client and server.
+    client_clock_skew: Option<i64>,
 }
 
 enum QueryResult {
@@ -300,6 +304,7 @@ impl<RT: Runtime> SyncWorker<RT> {
             modify_query_to_transition_timers: BTreeMap::new(),
             on_connect: Some((connect_timer(partition_id), on_connect)),
             partition_id,
+            client_clock_skew: None,
         }
     }
 
@@ -470,11 +475,18 @@ impl<RT: Runtime> SyncWorker<RT> {
                 last_close_reason,
                 max_observed_timestamp,
                 connection_count,
+                client_ts,
             } => {
                 if let Some((timer, on_connect)) = self.on_connect.take() {
                     timer.finish();
                     on_connect(session_id);
                 }
+
+                if let Some(ts) = client_ts {
+                    self.client_clock_skew =
+                        Some(ts as i64 - self.rt.unix_timestamp().as_ms_since_epoch()? as i64);
+                }
+
                 self.state.set_session_id(session_id);
                 if let Some(max_observed_timestamp) = max_observed_timestamp {
                     let latest_timestamp = *self
@@ -1001,6 +1013,8 @@ impl<RT: Runtime> SyncWorker<RT> {
             start_version: current_version,
             end_version: new_version,
             modifications: state_modifications.into_values().collect(),
+            client_clock_skew: self.client_clock_skew,
+            server_ts: None,
         };
         timer.finish();
         metrics::log_query_set_size(self.partition_id, self.state.num_queries());

crates/value/src/heap_size.rs

@@ -928,6 +928,8 @@ impl<V: HeapSize> HeapSize for ServerMessage<V> {
                 start_version,
                 end_version,
                 modifications,
+                client_clock_skew: _,
+                server_ts: _,
             } => {
                 start_version.heap_size()
                     + end_version.heap_size()