Introducer every-n-messages ReducePolicy

Summary:
Publish to collect early feedback.

The implementation details will change because the previous diffs in this stack might change. But the overall idea is the same:

We only coalesce updates when the buffer size is bigger than the policy.

Reviewed By: mariusae

Differential Revision: D74691188

fbshipit-source-id: 11514551f8ec3bd286e5f89b7805560774987f20

Author: pzhan9

hyperactor/src/mailbox/mod.rs

@@ -1277,6 +1277,30 @@ impl cap::sealed::CanOpenPort for Mailbox {
     }
 }
 
+#[derive(Default)]
+struct SplitPortBuffer(Vec<Serialized>);
+
+impl SplitPortBuffer {
+    /// Push a new item to the buffer, and optionally return any items that should
+    /// be flushed.
+    fn push(&mut self, serialized: Serialized) -> Option<Vec<Serialized>> {
+        static HYPERACTOR_SPLIT_MAX_BUFFER_SIZE: OnceLock<usize> = OnceLock::new();
+        let limit = HYPERACTOR_SPLIT_MAX_BUFFER_SIZE.get_or_init(|| {
+            std::env::var("HYPERACTOR_SPLIT_MAX_BUFFER_SIZE")
+                .ok()
+                .and_then(|val| val.parse::<usize>().ok())
+                .unwrap_or(5)
+        });
+
+        self.0.push(serialized);
+        if &self.0.len() >= limit {
+            Some(std::mem::take(&mut self.0))
+        } else {
+            None
+        }
+    }
+}
+
 impl cap::sealed::CanSplitPort for Mailbox {
     fn split(&self, port_id: PortId, reducer_typehash: Option<u64>) -> PortId {
         fn post(mailbox: &Mailbox, port_id: PortId, msg: Serialized) {
@@ -1302,23 +1326,22 @@ impl cap::sealed::CanSplitPort for Mailbox {
                 Ok(())
             }),
             Some(r) => {
-                let buffer = Arc::new(Mutex::new(Vec::<Serialized>::new()));
+                let buffer = Mutex::new(SplitPortBuffer::default());
                 Box::new(move |serialized: Serialized| {
                     // Hold the lock until messages are sent. This is to avoid another
                     // invocation of this method trying to send message concurrently and
                     // cause messages delivered out of order.
                     let mut buf = buffer.lock().unwrap();
-                    buf.push(serialized);
-                    // TODO(pzhang) add policy and use this buffer
-                    let buffered = std::mem::take(&mut *buf);
-                    let reduced = r.reduce_updates(buffered).map_err(|(e, mut b)| {
-                        (
-                            b.pop()
-                                .expect("there should be at least one update from buffer"),
-                            e,
-                        )
-                    })?;
-                    post(&mailbox, port_id.clone(), reduced);
+                    if let Some(buffered) = buf.push(serialized) {
+                        let reduced = r.reduce_updates(buffered).map_err(|(e, mut b)| {
+                            (
+                                b.pop()
+                                    .expect("there should be at least one update from buffer"),
+                                e,
+                            )
+                        })?;
+                        post(&mailbox, port_id.clone(), reduced);
+                    }
                     Ok(())
                 })
             }
@@ -2208,6 +2231,7 @@ mod tests {
     use std::time::Duration;
 
     use timed_test::async_timed_test;
+    use tracing::Level;
 
     use super::*;
     use crate::Actor;
@@ -2224,6 +2248,7 @@ mod tests {
     use crate::proc::Proc;
     use crate::reference::ProcId;
     use crate::reference::WorldId;
+    use crate::test_utils::tracing::set_tracing_env_filter;
 
     #[test]
     fn test_error() {
@@ -2937,6 +2962,9 @@ mod tests {
 
     #[async_timed_test(timeout_secs = 30)]
     async fn test_split_port_id_sum_reducer() {
+        std::env::set_var("HYPERACTOR_SPLIT_MAX_BUFFER_SIZE", "1");
+        set_tracing_env_filter(Level::INFO);
+
         let sum_accumulator = accum::sum::<u64>();
         let reducer_typehash = sum_accumulator.reducer_typehash();
         let Setup {
@@ -2965,4 +2993,35 @@ mod tests {
         let msg = receiver.try_recv().unwrap();
         assert_eq!(msg, None);
     }
+
+    #[async_timed_test(timeout_secs = 30)]
+    async fn test_split_port_id_every_n_messages() {
+        set_tracing_env_filter(Level::INFO);
+        let actor = Mailbox::new(
+            id!(test[0].actor),
+            BoxedMailboxSender::new(PanickingMailboxSender),
+        );
+        let (port_handle, mut receiver) = actor.open_port::<u64>();
+        let port_id = port_handle.bind().port_id().clone();
+        // Split it
+        let reducer_typehash = accum::sum::<u64>().reducer_typehash();
+        let split_port_id = port_id.split(&actor, Some(reducer_typehash));
+
+        // Send 9 messages.
+        for msg in [1, 5, 3, 4, 2, 91, 92, 93, 94] {
+            post(&actor, split_port_id.clone(), msg);
+        }
+        // The first 5 should be batched and reduced once due
+        // to every_n_msgs = 5.
+        let messages = wait_for(&mut receiver, 1, Duration::from_secs(2))
+            .await
+            .unwrap();
+        assert_eq!(messages, vec![15]);
+
+        // the last message unfortranately will never come because they do not
+        // reach batch size.
+        RealClock.sleep(Duration::from_secs(2)).await;
+        let msg = receiver.try_recv().unwrap();
+        assert_eq!(msg, None);
+    }
 }