fix(moonbit): serialize stream writer teardown

Author: peter-jerry-ye

crates/moonbit/src/async_support.rs

@@ -1984,7 +1984,11 @@ fn wasm{symbol_name}StreamCommit(handle: Int) -> Unit {{
                 wasmImport{symbol_name}DropWritable(writer)
             }}
         }}
-        defer close_writer()
+        let close_writer_serialized = async fn() -> Unit {{
+            writer_lock.acquire()
+            defer writer_lock.release()
+            close_writer()
+        }}
         let cleanup_value = fn(value : {result}) -> Unit {{
             let ptr = wasm{symbol_name}Malloc(1)
             wasm{symbol_name}Lower(value, ptr)
@@ -2054,11 +2058,7 @@ fn wasm{symbol_name}StreamCommit(handle: Int) -> Unit {{
                 // consumed the rest of this staging window.
                 data_len
             }},
-            () => {{
-                writer_lock.acquire()
-                defer writer_lock.release()
-                close_writer()
-            }},
+            () => close_writer_serialized(),
             () => !writer_closed.val,
             Some(cleanup_value),
         )
@@ -2096,6 +2096,13 @@ fn wasm{symbol_name}StreamCommit(handle: Int) -> Unit {{
                     }}
                 }}
         }}
+        // A retained Sink may still own an in-flight canonical write after its
+        // producer returns. Do not drop the writable endpoint until that write
+        // has reached a terminal event and released its staging buffer.
+        {ffi}protect_from_cancel(
+            () => close_writer_serialized(),
+            resume_on_cancel=true,
+        )
     }})
 }}
 

crates/moonbit/src/lib.rs

@@ -3585,8 +3585,13 @@ mod tests {
         assert!(ffi.contains("suspend_for_future_write_terminal"));
         assert!(ffi.contains("let data_len = if data.length() < 1"));
         assert!(ffi.contains("let writer_lock = @async-core.Mutex()"));
+        assert!(ffi.contains("let close_writer_serialized = async fn()"));
         assert!(ffi.contains("writer_lock.acquire()"));
         assert!(ffi.contains("defer writer_lock.release()"));
+        assert!(
+            ffi.contains("() => close_writer_serialized(),\n            resume_on_cancel=true,")
+        );
+        assert!(!ffi.contains("defer close_writer()"));
         assert!(ffi.contains("read_cleanup : @async-core.CondVar"));
         assert!(ffi.contains("self.read_cleanup.broadcast()"));
         assert!(ffi.contains("@async-core.cancel_future_read("));

tests/runtime/moonbit/stream-write-cancel/holder-service.rs

@@ -60,6 +60,13 @@ impl Guest for Component {
         result == StreamResult::Dropped && values.is_empty()
     }
 
+    async fn read_one_and_keep() -> bool {
+        let mut stream = HELD_STREAM.with(|stream| stream.borrow_mut().take().unwrap());
+        let (result, values) = stream.read(Vec::with_capacity(1)).await;
+        HELD_STREAM.with(|held| assert!(held.borrow_mut().replace(stream).is_none()));
+        result == StreamResult::Complete(1) && values.len() == 1
+    }
+
     async fn read_one_and_drop() -> bool {
         let mut stream = HELD_STREAM.with(|stream| stream.borrow_mut().take().unwrap());
         let (result, values) = stream.read(Vec::with_capacity(1)).await;

tests/runtime/moonbit/stream-write-cancel/runner.rs

@@ -24,20 +24,21 @@ impl Guest for Component {
             .is_pending());
         holder::wait_write_started().await;
         assert!(holder::write_started());
+        assert!(holder::read_one_and_keep().await);
 
         drop(task);
         wait_cancellation_observed().await;
         assert!(cancellation_observed());
         assert!(holder::writable_dropped().await);
         assert_eq!(holder::leaf_live_count(), 0);
-        assert_eq!(holder::leaf_drop_count(), 1);
+        assert_eq!(holder::leaf_drop_count(), 2);
 
         start_peer_drop().await;
         assert!(holder::read_one_and_drop().await);
         wait_producer_finished().await;
         assert!(second_write_started());
         assert!(producer_finished());
         assert_eq!(holder::leaf_live_count(), 0);
-        assert_eq!(holder::leaf_drop_count(), 5);
+        assert_eq!(holder::leaf_drop_count(), 6);
     }
 }

tests/runtime/moonbit/stream-write-cancel/test.mbt

@@ -25,17 +25,28 @@ let producer_done : (@async-core.Future[Unit], @async-core.Promise[Unit]) =
 
 ///|
 pub async fn run_until_cancelled(
-  _background_group : @async-core.TaskGroup[Unit]
+  background_group : @async-core.TaskGroup[Unit]
 ) -> Unit {
+  let writer_started : (@async-core.Future[Unit], @async-core.Promise[Unit]) =
+    @async-core.Future::new()
   let stream = @async-core.Stream::produce(async fn(sink) {
-    defer {
-      cancellation_seen.val = true
-      ignore(cancellation_done.1.complete(()))
-    }
-    @holder.mark_write_started()
-    let data : FixedArray[@holder.Leaf] = [@holder.Leaf::leaf()]
-    let _ = sink.write_all(data[:])
-  })
+    // Return while a child owns a partially accepted write, forcing generated
+    // finalization to wait for the child's canonical buffer cleanup.
+    background_group.spawn_bg(async fn() {
+      defer {
+        cancellation_seen.val = true
+        ignore(cancellation_done.1.complete(()))
+      }
+      @holder.mark_write_started()
+      ignore(writer_started.1.complete(()))
+      let data : FixedArray[@holder.Leaf] = [
+        @holder.Leaf::leaf(),
+        @holder.Leaf::leaf(),
+      ]
+      let _ = sink.write_all(data[:])
+    })
+    writer_started.0.get()
+  }, cleanup=fn(value) { value.drop() })
   @holder.hold(stream)
   never.0.get()
 }

tests/runtime/moonbit/stream-write-cancel/test.rs

@@ -44,7 +44,9 @@ impl Guest for Component {
         let (mut writer, reader) = wit_stream::new::<holder::Leaf>();
         holder::hold(reader).await;
         holder::mark_write_started();
-        let _ = writer.write_one(holder::Leaf::new()).await;
+        let _ = writer
+            .write(vec![holder::Leaf::new(), holder::Leaf::new()])
+            .await;
         std::future::pending::<()>().await;
     }
 

tests/runtime/moonbit/stream-write-cancel/test.wit

@@ -13,6 +13,7 @@ interface holder {
   write-started: func() -> bool;
   wait-write-started: async func();
   writable-dropped: async func() -> bool;
+  read-one-and-keep: async func() -> bool;
   read-one-and-drop: async func() -> bool;
   leaf-live-count: func() -> u32;
   leaf-drop-count: func() -> u32;