Cap outbound batches by message_batch_size

Author: shisoft

src/tests/compute/mod.rs

@@ -268,6 +268,65 @@ async fn test_compute_worker_scatter_and_apply() {
     assert_eq!(state_v3, 1);
 }
 
+#[tokio::test]
+async fn test_compute_worker_splits_outbound_batches_by_message_batch_size() {
+    let program = Arc::new(TestProgram::new());
+    let store = InMemoryStateStore::<i64, i64>::new();
+
+    let mut adjacency = LocalAdjacency::new();
+    let source = Id::new(1, 11);
+    let neighbors = [Id::new(1, 12), Id::new(1, 13), Id::new(1, 14), Id::new(1, 15), Id::new(1, 16)];
+    for neighbor in neighbors {
+        adjacency.add_edge(source, neighbor);
+    }
+
+    let transport = Arc::new(LocalTransport::new());
+
+    let mut worker = ComputeWorker::new(
+        1,
+        JobId {
+            coordinator_server_id: 1,
+            coordinator_job_id: 11,
+        },
+        program,
+        store,
+        adjacency,
+        transport,
+        AdjacencyConfig {
+            edge_schemas: vec![],
+            direction: Direction::Outbound,
+            type_list: None,
+            weight: WeightSpec::Hops,
+        },
+        JobOptions {
+            max_active_per_step: usize::MAX,
+            message_batch_size: 2,
+        },
+    );
+
+    let seed_msg = worker.program.codec().msg_encode(&1);
+    worker
+        .deliver_messages(vec![(source, seed_msg)])
+        .expect("deliver seed");
+
+    let output = worker
+        .run_superstep(0, &GlobalStepParams { step: 0 })
+        .expect("step 0");
+
+    assert_eq!(output.report.messages_sent, 5);
+    assert_eq!(output.report.batches_sent, 3);
+    assert_eq!(output.outbound.len(), 3);
+    assert!(output.outbound.iter().all(|batch| batch.server_id == 1));
+    assert_eq!(
+        output
+            .outbound
+            .iter()
+            .map(|batch| batch.messages.len())
+            .collect::<Vec<_>>(),
+        vec![2, 2, 1]
+    );
+}
+
 #[test]
 fn test_vertex_set_manifest_roundtrip() {
     let shard = VertexSetShard {

src/traversal/compute/engine.rs

@@ -14,7 +14,14 @@ use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct JobOptions {
+    /// Maximum number of active vertices processed in a single superstep.
     pub max_active_per_step: usize,
+    /// Maximum number of messages carried by a single `OutboundBatch` for one
+    /// destination server.
+    ///
+    /// When a server receives more emitted messages than this in one
+    /// superstep, the worker splits them into multiple `OutboundBatch`
+    /// records. Values less than 1 are clamped to 1 by the worker.
     pub message_batch_size: usize,
 }
 
@@ -37,6 +44,8 @@ pub struct SuperstepReport {
     pub active_vertices: usize,
     pub messages_sent: usize,
     pub bytes_sent: usize,
+    /// Number of outbound `OutboundBatch` chunks produced after grouping by
+    /// destination server and splitting by `JobOptions::message_batch_size`.
     pub batches_sent: usize,
     pub misrouted: usize,
 }
@@ -50,6 +59,11 @@ pub struct DeliverReport {
 #[derive(Debug, Default, Clone, Serialize, Deserialize)]
 pub struct SuperstepOutput {
     pub report: SuperstepReport,
+    /// Outbound message chunks ready for delivery.
+    ///
+    /// There may be multiple entries for the same `server_id` when one
+    /// superstep emits more than `JobOptions::message_batch_size` messages to
+    /// that destination server.
     pub outbound: Vec<OutboundBatch>,
 }
 
@@ -497,16 +511,36 @@ where
             .collect();
         self.retry_adjacency.extend(reduced.retry_adjacency.iter().copied());
 
+        let batch_size = self.options.message_batch_size.max(1);
         let outbound: Vec<OutboundBatch> = reduced
             .outbound
             .into_iter()
-            .map(|(server_id, messages)| {
+            .flat_map(|(server_id, messages)| {
                 report.messages_sent += messages.len();
                 report.bytes_sent += messages.iter().map(|(_, m)| m.len()).sum::<usize>();
-                OutboundBatch {
-                    server_id,
-                    messages,
+
+                let mut batches = Vec::new();
+                let mut current = Vec::with_capacity(batch_size.min(messages.len()));
+
+                for message in messages {
+                    current.push(message);
+                    if current.len() == batch_size {
+                        batches.push(OutboundBatch {
+                            server_id,
+                            messages: current,
+                        });
+                        current = Vec::with_capacity(batch_size);
+                    }
                 }
+
+                if !current.is_empty() {
+                    batches.push(OutboundBatch {
+                        server_id,
+                        messages: current,
+                    });
+                }
+
+                batches
             })
             .collect();
         report.batches_sent = outbound.len();