Skip to content

Commit 6a64085

Browse files
adamrkadamrk
authored
wasi-http: Restore shared ownership of I/O task (#14229)
* add regression test * Revert "Remove `BodyWithState`" This reverts commit e13936a.
1 parent e8890cc commit 6a64085

4 files changed

Lines changed: 131 additions & 9 deletions

File tree

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

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
use anyhow::Context as _;
2+
use futures::join;
3+
use test_programs::p3::wasi::http::client;
4+
use test_programs::p3::wasi::http::types::{
5+
Headers, Method, Request, RequestOptions, Response, Scheme,
6+
};
7+
use test_programs::p3::{wit_future, wit_stream};
8+
9+
struct Component;
10+
11+
test_programs::p3::export!(Component);
12+
13+
impl test_programs::p3::exports::wasi::cli::run::Guest for Component {
14+
async fn run() -> Result<(), ()> {
15+
const LEN: usize = 1 << 20;
16+
let body = vec![1; LEN];
17+
18+
let addr = test_programs::p3::wasi::cli::environment::get_environment()
19+
.into_iter()
20+
.find_map(|(k, v)| k.eq("HTTP_SERVER").then_some(v))
21+
.unwrap();
22+
23+
let headers = Headers::from_list(&[]).unwrap();
24+
let (mut contents_tx, contents_rx) = wit_stream::new();
25+
let (trailers_tx, trailers_rx) = wit_future::new(|| Ok(None));
26+
drop(trailers_tx);
27+
let options = RequestOptions::new();
28+
let (request, transmit) =
29+
Request::new(headers, Some(contents_rx), trailers_rx, Some(options));
30+
request.set_method(&Method::Post).unwrap();
31+
request.set_scheme(Some(&Scheme::Http)).unwrap();
32+
request.set_authority(Some(&addr)).unwrap();
33+
request.set_path_with_query(Some("/post")).unwrap();
34+
35+
drop(transmit);
36+
37+
let ((), echoed) = join!(
38+
async {
39+
let remaining = contents_tx.write_all((&body[..]).into()).await;
40+
assert!(remaining.is_empty());
41+
drop(contents_tx);
42+
},
43+
async {
44+
let response = client::send(request).await.context("send failed").unwrap();
45+
let status = response.get_status_code();
46+
assert_eq!(status, 200);
47+
let (_, result_rx) = wit_future::new(|| Ok(()));
48+
let (body_rx, _trailers_rx) = Response::consume_body(response, result_rx);
49+
body_rx.collect().await
50+
},
51+
);
52+
53+
assert_eq!(
54+
echoed.len(),
55+
LEN,
56+
"response body was truncated after dropping the transmit future"
57+
);
58+
Ok(())
59+
}
60+
}
61+
62+
fn main() {}

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

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -585,6 +585,39 @@ where
585585
}
586586
}
587587

588+
/// A wrapper around [http_body::Body], which allows attaching arbitrary state to it
589+
pub(crate) struct BodyWithState<T, U> {
590+
body: T,
591+
_state: U,
592+
}
593+
594+
impl<T, U> http_body::Body for BodyWithState<T, U>
595+
where
596+
T: http_body::Body + Unpin,
597+
U: Unpin,
598+
{
599+
type Data = T::Data;
600+
type Error = T::Error;
601+
602+
#[inline]
603+
fn poll_frame(
604+
self: Pin<&mut Self>,
605+
cx: &mut Context<'_>,
606+
) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
607+
Pin::new(&mut self.get_mut().body).poll_frame(cx)
608+
}
609+
610+
#[inline]
611+
fn is_end_stream(&self) -> bool {
612+
self.body.is_end_stream()
613+
}
614+
615+
#[inline]
616+
fn size_hint(&self) -> http_body::SizeHint {
617+
self.body.size_hint()
618+
}
619+
}
620+
588621
/// A wrapper around [http_body::Body], which validates `Content-Length`
589622
pub(crate) struct BodyWithContentLength<T, E> {
590623
body: T,
@@ -667,6 +700,16 @@ where
667700
}
668701

