p3-http: tie lifetime of the spawned task to the bodies

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

Author: rvolosatovs

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

@@ -4,13 +4,71 @@ use crate::p3::body::{Body, ConsumedBody, GuestBody, GuestBodyKind};
 use crate::p3::{HttpError, HttpResult, WasiHttp, WasiHttpCtxView, get_content_length};
 use anyhow::Context as _;
 use core::pin::Pin;
+use core::task::{Context, Poll, Waker};
 use http::header::HOST;
 use http::{HeaderValue, Uri};
 use http_body_util::BodyExt as _;
 use std::sync::Arc;
 use tokio::sync::oneshot;
 use tracing::debug;
-use wasmtime::component::{Accessor, AccessorTask, Resource};
+use wasmtime::component::{Accessor, AccessorTask, JoinHandle, Resource};
+
+/// A wrapper around [`JoinHandle`], which will [`JoinHandle::abort`] the task
+/// when dropped
+struct AbortOnDropJoinHandle(JoinHandle);
+
+impl Drop for AbortOnDropJoinHandle {
+    fn drop(&mut self) {
+        self.0.abort();
+    }
+}
+
+/// A wrapper around [http_body::Body], which allows attaching arbitrary state to it
+struct BodyWithState<T, U> {
+    body: T,
+    _state: U,
+}
+
+impl<T, U> http_body::Body for BodyWithState<T, U>
+where
+    T: http_body::Body + Unpin,
+    U: Unpin,
+{
+    type Data = T::Data;
+    type Error = T::Error;
+
+    #[inline]
+    fn poll_frame(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
+        Pin::new(&mut self.get_mut().body).poll_frame(cx)
+    }
+
+    #[inline]
+    fn is_end_stream(&self) -> bool {
+        self.body.is_end_stream()
+    }
+
+    #[inline]
+    fn size_hint(&self) -> http_body::SizeHint {
+        self.body.size_hint()
+    }
+}
+
+trait BodyExt {
+    fn with_state<T>(self, state: T) -> BodyWithState<Self, T>
+    where
+        Self: Sized,
+    {
+        BodyWithState {
+            body: self,
+            _state: state,
+        }
+    }
+}
+
+impl<T> BodyExt for T {}
 
 struct SendRequestTask {
     io: Pin<Box<dyn Future<Output = Result<(), ErrorCode>> + Send>>,
@@ -26,14 +84,35 @@ impl<T> AccessorTask<T, WasiHttp, wasmtime::Result<()>> for SendRequestTask {
     }
 }
 
+async fn io_task_result(
+    rx: oneshot::Receiver<(
+        Arc<AbortOnDropJoinHandle>,
+        oneshot::Receiver<Result<(), ErrorCode>>,
+    )>,
+) -> Result<(), ErrorCode> {
+    let Ok((_io, io_result_rx)) = rx.await else {
+        return Ok(());
+    };
+    io_result_rx.await.unwrap_or(Ok(()))
+}
+
 impl HostWithStore for WasiHttp {
     async fn handle<T>(
         store: &Accessor<T, Self>,
         req: Resource<Request>,
     ) -> HttpResult<Resource<Response>> {
-        let getter = store.getter();
+        // A handle to the I/O task, if spawned, will be sent on this channel
+        // and kept as part of request body state
+        let (io_task_tx, io_task_rx) = oneshot::channel();
+
+        // A handle to the I/O task and, if spawned, will be sent on this channel
+        // along with the result receiver
         let (io_result_tx, io_result_rx) = oneshot::channel();
+
+        // Response processing result will be sent on this channel
         let (res_result_tx, res_result_rx) = oneshot::channel();
+
+        let getter = store.getter();
         let fut = store.with(|mut store| {
             let WasiHttpCtxView { table, .. } = store.get();
             let Request {
@@ -62,7 +141,7 @@ impl HostWithStore for WasiHttp {
                         if let Ok(Err(err)) = http_result_rx.await {
                             return Err(err);
                         };
-                        io_result_rx.await.unwrap_or(Ok(()))
+                        io_task_result(io_result_rx).await
                     }));
                     GuestBody::new(
                         &mut store,
@@ -73,13 +152,12 @@ impl HostWithStore for WasiHttp {
                         GuestBodyKind::Request,
                         getter,
                     )
+                    .with_state(io_task_rx)
                     .boxed()
                 }
                 Body::Host { body, result_tx } => {
-                    _ = result_tx.send(Box::new(
-                        async move { io_result_rx.await.unwrap_or(Ok(())) },
-                    ));
-                    body
+                    _ = result_tx.send(Box::new(io_task_result(io_result_rx)));
+                    body.with_state(io_task_rx).boxed()
                 }
                 Body::Consumed => ConsumedBody.boxed(),
             };
@@ -129,16 +207,26 @@ impl HostWithStore for WasiHttp {
             ))
         })?;
         let (res, io) = Box::into_pin(fut).await?;
-        store.spawn(SendRequestTask {
-            io: Box::into_pin(io),
-            result_tx: io_result_tx,
-        });
         let (
             http::response::Parts {
                 status, headers, ..
             },
             body,
         ) = res.into_parts();
+
+        let mut io = Box::into_pin(io);
+        let body = match io.as_mut().poll(&mut Context::from_waker(Waker::noop()))? {
+            Poll::Ready(()) => body,
+            Poll::Pending => {
+                // I/O driver still needs to be polled, spawn a task and send handles to it
+                let (tx, rx) = oneshot::channel();
+                let io = store.spawn(SendRequestTask { io, result_tx: tx });
+                let io = Arc::new(AbortOnDropJoinHandle(io));
+                _ = io_result_tx.send((Arc::clone(&io), rx));
+                _ = io_task_tx.send(Arc::clone(&io));
+                body.with_state(io).boxed()
+            }
+        };
         let res = Response {
             status,
             headers: Arc::new(headers),