Reduce some complexity in wasi-http p3 paths (#13807)

* Use futures when creating a `FutureReader` rather than implementing
  `FutureProducer` directly.
* Use a helper of run-with-callback to avoid manual impls of
  `FutureConsumer`.

Author: alexcrichton

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

@@ -1,5 +1,6 @@
 use crate::FieldMap;
 use crate::p3::bindings::http::types::{ErrorCode, Trailers};
+use crate::p3::helpers::FutureReaderExt;
 use crate::p3::{WasiHttp, WasiHttpCtxView};
 use bytes::Bytes;
 use core::iter;
@@ -13,8 +14,8 @@ use std::sync::Arc;
 use tokio::sync::{mpsc, oneshot};
 use tokio_util::sync::PollSender;
 use wasmtime::component::{
-    Access, Destination, FutureConsumer, FutureReader, Resource, Source, StreamConsumer,
-    StreamProducer, StreamReader, StreamResult,
+    Access, Destination, FutureReader, Resource, Source, StreamConsumer, StreamProducer,
+    StreamReader, StreamResult,
 };
 use wasmtime::error::Context as _;
 use wasmtime::{AsContextMut, StoreContextMut};
@@ -39,33 +40,6 @@ pub(crate) enum Body {
     },
 }
 
-/// [FutureConsumer] implementation for future passed to `consume-body`.
-struct BodyResultConsumer(
-    Option<oneshot::Sender<Box<dyn Future<Output = Result<(), ErrorCode>> + Send>>>,
-);
-
-impl<D> FutureConsumer<D> for BodyResultConsumer
-where
-    D: 'static,
-{
-    type Item = Result<(), ErrorCode>;
-
-    fn poll_consume(
-        mut self: Pin<&mut Self>,
-        _: &mut Context<'_>,
-        store: StoreContextMut<D>,
-        mut src: Source<'_, Self::Item>,
-        _: bool,
-    ) -> Poll<wasmtime::Result<()>> {
-        let mut res = None;
-        src.read(store, &mut res).context("failed to read result")?;
-        let res = res.context("result value missing")?;
-        let tx = self.0.take().context("polled after returning `Ready`")?;
-        _ = tx.send(Box::new(async { res }));
-        Poll::Ready(Ok(()))
-    }
-}
-
 impl Body {
     /// Implementation of `consume-body` shared between requests and responses
     pub(crate) fn consume<T>(
@@ -77,25 +51,22 @@ impl Body {
         StreamReader<u8>,
         FutureReader<Result<Option<Resource<Trailers>>, ErrorCode>>,
     )> {
-        Ok(match self {
+        let (contents_rx, trailers_rx, result_tx) = match self {
             Body::Guest {
                 contents_rx: Some(contents_rx),
                 trailers_rx,
                 result_tx,
-            } => {
-                fut.pipe(&mut store, BodyResultConsumer(Some(result_tx)))?;
-                (contents_rx, trailers_rx)
-            }
+            } => (contents_rx, trailers_rx, result_tx),
             Body::Guest {
                 contents_rx: None,
                 trailers_rx,
                 result_tx,
-            } => {
-                fut.pipe(&mut store, BodyResultConsumer(Some(result_tx)))?;
-                (StreamReader::new(&mut store, iter::empty())?, trailers_rx)
-            }
+            } => (
+                StreamReader::new(&mut store, iter::empty())?,
+                trailers_rx,
+                result_tx,
+            ),
             Body::Host { body, result_tx } => {
-                fut.pipe(&mut store, BodyResultConsumer(Some(result_tx)))?;
                 let (trailers_tx, trailers_rx) = oneshot::channel();
                 (
                     StreamReader::new(
@@ -107,9 +78,16 @@ impl Body {
                         },
                     )?,
                     FutureReader::new(&mut store, trailers_rx)?,
+                    result_tx,
                 )
             }
-        })
+        };
+
+        fut.pipe_cb(&mut store, |_, res| {
+            _ = result_tx.send(Box::new(async { res }));
+            Ok(())
+        })?;
+        Ok((contents_rx, trailers_rx))
     }
 
     /// Implementation of `drop` shared between requests and responses
