p3: refactor future producers/consumers

Signed-off-by: Roman Volosatovs <rvolosatovs@riseup.net>

Author: rvolosatovs

crates/wasi-http/src/p3/body.rs

@@ -372,9 +372,9 @@ impl http_body::Body for ConsumedBody {
     }
 }
 
-pub(crate) struct GuestTrailerConsumer<T> {
-    pub(crate) tx: Option<oneshot::Sender<Result<Option<Arc<HeaderMap>>, ErrorCode>>>,
-    pub(crate) getter: fn(&mut T) -> WasiHttpCtxView<'_>,
+struct GuestTrailerConsumer<T> {
+    tx: Option<oneshot::Sender<Result<Option<Arc<HeaderMap>>, ErrorCode>>>,
+    getter: fn(&mut T) -> WasiHttpCtxView<'_>,
 }
 
 impl<D> FutureConsumer<D> for GuestTrailerConsumer<D>
@@ -387,12 +387,13 @@ where
         mut self: Pin<&mut Self>,
         _: &mut Context<'_>,
         mut store: StoreContextMut<D>,
-        mut source: Source<'_, Self::Item>,
+        mut src: Source<'_, Self::Item>,
         _: bool,
     ) -> Poll<wasmtime::Result<()>> {
-        let value = &mut None;
-        source.read(store.as_context_mut(), value)?;
-        let res = match value.take().unwrap() {
+        let mut result = None;
+        src.read(store.as_context_mut(), &mut result)
+            .context("failed to read result")?;
+        let res = match result.context("result value missing")? {
             Ok(Some(trailers)) => {
                 let WasiHttpCtxView { table, .. } = (self.getter)(store.data_mut());
                 let trailers = table

crates/wasi-http/src/p3/host/types.rs

@@ -129,34 +129,41 @@ fn parse_header_value(
     }
 }
 
-struct GuestBodyResultProducer(
-    Pin<Box<dyn Future<Output = wasmtime::Result<Result<(), ErrorCode>>> + Send>>,
-);
-
-impl GuestBodyResultProducer {
-    fn new(rx: oneshot::Receiver<Box<dyn Future<Output = Result<(), ErrorCode>> + Send>>) -> Self {
-        Self(Box::pin(async move {
-            let Ok(fut) = rx.await else {
-                return Ok(Ok(()));
-            };
-            Ok(Box::into_pin(fut).await)
-        }))
-    }
+enum GuestBodyResultProducer {
+    Receiver(oneshot::Receiver<Box<dyn Future<Output = Result<(), ErrorCode>> + Send>>),
+    Future(Pin<Box<dyn Future<Output = Result<(), ErrorCode>> + Send>>),
 }
 
 impl<D> FutureProducer<D> for GuestBodyResultProducer {
     type Item = Result<(), ErrorCode>;
 
     fn poll_produce(
-        self: Pin<&mut Self>,
+        mut self: Pin<&mut Self>,
         cx: &mut Context<'_>,
         _: StoreContextMut<D>,
         finish: bool,
     ) -> Poll<wasmtime::Result<Option<Self::Item>>> {
-        match Pin::new(&mut self.get_mut().0).poll(cx) {
-            Poll::Pending if finish => Poll::Ready(Ok(None)),
-            Poll::Pending => Poll::Pending,
-            Poll::Ready(result) => Poll::Ready(Ok(Some(result?))),
+        match &mut *self {
+            Self::Receiver(rx) => match Pin::new(rx).poll(cx) {
+                Poll::Ready(Ok(fut)) => {
+                    let mut fut = Box::into_pin(fut);
+                    match fut.as_mut().poll(cx) {
+                        Poll::Ready(res) => Poll::Ready(Ok(Some(res))),
+                        Poll::Pending => {
+                            *self = Self::Future(fut);
+                            Poll::Pending
+                        }
+                    }
+                }
+                Poll::Ready(Err(..)) => Poll::Ready(Ok(Some(Ok(())))),
+                Poll::Pending if finish => Poll::Ready(Ok(None)),
+                Poll::Pending => Poll::Pending,
+            },
+            Self::Future(fut) => match fut.as_mut().poll(cx) {
+                Poll::Ready(res) => Poll::Ready(Ok(Some(res))),
+                Poll::Pending if finish => Poll::Ready(Ok(None)),
+                Poll::Pending => Poll::Pending,
+            },
         }
     }
 }
@@ -332,7 +339,7 @@ impl HostRequestWithStore for WasiHttp {
                 FutureReader::new(
                     instance,
                     &mut store,
-                    GuestBodyResultProducer::new(result_rx),
+                    GuestBodyResultProducer::Receiver(result_rx),
                 ),
             ))
         })
@@ -609,7 +616,7 @@ impl HostResponseWithStore for WasiHttp {
                 FutureReader::new(
                     instance,
                     &mut store,
-                    GuestBodyResultProducer::new(result_rx),
+                    GuestBodyResultProducer::Receiver(result_rx),
                 ),
             ))
         })

crates/wasi/src/p3/mod.rs

@@ -67,7 +67,12 @@ where
         _: StoreContextMut<D>,
         _: bool,
     ) -> Poll<wasmtime::Result<Option<T>>> {
-        Poll::Ready(Ok(Some(self.get_mut().0.take().unwrap())))
+        let v = self
+            .get_mut()
+            .0
+            .take()
+            .context("polled after returning `Ready`")?;
+        Poll::Ready(Ok(Some(v)))
     }
 }
 
@@ -92,9 +97,10 @@ where
         finish: bool,
     ) -> Poll<wasmtime::Result<Option<T>>> {
         match Pin::new(&mut self.get_mut().0).poll(cx) {
+            Poll::Ready(Ok(v)) => Poll::Ready(Ok(Some(v))),
+            Poll::Ready(Err(err)) => Poll::Ready(Err(err).context("oneshot sender dropped")),
             Poll::Pending if finish => Poll::Ready(Ok(None)),
             Poll::Pending => Poll::Pending,
-            Poll::Ready(result) => Poll::Ready(Ok(Some(result.context("oneshot sender dropped")?))),
         }
     }
 }