fix a couple of partial read/write bugs (#11981)

* reset read/write state back to `Open` on event delivery

If one end of a stream does a partial read or write, we leave the other end in a
`GuestReady` state, allowing further reads or writes to proceed until the buffer
has been drained or filled, respectively.  However, once we've delivered the
event regarding the partial operation, we need to set the state back to `Open`,
since we'll have released the buffer back to the guest at that point.

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

* delay returning `Dropped` until producer buffer drained

If the `StreamProducer` calls `Destination::set_buffer`, we need to make sure
all the items in that buffer have been delivered to the receiver (or the
receiver closes its end) before telling it the write end has been dropped.

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

* add short reads tests

These cover a couple of scenarios where the guest and/or host read owned
resource items one-at-a-time from writes of more than one item, forcing the
writer to re-take ownership of the unwritten items between writes.

This also covers the case where the host's `StreamProducer` returns
`StreamResult::Dropped` after calling `Destination::set_buffer`, in which case
Wasmtime must delay telling the other end about the dropped stream until that
buffer has been drained.

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

---------

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

Author: dicej

crates/misc/component-async-tests/tests/scenario/streams.rs

@@ -6,20 +6,23 @@ use {
         util::{OneshotConsumer, OneshotProducer, PipeConsumer, PipeProducer},
     },
     futures::{
-        Sink, SinkExt, Stream, StreamExt,
+        FutureExt, Sink, SinkExt, Stream, StreamExt,
         channel::{mpsc, oneshot},
         future,
     },
     std::{
+        mem,
+        ops::DerefMut,
         pin::Pin,
         sync::{Arc, Mutex},
-        task::{Context, Poll},
+        task::{self, Context, Poll},
+        time::Duration,
     },
     wasmtime::{
         Engine, Store, StoreContextMut,
         component::{
-            Destination, FutureReader, Linker, ResourceTable, Source, StreamConsumer,
-            StreamProducer, StreamReader, StreamResult,
+            Destination, FutureReader, Lift, Linker, ResourceTable, Source, StreamConsumer,
+            StreamProducer, StreamReader, StreamResult, VecBuffer,
         },
     },
     wasmtime_wasi::WasiCtxBuilder,
@@ -296,3 +299,172 @@ pub async fn async_closed_stream() -> Result<()> {
         })
         .await?
 }
