wasmtime_wasi_http/p3/host/
handler.rs

1use crate::FieldMap;
2use crate::p3::bindings::http::client::{Host, HostWithStore};
3use crate::p3::bindings::http::types::{Request, Response};
4use crate::p3::body::{Body, BodyExt as _};
5use crate::p3::{HttpError, HttpResult};
6use crate::{Error, WasiHttp, WasiHttpCtxView};
7use core::task::{Context, Poll, Waker};
8use http_body_util::BodyExt as _;
9use std::sync::Arc;
10use tokio::sync::oneshot;
11use tokio::task::{self, JoinHandle};
12use tracing::debug;
13use wasmtime::component::{Accessor, Resource};
14use wasmtime::error::Context as _;
15
16/// A wrapper around [`JoinHandle`], which will [`JoinHandle::abort`] the task
17/// when dropped
18struct AbortOnDropJoinHandle(JoinHandle<()>);
19
20impl Drop for AbortOnDropJoinHandle {
21    fn drop(&mut self) {
22        self.0.abort();
23    }
24}
25
26async fn io_task_result(
27    rx: oneshot::Receiver<(
28        Arc<AbortOnDropJoinHandle>,
29        oneshot::Receiver<Result<(), Error>>,
30    )>,
31) -> Result<(), Error> {
32    let Ok((_io, io_result_rx)) = rx.await else {
33        return Ok(());
34    };
35    io_result_rx.await.unwrap_or(Ok(()))
36}
37
38impl<T> HostWithStore<T> for WasiHttp {
39    async fn send(
40        store: &Accessor<T, Self>,
41        req: Resource<Request>,
42    ) -> HttpResult<Resource<Response>> {
43        // A handle to the I/O task, if spawned, will be sent on this channel
44        // and kept as part of request body state
45        let (io_task_tx, io_task_rx) = oneshot::channel();
46
47        // A handle to the I/O task, if spawned, will be sent on this channel
48        // along with the result receiver
49        let (io_result_tx, io_result_rx) = oneshot::channel();
50
51        // Response processing result will be sent on this channel
52        let (res_result_tx, res_result_rx) = oneshot::channel();
53
54        let getter = store.getter();
55        let fut = store.with(|mut store| {
56            let WasiHttpCtxView { table, .. } = store.get();
57            let req = table
58                .delete(req)
59                .context("failed to delete request from table")
60                .map_err(HttpError::trap)?;
61            let (req, options) =
62                req.into_http_with_getter(&mut store, io_task_result(io_result_rx), getter)?;
63            HttpResult::Ok(store.get().hooks.send_request(
64                req.map(|body| body.with_state(io_task_rx).boxed_unsync()),
65                options.as_deref().copied(),
66                Box::new(async {
67                    // Forward the response processing result to `WasiHttpCtx` implementation
68                    let Ok(fut) = res_result_rx.await else {
69                        return Ok(());
70                    };
71                    Box::into_pin(fut).await
72                }),
73            ))
74        })?;
75        let (res, io) = Box::into_pin(fut)
76            .await
77            .map_err(|e| store.with(|mut store| store.get().error_to_p3(&e)))?;
78        let (
79            http::response::Parts {
80                status, headers, ..
81            },
82            body,
83        ) = res.into_parts();
84
85        let mut io = Box::into_pin(io);
86        let body = match io.as_mut().poll(&mut Context::from_waker(Waker::noop())) {
87            Poll::Ready(Ok(())) => body,
88            Poll::Ready(Err(e)) => {
89                return Err(store.with(|mut store| store.get().error_to_p3(&e)).into());
90            }
91            Poll::Pending => {
92                // I/O driver still needs to be polled, spawn a task and send handles to it
93                let (tx, rx) = oneshot::channel();
94                let io = task::spawn(async move {
95                    let res = io.await;
96                    debug!(?res, "`send_request` I/O future finished");
97                    _ = tx.send(res);
98                });
99                let io = Arc::new(AbortOnDropJoinHandle(io));
100                _ = io_result_tx.send((Arc::clone(&io), rx));
101                _ = io_task_tx.send(Arc::clone(&io));
102                body.with_state(io).boxed_unsync()
103            }
104        };
105        store.with(|mut store| {
106            let res = Response {
107                status,
108                headers: FieldMap::new_immutable(store.get().hooks, headers),
109                body: Body::Host {
110                    body,
111                    result_tx: res_result_tx,
112                },
113            };
114            store
115                .get()
116                .table
117                .push(res)
118                .context("failed to push response to table")
119                .map_err(HttpError::trap)
120        })
121    }
122}
123
124impl Host for WasiHttpCtxView<'_> {}
wasmtime_wasi_http/p3/host/handler.rs

wasmtime_wasi_http/p3/host/
handler.rs
