From 98336c6445b0079324de6d55d34e4c9a29f4451d Mon Sep 17 00:00:00 2001 From: yezihang Date: Fri, 24 Jul 2026 12:19:07 +0800 Subject: [PATCH] perf(moonbit): tune primitive stream buffering --- crates/moonbit/src/async/trait.mbt | 41 +++++++++++------ crates/moonbit/src/async_support.rs | 36 ++++++++++++++- crates/moonbit/src/lib.rs | 70 +++++++++++++++++++++++++++++ 3 files changed, 131 insertions(+), 16 deletions(-) diff --git a/crates/moonbit/src/async/trait.mbt b/crates/moonbit/src/async/trait.mbt index e219e6695..762c37c44 100644 --- a/crates/moonbit/src/async/trait.mbt +++ b/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, } } diff --git a/crates/moonbit/src/async_support.rs b/crates/moonbit/src/async_support.rs index 02c25a040..30988b245 100644 --- a/crates/moonbit/src/async_support.rs +++ b/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 {{ diff --git a/crates/moonbit/src/lib.rs b/crates/moonbit/src/lib.rs index 802c804c2..782c559da 100644 --- a/crates/moonbit/src/lib.rs +++ b/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) -> stream; + }} + + world runner {{ import api; }} + "# + ); + let files = generate(&wit, "runner"); + let generated = files + .iter() + .map(|(_, contents)| String::from_utf8_lossy(contents)) + .collect::>() + .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::>().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) -> stream; + } + "#, + "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(