+
+struct VecProducer<T> {
+    source: Vec<T>,
+    sleep: Pin<Box<dyn Future<Output = ()> + Send>>,
+}
+
+impl<T> VecProducer<T> {
+    fn new(source: Vec<T>, delay: bool) -> Self {
+        Self {
+            source,
+            sleep: if delay {
+                tokio::time::sleep(Duration::from_millis(10)).boxed()
+            } else {
+                async {}.boxed()
+            },
+        }
+    }
+}
+
+impl<D, T: Lift + Unpin + 'static> StreamProducer<D> for VecProducer<T> {
+    type Item = T;
+    type Buffer = VecBuffer<T>;
+
+    fn poll_produce(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        _: StoreContextMut<D>,
+        mut destination: Destination<Self::Item, Self::Buffer>,
+        _: bool,
+    ) -> Poll<Result<StreamResult>> {
+        let sleep = &mut self.as_mut().get_mut().sleep;
+        task::ready!(sleep.as_mut().poll(cx));
+        *sleep = async {}.boxed();
+
+        destination.set_buffer(mem::take(&mut self.get_mut().source).into());
+        Poll::Ready(Ok(StreamResult::Dropped))
+    }
+}
+
+struct OneAtATime<T> {
+    destination: Arc<Mutex<Vec<T>>>,
+    sleep: Pin<Box<dyn Future<Output = ()> + Send>>,
+}
+
+impl<T> OneAtATime<T> {
+    fn new(destination: Arc<Mutex<Vec<T>>>, delay: bool) -> Self {
+        Self {
+            destination,
+            sleep: if delay {
+                tokio::time::sleep(Duration::from_millis(10)).boxed()
+            } else {
+                async {}.boxed()
+            },
+        }
+    }
+}
+
+impl<D, T: Lift + 'static> StreamConsumer<D> for OneAtATime<T> {
+    type Item = T;
+
+    fn poll_consume(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        store: StoreContextMut<D>,
+        mut source: Source<Self::Item>,
+        _: bool,
+    ) -> Poll<Result<StreamResult>> {
+        let sleep = &mut self.as_mut().get_mut().sleep;
+        task::ready!(sleep.as_mut().poll(cx));
+        *sleep = async {}.boxed();
+
+        let value = &mut None;
+        source.read(store, value)?;
+        self.destination.lock().unwrap().push(value.take().unwrap());
+        Poll::Ready(Ok(StreamResult::Completed))
+    }
+}
+
+mod short_reads {
+    wasmtime::component::bindgen!({
+        path: "wit",
+        world: "short-reads-guest",
+        exports: { default: async | task_exit },
+    });
+}
+
+#[tokio::test]
+pub async fn async_short_reads() -> Result<()> {
+    test_async_short_reads(false).await
+}
+
+#[tokio::test]
+async fn async_short_reads_with_delay() -> Result<()> {
+    test_async_short_reads(true).await
+}
+
+async fn test_async_short_reads(delay: bool) -> Result<()> {
+    use short_reads::exports::local::local::short_reads::Thing;
+
+    let engine = Engine::new(&config())?;
+
+    let component = make_component(
+        &engine,
+        &[test_programs_artifacts::ASYNC_SHORT_READS_COMPONENT],
+    )
+    .await?;
+
+    let mut linker = Linker::new(&engine);
+
+    wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
+
+    let mut store = Store::new(
+        &engine,
+        Ctx {
+            wasi: WasiCtxBuilder::new().inherit_stdio().build(),
+            table: ResourceTable::default(),
+            continue_: false,
+            wakers: Arc::new(Mutex::new(None)),
+        },
+    );
+
+    let guest =
+        short_reads::ShortReadsGuest::instantiate_async(&mut store, &component, &linker).await?;
+    let thing = guest.local_local_short_reads().thing();
+
+    let strings = ["a", "b", "c", "d", "e"];
+    let mut things = Vec::with_capacity(strings.len());
+    for string in strings {
+        things.push(thing.call_constructor(&mut store, string).await?);
+    }
+
+    store
+        .run_concurrent(async |store| {
+            let count = things.len();
+            let stream =
+                store.with(|store| StreamReader::new(store, VecProducer::new(things, delay)));
+
+            let (stream, task) = guest
+                .local_local_short_reads()
+                .call_short_reads(store, stream)
+                .await?;
+
+            let received_things = Arc::new(Mutex::new(Vec::<Thing>::with_capacity(count)));
+            // Read just one item at a time from the guest, forcing it to
+            // re-take ownership of any unwritten items.
+            store.with(|store| stream.pipe(store, OneAtATime::new(received_things.clone(), delay)));
+
+            task.block(store).await;
+
+            assert_eq!(count, received_things.lock().unwrap().len());
+
+            let mut received_strings = Vec::with_capacity(strings.len());
+            let received_things = mem::take(received_things.lock().unwrap().deref_mut());
+            for it in received_things {
+                received_strings.push(thing.call_get(store, it).await?.0);
+            }
+
+            assert_eq!(
+                &strings[..],
+                &received_strings
+                    .iter()
+                    .map(|s| s.as_str())
+                    .collect::<Vec<_>>()
+            );
+
+            anyhow::Ok(())
+        })
+        .await?
+}

crates/misc/component-async-tests/tests/test_all.rs

@@ -30,7 +30,7 @@ use scenario::round_trip_many::{
     async_round_trip_many_stackful, async_round_trip_many_stackless,
     async_round_trip_many_synchronous, async_round_trip_many_wait,
 };
