wasi-http: Restore shared ownership of I/O task

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

@@ -0,0 +1,62 @@
+use anyhow::Context as _;
+use futures::join;
+use test_programs::p3::wasi::http::client;
+use test_programs::p3::wasi::http::types::{
+    Headers, Method, Request, RequestOptions, Response, Scheme,
+};
+use test_programs::p3::{wit_future, wit_stream};
+
+struct Component;
+
+test_programs::p3::export!(Component);
+
+impl test_programs::p3::exports::wasi::cli::run::Guest for Component {
+    async fn run() -> Result<(), ()> {
+        const LEN: usize = 1 << 20;
+        let body = vec![1; LEN];
+
+        let addr = test_programs::p3::wasi::cli::environment::get_environment()
+            .into_iter()
+            .find_map(|(k, v)| k.eq("HTTP_SERVER").then_some(v))
+            .unwrap();
+
+        let headers = Headers::from_list(&[]).unwrap();
+        let (mut contents_tx, contents_rx) = wit_stream::new();
+        let (trailers_tx, trailers_rx) = wit_future::new(|| Ok(None));
+        drop(trailers_tx);
+        let options = RequestOptions::new();
+        let (request, transmit) =
+            Request::new(headers, Some(contents_rx), trailers_rx, Some(options));
+        request.set_method(&Method::Post).unwrap();
+        request.set_scheme(Some(&Scheme::Http)).unwrap();
+        request.set_authority(Some(&addr)).unwrap();
+        request.set_path_with_query(Some("/post")).unwrap();
+
+        drop(transmit);
+
+        let ((), echoed) = join!(
+            async {
+                let remaining = contents_tx.write_all((&body[..]).into()).await;
+                assert!(remaining.is_empty());
+                drop(contents_tx);
+            },
+            async {
+                let response = client::send(request).await.context("send failed").unwrap();
+                let status = response.get_status_code();
+                assert_eq!(status, 200);
+                let (_, result_rx) = wit_future::new(|| Ok(()));
+                let (body_rx, _trailers_rx) = Response::consume_body(response, result_rx);
+                body_rx.collect().await
+            },
+        );
+
+        assert_eq!(
+            echoed.len(),
+            LEN,
+            "response body was truncated after dropping the transmit future"
+        );
+        Ok(())
+    }
+}
+
+fn main() {}

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

@@ -585,6 +585,39 @@ where
     }
 }
 
+/// A wrapper around [http_body::Body], which allows attaching arbitrary state to it
+pub(crate) 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()
+    }
+}
+
 /// A wrapper around [http_body::Body], which validates `Content-Length`
 pub(crate) struct BodyWithContentLength<T, E> {
     body: T,
@@ -667,6 +700,16 @@ where
 }
 
 pub(crate) trait BodyExt {
+    fn with_state<T>(self, state: T) -> BodyWithState<Self, T>
+    where
+        Self: Sized,
+    {
+        BodyWithState {
+            body: self,
+            _state: state,
+        }
+    }
+
     fn with_content_length<E>(
         self,
         limit: u64,

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

@@ -1,10 +1,12 @@
 use crate::FieldMap;
 use crate::p3::bindings::http::client::{Host, HostWithStore};
 use crate::p3::bindings::http::types::{Request, Response};
-use crate::p3::body::Body;
+use crate::p3::body::{Body, BodyExt as _};
 use crate::p3::{HttpError, HttpResult};
 use crate::{Error, WasiHttp, WasiHttpCtxView};
 use core::task::{Context, Poll, Waker};
+use http_body_util::BodyExt as _;
+use std::sync::Arc;
 use tokio::sync::oneshot;
 use tokio::task::{self, JoinHandle};
 use tracing::debug;
@@ -26,7 +28,7 @@ const DROPPED_FUTURE_ERROR: &str =
 
 async fn io_task_result(
     rx: oneshot::Receiver<(
-        Option<AbortOnDropJoinHandle>,
+        Option<Arc<AbortOnDropJoinHandle>>,
         oneshot::Receiver<Result<(), Error>>,
     )>,
 ) -> Result<(), Error> {
@@ -41,7 +43,7 @@ async fn io_task_result(
 fn send_dummy_io(
     result: Result<(), Error>,
     io_result_tx: oneshot::Sender<(
-        Option<AbortOnDropJoinHandle>,
+        Option<Arc<AbortOnDropJoinHandle>>,
         oneshot::Receiver<Result<(), Error>>,
     )>,
 ) {
@@ -54,7 +56,7 @@ fn send_dummy_io_err<T>(
     store: &Accessor<T, WasiHttp>,
     e: Error,
     io_result_tx: oneshot::Sender<(
-        Option<AbortOnDropJoinHandle>,
+        Option<Arc<AbortOnDropJoinHandle>>,
         oneshot::Receiver<Result<(), Error>>,
     )>,
 ) -> HttpError {
@@ -68,6 +70,10 @@ impl<T> HostWithStore<T> for WasiHttp {
         store: &Accessor<T, Self>,
         req: Resource<Request>,
     ) -> HttpResult<Resource<Response>> {
+        // 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, if spawned, will be sent on this channel
         // along with the result receiver
         let (io_result_tx, io_result_rx) = oneshot::channel();
@@ -85,7 +91,9 @@ impl<T> HostWithStore<T> for WasiHttp {
             let (req, options) =
                 req.into_http_with_getter(&mut store, io_task_result(io_result_rx), getter)?;
             HttpResult::Ok(store.get().hooks.send_request(
-                req,
+                // Attach a reference to the io task to the body so that it
+                // isn't cancelled if the body is dropped.
+                req.map(|body| body.with_state(io_task_rx).boxed_unsync()),
                 options.as_deref().copied(),
                 Box::new(async {
                     // Forward the response processing result to `WasiHttpCtx` implementation
@@ -134,13 +142,16 @@ impl<T> HostWithStore<T> for WasiHttp {
             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 = AbortOnDropJoinHandle(task::spawn(async move {
+                let io = Arc::new(AbortOnDropJoinHandle(task::spawn(async move {
                     let res = io.await;
                     debug!(?res, "`send_request` I/O future finished");
                     _ = tx.send(res);
-                }));
-                _ = io_result_tx.send((Some(io), rx));
-                body
+                })));
+                _ = io_result_tx.send((Some(Arc::clone(&io)), rx));
+                _ = io_task_tx.send(Arc::clone(&io));
+                // Attach a reference to the io task to the body so that it
+                // isn't cancelled if the body is dropped.
+                body.with_state(io).boxed_unsync()
             }
         };
         store.with(|mut store| {

crates/wasi-http/tests/all/p3/mod.rs

@@ -946,3 +946,9 @@ async fn p3_http_outbound_request_chunk_size() -> Result<()> {
     let server = Server::http1(1)?;
     run_cli(P3_HTTP_OUTBOUND_REQUEST_CHUNK_SIZE_COMPONENT, &server).await
 }
+
+#[test_log::test(tokio::test(flavor = "multi_thread"))]
+async fn p3_http_drop_transmit() -> Result<()> {
+    let server = Server::http1(1)?;
+    run_cli(P3_HTTP_DROP_TRANSMIT_COMPONENT, &server).await
+}