669702
pub(crate) trait BodyExt {
703+
fn with_state<T>(self, state: T) -> BodyWithState<Self, T>
704+
where
705+
Self: Sized,
706+
{
707+
BodyWithState {
708+
body: self,
709+
_state: state,
710+
}
711+
}
712+
670713
fn with_content_length<E>(
671714
self,
672715
limit: u64,

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

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
use crate::FieldMap;
22
use crate::p3::bindings::http::client::{Host, HostWithStore};
33
use crate::p3::bindings::http::types::{Request, Response};
4-
use crate::p3::body::Body;
4+
use crate::p3::body::{Body, BodyExt as _};
55
use crate::p3::{HttpError, HttpResult};
66
use crate::{Error, WasiHttp, WasiHttpCtxView};
77
use core::task::{Context, Poll, Waker};
8+
use http_body_util::BodyExt as _;
9+
use std::sync::Arc;
810
use tokio::sync::oneshot;
911
use tokio::task::{self, JoinHandle};
1012
use tracing::debug;
@@ -26,7 +28,7 @@ const DROPPED_FUTURE_ERROR: &str =
2628

2729
async fn io_task_result(
2830
rx: oneshot::Receiver<(
29-
Option<AbortOnDropJoinHandle>,
31+
Option<Arc<AbortOnDropJoinHandle>>,
3032
oneshot::Receiver<Result<(), Error>>,
3133
)>,
3234
) -> Result<(), Error> {
@@ -41,7 +43,7 @@ async fn io_task_result(
4143
fn send_dummy_io(
4244
result: Result<(), Error>,
4345
io_result_tx: oneshot::Sender<(
44-
Option<AbortOnDropJoinHandle>,
46+
Option<Arc<AbortOnDropJoinHandle>>,
4547
oneshot::Receiver<Result<(), Error>>,
4648
)>,
4749
) {
@@ -54,7 +56,7 @@ fn send_dummy_io_err<T>(
5456
store: &Accessor<T, WasiHttp>,
5557
e: Error,
5658
io_result_tx: oneshot::Sender<(
57-
Option<AbortOnDropJoinHandle>,
59+
Option<Arc<AbortOnDropJoinHandle>>,
5860
oneshot::Receiver<Result<(), Error>>,
5961
)>,
6062
) -> HttpError {
@@ -68,6 +70,10 @@ impl<T> HostWithStore<T> for WasiHttp {
6870
store: &Accessor<T, Self>,
6971
req: Resource<Request>,
7072
) -> HttpResult<Resource<Response>> {
73+
// A handle to the I/O task, if spawned, will be sent on this channel
74+
// and kept as part of request body state
75+
let (io_task_tx, io_task_rx) = oneshot::channel();
76+
7177
// A handle to the I/O task, if spawned, will be sent on this channel
7278
// along with the result receiver
7379
let (io_result_tx, io_result_rx) = oneshot::channel();
@@ -85,7 +91,9 @@ impl<T> HostWithStore<T> for WasiHttp {
8591
let (req, options) =
8692
req.into_http_with_getter(&mut store, io_task_result(io_result_rx), getter)?;
8793
HttpResult::Ok(store.get().hooks.send_request(
88-
req,
94+
// Attach a reference to the io task to the body so that it
95+
// isn't cancelled if the body is dropped.
96+
req.map(|body| body.with_state(io_task_rx).boxed_unsync()),
8997
options.as_deref().copied(),
9098
Box::new(async {
9199
// Forward the response processing result to `WasiHttpCtx` implementation
@@ -134,13 +142,16 @@ impl<T> HostWithStore<T> for WasiHttp {
134142
Poll::Pending => {
135143
// I/O driver still needs to be polled, spawn a task and send handles to it
136144
let (tx, rx) = oneshot::channel();
137-
let io = AbortOnDropJoinHandle(task::spawn(async move {
145+
let io = Arc::new(AbortOnDropJoinHandle(task::spawn(async move {
138146
let res = io.await;
139147
debug!(?res, "`send_request` I/O future finished");
140148
_ = tx.send(res);
141-
}));
142-
_ = io_result_tx.send((Some(io), rx));
143-
body
149+
})));
150+
_ = io_result_tx.send((Some(Arc::clone(&io)), rx));
151+
_ = io_task_tx.send(Arc::clone(&io));
152+
// Attach a reference to the io task to the body so that it
153+
// isn't cancelled if the body is dropped.
154+
body.with_state(io).boxed_unsync()
144155
}
145156
};
146157
store.with(|mut store| {

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -946,3 +946,9 @@ async fn p3_http_outbound_request_chunk_size() -> Result<()> {
946946
let server = Server::http1(1)?;
947947
run_cli(P3_HTTP_OUTBOUND_REQUEST_CHUNK_SIZE_COMPONENT, &server).await
948948
}
949+
950+
#[test_log::test(tokio::test(flavor = "multi_thread"))]
951+
async fn p3_http_drop_transmit() -> Result<()> {
952+
let server = Server::http1(1)?;
953+
run_cli(P3_HTTP_DROP_TRANSMIT_COMPONENT, &server).await
954+
}

0 commit comments

Comments
 (0)