-use scenario::streams::{async_closed_stream, async_closed_streams};
+use scenario::streams::{async_closed_stream, async_closed_streams, async_short_reads};
 use scenario::transmit::{
     async_cancel_callee, async_cancel_caller, async_cancel_transmit, async_intertask_communication,
     async_poll_stackless, async_poll_synchronous, async_readiness, async_synchronous_transmit,

crates/misc/component-async-tests/wit/test.wit

@@ -180,6 +180,15 @@ interface closed-stream {
   get: func() -> stream<u8>;
 }
 
+interface short-reads {
+  resource thing {
+    constructor(s: string);
+    get: async func() -> string;
+  }
+
+  short-reads: async func(s: stream<thing>) -> stream<thing>;
+}
+
 world yield-caller {
   import continue;
   import ready;
@@ -348,3 +357,7 @@ world sleep-post-return-caller {
 world closed-stream-guest {
   export closed-stream;
 }
+
+world short-reads-guest {
+  export short-reads;
+}

crates/test-programs/src/bin/async_short_reads.rs

@@ -0,0 +1,69 @@
+use {
+    bindings::{
+        exports::local::local::short_reads::{self, Guest, GuestThing},
+        wit_stream,
+    },
+    wit_bindgen::{StreamReader, StreamResult, rt::async_support},
+};
+
+mod bindings {
+    wit_bindgen::generate!({
+        path: "../misc/component-async-tests/wit",
+        world: "short-reads-guest",
+    });
+
+    use super::Component;
+    export!(Component);
+}
+
+struct Thing {
+    value: String,
+}
+
+impl GuestThing for Thing {
+    fn new(value: String) -> Self {
+        Self { value }
+    }
+
+    async fn get(&self) -> String {
+        self.value.clone()
+    }
+}
+
+struct Component;
+
+impl Guest for Component {
+    type Thing = Thing;
+
+    async fn short_reads(
+        mut stream: StreamReader<short_reads::Thing>,
+    ) -> StreamReader<short_reads::Thing> {
+        let (mut tx, rx) = wit_stream::new();
+
+        async_support::spawn(async move {
+            // Read the things one at a time, forcing the host to re-take
+            // ownership of any unwritten items between writes.
+            let mut things = Vec::new();
+            loop {
+                let (status, buffer) = stream.read(Vec::with_capacity(1)).await;
+                match status {
+                    StreamResult::Complete(_) => {
+                        things.extend(buffer);
+                    }
+                    StreamResult::Dropped => break,
+                    StreamResult::Cancelled => unreachable!(),
+                }
+            }
+            // Write the things all at once.  The host will read them only one
+            // at a time, forcing us to re-take ownership of any unwritten
+            // items between writes.
+            things = tx.write_all(things).await;
+            assert!(things.is_empty());
+        });
+
+        rx
+    }
+}
+
+// Unused function; required since this file is built as a `bin`:
+fn main() {}

crates/wasmtime/src/runtime/component/concurrent.rs

@@ -54,9 +54,7 @@ use crate::component::func::{self, Func, Options};
 use crate::component::{HasData, HasSelf, Instance, Resource, ResourceTable, ResourceTableError};
 use crate::fiber::{self, StoreFiber, StoreFiberYield};
 use crate::store::{Store, StoreId, StoreInner, StoreOpaque, StoreToken};
-use crate::vm::component::{
-    CallContext, ComponentInstance, InstanceFlags, ResourceTables, TransmitLocalState,
-};
+use crate::vm::component::{CallContext, ComponentInstance, InstanceFlags, ResourceTables};
 use crate::vm::{AlwaysMut, SendSyncPtr, VMFuncRef, VMMemoryDefinition, VMStore};
 use crate::{AsContext, AsContextMut, FuncType, StoreContext, StoreContextMut, ValRaw, ValType};
 use anyhow::{Context as _, Result, anyhow, bail};
@@ -1568,7 +1566,7 @@ impl Instance {
                     "deliver event {event:?} to {guest_task:?} for {waitable:?} (handle {handle}); set {set:?}"
                 );
 
-                waitable.on_delivery(self.id().get_mut(store), event);
+                waitable.on_delivery(store, self, event);
 
                 Some((event, Some((waitable, handle))))
             } else {
@@ -4356,55 +4354,6 @@ impl Waitable {
         Ok(())
     }
 
-    /// Handle the imminent delivery of the specified event, e.g. by updating
-    /// the state of the stream or future.
-    fn on_delivery(&self, instance: Pin<&mut ComponentInstance>, event: Event) {
-        match event {
-            Event::FutureRead {
-                pending: Some((ty, handle)),
-                ..
-            }
-            | Event::FutureWrite {
-                pending: Some((ty, handle)),
-                ..
-            } => {
-                let runtime_instance = instance.component().types()[ty].instance;
-                let (rep, state) = instance.guest_tables().0[runtime_instance]
-                    .future_rep(ty, handle)
-                    .unwrap();
-                assert_eq!(rep, self.rep());
-                assert_eq!(*state, TransmitLocalState::Busy);
-                *state = match event {
-                    Event::FutureRead { .. } => TransmitLocalState::Read { done: false },
-                    Event::FutureWrite { .. } => TransmitLocalState::Write { done: false },
-                    _ => unreachable!(),
-                };
-            }
-            Event::StreamRead {
-                pending: Some((ty, handle)),
-                code,
-            }
-            | Event::StreamWrite {
-                pending: Some((ty, handle)),
-                code,
-            } => {
-                let runtime_instance = instance.component().types()[ty].instance;
-                let (rep, state) = instance.guest_tables().0[runtime_instance]
-                    .stream_rep(ty, handle)
-                    .unwrap();
-                assert_eq!(rep, self.rep());
-                assert_eq!(*state, TransmitLocalState::Busy);
-                let done = matches!(code, ReturnCode::Dropped(_));
-                *state = match event {
-                    Event::StreamRead { .. } => TransmitLocalState::Read { done },
-                    Event::StreamWrite { .. } => TransmitLocalState::Write { done },
-                    _ => unreachable!(),
-                };
-            }
-            _ => {}
-        }
-    }
-
     /// Remove this waitable from the instance's rep table.
     fn delete_from(&self, state: &mut ConcurrentState) -> Result<()> {
         match self {