Support "linear writes" in Rust future bindings (#1306)

To account for WebAssembly/wasip3-prototyping#172 and
WebAssembly/component-model#521 the Rust bindings need to take into
account that a future writer cannot be dropped without actually writing
something. To model this futures now must be created with a thunk that
creates the default value to be sent. This value is only used if another
value was not otherwise written.

Closes #1304

Author: alexcrichton

crates/guest-rust/rt/src/async_support/future_support.rs

@@ -86,6 +86,30 @@
 //! represented the same in Rust as it is in the canonical ABI. That means that
 //! sending `list<T>` into a future might require copying the entire list and
 //! changing its layout. Currently this is par-for-the-course with bindings.
+//!
+//! ## Linear (exactly once) Writes
+//!
+//! The component model requires that a writable end of a future must be written
+//! to before closing, otherwise the close operation traps. Ideally usage of
+//! this API shouldn't result in traps so this is modeled in the Rust-level API
+//! to prevent this trap from occurring. Rust does not support linear types
+//! (types that must be used exactly once), instead it only has affine types
+//! (types which must be used at most once), meaning that this requires some
+//! runtime support.
+//!
+//! Specifically the `FutureWriter` structure stores two auxiliary Rust-specific
+//! pieces of information:
+//!
+//! * A `should_write_default_value` boolean - if `true` on destruction then a
+//!   value has not yet been written and something must be written.
+//! * A `default: fn() -> T` constructor to lazily create the default value to
+//!   be sent in this situation.
+//!
+//! This `default` field is provided by the user when the future is initially
+//! created. Additionally during `Drop` a new Rust-level task will be spawned to
+//! perform the write in the background. That'll keep the component-level task
+//! alive until that write completes but otherwise shouldn't hinder anything
+//! else.
 
 use {
     super::waitable::{WaitableOp, WaitableOperation},
@@ -162,6 +186,7 @@ pub struct FutureVtable<T> {
 /// This function is unsafe as it requires the functions within `vtable` to
 /// correctly uphold the contracts of the component model.
 pub unsafe fn future_new<T>(
+    default: fn() -> T,
     vtable: &'static FutureVtable<T>,
 ) -> (FutureWriter<T>, FutureReader<T>) {
     unsafe {
@@ -170,7 +195,7 @@ pub unsafe fn future_new<T>(
         let writer = (handles >> 32) as u32;
         rtdebug!("future.new() = [{writer}, {reader}]");
         (
-            FutureWriter::new(writer, vtable),
+            FutureWriter::new(writer, default, vtable),
             FutureReader::new(reader, vtable),
         )
     }
@@ -183,6 +208,19 @@ pub unsafe fn future_new<T>(
 pub struct FutureWriter<T: 'static> {
     handle: u32,
     vtable: &'static FutureVtable<T>,
+
+    /// Whether or not a value should be written during `drop`.
+    ///
+    /// This is set to `false` when a value is successfully written or when a
+    /// value is written but the future is witnessed as being closed.
+    ///
+    /// Note that this is set to `true` on construction to ensure that only
+    /// location which actually witness a completed write set it to `false`.
+    should_write_default_value: bool,
+
+    /// Constructor for the default value to write during `drop`, should one
+    /// need to be written.
+    default: fn() -> T,
 }
 
 impl<T> FutureWriter<T> {
@@ -193,8 +231,13 @@ impl<T> FutureWriter<T> {
     /// This function is unsafe as it requires the functions within `vtable` to
     /// correctly uphold the contracts of the component model.
     #[doc(hidden)]
-    pub unsafe fn new(handle: u32, vtable: &'static FutureVtable<T>) -> Self {
-        Self { handle, vtable }
+    pub unsafe fn new(handle: u32, default: fn() -> T, vtable: &'static FutureVtable<T>) -> Self {
+        Self {
+            handle,
+            default,
+            should_write_default_value: true,
+            vtable,
+        }
     }
 
     /// Write the specified `value` to this `future`.
@@ -241,9 +284,33 @@ impl<T> fmt::Debug for FutureWriter<T> {
 
 impl<T> Drop for FutureWriter<T> {
     fn drop(&mut self) {
-        unsafe {
-            rtdebug!("future.close-writable({})", self.handle);
-            (self.vtable.close_writable)(self.handle);
+        // If a value has not yet been written into this writer than that must
+        // be done so now. Perform a "clone" of `self` by moving our data into a
+        // subtask, but ensure that `should_write_default_value` is set to
+        // `false` to avoid infinite loops by accident. Once the task is spawned
+        // we're done and the subtask's destructor of the closed-over
+        // `FutureWriter` will be responsible for performing the
+        // `close-writable` call below.
+        //
+        // Note, though, that if `should_write_default_value` is `false` then a
+        // write has already happened and we can go ahead and just synchronously
+        // drop this writer as we would any other handle.
+        if self.should_write_default_value {
+            let clone = FutureWriter {
+                handle: self.handle,
+                default: self.default,
+                should_write_default_value: false,
+                vtable: self.vtable,
+            };
+            crate::async_support::spawn(async move {
+                let value = (clone.default)();
+                let _ = clone.write(value).await;
+            });
+        } else {
+            unsafe {
+                rtdebug!("future.close-writable({})", self.handle);
+                (self.vtable.close_writable)(self.handle);
+            }
         }
     }
 }
@@ -299,7 +366,7 @@ where
     }
 
     fn in_progress_update(
-        (writer, cleanup): Self::InProgress,
+        (mut writer, cleanup): Self::InProgress,
         code: u32,
     ) -> Result<Self::Result, Self::InProgress> {
         let ptr = cleanup
@@ -321,6 +388,11 @@ where
                 // pass here.
                 let value = unsafe { (writer.vtable.lift)(ptr) };
                 let status = if c == ReturnCode::Closed(0) {
+                    // This writer has been witnessed to be closed, meaning that
+                    // `writer` is going to get destroyed soon as this return
+                    // value propagates up the stack. There's no need to write
+                    // the default value, so set this to `false`.
+                    writer.should_write_default_value = false;
                     WriteComplete::Closed(value)
                 } else {
                     WriteComplete::Cancelled(value)
@@ -338,6 +410,9 @@ where
             // Afterwards the `cleanup` itself is naturally dropped and cleaned
             // up.
             ReturnCode::Completed(1) | ReturnCode::Closed(1) | ReturnCode::Cancelled(1) => {
+                // A value was written, so no need to write the default value.
+                writer.should_write_default_value = false;
+
                 // SAFETY: we're the ones managing `ptr` so we know it's safe to
                 // pass here.
                 unsafe {
@@ -456,7 +531,7 @@ pub enum FutureWriteCancel<T: 'static> {
     Cancelled(T, FutureWriter<T>),
 }
 
-/// Represents the readable end of a Component Model `future`.
+/// Represents the readable end of a Component Model `future<T>`.
 pub struct FutureReader<T: 'static> {
     handle: AtomicU32,
     vtable: &'static FutureVtable<T>,
@@ -499,12 +574,11 @@ impl<T> FutureReader<T> {
 }
 
 impl<T> IntoFuture for FutureReader<T> {
-    type Output = Option<T>;
+    type Output = T;
     type IntoFuture = FutureRead<T>;
 
     /// Convert this object into a `Future` which will resolve when a value is
-    /// written to the writable end of this `future` (yielding a `Some` result)
-    /// or when the writable end is dropped (yielding a `None` result).
+    /// written to the writable end of this `future`.
     fn into_future(self) -> Self::IntoFuture {
         FutureRead {
             op: WaitableOperation::new(self),
@@ -536,7 +610,6 @@ struct FutureReadOp<T>(marker::PhantomData<T>);
 
 enum ReadComplete<T> {
     Value(T),
-    Closed,
     Cancelled,
 }
 
@@ -547,7 +620,7 @@ where
     type Start = FutureReader<T>;
     type InProgress = (FutureReader<T>, Option<Cleanup>);
     type Result = (ReadComplete<T>, FutureReader<T>);
-    type Cancel = Result<Option<T>, FutureReader<T>>;
+    type Cancel = Result<T, FutureReader<T>>;
 
     fn start(reader: Self::Start) -> (u32, Self::InProgress) {
         let (ptr, cleanup) = Cleanup::new(reader.vtable.layout);
@@ -570,13 +643,9 @@ where
         match ReturnCode::decode(code) {
             ReturnCode::Blocked => Err((reader, cleanup)),
 
-            // The read didn't complete, so `cleanup` is still uninitialized, so
-            // let it fall out of scope.
-            ReturnCode::Closed(0) => Ok((ReadComplete::Closed, reader)),
-
-            // Like `in_progress_closed` the read operation has finished but
-            // without a value, so let `cleanup` fall out of scope to clean up
-            // its allocation.
+            // Let `cleanup` fall out of scope to clean up its allocation here,
+            // and otherwise tahe reader is plumbed through to possibly restart
+            // the read in the future.
             ReturnCode::Cancelled(0) => Ok((ReadComplete::Cancelled, reader)),
 
             // The read has completed, so lift the value from the stored memory and
@@ -613,28 +682,27 @@ where
     fn result_into_cancel((value, reader): Self::Result) -> Self::Cancel {
         match value {
             // The value was actually read, so thread that through here.
-            ReadComplete::Value(value) => Ok(Some(value)),
+            ReadComplete::Value(value) => Ok(value),
 
             // The read was successfully cancelled, so thread through the
             // `reader` to possibly restart later on.
             ReadComplete::Cancelled => Err(reader),
-
-            // The other end was closed, so this can't possibly ever complete
-            // again, so thread that through.
-            ReadComplete::Closed => Ok(None),
         }
     }
 }
 
 impl<T: 'static> Future for FutureRead<T> {
-    type Output = Option<T>;
+    type Output = T;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         self.pin_project()
             .poll_complete(cx)
             .map(|(result, _reader)| match result {
-                ReadComplete::Value(val) => Some(val),
-                ReadComplete::Cancelled | ReadComplete::Closed => None,
+                ReadComplete::Value(val) => val,
+                // This is only possible if, after calling `FutureRead::cancel`,
+                // the future is polled again. The `cancel` method is documented
+                // as "don't do that" so this is left to panic.
+                ReadComplete::Cancelled => panic!("cannot poll after cancelling"),
             })
     }
 }
@@ -650,21 +718,17 @@ impl<T> FutureRead<T> {
     ///
     /// Return values include:
     ///
-    /// * `Ok(Some(value))` - future completed before this cancellation request
+    /// * `Ok(value)` - future completed before this cancellation request
     ///   was received.
-    /// * `Ok(None)` - future closed before this cancellation request was
-    ///   received.
     /// * `Err(reader)` - read operation was cancelled and it can be retried in
     ///   the future if desired.
     ///
-    /// Note that if this method is called after the write was already cancelled
-    /// then `Ok(None)` will be returned.
-    ///
     /// # Panics
     ///
     /// Panics if the operation has already been completed via `Future::poll`,
-    /// or if this method is called twice.
-    pub fn cancel(self: Pin<&mut Self>) -> Result<Option<T>, FutureReader<T>> {
+    /// or if this method is called twice. Additionally if this method completes
+    /// then calling `poll` again on `self` will panic.
+    pub fn cancel(self: Pin<&mut Self>) -> Result<T, FutureReader<T>> {
         self.pin_project().cancel()
     }
 }

crates/rust/src/lib.rs

@@ -452,8 +452,11 @@ pub mod wit_future {{
             self.src.push_str(&format!(
                 "\
     /// Creates a new Component Model `future` with the specified payload type.
-    pub fn new<T: FuturePayload>() -> ({async_support}::FutureWriter<T>, {async_support}::FutureReader<T>) {{
-        unsafe {{ {async_support}::future_new::<T>(T::VTABLE) }}
+    ///
+    /// The `default` function provided computes the default value to be sent in
+    /// this future if no other value was otherwise sent.
+    pub fn new<T: FuturePayload>(default: fn() -> T) -> ({async_support}::FutureWriter<T>, {async_support}::FutureReader<T>) {{
+        unsafe {{ {async_support}::future_new::<T>(default, T::VTABLE) }}
     }}
 }}
                 ",

tests/runtime-async/async/cancel-import/runner.rs

@@ -9,7 +9,7 @@ use wit_bindgen::yield_async;
 fn main() {
     println!("test cancelling an import in progress");
     wit_bindgen::block_on(async {
-        let (tx, rx) = wit_future::new();
+        let (tx, rx) = wit_future::new(|| unreachable!());
         let mut import = Box::pin(pending_import(rx));
         assert!(import
             .as_mut()
@@ -21,16 +21,16 @@ fn main() {
 
     println!("test cancelling an import before it starts");
     wit_bindgen::block_on(async {
-        let (tx, rx) = wit_future::new();
+        let (tx, rx) = wit_future::new(|| unreachable!());
         let import = Box::pin(pending_import(rx));
         drop(import);
         tx.write(()).await.unwrap_err();
     });
 
     println!("test cancelling an import in the started state");
     wit_bindgen::block_on(async {
-        let (tx1, rx1) = wit_future::new();
-        let (tx2, rx2) = wit_future::new();
+        let (tx1, rx1) = wit_future::new(|| unreachable!());
+        let (tx2, rx2) = wit_future::new(|| unreachable!());
 
         // create a task in the "started" state, but don't complete it yet
         let mut started_import = Box::pin(pending_import(rx1));
@@ -66,7 +66,7 @@ fn main() {
     println!("test cancellation with a status code saying it's done");
     wit_bindgen::block_on(async {
         // Start a subtask and get it into the "started" state
-        let (tx, rx) = wit_future::new();
+        let (tx, rx) = wit_future::new(|| unreachable!());
         let mut import = Box::pin(pending_import(rx));
         assert!(import
             .as_mut()
@@ -92,7 +92,7 @@ fn main() {
     println!("race cancellation with pending status code");
     wit_bindgen::block_on(async {
         // Start a subtask and get it into the "started" state
-        let (tx1, rx1) = wit_future::new();
+        let (tx1, rx1) = wit_future::new(|| unreachable!());
         let mut started_import = Box::pin(pending_import(rx1));
         assert!(started_import
             .as_mut()
@@ -102,7 +102,7 @@ fn main() {
         // force the next subtask to start out in the "starting" state, not the
         // "started" state.
         backpressure_set(true);
-        let (tx2, rx2) = wit_future::new();
+        let (tx2, rx2) = wit_future::new(|| unreachable!());
         let mut starting_import = Box::pin(pending_import(rx2));
         assert!(starting_import
             .as_mut()

tests/runtime-async/async/cancel-import/test.rs

@@ -8,7 +8,7 @@ export!(Component);
 
 impl crate::exports::my::test::i::Guest for Component {
     async fn pending_import(x: FutureReader<()>) {
-        x.await.unwrap();
+        x.await
     }
 
     fn backpressure_set(x: bool) {