Add a start_request function for requests with streaming output bodies.

Author: sunfishcode

src/http/client.rs

@@ -1,9 +1,16 @@
-use super::{body::IncomingBody, Body, Error, Request, Response, Result};
+use super::{
+    body::{BodyForthcoming, IncomingBody, OutgoingBody},
+    fields::header_map_to_wasi,
+    Body, Error, HeaderMap, Request, Response, Result,
+};
 use crate::http::request::try_into_outgoing;
 use crate::http::response::try_from_incoming;
 use crate::io::{self, AsyncOutputStream, AsyncPollable};
 use crate::time::Duration;
-use wasi::http::types::{OutgoingBody, RequestOptions as WasiRequestOptions};
+use wasi::http::types::{
+    FutureIncomingResponse as WasiFutureIncomingResponse, OutgoingBody as WasiOutgoingBody,
+    RequestOptions as WasiRequestOptions,
+};
 
 /// An HTTP client.
 // Empty for now, but permits adding support for RequestOptions soon:
@@ -19,22 +26,27 @@ impl Client {
     }
 
     /// Send an HTTP request.
+    ///
+    /// TODO: Should this automatically add a "Content-Length" header if the
+    /// body size is known?
+    ///
+    /// To respond with trailers, use [`Client::start_request`] instead.
     pub async fn send<B: Body>(&self, req: Request<B>) -> Result<Response<IncomingBody>> {
         // We don't use `body::OutputBody` here because we can report I/O
         // errors from the `copy` directly.
         let (wasi_req, body) = try_into_outgoing(req)?;
         let wasi_body = wasi_req.body().unwrap();
-        let body_stream = wasi_body.write().unwrap();
+        let wasi_stream = wasi_body.write().unwrap();
 
         // 1. Start sending the request head
         let res = wasi::http::outgoing_handler::handle(wasi_req, self.wasi_options()?).unwrap();
 
         // 2. Start sending the request body
-        io::copy(body, AsyncOutputStream::new(body_stream)).await?;
+        io::copy(body, AsyncOutputStream::new(wasi_stream)).await?;
 
         // 3. Finish sending the request body
         let trailers = None;
-        OutgoingBody::finish(wasi_body, trailers).unwrap();
+        WasiOutgoingBody::finish(wasi_body, trailers).unwrap();
 
         // 4. Receive the response
         AsyncPollable::new(res.subscribe()).wait_for().await;
@@ -46,6 +58,55 @@ impl Client {
         try_from_incoming(res)
     }
 
+    /// Start sending an HTTP request, and return an `OutgoingBody` stream to
+    /// write the body to.
+    ///
+    /// The returned `OutgoingBody` must be consumed by [`Client::finish`] or
+    /// [`Client::fail`].
+    pub async fn start_request(
+        &self,
+        req: Request<BodyForthcoming>,
+    ) -> Result<(OutgoingBody, FutureIncomingResponse)> {
+        let (wasi_req, _body_forthcoming) = try_into_outgoing(req)?;
+        let wasi_body = wasi_req.body().unwrap();
+        let wasi_stream = wasi_body.write().unwrap();
+
+        // Start sending the request head.
+        let res = wasi::http::outgoing_handler::handle(wasi_req, self.wasi_options()?).unwrap();
+
+        let outgoing_body = OutgoingBody::new(AsyncOutputStream::new(wasi_stream), wasi_body);
+
+        Ok((outgoing_body, FutureIncomingResponse(res)))
+    }
+
+    /// Finish the body, optionally with trailers.
+    ///
+    /// This is used with [`Client::start_request`].
+    pub fn finish(body: OutgoingBody, trailers: Option<HeaderMap>) -> Result<()> {
+        let (stream, body) = body.consume();
+
+        // The stream is a child resource of the `OutgoingBody`, so ensure that
+        // it's dropped first.
+        drop(stream);
+
+        let wasi_trailers = match trailers {
+            Some(trailers) => Some(header_map_to_wasi(&trailers)?),
+            None => None,
+        };
+
+        wasi::http::types::OutgoingBody::finish(body, wasi_trailers)
+            .expect("body length did not match Content-Length header value");
+        Ok(())
+    }
+
+    /// Consume the `OutgoingBody` and indicate that the body was not
+    /// completed.
+    ///
+    /// This is used with [`Client::start_request`].
+    pub fn fail(body: OutgoingBody) {
+        let (_stream, _body) = body.consume();
+    }
+
     /// Set timeout on connecting to HTTP server
     pub fn set_connect_timeout(&mut self, d: impl Into<Duration>) {
         self.options_mut().connect_timeout = Some(d.into());
@@ -76,6 +137,25 @@ impl Client {
     }
 }
 
+/// Returned from [`Client::start_request`], this represents a handle to a
+/// response which has not arrived yet. Call [`FutureIncomingResponse::get`]
+/// to wait for the response.
+pub struct FutureIncomingResponse(WasiFutureIncomingResponse);
+
+impl FutureIncomingResponse {
+    /// Consume this `FutureIncomingResponse`, wait, and return the `Response`.
+    pub async fn get(self) -> Result<Response<IncomingBody>> {
+        // Wait for the response.
+        AsyncPollable::new(self.0.subscribe()).wait_for().await;
+
+        // NOTE: the first `unwrap` is to ensure readiness, the second `unwrap`
+        // is to trap if we try and get the response more than once. The final
+        // `?` is to raise the actual error if there is one.
+        let res = self.0.get().unwrap().unwrap()?;
+        try_from_incoming(res)
+    }
+}
+
 #[derive(Default, Debug)]
 struct RequestOptions {
     connect_timeout: Option<Duration>,