perf(moonbit): tune primitive stream buffering (#1664)

Author: peter-jerry-ye

crates/moonbit/src/async/trait.mbt

@@ -5,10 +5,11 @@ pub struct Sink[X] {
   priv is_open : () -> Bool
   priv has_cleanup : () -> Bool
   priv cleanup : ((X) -> Unit)?
+  priv write_window_size : Int
 }
 
 ///|
-let sink_write_window_size : Int = 64
+let default_sink_write_window_size : Int = 64
 
 ///|
 /// The write callback must either raise before consuming any value or return
@@ -21,8 +22,17 @@ pub fn[X] Sink::from_callbacks(
   close : async () -> Unit,
   is_open : () -> Bool,
   cleanup : ((X) -> Unit)?,
+  write_window_size : Int,
 ) -> Sink[X] {
-  { write, close, is_open, has_cleanup: () => cleanup is Some(_), cleanup }
+  guard write_window_size > 0
+  {
+    write,
+    close,
+    is_open,
+    has_cleanup: () => cleanup is Some(_),
+    cleanup,
+    write_window_size,
+  }
 }
 
 ///|
@@ -34,13 +44,13 @@ pub async fn[X] Sink::write(self : Sink[X], data : ArrayView[X]) -> Int {
   if data.length() == 0 {
     return 0
   }
-  let length = if data.length() < sink_write_window_size {
+  let length = if data.length() < self.write_window_size {
     data.length()
   } else {
-    sink_write_window_size
+    self.write_window_size
   }
   // The callback may suspend, so it must not retain the caller's borrowed view.
-  let buffer = FixedArray::makei(length, i => data[i])
+  let buffer = FixedArray::from_array(data[:length])
   let written = (self.write)(buffer[:])
   guard written >= 0 && written <= length
   if !(self.is_open)() && written < length && (self.has_cleanup)() {
@@ -60,12 +70,12 @@ pub async fn Sink::write_bytes(self : Sink[Byte], data : BytesView) -> Int {
   if data.length() == 0 {
     return 0
   }
-  let length = if data.length() < sink_write_window_size {
+  let length = if data.length() < self.write_window_size {
     data.length()
   } else {
-    sink_write_window_size
+    self.write_window_size
   }
-  let buffer = FixedArray::makei(length, i => data[i])
+  let buffer = data[:length].to_fixedarray()
   let written = (self.write)(buffer[:])
   guard written >= 0 && written <= length
   if !(self.is_open)() && written < length && (self.has_cleanup)() {
@@ -314,7 +324,7 @@ fn[X] fill_stream_pipe_from_writers(pipe : Ref[StreamPipe[X]]) -> Unit {
     guard take_stream_pipe_writer(pipe) is Some(writer) else { return }
     guard writer.value is Some(data) else { continue }
     let take = if available < data.length() { available } else { data.length() }
-    let chunk = FixedArray::makei(take, i => data[i])
+    let chunk = FixedArray::from_array(data[:take])
     pipe.val.chunks.push_back(chunk)
     pipe.val.buffered = pipe.val.buffered + take
     wake_stream_pipe_writer(writer, take)
@@ -415,7 +425,7 @@ async fn[X] stream_pipe_write(
       } else {
         data.length()
       }
-      let chunk = FixedArray::makei(take, i => data[i])
+      let chunk = FixedArray::from_array(data[:take])
       wake_stream_pipe_reader(reader, Some(chunk), false)
       return take
     }
@@ -424,13 +434,13 @@ async fn[X] stream_pipe_write(
   if pipe.val.capacity > 0 && pipe.val.buffered < pipe.val.capacity {
     let available = pipe.val.capacity - pipe.val.buffered
     let take = if available < data.length() { available } else { data.length() }
-    let chunk = FixedArray::makei(take, i => data[i])
+    let chunk = FixedArray::from_array(data[:take])
     pipe.val.chunks.push_back(chunk)
     pipe.val.buffered = pipe.val.buffered + take
     return take
   }
   let writer = StreamPipeWriter::{
-    value: Some(FixedArray::makei(data.length(), i => data[i])),
+    value: Some(FixedArray::from_array(data)),
     accepted: 0,
     coro: Some(current_coroutine()),
   }
@@ -463,7 +473,9 @@ async fn[X] stream_pipe_read(
       Some(head) => {
         let available = head.length() - pipe.val.head_pos
         let take = if count < available { count } else { available }
-        let result = FixedArray::makei(take, i => head[pipe.val.head_pos + i])
+        let result = FixedArray::from_array(
+          head[pipe.val.head_pos:pipe.val.head_pos + take],
+        )
         pipe.val.head_pos = pipe.val.head_pos + take
         pipe.val.buffered = pipe.val.buffered - take
         if pipe.val.head_pos >= head.length() {
@@ -484,7 +496,7 @@ async fn[X] stream_pipe_read(
       Some(writer) => {
         guard writer.value is Some(data) else { continue }
         let take = if count < data.length() { count } else { data.length() }
-        let result = FixedArray::makei(take, i => data[i])
+        let result = FixedArray::from_array(data[:take])
         wake_stream_pipe_writer(writer, take)
         return Some(result)
       }
@@ -740,6 +752,7 @@ fn[X] stream_pipe_sink(pipe : Ref[StreamPipe[X]]) -> Sink[X] {
         None => ()
       }
     }),
+    write_window_size: default_sink_write_window_size,
   }
 }
 

crates/moonbit/src/async_support.rs

@@ -21,6 +21,30 @@ use super::FunctionBindgen;
 use super::InterfaceGenerator;
 use super::wasm_type;
 
+const DEFAULT_STREAM_WINDOW_ELEMENTS: usize = 64;
+const PRIMITIVE_STREAM_BUFFER_BYTES: usize = 4 * 1024;
+
+fn is_fixed_primitive(resolve: &Resolve, ty: &Type) -> bool {
+    match ty {
+        Type::U8
+        | Type::S8
+        | Type::U16
+        | Type::S16
+        | Type::U32
+        | Type::S32
+        | Type::U64
+        | Type::S64
+        | Type::F32
+        | Type::F64
+        | Type::Char => true,
+        Type::Id(id) => match &resolve.types[*id].kind {
+            TypeDefKind::Type(ty) => is_fixed_primitive(resolve, ty),
+            _ => false,
+        },
+        _ => false,
+    }
+}
+
 #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
 enum PayloadFor {
     Future,
@@ -1233,8 +1257,15 @@ impl<'a> InterfaceGenerator<'a> {
             .map(|ty| self.world_gen.sizes.size(ty).size_wasm32())
             .unwrap_or(0);
         let read_chunk_owns_buffer = result_type.is_some_and(|ty| self.is_list_canonical(ty));
-        let staging_window = if payload_sites.is_empty() { 64 } else { 1 };
-        let max_read_count = 64;
+        let primitive_window = result_type
+            .filter(|ty| is_fixed_primitive(self.resolve, ty))
+            .map(|_| (PRIMITIVE_STREAM_BUFFER_BYTES / elem_size).max(1));
+        let staging_window = if payload_sites.is_empty() {
+            primitive_window.unwrap_or(DEFAULT_STREAM_WINDOW_ELEMENTS)
+        } else {
+            1
+        };
+        let max_read_count = primitive_window.unwrap_or(DEFAULT_STREAM_WINDOW_ELEMENTS);
 
         let EndpointPayloadFragments {
             lift,
@@ -2104,6 +2135,7 @@ fn wasm{symbol_name}StreamCommit(handle : Int) -> Unit {{
             () => close_writer_serialized(),
             () => !writer_closed.val,
             Some(cleanup_value),
+            {staging_window},
         )
         let relay_source = producer is None
         let run_producer = async fn() -> Unit {{

crates/moonbit/src/lib.rs

@@ -3445,6 +3445,76 @@ mod tests {
         );
     }
 
+    #[test]
+    fn fixed_primitive_streams_use_four_kibibyte_buffer_budget() {
+        for (wit_type, expected_window) in [
+            ("u8", 4096),
+            ("s8", 4096),
+            ("u16", 2048),
+            ("s16", 2048),
+            ("u32", 1024),
+            ("s32", 1024),
+            ("f32", 1024),
+            ("char", 1024),
+            ("u64", 512),
+            ("s64", 512),
+            ("f64", 512),
+            ("bool", 64),
+            ("string", 64),
+        ] {
+            let wit = format!(
+                r#"
+                package a:b;
+
+                interface api {{
+                    type scalar = {wit_type};
+                    exchange: func(input: stream<scalar>) -> stream<scalar>;
+                }}
+
+                world runner {{ import api; }}
+                "#
+            );
+            let files = generate(&wit, "runner");
+            let generated = files
+                .iter()
+                .map(|(_, contents)| String::from_utf8_lossy(contents))
+                .collect::<Vec<_>>()
+                .join("\n");
+            assert!(
+                generated.contains(&format!("let read_count = if count < {expected_window}")),
+                "unexpected read window for {wit_type}: {generated}"
+            );
+            assert!(
+                generated.contains(&format!(
+                    "let data_len = if data.length() < {expected_window}"
+                )),
+                "unexpected write window for {wit_type}: {generated}"
+            );
+            let compact = generated.split_whitespace().collect::<Vec<_>>().join(" ");
+            assert!(
+                compact.contains(&format!("Some(cleanup_value), {expected_window}, )")),
+                "sink did not receive the {expected_window}-element window for {wit_type}: \
+                 {generated}"
+            );
+        }
+
+        let runtime = generate(
+            r#"
+            package a:b;
+            world runner {
+                import exchange: func(input: stream<u8>) -> stream<u8>;
+            }
+            "#,
+            "runner",
+        );
+        let sink = file(&runtime, "async-core/async_trait.mbt");
+        assert!(sink.contains("priv write_window_size : Int"));
+        assert!(sink.contains("data.length() < self.write_window_size"));
+        assert!(sink.contains("FixedArray::from_array(data[:length])"));
+        assert!(sink.contains("data[:length].to_fixedarray()"));
+        assert!(!sink.contains("FixedArray::makei"));
+    }
+
     #[test]
     fn async_export_background_group_name_is_deconflicted() {
         let files = generate(