@@ -275,13 +253,21 @@ impl GuestBody {
         getter: fn(&mut T) -> WasiHttpCtxView<'_>,
     ) -> wasmtime::Result<Self> {
         let (trailers_http_tx, trailers_http_rx) = oneshot::channel();
-        trailers_rx.pipe(
-            &mut store,
-            GuestTrailerConsumer {
-                tx: Some(trailers_http_tx),
-                getter,
-            },
-        )?;
+        trailers_rx.pipe_cb(&mut store, move |data, res| {
+            let res = match res {
+                Ok(Some(trailers)) => {
+                    let WasiHttpCtxView { table, .. } = getter(data);
+                    let trailers = table
+                        .delete(trailers)
+                        .context("failed to delete trailers")?;
+                    Ok(Some(Arc::from(trailers)))
+                }
+                Ok(None) => Ok(None),
+                Err(err) => Err(err),
+            };
+            _ = trailers_http_tx.send(res);
+            Ok(())
+        })?;
 
         let contents_rx = if let Some(rx) = contents_rx {
             let (http_tx, http_rx) = mpsc::channel(1);
@@ -401,44 +387,6 @@ impl http_body::Body for GuestBody {
     }
 }
 
-/// [FutureConsumer] implementation for trailers originating in the guest.
-struct GuestTrailerConsumer<T> {
-    tx: Option<oneshot::Sender<Result<Option<Arc<FieldMap>>, ErrorCode>>>,
-    getter: fn(&mut T) -> WasiHttpCtxView<'_>,
-}
-
-impl<D> FutureConsumer<D> for GuestTrailerConsumer<D>
-where
-    D: 'static,
-{
-    type Item = Result<Option<Resource<Trailers>>, ErrorCode>;
-
-    fn poll_consume(
-        mut self: Pin<&mut Self>,
-        _: &mut Context<'_>,
-        mut store: StoreContextMut<D>,
-        mut src: Source<'_, Self::Item>,
-        _: bool,
-    ) -> Poll<wasmtime::Result<()>> {
-        let mut res = None;
-        src.read(&mut store, &mut res)
-            .context("failed to read result")?;
-        let res = match res.context("result value missing")? {
-            Ok(Some(trailers)) => {
-                let WasiHttpCtxView { table, .. } = (self.getter)(store.data_mut());
-                let trailers = table
-                    .delete(trailers)
-                    .context("failed to delete trailers")?;
-                Ok(Some(Arc::from(trailers)))
-            }
-            Ok(None) => Ok(None),
-            Err(err) => Err(err),
-        };
-        _ = self.tx.take().unwrap().send(res);
-        Poll::Ready(Ok(()))
-    }
-}
-
 /// [StreamProducer] implementation for bodies originating in the host.
 pub(crate) struct HostBodyStreamProducer<T> {
     pub(crate) body: UnsyncBoxBody<Bytes, ErrorCode>,

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

@@ -0,0 +1,70 @@
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use wasmtime::component::{FutureConsumer, FutureReader, Lift, Source};
+use wasmtime::error::Context as _;
+use wasmtime::{AsContextMut, StoreContextMut};
+
+/// Extension methosd for `FutureReader`
+pub trait FutureReaderExt<T> {
+    /// Get the underlying `FutureReader`.
+    fn as_future_reader(self) -> FutureReader<T>;
+
+    /// Run `cb` with the result of this future when it's ready.
+    ///
+    /// The `cb` is given the store's data-at-the-time, the result of the
+    /// future, and can produce a trapping error if so desirable.
+    fn pipe_cb<S>(
+        self,
+        store: S,
+        cb: impl FnOnce(&mut S::Data, T) -> wasmtime::Result<()> + Unpin + Send + 'static,
+    ) -> wasmtime::Result<()>
+    where
+        Self: Sized,
+        S: AsContextMut,
+        T: Lift + 'static,
+    {
+        struct Consumer<F, D, T> {
+            cb: Option<F>,
+            _marker: std::marker::PhantomData<fn(D, T)>,
+        }
+
+        impl<T, D, F> FutureConsumer<D> for Consumer<F, D, T>
+        where
+            T: Lift + 'static,
+            F: FnOnce(&mut D, T) -> wasmtime::Result<()> + Send + Unpin + 'static,
+            D: 'static,
+        {
+            type Item = T;
+
+            fn poll_consume(
+                mut self: Pin<&mut Self>,
+                _: &mut Context<'_>,
+                mut store: StoreContextMut<D>,
+                mut src: Source<'_, Self::Item>,
+                _: bool,
+            ) -> Poll<wasmtime::Result<()>> {
+                let mut res = None;
+                src.read(&mut store, &mut res)
+                    .context("failed to read result")?;
+                let res = res.context("result value missing")?;
+                let cb = self.cb.take().context("polled after returning `Ready`")?;
+                cb(store.data_mut(), res)?;
+                Poll::Ready(Ok(()))
+            }
+        }
+
+        self.as_future_reader().pipe(
+            store,
+            Consumer {
+                cb: Some(cb),
+                _marker: std::marker::PhantomData,
+            },
+        )
+    }
+}
+
+impl<T> FutureReaderExt<T> for FutureReader<T> {
+    fn as_future_reader(self) -> FutureReader<T> {
+        self
+    }
+}

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

@@ -9,15 +9,12 @@ use crate::p3::body::{Body, HostBodyStreamProducer};
 use crate::p3::{HeaderResult, HttpError, RequestOptionsResult, WasiHttp, WasiHttpCtxView};
 use core::mem;
 use core::pin::Pin;
-use core::task::{Context, Poll, ready};
 use http::header::CONTENT_LENGTH;
 use std::sync::Arc;
 use tokio::sync::oneshot;
-use wasmtime::component::{
-    Access, FutureProducer, FutureReader, Resource, ResourceTable, StreamReader,
-};
+use wasmtime::AsContextMut;
+use wasmtime::component::{Access, FutureReader, Resource, ResourceTable, StreamReader};
 use wasmtime::error::Context as _;
-use wasmtime::{AsContextMut, StoreContextMut};
 
 fn get_fields<'a>(
     table: &'a ResourceTable,
@@ -161,51 +158,13 @@ fn parse_header_value(
     }
 }
 
-enum GuestBodyResultProducer {
-    Receiver(oneshot::Receiver<Box<dyn Future<Output = Result<(), ErrorCode>> + Send>>),
-    Future(Pin<Box<dyn Future<Output = Result<(), ErrorCode>> + Send>>),
-}
-
-fn poll_future<T>(
-    cx: &mut Context<'_>,
-    fut: Pin<&mut (impl Future<Output = T> + ?Sized)>,
-    finish: bool,
-) -> Poll<Option<T>> {
-    match fut.poll(cx) {
-        Poll::Ready(v) => Poll::Ready(Some(v)),
-        Poll::Pending if finish => Poll::Ready(None),
-        Poll::Pending => Poll::Pending,
-    }
-}
-
-impl<D> FutureProducer<D> for GuestBodyResultProducer {
-    type Item = Result<(), ErrorCode>;
-
-    fn poll_produce(
-        mut self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-        _: StoreContextMut<D>,
-        finish: bool,
-    ) -> Poll<wasmtime::Result<Option<Self::Item>>> {
-        match &mut *self {
-            Self::Receiver(rx) => {
-                match ready!(poll_future(cx, Pin::new(rx), finish)) {
-                    Some(Ok(fut)) => {
-                        let mut fut = Box::into_pin(fut);
-                        // poll the received future once and update state
-                        let res = poll_future(cx, fut.as_mut(), finish);
-                        *self = Self::Future(fut);
-                        res.map(Ok)
-                    }
-                    Some(Err(..)) => {
-                        // oneshot sender dropped, treat as success
-                        Poll::Ready(Ok(Some(Ok(()))))
-                    }
-                    None => Poll::Ready(Ok(None)),
-                }
-            }
-            Self::Future(fut) => poll_future(cx, fut.as_mut(), finish).map(Ok),
-        }
+async fn guest_body_result(
+    rx: oneshot::Receiver<Box<dyn Future<Output = Result<(), ErrorCode>> + Send>>,
+) -> wasmtime::Result<Result<(), ErrorCode>> {
+    match rx.await {
+        Ok(fut) => Ok(Pin::from(fut).await),
+        // oneshot sender dropped, treat as success
+        Err(..) => Ok(Ok(())),
     }
 }
 
@@ -375,7 +334,7 @@ impl<T> HostRequestWithStore<T> for WasiHttp {
         let req = table.push(req).context("failed to push request to table")?;
         Ok((
             req,
-            FutureReader::new(&mut store, GuestBodyResultProducer::Receiver(result_rx))?,
+            FutureReader::new(&mut store, guest_body_result(result_rx))?,
         ))
     }
 
@@ -648,7 +607,7 @@ impl<T> HostResponseWithStore<T> for WasiHttp {
             .context("failed to push response to table")?;
         Ok((
             res,
-            FutureReader::new(&mut store, GuestBodyResultProducer::Receiver(result_rx))?,
+            FutureReader::new(&mut store, guest_body_result(result_rx))?,
         ))
     }
 

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

@@ -11,6 +11,7 @@
 pub mod bindings;
 mod body;
 mod conv;
+mod helpers;
 mod host;
 mod proxy;
 mod request;