Metrics for ClientMessage args size (#40989)

following up on #40977, adds metrics for argument size of ClientMessage variants

GitOrigin-RevId: 3e54728d67187e811021fa8d2f0af5fb92e394c8

Author: emmaling27

crates/sync/src/metrics.rs

@@ -99,6 +99,54 @@ pub fn mutation_queue_timer(partition_id: u64) -> StatusTimer {
     timer
 }
 
+register_convex_histogram!(
+    SYNC_QUERY_MODIFICATION_ARGS_BYTES,
+    "Size of query modification args in ClientMessages",
+    &["partition_id"]
+);
+pub fn log_query_modification_args_size(partition_id: u64, size: usize) {
+    log_distribution_with_labels(
+        &SYNC_QUERY_MODIFICATION_ARGS_BYTES,
+        size as f64,
+        vec![StaticMetricLabel::new(
+            "partition_id",
+            partition_id.to_string(),
+        )],
+    );
+}
+
+register_convex_histogram!(
+    SYNC_MUTATION_ARGS_BYTES,
+    "Size of mutation args in ClientMessages",
+    &["partition_id"]
+);
+pub fn log_mutation_args_size(partition_id: u64, size: usize) {
+    log_distribution_with_labels(
+        &SYNC_MUTATION_ARGS_BYTES,
+        size as f64,
+        vec![StaticMetricLabel::new(
+            "partition_id",
+            partition_id.to_string(),
+        )],
+    );
+}
+
+register_convex_histogram!(
+    SYNC_ACTION_ARGS_BYTES,
+    "Size of action args in ClientMessages",
+    &["partition_id"]
+);
+pub fn log_action_args_size(partition_id: u64, size: usize) {
+    log_distribution_with_labels(
+        &SYNC_ACTION_ARGS_BYTES,
+        size as f64,
+        vec![StaticMetricLabel::new(
+            "partition_id",
+            partition_id.to_string(),
+        )],
+    );
+}
+
 register_convex_counter!(
     SYNC_QUERY_FAILED_TOTAL,
     "Number of query failures",

crates/sync/src/worker.rs

@@ -101,6 +101,9 @@ use crate::{
     metrics::{
         self,
         connect_timer,
+        log_action_args_size,
+        log_mutation_args_size,
+        log_query_modification_args_size,
         modify_query_to_transition_timer,
         mutation_queue_timer,
         TypedClientEvent,
@@ -517,6 +520,16 @@ impl<RT: Runtime> SyncWorker<RT> {
                 new_version,
                 modifications,
             } => {
+                let total_args_size = modifications
+                    .iter()
+                    .filter_map(|m| match m {
+                        QuerySetModification::Add(q) => {
+                            Some(q.args.iter().map(|v| v.heap_size()).sum::<usize>())
+                        },
+                        QuerySetModification::Remove { .. } => None,
+                    })
+                    .sum();
+                log_query_modification_args_size(self.partition_id, total_args_size);
                 self.state
                     .modify_query_set(base_version, new_version, modifications)?;
                 self.schedule_update();
@@ -531,6 +544,7 @@ impl<RT: Runtime> SyncWorker<RT> {
                 args,
                 component_path,
             } => {
+                log_mutation_args_size(self.partition_id, args.iter().map(|v| v.heap_size()).sum());
                 let identity = self.state.identity(self.rt.system_time())?;
                 let mutation_identifier =
                     self.state.session_id().map(|id| SessionRequestIdentifier {
@@ -638,6 +652,7 @@ impl<RT: Runtime> SyncWorker<RT> {
                 args,
                 component_path,
             } => {
+                log_action_args_size(self.partition_id, args.iter().map(|v| v.heap_size()).sum());
                 let identity = self.state.identity(self.rt.system_time())?;
 
                 let api = self.api.clone();